#!/usr/bin/env bash
# Proves pids-guard kills a REAL fork bomb, in the danger lab only.
#
# DANGEROUS: this script creates a fork bomb. It refuses to run anywhere but the
# disposable VM. Host: /home/user/dangerlab/dangerlab-run --payload <repo>/modules/monitor \
#   --timeout 240 -- bash monitor/tests/pids-guard-kill.dangerlab.sh
#
# Scopes are made the way production makes them — systemd-run --user --scope into
# agent.slice — so the guard sees the real cgroup layout, not a mock.
#
# What it proves:
#   1. an agent scope crossing the kill threshold is destroyed, in bounded time
#   2. a sibling agent scope keeps forking throughout — the wedge never spreads
#   3. human-owned cgroups at the same saturation are NOT killed
set -uo pipefail

[[ "$(hostname)" == dangerlab ]] || {
  echo "REFUSED: this script fork-bombs the machine it runs on; it runs in the danger lab only." >&2
  exit 78
}

HERE="$(cd "$(dirname "$0")" && pwd)"
GUARD="$HERE/../bin/pids-guard"
PASS=0; FAIL=0
ok()  { PASS=$((PASS+1)); printf 'PASS %s\n' "$1"; }
bad() { FAIL=$((FAIL+1)); printf 'FAIL %s\n     %s\n' "$1" "$2"; }

UID_NUM="$(id -u)"
CG=/sys/fs/cgroup/user.slice/user-${UID_NUM}.slice/user@${UID_NUM}.service
SCOPE_CAP=512
SLICE_CAP=3072
DECOY_TASKS=340          # 66% of SCOPE_CAP: past the 60% kill threshold

cur() { cat "$1/pids.current" 2>/dev/null || echo -1; }
denied() { awk '/^max /{print $2}' "$1/pids.events.local" 2>/dev/null || echo 0; }

scope() { # unit slice tasksmax -- command...
  local unit=$1 slice=$2 tasks=$3; shift 3
  systemd-run --user --scope --quiet --collect --unit="$unit" --slice="$slice" \
    -p TasksMax="$tasks" -- "$@"
}
sleepers() { printf 'for i in $(seq 1 %s); do sleep 600 & done; wait' "$1"; }

systemctl --user set-property agent.slice TasksMax="$SLICE_CAP" 2>/dev/null

BOMB_UNIT=confine-agent-1-2.scope
SIB_UNIT=confine-agent-3-4.scope
HUMAN_UNIT=human-decoy.scope
TERM_UNIT=vte-spawn-11111111-2222-3333-4444-555555555555.scope
BOMB="$CG/agent.slice/$BOMB_UNIT"
SIBLING="$CG/agent.slice/$SIB_UNIT"
HUMAN="$CG/human.slice/$HUMAN_UNIT"
TERM_CG="$CG/app.slice/$TERM_UNIT"

scope "$HUMAN_UNIT" human.slice "$SCOPE_CAP" bash -c "$(sleepers $DECOY_TASKS)" &
disown
scope "$TERM_UNIT" app.slice "$SCOPE_CAP" bash -c "$(sleepers $DECOY_TASKS)" &
disown
for _ in $(seq 1 30); do
  [[ "$(cur "$HUMAN")" -gt 300 && "$(cur "$TERM_CG")" -gt 300 ]] && break
  sleep 0.5
done
if [[ "$(cur "$HUMAN")" -gt 300 && "$(cur "$TERM_CG")" -gt 300 ]]; then
  ok "setup: human decoy=$(cur "$HUMAN") vte decoy=$(cur "$TERM_CG") tasks, both past the kill threshold"
else
  bad "setup: decoys populated" "human=$(cur "$HUMAN") vte=$(cur "$TERM_CG")"
  exit 1
fi

