#!/usr/bin/env bash

# timeout(1) is GNU coreutils and absent on macOS; where it's missing, run the
# command unbounded (the flaky hang this guards against is Linux-only CI).
# A prefix variable rather than a shell function: backgrounding a function
# makes $! the forked subshell, so a later "kill $server_pid" would stop the
# wrapper and orphan the server it was meant to kill.
if command -v timeout >/dev/null 2>&1; then
    TIMEOUT_KILL_1M="timeout -s KILL 1m"
else
    TIMEOUT_KILL_1M=""
fi

# 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
# Prerequisites, each of which makes the script skip (exit 77) when missing:
# the udp_proxy helper, python3, and a server build that can report an
# ephemeral port through -p 0 with -R.
# 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

[ ! -x ./examples/client/client ] && printf '\n\n%s\n' "Client doesn't exist" \
                                  && exit 1

if ./examples/client/client -? 2>&1 | grep "Client not compiled in!" ; then
    echo 'skipping dtls.test because client not compiled in.' 1>&2
    exit 77
fi

if ./examples/server/server -? 2>&1 | grep "Server not compiled in!" ; then
    echo 'skipping dtls.test because server not compiled in.' 1>&2
    exit 77
fi

# Every test runs through the udp_proxy helper. Without it each one fails at the
# client, so skip the script rather than report a wall of failures.
# command -v matches how the proxy is actually launched below: it searches
# $PATH for a bare name and tests executability for a path with a slash. An
# extra -x test would accept a ./udp_proxy that the launch could never find.
if ! command -v "$UDP_PROXY_BIN" >/dev/null 2>&1; then
    echo "skipping dtls.test because $UDP_PROXY_BIN was not found." 1>&2
    exit 77
fi

# python3 picks the proxy port for every test, on top of generating the
# permutation and delay lists, so it is a hard requirement of the whole script.
if ! command -v python3 >/dev/null 2>&1; then
    echo 'skipping dtls.test because python3 was not found.' 1>&2
    exit 77
fi

# 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

kill_server() {
    sleepseconds=1
    for i in $(jobs -pr); do
        if [ "$i" != "$TCPDUMP_PID" ]; then
            # sleep to give sanitizers time to dump backtraces.
            sleep $sleepseconds
            sleepseconds=0
            kill -9 $i
        fi
    done
    # empty print to show which backgrounded processes were killed
    sleep 0.2 && echo
}

cleanup () {
    echo
    echo "Cleaning up..."
    kill_server

    if [ ! -z "$TCPDUMP_PID" ];then
        echo "Killing tcpdump $TCPDUMP_PID"
        sleep 1
        kill $TCPDUMP_PID
    fi

    rm -f "$READY_FILE"
}

trap cleanup err exit

# Nothing ever binds the marker port - the three datagrams sent to it only tag
# the capture - so a fixed number is safe here. The server and proxy ports are
# chosen per test below.
MARKER_PORT=4321
KEY_UPDATE_SIZE=35
NUM_TESTS_FAILED=0
NUM_TESTS_RUN=0
# The server writes the port it bound to this file. The pid keeps concurrent
# runs from sharing one.
READY_FILE="$WOLFSSL_ROOT/wolfssl_dtls_ready$$"

HAVE_SS=no
command -v ss >/dev/null 2>&1 && HAVE_SS=yes

# A free UDP port for the proxy. udp_proxy has no equivalent of the server's -R
# ready file, so the port is picked here and the proxy polled below until it is
# actually bound.
free_udp_port() {
    python3 -c 'import socket
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.bind(("127.0.0.1", 0))
print(s.getsockname()[1])
s.close()'
}

# A usable port is a non-zero decimal below 65536. Worth checking rather than
# assuming: a SINGLE_THREADED build never resolves the ephemeral port in
# udp_accept(), so the ready file legitimately contains 0 there, and pointing
# the proxy at port 0 would report every test as a DTLS failure instead.
valid_port() {
    case "$1" in
        ''|*[!0-9]*) return 1 ;;
    esac
    [ "$1" -gt 0 ] && [ "$1" -lt 65536 ]
}

# Wait for $1 to be bound for UDP. Where ss is not available to be asked this
# cannot be checked, so say so rather than let a blind sleep pass for a check.
wait_udp_bound() {
    local counter=0

    if [ "$HAVE_SS" != "yes" ]; then
        echo "[warning] ss not available - cannot confirm the proxy bound $1"
        sleep 0.5
        return 0
    fi

    while [ "$counter" -lt 50 ]; do
        # Match on the local port specifically - a bare grep for the number can
        # be satisfied by the peer column or an unrelated socket.
        if [ -n "$(ss -lunH "sport = :$1" 2>/dev/null)" ]; then
            return 0
        fi
        sleep 0.1
        counter=$((counter + 1))
    done
    return 1
}

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

