#! /bin/bash
# precise shaper script v0.02 (c) Dzianis Kahanovich, 2007, GPLv2
# This script are channel shaper for Linux router box with 2 network cards
# and PSPacer (http://www.gridmpi.org/gridtcp.en.jsp)
# I am repeat ~99% of uplink in/out speeds with precise pacing and queue to
# avoid packet dropping on busy channels.
#
# Required: sch_psp and U32 (w/hashing to speed-up lease non-shaped traffic).
# Additional - see options.
#
# and IGNORE some errors ("file exists") in advanced modes

x=${LATENCY:=900}
x=${MAX_RATE:=100mbit}
x=${LAN:=eth0}
x=${WAN:=eth1}
x=${IN_RATE:=$[(1*1024)*99/100]Kibit}
x=${OUT_RATE:=$[(1*1024)*99/100]Kibit}
x=${INVERT_CLASSES:=0}
x=${HASHKEY:=0 1} # IP's bytes or none, 2 keys experemental, 3+ in theory (shared)
x=${PRIO:=100}
x=${INTRANET:=10.0.0.0/8 172.16.0.0/12 192.168.0.0/16}
x=${BATCH:=0} # fast script run, with 1|2 - ignore errors from "tc -force -batch"
x=${V:=1} # 0-quiet, 1-errors, 2-messages, 3-full (BATCH=0 to safe)
#x=${ACK_IN:=} # % of rate or none
#x=${ACK_OUT:=}
x=${FIREWALL:=0} # required actions drop/pass & ingress
x=${PORT_CLOSE:=135 136 137 138 139 445 9996}
## if defined next:
## - traffic will previously (ingress/mangle/prerouting) NFMARK-ed
## and finally scheduled via NFMARK, may be used with netfilter/iptables
x=${NFMARK_IN:=} # required: ingress, cls_u32_mark, cls_act/act_ipt, iptables (libs)
x=${NFMARK_OUT:=}
x=${NFMARK_INTRANET:=} # "or" to NFMARK_IN/NFMARK_OUT

##############################################################################

