#!/bin/bash

# This script can be run with several environment variables set dictating its
# run. You can set the following to what you like:
WOLFSSL_ROOT=${WOLFSSL_ROOT:-$(pwd)}
UDP_PROXY_BIN=${UDP_PROXY_BIN:-"udp_proxy"}
DTLS_VERSION=${DTLS_VERSION:-"-v4"}
PCAP_FILENAME=${PCAP_FILENAME:-"dtls${DTLS_VERSION}.pcap"}
# Additionally, you can add the following tests by setting it to '1':
DO_EXTENDED_SERVER_PERMUTATION_TEST=${DO_EXTENDED_SERVER_PERMUTATION_TEST:-0}
DO_DELAY_TEST=${DO_DELAY_TEST:-0}

# An example use would be: DTLS_VERSION=-v3 scripts/dtls.test
# Note the output also consists of a single PCAP file which has a set of
# three packets (1-byte, strlen()-byte, 1-byte payload) deliniating each test.

#set -x # enable debug output

# bwrap execution environment to avoid port conflicts
if [ "${AM_BWRAPPED-}" != "yes" ]; then
    bwrap_path="$(command -v bwrap)"
    if [ -n "$bwrap_path" ]; then
        export AM_BWRAPPED=yes
        exec "$bwrap_path" --cap-add ALL --unshare-net --dev-bind / / "$0" "$@"
    fi
fi

cleanup () {
    echo
    echo "Cleaning up..."
    if [ ! -z "$UDP_PROXY_PID" ];then
        echo "Killing udp_proxy $UDP_PROXY_PID"
        kill $UDP_PROXY_PID
    fi
    if [ ! -z "$SERVER_PID" ];then
        echo "Killing server $SERVER_PID"
        kill $SERVER_PID
    fi
    if [ ! -z "$TCPDUMP_PID" ];then
        echo "Killing tcpdump $TCPDUMP_PID"
        sleep 1
        kill $TCPDUMP_PID
    fi
}

trap cleanup err exit

PROXY_PORT=1234
SERVER_PORT=4321
KEY_UPDATE_SIZE=35
NUM_TESTS_FAILED=0
NUM_TESTS_RUN=0

if [ "$DTLS_VERSION" = "-v4" ]; then
    UDP_PROXY_EXTRA_ARGS="-u"
fi

# $WOLFSSL_ROOT/tests/unit.test tests/test-dtls13.conf

set -o pipefail
prepend() { # Usage: cmd 2>&1 | prepend "sometext "
    while read line; do echo "${1}${line}"; done
}

run_test() { # usage: run_test "<testName>" "<udp-proxy args>" "<server args>" "<client args>"
    ((NUM_TESTS_RUN++))
    echo "" | nc -u 127.0.0.1 $SERVER_PORT # This is a marker for the PCAP file
    echo "$1" | nc -u 127.0.0.1 $SERVER_PORT # This is a marker for the PCAP file
    echo "" | nc -u 127.0.0.1 $SERVER_PORT # This is a marker for the PCAP file
    echo -e "\n${1}\n"
    stdbuf -oL -eL $WOLFSSL_ROOT/examples/server/server -u -p$SERVER_PORT $DTLS_VERSION $3 2>&1 | prepend "[server] " &
    SERVER_PID=$(($! - 1))
    stdbuf -oL -eL $UDP_PROXY_BIN -p $PROXY_PORT -s 127.0.0.1:$SERVER_PORT $UDP_PROXY_EXTRA_ARGS $2 2>&1 | prepend "[udp-proxy] " &
    UDP_PROXY_PID=$(($! - 1))
    sleep 0.2
    # Wrap this command in a timeout so that a deadlock won't bring down the entire test
    timeout -s KILL 1m stdbuf -oL -eL $WOLFSSL_ROOT/examples/client/client -u -p$PROXY_PORT $DTLS_VERSION $4 2>&1 | prepend "[client] "
    if [ $? != 0 ]; then
        echo "***Test failed***"
        ((NUM_TESTS_FAILED++))
    fi
    kill $SERVER_PID >&/dev/null # make sure the server is no longer running
    SERVER_PID=
    kill $UDP_PROXY_PID
    UDP_PROXY_PID=
}

test_dropping_packets () {
    for i in $(seq 0 11);do
        run_test "Dropping ${i}th packet" "-f $i" "-Ta" ""
    done

    # dropping last ack would be client error as wolfssl_read doesn't support WANT_WRITE as returned error
    for i in $(seq 0 10);do
        run_test "Testing WANT_WRITE: dropping packet $i" "-f $i" "-Ta -6" "-6"
    done
}

# this test is based on detecting newSessionTicket message by its size. This is rather fragile.
test_dropping_new_session_ticket() { # usage: test_dropping_new_session_ticket <size>
    run_test "Dropping new session ticket packet of size $1" "-F $1" "-w" "-w --waitTicket"
}

test_permutations () {
    SIDE=$1
    PERMUTATIONS=$(python3 << EOF
import itertools
for p in itertools.permutations("$2"):
    print(''.join(p))
EOF
    )
    for i in $PERMUTATIONS;do
        UDP_LOGFILE=$(mktemp)
        run_test "Testing $SIDE permutations order $i" "-r $i -S $SIDE -l $UDP_LOGFILE" "-Ta -w" "-w"
        echo "...produced $(grep -P 'client:|server:' $UDP_LOGFILE | wc -l) messages"
        rm -f $UDP_LOGFILE
    done
    echo "All $SIDE msg permutations succeeded"
}


test_time_delays () {
    DELAYS=$(python3 << EOF
import itertools
t = [0.1, 0.5, 1.1]
tt = []
for i in itertools.product(t, t, t):
    tt.append(i * 15)
for i in tt:
    print(','.join(map(lambda x: str(x) , i)))
EOF
    )
    for DELAY in $DELAYS;do
        UDP_LOGFILE=$(mktemp)
        run_test "Testing delay $DELAY" "-l $UDP_LOGFILE -t $DELAY" "-Ta -w" "-w"
        echo "...produced $(grep -P 'client:|server:' $UDP_LOGFILE | wc -l) messages"
        rm -f $UDP_LOGFILE
    done
}

echo "Starting capture"
tcpdump -i lo -n port ${SERVER_PORT} -w ${PCAP_FILENAME} -U &
TCPDUMP_PID=$!
sleep 0.5

test_dropping_packets
test_permutations client 012

if [ "$DO_EXTENDED_SERVER_PERMUTATION_TEST" = "1" ];then
   test_permutations server 0123456
else
   test_permutations server 012
fi

test_dropping_new_session_ticket 200

# TODO: fix udp_proxy to not re-order close alert before app data
if [ "$DO_DELAY_TEST" = "1" ];then
   test_time_delays
fi

if [ $NUM_TESTS_FAILED == 0 ]; then
    echo -e "\nAll $NUM_TESTS_RUN tests SUCCEEDED!!!\n"
else
    echo -e "\nThere were $NUM_TESTS_FAILED failures out of $NUM_TESTS_RUN tests\n"
fi

echo "The script ran for $SECONDS seconds"
exit $NUM_TESTS_FAILED