#!/bin/bash -e
# Adapted from http://lartc.org/wondershaper/

## CONFIGURATION ##
# Unless specified in the list below, all UPLINK traffic is medium priority.  
# Specify source MAC addresses for device which should be subject to other 
# priorities.

HIGH_PRIORITY=(
    '74:5e:1c:56:c2:97' # Pioneer SC-LX57
    )

# NB: Has *no* affect, just a convenient bucket for MAC address I may want to 
# cut-and-paste into HIGH or LOW in the future
MEDIUM_PRIORITY=(
    '78:31:c1:be:0c:34' # Macbook Pro, Robert
    '00:50:8d:9d:75:83' # Ikari (top)
    '00:50:8d:9d:75:84' # Ikari (bottom)
    '20:02:af:18:dc:15' # Samsung Galaxy Tab 10.1, Vicky
    '20:02:af:44:52:79' # Samsung Galaxy Tab 10.1, Regina
    '60:45:bd:9c:9d:bc' # Xbox 360 Slim
    )

LOW_PRIORITY=(
    '34:af:2c:9c:ff:fd' # Nintendo WiiU
    'bc:8c:cd:fb:c0:71' # Samsung Smart TV, Vicky    
    'b8:e8:56:3a:cb:34' # Macbook Pro, Vicky
    '88:30:8a:77:03:70' # Samsung Galaxy S3, Vicky
    '88:32:9b:02:6a:ce' # Samsung Galaxy S3, Regina
    'b8:27:eb:8f:c0:f3' # Raspberry Pi, Flux
    'd0:22:be:2c:26:ef' # Samsung Galaxy Note 3, Rob
    )

DEV=$1
DOWNLINK=$2
UPLINK=$3

# The following fudge factors allow you to express the usable % of your link.  
#
# Experimentation has shown that ~75% of the author's ADSL downlink can be used 
# before upstream congestion starts to affect round-trip times. In other words, 
# by throttling our download speeds we can ensure that our ISP does not queue 
# any packs on our behalf, giving us full control over congestion.
#
# Methodology used: ping robmeerman.co.uk from home in one window, ping home 
# from robmeerman.co.uk in another. Set both factors to 100/100. Run a speed 
# test (e.g. speedtest.net), and observe if either shows a *consistent* latency 
# above 200ms.  It's usually pretty drastic - from 22ms when idle -> 1019ms 
# during test. Bisect DOWNFACTOR and repeat until you find the highest value 
# which maintains acceptable latency during a speedtest. For me, the factors 
# tried were: 100 (bad), 50 (good), 75 (good), 87.5 (bad), 81.25 (bad), 78.125 
# (bad) and then I gave up and left it at 75 (good). Repeat process for 
# UPFACTOR.
#
# It is important to do the DOWNFACTOR first because if your ISP holds a queue 
# of traffic for us, then we cannot control round-trip time. So first we find a 
# factor that eliminates ISP queuing, and then we can find a factor to control 
# our own egress queuing.
DOWNFACTOR='285/350'
UPFACTOR='390/448'

if [ "x$DEV" = "x" ]
then
    echo "Usage: $0 (DEV) [ 'clear' | (DOWNLINK kbit/s) (UPLINK kbit/s) ]"
    exit 0
fi

# Display status when DOWNLINK/UPLINK are ommitted
if [ "x$DOWNLINK" = "x" ]
then
    echo "--------------------------------------------------------------------------------"
    iptables -nvL -t mangle
    # echo "--------------------------------------------------------------------------------"
    # tc -s filter ls dev $DEV
    # echo "--------------------------------------------------------------------------------"
    # tc -s qdisc ls dev $DEV
    # echo "--------------------------------------------------------------------------------"
    # tc -s class ls dev $DEV
    exit 0
fi

# Clear both IN and OUT
tc qdisc del dev $DEV root    2> /dev/null > /dev/null || true
tc qdisc del dev $DEV ingress 2> /dev/null > /dev/null || true

# Flush and delete all mangle rules
iptables -F 
iptables -X 
iptables -t mangle -F 
iptables -t mangle -X 

