# #-- ip_ratelimit.test --#
# source the master var file when it's there
[ -f ../.tpkg.var.master ] && source ../.tpkg.var.master
# use .tpkg.var.test for in test variable passing
[ -f .tpkg.var.test ] && source .tpkg.var.test

PRE="../.."
. ../common.sh

get_make
(cd $PRE; $MAKE streamtcp)

# These tests rely on second time precision. To combat false negatives the
# tests run multiple times and we allow 1/3 of the runs to fail.
total_runs=6
success_threshold=4  # 2/3*total_runs

if dig -h 2>&1 | grep "cookie" >/dev/null; then
	nocookie="+nocookie"
else
	nocookie=""
fi

echo "> First get a valid cookie"
dig @127.0.0.1 -p $UNBOUND_PORT +ednsopt=10:0102030405060708 $nocookie +tcp +retry=0 +time=1 test. TXT >outfile 2>&1
if test "$?" -ne 0; then
	echo "exit status not OK"
	echo "> cat logfiles"
	cat outfile
	cat unbound.log
	echo "Not OK"
	exit 1
fi
if test `grep "COOKIE: " outfile | wc -l` -ne 1; then
	echo "Could not get cookie"
	echo "> cat logfiles"
	cat outfile
	cat unbound.log
	echo "Not OK"
	exit 1
fi
cookie=`grep "COOKIE: " outfile | cut -d ' ' -f 3`

successes=0
echo "> Three parallel queries with backoff and cookie"
# For this test we send three parallel queries. The ratelimit should be reached
# for that second. We send a query to verify that there is no reply.
# Then for the next second we again send three parallel queries and we expect
# none of them to be allowed through because of the backoff logic that keeps
# rolling the RATE_WINDOW based on demand.
# Again we send another query but with a valid cookie and we expect to receive
# an answer.
for i in $(seq 1 $total_runs); do
	# Try to hit limit
	$PRE/streamtcp -nu -f 127.0.0.1@$UNBOUND_PORT test. TXT IN test. TXT IN test. TXT IN >outfile 2>&1
	if test "$?" -ne 0; then
		echo "exit status not OK"
		echo "> cat logfiles"
		cat outfile
		cat unbound.log
		echo "Not OK"
		exit 1
	fi
	# Expect no answer because of limit
	dig @127.0.0.1 -p $UNBOUND_PORT $nocookie +retry=0 +time=1 test. TXT >outfile 2>&1
	if test "$?" -eq 0; then
		continue
	fi
	# Try to keep limit
	$PRE/streamtcp -nu -f 127.0.0.1@$UNBOUND_PORT test. TXT IN test. TXT IN test. TXT IN >outfile 2>&1
	if test "$?" -ne 0; then
		echo "exit status not OK"
		echo "> cat logfiles"
		cat outfile
		cat unbound.log
		echo "Not OK"
		exit 1
	fi
	# Expect answer because of DNS cookie
	dig @127.0.0.1 -p $UNBOUND_PORT +ednsopt=10:$cookie $nocookie +retry=0 +time=1 test. TXT >outfile 2>&1
	if test "$?" -ne 0; then
		continue
	fi
	((successes++))
	# We don't have to wait for all the runs to complete if we know
	# we passed the threshold.
	if test $successes -ge $success_threshold; then
		break
	fi
done

if test $successes -ge $success_threshold; then
	echo "Three parallel queries with backoff and cookie OK"
else
	echo "Three parallel queries with backoff and cookie NOT OK"
	echo "> cat logfiles"
	cat outfile
	cat unbound.log
	echo "Three parallel queries with backoff and cookie NOT OK"
	exit 1
fi

echo "> Activating ip-ratelimit-cookie"
echo "$PRE/unbound-control -c ub.conf set_option ip-ratelimit-cookie: 1"
$PRE/unbound-control -c ub.conf set_option ip-ratelimit-cookie: 1
if test $? -ne 0; then
	echo "wrong exit value after success"
	exit 1
fi

successes=0
echo "> Three parallel queries with backoff and cookie with ip-ratelimit-cookie"
# This is the exact same test as above with the exception that we don't expect
# an answer on the last query because ip-ratelimit-cookie is now enabled.
for i in $(seq 1 $total_runs); do
	# Try to hit limit
	$PRE/streamtcp -nu -f 127.0.0.1@$UNBOUND_PORT test. TXT IN test. TXT IN test. TXT IN >outfile 2>&1
	if test "$?" -ne 0; then
		echo "exit status not OK"
		echo "> cat logfiles"
		cat outfile
		cat unbound.log
		echo "Not OK"
		exit 1
	fi
	# Expect no answer because of limit
	dig @127.0.0.1 -p $UNBOUND_PORT $nocookie +retry=0 +time=1 test. TXT >outfile 2>&1
	if test "$?" -eq 0; then
		continue
	fi
	# Try to keep limit
	$PRE/streamtcp -nu -f 127.0.0.1@$UNBOUND_PORT test. TXT IN test. TXT IN test. TXT IN >outfile 2>&1
	if test "$?" -ne 0; then
		echo "exit status not OK"
		echo "> cat logfiles"
		cat outfile
		cat unbound.log
		echo "Not OK"
		exit 1
	fi
	# Expect no answer because of ip-ratelimit-cookie
	dig @127.0.0.1 -p $UNBOUND_PORT +ednsopt=10:$cookie $nocookie +retry=0 +time=1 test. TXT >outfile 2>&1
	if test "$?" -eq 0; then
		continue
	fi
	((successes++))
	# We don't have to wait for all the runs to complete if we know
	# we passed the threshold.
	if test $successes -ge $success_threshold; then
		break
	fi
done

if test $successes -ge $success_threshold; then
	echo "Three parallel queries with backoff and cookie with ip-ratelimit-cookie OK"
else
	echo "Three parallel queries with backoff and cookie with ip-ratelimit-cookie NOT OK"
	echo "> cat logfiles"
	cat outfile
	cat unbound.log
	echo "Three parallel queries with backoff and cookie with ip-ratelimit-cookie NOT OK"
	exit 1
fi

exit 0