#!/bin/bash

setglobal MY_DIR = $[dirname $0]
# Details on the bpf prog
setglobal BPF_CGRP2_ARRAY_NAME = ''test_cgrp2_array_pin''
setglobal BPF_PROG = ""$MY_DIR/test_cgrp2_tc_kern.o""
setglobal BPF_SECTION = ''filter''

test -z $TC && setglobal TC = ''tc''
test -z $IP && setglobal IP = ''ip''

# Names of the veth interface, net namespace...etc.
setglobal HOST_IFC = ''ve''
setglobal NS_IFC = ''vens''
setglobal NS = ''ns''

proc find_mnt {
    cat /proc/mounts | \
	awk '{ if ($3 == "'$1'" && mnt == "") { mnt = $2 }} END { print mnt }'
}

# Init cgroup2 vars
proc init_cgrp2_vars {
    setglobal CGRP2_ROOT = $[find_mnt cgroup2]
    if test -z $CGRP2_ROOT
    {
	setglobal CGRP2_ROOT = ''/mnt/cgroup2''
	setglobal MOUNT_CGRP2 = '"yes'"
    }
    setglobal CGRP2_TC = ""$CGRP2_ROOT/tc""
    setglobal CGRP2_TC_LEAF = ""$CGRP2_TC/leaf""
}

# Init bpf fs vars
proc init_bpf_fs_vars {
    var bpf_fs_root = $[find_mnt bpf]
    test -n $bpf_fs_root || return -1
    setglobal BPF_FS_TC_SHARE = ""$bpf_fs_root/tc/globals""
}

proc setup_cgrp2 {
    match $1 {
	with start
	    if test $MOUNT_CGRP2 == 'yes'
	    {
		test -d $CGRP2_ROOT || mkdir -p $CGRP2_ROOT
		mount -t cgroup2 none $CGRP2_ROOT || return $?
	    }
	    mkdir -p $CGRP2_TC_LEAF
	    
	with *
	    rmdir $CGRP2_TC_LEAF && rmdir $CGRP2_TC
	    test $MOUNT_CGRP2 == 'yes' && umount $CGRP2_ROOT
	    
    }
}

proc setup_bpf_cgrp2_array {
    var bpf_cgrp2_array = ""$BPF_FS_TC_SHARE/$BPF_CGRP2_ARRAY_NAME""
    match $1 {
	with start
	    $MY_DIR/test_cgrp2_array_pin -U $bpf_cgrp2_array -v $CGRP2_TC
	    
	with *
	    test -d $BPF_FS_TC_SHARE && rm -f $bpf_cgrp2_array
	    
    }
}

proc setup_net {
    match $1 {
	with start
	    $IP link add $HOST_IFC type veth peer name $NS_IFC || return $?
	    $IP link set dev $HOST_IFC up || return $?
	    sysctl -q net.ipv6.conf.$HOST_IFC.accept_dad=0

	    $IP netns add ns || return $?
	    $IP link set dev $NS_IFC netns ns || return $?
	    $IP -n $NS link set dev $NS_IFC up || return $?
	    $IP netns exec $NS sysctl -q net.ipv6.conf.$NS_IFC.accept_dad=0
	    $TC qdisc add dev $HOST_IFC clsact || return $?
	    $TC filter add dev $HOST_IFC egress bpf da obj $BPF_PROG sec $BPF_SECTION || return $?
	    
	with *
	    $IP netns del $NS
	    $IP link del $HOST_IFC
	    
    }
}

proc run_in_cgrp {
    # Fork another bash and move it under the specified cgroup.
    # It makes the cgroup cleanup easier at the end of the test.
    setglobal cmd = ''echo $$ > ''
    setglobal cmd = ""$cmd $1/cgroup.procs; exec $2""
    bash -c $cmd
}

proc do_test {
    run_in_cgrp $CGRP2_TC_LEAF "ping -6 -c3 ff02::1%$HOST_IFC >& /dev/null"
    var dropped = $[$TC -s qdisc show dev $HOST_IFC | tail -3 | \
			   awk '/drop/{print substr($7, 0, index($7, ",")-1)}]
    if [[ $dropped -eq 0 ]]
    {
	echo "FAIL"
	return 1
    } else {
	echo "Successfully filtered $dropped packets"
	return 0
    }
}

proc do_exit {
    if test $DEBUG == "yes" && test $MODE != 'cleanuponly'
    {
	echo "------ DEBUG ------"
	echo "mount: "; mount | egrep '(cgroup2|bpf)'; echo
	echo "$CGRP2_TC_LEAF: "; ls -l $CGRP2_TC_LEAF; echo
	if test -d $BPF_FS_TC_SHARE
	{
	    echo "$BPF_FS_TC_SHARE: "; ls -l $BPF_FS_TC_SHARE; echo
	}
	echo "Host net:"
	$IP netns
	$IP link show dev $HOST_IFC
	$IP -6 a show dev $HOST_IFC
	$TC -s qdisc show dev $HOST_IFC
	echo
	echo "$NS net:"
	$IP -n $NS link show dev $NS_IFC
	$IP -n $NS -6 link show dev $NS_IFC
	echo "------ DEBUG ------"
	echo
    }

    if test $MODE != 'nocleanup'
    {
	setup_net stop
	setup_bpf_cgrp2_array stop
	setup_cgrp2 stop
    }
}

init_cgrp2_vars
init_bpf_fs_vars

while [[ $# -ge 1 ]]
{
    setglobal a = $1
    match $a {
	with debug
	    setglobal DEBUG = ''yes''
	    shift 1
	    
	with cleanup-only
	    setglobal MODE = ''cleanuponly''
	    shift 1
	    
	with no-cleanup
	    setglobal MODE = ''nocleanup''
	    shift 1
	    
	with *
	    echo "test_cgrp2_tc [debug] [cleanup-only | no-cleanup]"
	    echo "  debug: Print cgrp and network setup details at the end of the test"
	    echo "  cleanup-only: Try to cleanup things from last test.  No test will be run"
	    echo "  no-cleanup: Run the test but don't do cleanup at the end"
	    echo "[Note: If no arg is given, it will run the test and do cleanup at the end]"
	    echo
	    exit -1
	    
    }
}

trap do_exit 0

test $MODE == 'cleanuponly' && exit

setup_cgrp2 start || exit $?
setup_net start || exit $?
init_bpf_fs_vars || exit $?
setup_bpf_cgrp2_array start || exit $?
do_test
echo