if [ "x$DOWNLINK" = "xclear" ]
then
    echo "Cleared traffic rules on $DEV"
    exit 0
fi

trap "$0 $1 clear" ERR

# Calculations
#
# Target latency is < 50ms. This means max burst length should be limited to 
# 1/20th the queue's rate.

# =============================================================================
# Queues and Classes
# =============================================================================
# 1: ROOT
# |-- 1:ff LOCAL_TRAFFIC (to/from this host itself)
# | `-- ff: (sfq)
# |-- 1:1 INTERNET->LAN (downlink)
# | `-- 10: (red) Drop traffic as link approaches congestion
# `-- 1:2 LAN->INTERNET (uplink)
#   |-- 1:21: High priority
#   | `-- 21: (sfq)
#   |-- 1:22: Medium priority
#   | `-- 22: (sfq)
#   `-- 1:23: Low priority
#     `-- 23: (sfq)

# ROOT
tc qdisc add dev $DEV root handle 1: htb

    # LOCAL TRAFFIC
    tc class add dev $DEV parent 1: classid 1:ff htb \
        rate 100mbit \
        burst $((100/20))mbit \
        cburst $((100/20))mbit \
        prio 1

        # .. and its actual queue that holds the packets
		tc qdisc add dev $DEV parent 1:ff handle ff: sfq perturb 10

    # INTERNET->LAN (downlink)
    #
    # NOTES: Target is 50ms (1/20th of a second), so our burst and min sizes 
    # must be smaller than this or we're too late to influence latency.  I've 
    # chosen to use 1/40th as the lower bound.
    #
    # cburst should be set the network's Maximum Transmission Unit (MTU), or 
    # bad things happen - presumably the packets are fragmented, increasing 
    # overheads and processing cost. Regardless, performance was poor when this 
    # was less than MTU.
    tc class add dev $DEV parent 1: classid 1:1 htb \
        rate $(($DOWNLINK*$DOWNFACTOR))kbit \
        ceil $(($DOWNLINK*$DOWNFACTOR))kbit \
        burst $(($DOWNLINK*$DOWNFACTOR/40))kbit \
        cburst 1500 \
        prio 10

        # .. and its actual queue that holds the packets
        # Note: All values are in BYTES. It doesn't seem to accept "kbit"
        #
        # The burst calculation needs to be increased by one so as to avoid an 
        # internal assert in the qdisc (seems our target and their min 
        # acceptable burst are one and the same)
        tc qdisc add dev $DEV parent 1:1 handle 10: red \
            limit $(($DOWNLINK*$DOWNFACTOR*1000/8)) \
            avpkt 1500 \
            burst $((($DOWNLINK*1000/8/40/1500)+1)) \
            min   $(($DOWNLINK*1000/8/40)) \
            max   $(($DOWNLINK*1000/8/20)) \
            probability 1

    # LAN->INTERNET (uplink)
    tc class add dev $DEV parent 1: classid 1:2 htb \
        rate $(($UPLINK*$UPFACTOR))kbit \
        ceil $(($UPLINK*$UPFACTOR))kbit \
        burst $(($UPLINK/20))kbit \
        cburst $(($UPLINK/20))kbit \
        prio 20

        # High priority
        tc class add dev $DEV parent 1:2 classid 1:21 htb \
            rate $(($UPLINK*$UPFACTOR*3/9))kbit \
            ceil $(($UPLINK*$UPFACTOR))kbit \
            prio 0

        # Medium priority
        tc class add dev $DEV parent 1:2 classid 1:22 htb \
            rate $(($UPLINK*$UPFACTOR*3/9))kbit \
            ceil $(($UPLINK*$UPFACTOR))kbit \
            prio 1

        # Low priority
        tc class add dev $DEV parent 1:2 classid 1:23 htb \
            rate $(($UPLINK*$UPFACTOR*3/9))kbit \
            prio 2

        # .. and their actual queues that hold the packets
        for ID in 21 22 23
        do
            tc qdisc add dev $DEV parent 1:$ID handle $ID: sfq
        done


# =============================================================================
# Filters
# =============================================================================