export XDG_STATE_HOME=/tmp/pg-state
export PIDS_KILL_PCT=60 PIDS_SLOW_EVERY=1000000
python3 "$GUARD" --watch &
GUARD_PID=$!
sleep 1
if kill -0 "$GUARD_PID" 2>/dev/null; then
  ok "guard: watching from $(awk -F: '{print $3}' /proc/$GUARD_PID/cgroup), outside every scope it kills"
else
  bad "guard: started" "the watcher exited immediately"
  exit 1
fi

# A sibling agent scope that keeps forking: if the wedge spreads past the bomb's own
# scope, these forks start failing.
scope "$SIB_UNIT" agent.slice "$SCOPE_CAP" bash -c \
  'for i in $(seq 1 200); do if (exec true) 2>/dev/null; then echo ok; else echo fail; fi; sleep 0.05; done' \
  > /tmp/sibling.log 2>&1 &
SIB=$!
disown

start=$(date +%s.%N)
scope "$BOMB_UNIT" agent.slice "$SCOPE_CAP" bash -c 'b() { b | b & }; b' >/dev/null 2>&1 &
disown

peak=0; killed_at=""
for _ in $(seq 1 400); do
  n=$(cur "$BOMB")
  (( n > peak )) && peak=$n
  if [[ -z "$killed_at" && "$n" -le 1 && "$peak" -gt 50 ]]; then
    killed_at=$(date +%s.%N)
    break
  fi
  sleep 0.05
done
for _ in $(seq 1 100); do kill -0 "$SIB" 2>/dev/null || break; sleep 0.2; done
sib_ok=$(grep -c '^ok$' /tmp/sibling.log 2>/dev/null); sib_ok=${sib_ok:-0}
sib_fail=$(grep -c '^fail$' /tmp/sibling.log 2>/dev/null); sib_fail=${sib_fail:-0}

if [[ -n "$killed_at" ]]; then
  ok "kill: bomb scope emptied $(python3 -c "print(f'{$killed_at-$start:.2f}s')") after launch, peak $peak tasks"
else
  bad "kill: bomb scope emptied" "still $(cur "$BOMB") tasks after 20s, peak $peak"
fi

if grep -q "KILL .*$BOMB_UNIT" "$XDG_STATE_HOME/pids-guard/events.log" 2>/dev/null; then
  ok "kill: recorded — $(grep -m1 "KILL .*$BOMB_UNIT" "$XDG_STATE_HOME/pids-guard/events.log")"
else
  bad "kill: recorded in events.log" "$(tail -3 "$XDG_STATE_HOME/pids-guard/events.log" 2>/dev/null)"
fi

if (( sib_fail == 0 && sib_ok > 0 )); then
  ok "containment: sibling agent scope forked $sib_ok/$((sib_ok+sib_fail)) times during the bomb"
else
  bad "containment: sibling kept forking" "ok=$sib_ok fail=$sib_fail"
fi

if [[ "$(denied "$CG/agent.slice")" == 0 ]]; then
  ok "containment: agent.slice itself never denied a fork (pids.events.local max=0)"
else
  bad "containment: agent.slice denied forks" "max=$(denied "$CG/agent.slice") — the wedge spread past the scope"
fi

for d in "$HUMAN" "$TERM_CG"; do
  n=$(cur "$d")
  if (( n > 300 )); then
    ok "human-safety: $(basename "$d") untouched at $n tasks, past the kill threshold"
  else
    bad "human-safety: $(basename "$d") untouched" "dropped to $n — the guard killed a human cgroup"
  fi
done

exec 2>/dev/null   # teardown SIGTERMs are job-control noise, not results
{ kill "$GUARD_PID"; wait "$GUARD_PID"; } 2>/dev/null
systemctl --user stop "$HUMAN_UNIT" "$TERM_UNIT" 2>/dev/null
wait 2>/dev/null
printf 'bomb scope peak: %s tasks of a %s cap\n' "$peak" "$SCOPE_CAP"
printf 'pids-guard-kill: pass=%d fail=%d\n' "$PASS" "$FAIL"
(( FAIL == 0 ))