#!/bin/bash

set -e

RATE="10mbit"

usage() {
    echo "Usage: $0 [start|stop] container-name"
    echo
    echo "Start or stop traffic policing for the given docker container"
}

get_pid() {
    local CONTAINER=$1
    docker inspect --format '{{.State.Pid}}' "$CONTAINER"
}

get_iface() {
    local CONTAINER=$1
    INDEX=$(docker exec "$CONTAINER" sh -c 'cat /sys/class/net/eth0/iflink')
    ip addr | grep "^${INDEX}:" | awk '{print $2}' | awk -F '@' '{print $1}'
}

start_policing() {
    local CONTAINER=$1

    PID=$(get_pid "$CONTAINER")
    nsenter --target "$PID" --net tc qdisc add dev eth0 handle ffff: ingress
    nsenter --target "$PID" --net \
        tc filter add dev eth0 parent ffff: u32 \
            match u32 0 0 \
            police rate "$RATE" burst 100k conform-exceed drop/ok

    IFACE=$(get_iface "$CONTAINER")
    tc qdisc add dev "$IFACE" handle ffff: ingress
    tc filter add dev "$IFACE" parent ffff: u32 \
        match u32 0 0 \
        police rate "$RATE" burst 100k conform-exceed drop/ok
}

stop_policing() {
    local CONTAINER=$1

    PID=$(get_pid "$CONTAINER")
    nsenter --target "$PID" --net tc qdisc del dev eth0 ingress

    IFACE=$(get_iface "$CONTAINER")
    tc qdisc del dev "$IFACE" ingress
}

if [ $# != 2 ]; then
    usage
    exit 2
fi

case "$1" in
    "--help" | "-h")
        usage
        exit 0
        ;;
    "start")
        start_policing "$2"
        exit 0
        ;;
    "stop")
        stop_policing "$2"
        exit 0
        ;;
    *)
        usage
        exit 2
        ;;
esac
