#!/usr/bin/env bash
# NAT-traversal simulation for Covalence direct (SYNC_PROTOCOL §3.3). Stands up two "home
# networks" behind NATs, joined by a middle "internet" namespace running a STUN server, and
# drives the REAL punch core (`nucleic-punch-harness`, which is `NucleicProtocol/Sync/Direct`)
# across them — so a green run is evidence about the shipping STUN/punch logic, not a mock.
#
# Topology (all Linux network namespaces + veth pairs):
#
#     peerA ── natA ──┐            ┌── natB ── peerB
#      10.0.1.2        │  inet ns  │        10.0.2.2
#                 198.51.100.x  (STUN)  203.0.113.x
#
# Each NAT namespace masquerades its home subnet toward `inet`. The masquerade *mode* is the
# whole experiment:
#   • port-preserving (default nftables masquerade)  → endpoint-independent  → CONE (punchable)
#   • `masquerade random`  (per-flow random source port) → endpoint-dependent → SYMMETRIC
#
# Usage:  sudo scripts/nat-sim/run.sh <cone|symmetric> [harness_path]
#   cone       both NATs port-preserving  → punch must SUCCEED
#   symmetric  both NATs random-port       → classification must detect symmetric, punch gives up
#
# Requires: root (netns/nft), iproute2, nftables, and a built harness (defaults to the debug
# build; pass an explicit path for a musl/release binary). Exercised inside a privileged
# container: `docker run --privileged … scripts/nat-sim/run.sh cone`.
set -euo pipefail

MODE="${1:-cone}"
HARNESS="${2:-.build/debug/nucleic-punch-harness}"

if [[ $EUID -ne 0 ]]; then echo "error: must run as root (netns + nftables)" >&2; exit 1; fi
if [[ ! -x "$HARNESS" ]]; then echo "error: harness not found at $HARNESS" >&2; exit 1; fi
case "$MODE" in cone|symmetric) ;; *) echo "error: mode must be cone|symmetric" >&2; exit 1;; esac

WORK="$(mktemp -d)"
cleanup() {
  for ns in peerA peerB natA natB inet; do ip netns del "$ns" 2>/dev/null || true; done
  rm -rf "$WORK"
}
trap cleanup EXIT

echo "==> building topology (mode: $MODE)"
for ns in peerA peerB natA natB inet; do ip netns add "$ns"; done

# veth: peerA<->natA, natA<->inet, inet<->natB, natB<->peerB
link() { # ns1 if1 addr1  ns2 if2 addr2
  ip link add "$2" netns "$1" type veth peer name "$5" netns "$4"
  ip -n "$1" addr add "$3" dev "$2"; ip -n "$1" link set "$2" up
  ip -n "$4" addr add "$6" dev "$5"; ip -n "$4" link set "$5" up
}
link peerA a-nat 10.0.1.2/24     natA nat-a 10.0.1.1/24
link natA  a-inet 198.51.100.2/24 inet inet-a 198.51.100.1/24
link inet  inet-b 203.0.113.1/24  natB nat-b 203.0.113.2/24
link natB  b-peer 10.0.2.1/24      peerB b-nat 10.0.2.2/24

for ns in peerA peerB natA natB inet; do ip -n "$ns" link set lo up; done

# Routing: peers default via their NAT; NATs reach the far *public* subnet via inet. The
# private home subnets (10.0.x) are deliberately NOT cross-routed — so the only path between
# the peers is through their public (reflexive) addresses, i.e. a genuine hole punch. (Cross-
# routing the private subnets would let the punch cheat via a directly-reachable private
# address, which would also make the symmetric test falsely "succeed".)
ip -n peerA route add default via 10.0.1.1
ip -n peerB route add default via 10.0.2.1
ip -n natA route add 203.0.113.0/24 via 198.51.100.1
ip -n natB route add 198.51.100.0/24 via 203.0.113.1