# Everything below depends on the server being able to report the port it bound.
# A SINGLE_THREADED or WOLFSSL_TIRTOS build never resolves an ephemeral port in
# udp_accept(), and a NO_FILESYSTEM or NETOS build compiles write_ready_file()
# to a no-op, so probe the capability once here rather than rediscovering it on
# every one of the tests below.
check_ready_file_support() {
    local probe_pid counter port

    rm -f "$READY_FILE"
    $WOLFSSL_ROOT/examples/server/server -u -p 0 -R "$READY_FILE" $DTLS_VERSION \
        >/dev/null 2>&1 &
    probe_pid=$!

    counter=0
    while [ ! -s "$READY_FILE" -a "$counter" -lt 50 ]; do
        if ! kill -0 "$probe_pid" 2>&- && [ ! -s "$READY_FILE" ]; then
            break
        fi
        sleep 0.1
        counter=$((counter + 1))
    done

    port=""
    if [ -s "$READY_FILE" ]; then
        sleep 0.1
        port=$(cat "$READY_FILE")
    fi

    kill -9 "$probe_pid" 2>&-
    wait "$probe_pid" 2>&-
    rm -f "$READY_FILE"

    valid_port "$port"
}

if ! check_ready_file_support; then
    echo 'skipping dtls.test because the server cannot report an ephemeral port' 1>&2
    echo 'with -p 0 -R (needs threads and a filesystem in the build).' 1>&2
    exit 77
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
}

# Record a failed test. Plain assignment rather than ((n++)) so the count is
# not also the command's exit status.
fail_test() {
    echo "***Test failed***${1:+ $1}"
    NUM_TESTS_FAILED=$((NUM_TESTS_FAILED + 1))
}

run_test() { # usage: run_test "<testName>" "<udp-proxy args>" "<server args>" "<client args>"
    ((NUM_TESTS_RUN++))
    echo "" | nc -u -w 1 127.0.0.1 $MARKER_PORT # This is a marker for the PCAP file
    echo "$1" | nc -u -w 1 127.0.0.1 $MARKER_PORT # This is a marker for the PCAP file
    echo "" | nc -u -w 1 127.0.0.1 $MARKER_PORT # This is a marker for the PCAP file
    echo -e "\n${1}\n"

    # Bind an ephemeral port and wait for the server to report it, rather than
    # taking a fixed port and assuming a fixed sleep covers the bind.
    rm -f "$READY_FILE"
    stdbuf -oL -eL $WOLFSSL_ROOT/examples/server/server -u -p 0 -R "$READY_FILE" $DTLS_VERSION $3 2>&1 | prepend "[server] " &
    # The pid is prepend's, but it exits on EOF when the server dies, so it
    # still tells us whether the server is alive - and saves the full poll when
    # it is not, which matters at 5040 tests with the extended permutations.
    SERVER_JOB_PID=$!

    counter=0
    while [ ! -s "$READY_FILE" -a "$counter" -lt 50 ]; do
        if ! kill -0 "$SERVER_JOB_PID" 2>&- && [ ! -s "$READY_FILE" ]; then
            fail_test "server exited before creating the ready file"
            kill_server
            return
        fi
        sleep 0.1
        counter=$((counter + 1))
    done
    if [ ! -s "$READY_FILE" ]; then
        fail_test "server did not become ready"
        kill_server
        return
    fi

    # sleep for an additional 0.1 to mitigate race on write/read of $READY_FILE:
    sleep 0.1

    SERVER_PORT=$(cat "$READY_FILE")
    if ! valid_port "$SERVER_PORT"; then
        fail_test "server reported unusable port '$SERVER_PORT'"
        kill_server
        return
    fi

    PROXY_PORT=$(free_udp_port)
    if ! valid_port "$PROXY_PORT"; then
        fail_test "could not obtain a free UDP port for the proxy"
        kill_server
        return
    fi
    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] " &
    if ! wait_udp_bound $PROXY_PORT; then
        fail_test "udp-proxy did not bind port $PROXY_PORT"
        kill_server
        return
    fi

    # Wrap this command in a timeout so that a deadlock won't bring down the entire test
    $TIMEOUT_KILL_1M stdbuf -oL -eL $WOLFSSL_ROOT/examples/client/client -u -p$PROXY_PORT $DTLS_VERSION $4 2>&1 | prepend "[client] "
    if [ $? != 0 ]; then
        fail_test
    fi
    kill_server
}

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"
# The server port is ephemeral and changes per test, so capture by protocol and
# address rather than one fixed port.
tcpdump -i lo -n udp and host 127.0.0.1 -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