zero(){
local x="$*"
x=${x// }
test "${x//0}" && return 1
return 0
}

substr(){
local s=" $1 "
local s1="$2";
test "${s// $s1 }" = "$s" && return 1
return 0
}

run(){
local ee
if zero $[V/2] ; then
 ee=`$* 2>&1`
else
 test $V == 3 && echo "$*"
 $*
fi
ret=$?
substr "$ee" "RTNETLINK answers: File exists" && ret=0
zero "$V" || zero "$ret" || echo "[$ret] $*
$ee"
return $ret
}

tc(){
if zero "$BATCH" ; then
 run /sbin/tc $*
 return $ret
fi
tc="$tc
$*"
zero "$[${#tc}/bufsize]" || flush
return 0
}

ifconfig(){
run /sbin/ifconfig $*
}

ethtool(){
run /usr/sbin/ethtool $*
}

modprobe(){
run /sbin/modprobe $*
}

int2base() {
local d="0123456789abcdef"
local i=$2
int2base=
while [[ $i != 0 ]] ; do
 int2base=${d:$[i%$1]:1}$int2base
 i=$[i/$1]
done
}

flush(){
test "$tc" && echo "$tc"|BATCH="" tc -force -batch
tc=""
}

rate(){
local r="$1"
local i
r=${r//Kibit/ 1024}
r=${r//kbit/ 1000}
r=${r//mibit/ 1024 1024}
r=${r//mbit/ 1000000}
r=${r//gibit/ 1024 1024 1024}
r=${r//gbit/ 1000000000}
r=${r//tibit/ 1024 1024 1024 1024}
r=${r//tbit/ 1000000000000}
r=${r//KiBps/ 8 1024}
r=${r//KBps/ 8000}
r=${r//MiBps/ 8 1024 1024}
r=${r//MBps/ 8000000}
r=${r//GiBps/ 8 1024 1024 1024}
r=${r//GBps/ 8000000000}
r=${r//TiBps/ 8 1024 1024 1024 1024}
r=${r//TBps/ 8000000000000}
r=${r//Bps/ 8}
r=${r//bit}
rate=1
for i in ${r[@]} ; do
 rate=$[rate*i]
done
}

wlog(){
if [[ $[V/3] != 0 ]] ; then
 echo "$*"
fi
}

# ethX parent_class rate [ack_rate%]
sched_tcp(){
rate $3
rate=$[rate*$4/100]
if zero $[rate/1000] ; then
 rate="${rate}bit"
else
 rate="$[rate/1000]kbit"
fi
tc class add dev $1 parent 1: classid 1:1$2 psp rate $3
tc class add dev $1 parent 1:1$2 classid 1:a$2 psp rate $rate
tc qdisc add dev $1 parent 1:a$2 $estimator handle a$2: pfifo
tc class add dev $1 parent 1:1$2 classid 1:$2 psp mode 0
tc filter add dev $1 parent 1: protocol all prio $[PRIO+3] u32 \
 match ip protocol 6 0xff \
 match u8 0x05 0x0f at 0 \
 match u8 0x10 0xff at 33 \
 match u16 0x0000 0xffc0 at 2 \
 flowid 1:a$2
}

# ethX rate [ack_rate%]
sched(){
local c1=$[INVERT_CLASSES+1]
local c2=$[2-INVERT_CLASSES]
ethtool -K tso off $1
ifconfig $1 txqueuelen 10000
V=0 tc qdisc del dev $1 ingress $Q
V=0 tc qdisc del dev $1 root $Q
tc qdisc add dev $1 root handle 1: psp default 1 || exit
# rate ${MAX_RATE}
if zero "$3" ; then
 tc class add dev $1 parent 1: classid 1:$c1 psp rate $2
else
 sched_tcp $1 $c1 $2 $3
fi
tc qdisc add dev $1 parent 1:$c1 $estimator handle 100$c1: pfifo
tc class add dev $1 parent 1: classid 1:$c2 psp mode 0
tc qdisc add dev $1 parent 1:$c2 handle 100$c2: pfifo
}

hkey(){
local x=(0xff000000 0x00ff0000 0x0000ff00 0x000000ff)
hkey="match u32 0 0 at $2 hashkey mask ${x[$1]} at $2"
}

# ethX ethY addr_offset [mark]
filters(){
local i
local p
local id=1
local eth=$1
local m=$4
test "$m" &&
 tc qdisc add dev $2 ingress &&
 tc filter add dev $1 protocol all parent 1: prio $PRIO u32 match mark $m $m flowid 1:2 &&
 id="ffff 1" &&
 eth=$2 &&
 test "$NFMARK_INTRANET" &&
 test "$FIREWALL" == 1 &&
 tc filter add dev $1 protocol all parent 1: prio $[PRIO-1] u32 match mark $[NFMARK_INTRANET|m] $[NFMARK_INTRANET|m] action drop
hkey ${HKEY[0]} $3
test "$HKEY" && for i in ${id[@]} ; do
 tc filter add dev $eth parent $i: prio $PRIO protocol all u32
 tc filter add dev $eth parent $i: prio $PRIO handle $hid: protocol all u32 divisor 256
 tc filter add dev $eth protocol all parent $i: prio $PRIO u32 ht 800:: $hkey link $hid: && break
done
}

firewall(){
local h=ff
local d="dev $LAN parent 1:"
local p
#if [[ $NFMARK_IN ]] ; then
 tc qdisc add dev $WAN ingress
 d="dev $WAN parent ffff:"
#fi
tc filter add $d prio $[PRIO-2] protocol all u32
tc filter add $d prio $[PRIO-2] handle $h: protocol all u32 divisor 256
#tc filter add $d protocol ip prio $[PRIO-1] u32 ht 800:: match ip dport 0 0 hashkey mask 0x000000ff at 20 link $h:
for p in {6,17} ; do
 tc filter add $d protocol all prio $[PRIO-1] u32 ht 800:: match ip protocol $p 0xff hashkey mask 0x000000ff at 20 link $h:
done
for p in ${PORT_CLOSE[@]} ; do
 int2base 16 $p
 tc filter add $d protocol all prio $[PRIO-1] u32 ht $h:${int2base:${#int2base}-2} match u32 0x$int2base 0x0000ffff at 20 action drop
done
}

tc2(){
local c="$*"
local i
for i in {"$LAN $WAN $NFMARK_IN","$WAN $LAN $NFMARK_OUT"} ; do
i=($i)
if [[ "${i[2]}" ]] ; then
 c=${c//${i[0]}/${i[1]}}
 c=${c//parent 1: /parent ffff: }
 c=${c//flowid 1:2/action ipt -j MARK --set-mark $[i[2]|om]}
else
 c=${c//${i[1]}/${i[0]}}
 if [[ "$om" && "$FIREWALL" == 1 ]] ; then
  c=${c//flowid 1:2/action drop}
 else
  zero "$FIREWALL" || c=${c//flowid 1:2/flowid 1:2 action ok}
 fi
fi
tc $c
c="$*"
c=${c// src / dst }
c=${c//at 12 /at 16 }
done
}

# "add" ip [or-mark]
noshape(){
local a=$2
om=$3
a=${a//./ }
a=(${a//\// })
local ht=""
local htt
local h="$hid"
local hk
local i=0
test "${a[4]}" || a[4]=32
local id=
test "$HKEY" && for hk in ${HKEY[@]} ; do
 test $[${a[4]}/((hk+1)*8)] == 0 && break
 if [[ "$ht" ]] ; then
  tc2 filter add dev $LAN parent 1: prio $PRIO handle $h: protocol all u32 $ht divisor 256
  hkey i 12
  tc2 filter $1 dev $LAN protocol all parent 1: prio $PRIO $htt u32 $ht $hkey link $h:
 fi 
 int2base 16 ${a[hk]}
 htt="handle $h:$int2base:1$hk"
 ht="ht $h:$int2base:"
 i=$[i+1]
 test ${#int2base} == 1 && int2base="0$int2base"
 h="$hk$int2base"
done
tc2 filter $1 dev $LAN protocol all parent 1: prio $PRIO u32 $ht match ip src $2 flowid 1:2
}

help(){
local c="psp-shaper"
echo "precise shaper script v0.02 (c) Dzianis Kahanovich, 2007, GPLv2
Usage: $c {-command_to_eval} [{CMD {[net|@ file]}}]
 CMD are add|del
Examples:
 $c -show				- show active variables
 $c -init				- start shaper
 $c -V=3 -IN_RATE=2mbit -init		- change variables and start
 HASHKEY=0 IN_RATE=10mbit $c -init	- same
 $c -V=3 add 199.1.1.0/24 198.1.1.1	- add unshaped networks
 $c -init add @ /etc/shaper/peering	- init & load from file"
exit
}

show(){
local vars
local i
local x
local r
while read i ; do
 test "${i:0:5}" == "#####" && break
 r=${i/*# }
 if [[ "$r" == "$i" ]] ; then
  r=""
 else
  r="## $r
"
 fi
 i=${i// }
 i=${i//	}
 i=(${i//\:\=/ })
 if [[ "${i[0]:0:4}" == "x=\${" && "${i[1]}" != "" ]] ; then
  i[0]="${i[0]:4}"
  eval "x=\"\$$i\""
  echo "$r $i=$x"
 fi
done <"$0"
}

modules(){
modprobe cls_u32
modprobe sch_psp
if ! zero $FIREWALL ; then
 modprobe sch_ingress
 modprobe act_gact
fi
if [[ "$FWMARK_IN$FWMARK_OUT" ]] ; then
 modprobe sch_ingress
 modprobe act_ipt
fi
}

init(){
local i
local nsh
HKEY=($HASHKEY)
hid="f00"
bufsize=$[32*1024]
tc=
V=0 modules
#estimator="est ${LATENCY}msec $[LATENCY*8]msec"
test "$NFMARK_IN$NFMARK_OUT" && test "$BATCH" == "1" && BATCH=
rate $MAX_RATE
local rmem=$[(rate/8)*LATENCY/500+1000000]
test "$V" == "3" && show
#echo 0 >/proc/sys/net/ipv4/ip_forward
sched $LAN $IN_RATE $ACK_IN
sched $WAN $OUT_RATE $ACK_OUT
filters $LAN $WAN 12 $NFMARK_IN
filters $WAN $LAN 16 $NFMARK_OUT
if [[ "$NFMARK_INTRANET" && "$NFMARK_OUT" ]] ; then
 local f="tc filter add dev $LAN protocol all parent 1: prio $[PRIO-1] u32 match mark $[NFMARK_INTRANET|NFMARK_OUT] $[NFMARK_INTRANET|NFMARK_OUT] flowid 1:2"
 zero "$FIREWALL" || f="$f action pass"
 $f
fi
for nsh in ${INTRANET[@]} ; do
 noshape add $nsh $NFMARK_INTRANET
done
zero $FIREWALL || firewall
test $V == 3 && echo rmem: $rmem
echo $rmem >/proc/sys/net/core/rmem_max
echo $rmem >/proc/sys/net/core/rmem_default
echo 2500 >/proc/sys/net/core/netdev_max_backlog
flush
echo 1 >/proc/sys/net/ipv4/ip_forward
}

cmd_del(){
echo under construction
exit
}

help1(){
test "$*" && return
echo "Say $0 -help"
exit
}

help1 $*
cmd="eval"
cmds="  "
while [[ "$1" ]] ; do
if substr "add del" "$1" ; then
 cmd="$1"
 shift
 help1 $*
fi
case "$cmd" in
add)
 if [[ "$1" == "@" ]] ; then
  shift
  while read nsh j ; do
   noshape add $nsh
  done <"$1"
 else
  noshape add $1
 fi
;;
del)
 if [[ "$1" == "@" ]] ; then
  shift
  while read nsh j ; do
   cmd_del $nsh
  done <"$1"
 else
  cmd_del $1
 fi
;;
eval)
 x=$1
 if [[ "${x:0:1}" == "-" ]] ; then
  eval "${x:1}"
 else
  help1
 fi
;;
esac
shift
done
exit