# Enable forwarding in the router namespaces. Prefer sysctl; fall back to the procfs knob
# directly (minimal container images often ship no `sysctl` binary).
enable_forwarding() {
  if ip netns exec "$1" sh -c 'command -v sysctl >/dev/null'; then
    ip netns exec "$1" sysctl -qw net.ipv4.ip_forward=1
  else
    ip netns exec "$1" sh -c 'echo 1 > /proc/sys/net/ipv4/ip_forward'
  fi
}
for ns in natA natB inet; do enable_forwarding "$ns"; done

# NAT masquerade. `random` makes the source-port mapping endpoint-dependent = symmetric.
RANDFLAG=""; [[ "$MODE" == "symmetric" ]] && RANDFLAG="random"
nat_rules() { # ns  wan-if
  ip netns exec "$1" nft -f - <<EOF
table ip nat {
  chain post {
    type nat hook postrouting priority 100;
    oifname "$2" masquerade $RANDFLAG
  }
}
EOF
}
# Masquerade on each NAT's *WAN* interface (toward inet), not its LAN side.
nat_rules natA a-inet
nat_rules natB nat-b

STUN_PORT=3478
echo "==> starting two STUN servers in inet ns (198.51.100.1, 203.0.113.1)"
# Two servers, each bound to a *specific* inet interface, so every reply leaves from the exact
# IP the client queried (a wildcard-bound server would reply from the kernel-default source,
# which the peer's NAT conntrack then drops). Two distinct STUN endpoints is what lets
# `StunClient.discover` tell cone (same mapped port both servers) from symmetric (differing).
ip netns exec inet "$HARNESS" stun-server --port "$STUN_PORT" --bind 198.51.100.1 &
STUN_PID1=$!
ip netns exec inet "$HARNESS" stun-server --port "$STUN_PORT" --bind 203.0.113.1 &
STUN_PID2=$!
trap 'kill $STUN_PID1 $STUN_PID2 2>/dev/null || true; cleanup' EXIT
sleep 1

export_stun() { echo "NUCLEIC_STUN_SERVERS=198.51.100.1:$STUN_PORT,203.0.113.1:$STUN_PORT"; }

echo "==> classifying NAT from peerA"
set +e
ip netns exec peerA env "$(export_stun)" "$HARNESS" classify
CLASSIFY_A=$?
set -e

echo "==> punching peerA <-> peerB"
TOKEN_A="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
TOKEN_B="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"
set +e
ip netns exec peerA env "$(export_stun)" "$HARNESS" punch \
  --role a --out "$WORK/a.cand" --peer-file "$WORK/b.cand" \
  --token "$TOKEN_A" --peer-token "$TOKEN_B" --budget 12 &
PA=$!
ip netns exec peerB env "$(export_stun)" "$HARNESS" punch \
  --role b --out "$WORK/b.cand" --peer-file "$WORK/a.cand" \
  --token "$TOKEN_B" --peer-token "$TOKEN_A" --budget 12 &
PB=$!
wait $PA; RA=$?
wait $PB; RB=$?
set -e

echo "==> results: classifyA=$CLASSIFY_A  punchA=$RA  punchB=$RB"

if [[ "$MODE" == "cone" ]]; then
  # Cone: classification is cone (exit 0) and both sides validate a round trip (exit 0).
  if [[ $RA -eq 0 && $RB -eq 0 ]]; then
    echo "PASS: cone NATs punched through"
    exit 0
  fi
  echo "FAIL: cone NATs should punch (punchA=$RA punchB=$RB)"; exit 1
else
  # Symmetric: classification must NOT be cone, and the punch must fail or give up
  # (non-zero) — validating that the policy doesn't waste effort where it can't win.
  if [[ $CLASSIFY_A -ne 0 && $RA -ne 0 && $RB -ne 0 ]]; then
    echo "PASS: symmetric NATs detected and punch abandoned"
    exit 0
  fi
  echo "FAIL: symmetric NATs should be detected + punch should give up (classifyA=$CLASSIFY_A punchA=$RA punchB=$RB)"; exit 1
fi