# -----------------------------------------------------------------------------
# LOCAL TRAFFIC
# Mark traffic generated by this host itself (INPUT + OUTPUT, but not FORWARD)
iptables -t mangle -A INPUT  -p all -i $DEV -j MARK --set-mark 0xff
iptables -t mangle -A OUTPUT -p all -o $DEV -j MARK --set-mark 0xff

# ("fw" means the handle refers to a MARK, rather than a qdisc)
tc filter add dev $DEV parent 1: protocol ip prio 1 handle 0xff fw classid 1:ff

# -----------------------------------------------------------------------------
# INTERNET->LAN (downlink)
#
# Note: We assume that LAN->LAN traffic is *not* forwarded through this host, 
# and so we need only check the destination of a given packet. We've already 
# taken care of this host's own traffic above.

iptables -t mangle -N DOWNLINK
iptables -t mangle -A DOWNLINK -p all -j MARK --set-mark 0x1
tc filter add dev $DEV parent 1: protocol ip prio 2 handle 0x1 fw classid 1:1

for SUBNET in 192.168.0.0/16 10.0.0.0/8 172.16.0.0/12 
do
    iptables -t mangle -A PREROUTING -p all -i $DEV ! -s $SUBNET -d $SUBNET -j DOWNLINK
done


# -----------------------------------------------------------------------------
# LAN->INTERNET (uplink)
#
# Note: Assumes that all downlink and private traffic have already been 
# classified, so no source checks are performed.

iptables -t mangle -N UPLINK
iptables -t mangle -A UPLINK -p all -j MARK --set-mark 0x22     # Default to medium priority

for SUBNET in 192.168.0.0/16 10.0.0.0/8 172.16.0.0/12 
do
    iptables -t mangle -A PREROUTING -p all -i $DEV -s $SUBNET ! -d $SUBNET -j UPLINK
done

# Note: Use $PRIORITY and $FLOWID in the rules, so they can cut-and-paste into 
# different priority bands without having to tediously modify all the 
# priorty-related numbers (I've forgotten to do this a number of times!)

##
## HIGH PRIORITY ##
##
PRIORITY=20
FLOWID=1:21
MARK=0x21

# TOS Minimum Delay (ssh, NOT scp)
tc filter add dev $DEV parent 1: protocol ip prio $PRIORITY u32 \
    match ip tos 0x10 0xff \
    flowid $FLOWID

# Traffic headed to robmeerman.co.uk (typically SSH proxying to else where)
iptables -t mangle -A UPLINK --destination 85.119.82.218/32 -j MARK --set-mark $MARK

# Traffic headed to vpn1.cambridge.arm.com (working from home)
iptables -t mangle -A UPLINK --destination 217.140.97.4/32 -j MARK --set-mark $MARK

for MAC in ${HIGH_PRIORITY[@]}
do
    iptables -t mangle -A UPLINK -m mac --mac-source $MAC -j MARK --set-mark $MARK
done


##
## LOW PRIORITY ##
##
PRIORITY=30
FLOWID=1:23
MARK=0x22

# TOS High Throughput
tc filter add dev $DEV parent 1: protocol ip prio $PRIORITY u32 \
    match ip tos 0x8 0xff \
    flowid $FLOWID

for MAC in ${LOW_PRIORITY[@]}
do
    iptables -t mangle -A UPLINK -m mac --mac-source $MAC -j MARK --set-mark $MARK
done



##
## MEDIUM PRIORITY ##
##

# If no other filter has classified the packet, then use FW markers (set by 
# iptables -j MARK). All UPLINK packets are marked as 0x22 by default (see 
# iptables command earlier)
tc filter add dev $DEV parent 1: protocol ip prio 40 handle 0x21 fw classid 1:21 # High priority
tc filter add dev $DEV parent 1: protocol ip prio 40 handle 0x22 fw classid 1:22 # Medium priority
tc filter add dev $DEV parent 1: protocol ip prio 40 handle 0x23 fw classid 1:23 # Low priority


# Reset counters, so that packet counts are in sync (it takes time to add 
# rules, and during that time the first rule added may be hit, leading to 
# confusing packet counts: "But these rules should always apply to the same 
# packets! How can their hit count be different?"
iptables -t mangle -Z