#!/usr/bin/env bash
# Hermetic launcher tests. ssh, rsync, systemd-run, systemctl, loginctl and node
# are stubbed on PATH, so nothing leaves this machine and no container starts.
set -u

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LAUNCHER="$SCRIPT_DIR/../bin/agent-sandbox"
ROOT="$(mktemp -d)"
trap 'rm -rf "$ROOT"' EXIT
PASS=0
FAIL=0
FAIL_LIST=()

ok()  { printf 'PASS %s\n' "$1"; PASS=$((PASS + 1)); }
bad() { printf 'FAIL %s: %s\n' "$1" "$2"; FAIL_LIST+=("$1"); FAIL=$((FAIL + 1)); }

# shellcheck disable=SC2016  # a literal sh program, expanded by the stub not here
STUB_BODY='#!/usr/bin/env bash
set -u
name="$(basename -- "$0")"
{ printf "%s" "$name"; for a in "$@"; do printf " %q" "$a"; done; printf "\n"; } >>"$STUB_LOG"
case "$name" in
  node) printf "%s" "${NODE_HOST:-stub-host}" ;;
  ssh)
    [ "${SSH_MODE:-ok}" = fail ] && exit "${SSH_RC:-99}"
    case "${*}" in *"printf %s \"\$HOME\""*) printf "%s" "${REMOTE_HOME:-/remote/home}" ;; esac
    ;;
  rsync)
    [ "${RSYNC_MODE:-ok}" = fail ] && exit "${RSYNC_RC:-99}"
    dest="${!#}"
    case "$dest" in *:*) ;; *) mkdir -p "$(dirname "$dest")" && : >"$dest" ;; esac
    ;;
  loginctl) [ "${1-}" = show-user ] && printf "%s\n" "${LINGER:-yes}" ;;
  systemctl) : ;;
  systemd-run)
    status_file=""; payload=(); setenv=(); in_payload=0
    for a in "$@"; do
      if [ "$in_payload" = 1 ]; then payload+=("$a"); continue; fi
      case "$a" in
        --setenv=AGENT_SANDBOX_REEXEC_STATUS_FILE=*)
          status_file="${a#--setenv=AGENT_SANDBOX_REEXEC_STATUS_FILE=}"; setenv+=("${a#--setenv=}") ;;
        --setenv=*) setenv+=("${a#--setenv=}") ;;
        --) in_payload=1 ;;
      esac
    done
    case "${SYSTEMD_RUN_MODE:-payload}" in
      fail) exit "${SYSTEMD_RUN_RC:-99}" ;;
      exit-7) [ -n "$status_file" ] && printf "7\n" >"$status_file"; exit 0 ;;
      lose-status)
        kept=()
        for e in "${setenv[@]}"; do
          case "$e" in AGENT_SANDBOX_REEXEC_STATUS_FILE=*) ;; *) kept+=("$e") ;; esac
        done
        setenv=(${kept[@]+"${kept[@]}"}) ;;
    esac
    # Real systemd-run applies --setenv to the payload; without it the launcher
    # would never see its own recursion marker. The spawned process is fresh, so
    # its NoNewPrivs is 0 — the fixture the payload reads is updated to match.
    for i in "${!payload[@]}"; do
      [ "${payload[$i]}" = --proc-status ] && printf "NoNewPrivs:\t0\n" >"${payload[$((i + 1))]}"
    done
    [ "${#payload[@]}" -gt 0 ] && { env "${setenv[@]}" "${payload[@]}"; exit $?; }
    ;;
esac
exit 0
'

new_case() { # name -> echoes the case dir
  local dir="$ROOT/$1"
  mkdir -p "$dir/bin" "$dir/home" "$dir/codex" "$dir/config" "$dir/runtime/$(id -u)"
  printf '%s' "$STUB_BODY" >"$dir/bin/stub"
  chmod +x "$dir/bin/stub"
  local c
  for c in node ssh rsync loginctl systemctl systemd-run; do ln -sf stub "$dir/bin/$c"; done
  : >"$dir/stub.log"
  printf 'NoNewPrivs:\t0\n' >"$dir/proc-status"
  printf '%s' "$dir"
}

# Leading VAR=value arguments become the launcher's environment.
launch() { # dir [VAR=value...] [launcher args...]
  local dir="$1"; shift
  local env_args=()
  while [ $# -gt 0 ] && [[ "$1" == *=* && "$1" != -* ]]; do env_args+=("$1"); shift; done
  RUN_OUT="$(
    env -i \
      HOME="$dir/home" CODEX_HOME="$dir/codex" XDG_CONFIG_HOME="$dir/config" \
      XDG_RUNTIME_DIR="$dir/runtime/$(id -u)" USER="${USER:-user}" \
      PATH="$dir/bin:$PATH" STUB_LOG="$dir/stub.log" \
      "${env_args[@]}" \
      timeout 20 bash "$LAUNCHER" "$@" 2>&1
  )"
  RUN_RC=$?
}

stub_called() { grep -q "^$2 " "$1/stub.log"; }

case_no_reexec_when_nonewprivs_zero() {
  local d; d="$(new_case no-reexec)"
  launch "$d" --host box --id c --proc-status "$d/proc-status" --print-plan
  [ "$RUN_RC" = 0 ] && ! stub_called "$d" systemd-run
}

case_reexec_when_nonewprivs_one() {
  local d; d="$(new_case reexec)"
  printf 'NoNewPrivs:\t1\n' >"$d/proc-status"
  touch "$d/runtime/$(id -u)/bus.plain"
  python3 -c "import socket,sys; s=socket.socket(socket.AF_UNIX); s.bind(sys.argv[1])" "$d/runtime/$(id -u)/bus"
  launch "$d" --host box --id c --proc-status "$d/proc-status" --preflight-root "$d/runtime" --print-plan
  local line; line="$(grep '^systemd-run ' "$d/stub.log")"
  [ "$RUN_RC" = 0 ] || return 1
  for flag in -- --user --wait --pipe --collect --expand-environment=no --setenv=PATH= --setenv=HOME=; do
    [[ "$line" == *"$flag"* ]] || { printf 'missing %s in: %s\n' "$flag" "$line" >&2; return 1; }
  done
}

case_recursion_guard() {
  local d; d="$(new_case recursion)"
  printf 'NoNewPrivs:\t1\n' >"$d/proc-status"
  launch "$d" AGENT_SANDBOX_REEXEC=1 --host box --id c --proc-status "$d/proc-status"
  [ "$RUN_RC" = 13 ] && ! stub_called "$d" systemd-run
}

case_recursion_guard_by_cgroup() {
  local d; d="$(new_case recursion-cgroup)"
  printf 'NoNewPrivs:\t1\n' >"$d/proc-status"
  printf '0::/user.slice/user-0.slice/user@0.service/agent-sandbox-1-2-3.service\n' >"$d/cgroup"
  launch "$d" --host box --id c --proc-status "$d/proc-status" --cgroup-path "$d/cgroup"
  [ "$RUN_RC" = 13 ] && ! stub_called "$d" systemd-run
}

case_payload_without_status_is_12() {
  local d; d="$(new_case systemd-no-status)"
  printf 'NoNewPrivs:\t1\n' >"$d/proc-status"
  python3 -c "import socket,sys; s=socket.socket(socket.AF_UNIX); s.bind(sys.argv[1])" "$d/runtime/$(id -u)/bus"
  launch "$d" SYSTEMD_RUN_MODE=lose-status \
    --host box --id c --proc-status "$d/proc-status" --preflight-root "$d/runtime" --print-plan
  [ "$RUN_RC" = 12 ]
}

case_print_plan_makes_no_remote_call() {
  local d; d="$(new_case plan-no-ssh)"
  printf 'cred' >"$d/codex/auth.json"
  launch "$d" --host box --id c --runtime codex --proc-status "$d/proc-status" --print-plan -- agent go
  [ "$RUN_RC" = 0 ] || return 1
  ! stub_called "$d" ssh && ! stub_called "$d" rsync
}

# The node expands the credential mount source itself; an unexpanded marker would
# mount a path that does not exist.
case_credential_mount_source_expands_on_the_node() {
  local d; d="$(new_case cred-src)"
  printf 'cred' >"$d/codex/auth.json"
  launch "$d" --host box --id c --runtime codex --proc-status "$d/proc-status" --print-plan
  [ "$RUN_RC" = 0 ] || return 1
  # shellcheck disable=SC2016  # $HOME is expanded by the node's shell, not here
  [[ "$RUN_OUT" == *'"$HOME"/.local/state/overdeck-sandbox/creds/'*'/codex/auth.json:/sandbox-secrets/codex/auth.json:ro'* ]] &&
    ! grep -q 'remote-home' <<<"$RUN_OUT"
}

case_systemd_run_failure_is_12() {
  local d; d="$(new_case systemd-fail)"
  printf 'NoNewPrivs:\t1\n' >"$d/proc-status"
  python3 -c "import socket,sys; s=socket.socket(socket.AF_UNIX); s.bind(sys.argv[1])" "$d/runtime/$(id -u)/bus"
  launch "$d" SYSTEMD_RUN_MODE=fail SYSTEMD_RUN_RC=99 \
    --host box --id c --proc-status "$d/proc-status" --preflight-root "$d/runtime"
  [ "$RUN_RC" = 12 ]
}

case_payload_exit_is_passed_through() {
  local d; d="$(new_case systemd-exit7)"
  printf 'NoNewPrivs:\t1\n' >"$d/proc-status"
  python3 -c "import socket,sys; s=socket.socket(socket.AF_UNIX); s.bind(sys.argv[1])" "$d/runtime/$(id -u)/bus"
  launch "$d" SYSTEMD_RUN_MODE=exit-7 \
    --host box --id c --proc-status "$d/proc-status" --preflight-root "$d/runtime"
  [ "$RUN_RC" = 7 ]
}

preflight_case() { # dir -> runs a launcher that needs the re-exec
  launch "$1" "${@:2}" --host box --id c --proc-status "$1/proc-status" --preflight-root "$1/runtime"
}

case_preflight_linger() {
  local d; d="$(new_case preflight-linger)"
  printf 'NoNewPrivs:\t1\n' >"$d/proc-status"
  python3 -c "import socket,sys; s=socket.socket(socket.AF_UNIX); s.bind(sys.argv[1])" "$d/runtime/$(id -u)/bus"
  preflight_case "$d" LINGER=no
  [ "$RUN_RC" = 11 ] && [[ "$RUN_OUT" == *"loginctl enable-linger"* ]] && [[ "$RUN_OUT" == *lingering* ]]
}

case_preflight_runtime_dir() {
  local d; d="$(new_case preflight-xdg)"
  printf 'NoNewPrivs:\t1\n' >"$d/proc-status"
  python3 -c "import socket,sys; s=socket.socket(socket.AF_UNIX); s.bind(sys.argv[1])" "$d/runtime/$(id -u)/bus"
  preflight_case "$d" XDG_RUNTIME_DIR=/tmp/wrong
  [ "$RUN_RC" = 11 ] && [[ "$RUN_OUT" == *XDG_RUNTIME_DIR* ]] && [[ "$RUN_OUT" == *"loginctl enable-linger"* ]]
}

case_preflight_missing_bus() {
  local d; d="$(new_case preflight-bus)"
  printf 'NoNewPrivs:\t1\n' >"$d/proc-status"
  preflight_case "$d"
  [ "$RUN_RC" = 11 ] && [[ "$RUN_OUT" == *"bus socket"* ]] && [[ "$RUN_OUT" == *"loginctl enable-linger"* ]]
}

case_preflight_skipped_without_nonewprivs() {
  local d; d="$(new_case preflight-skipped)"
  launch "$d" --host box --id c --proc-status "$d/proc-status" --print-plan
  [ "$RUN_RC" = 0 ] && ! stub_called "$d" loginctl
}

case_unreadable_proc_status_fails_closed() {
  local d; d="$(new_case proc-unreadable)"
  launch "$d" --host box --id c --proc-status "$d/does-not-exist" --preflight-root "$d/runtime"
  [ "$RUN_RC" = 11 ]
}

case_credential_ro_mount() {
  local d; d="$(new_case cred-mount)"
  printf 'cred' >"$d/codex/auth.json"
  launch "$d" --host box --id c --runtime codex --proc-status "$d/proc-status" --print-plan
  [ "$RUN_RC" = 0 ] && [[ "$RUN_OUT" == *":/sandbox-secrets/codex/auth.json:ro"* ]]
}

# Codex's first implementation generated the stage script against the node-side
# staging directory, so the container read from a path nothing was mounted at.
case_stage_script_reads_the_mount_destination() {
  local d; d="$(new_case cred-stage-dest)"
  printf 'cred' >"$d/codex/auth.json"
  launch "$d" --host box --id c --runtime codex --proc-status "$d/proc-status" --print-plan -- agent go
  [ "$RUN_RC" = 0 ] || return 1
  local dest=/sandbox-secrets/codex/auth.json
  [[ "$RUN_OUT" == *"$dest:ro"* ]] || return 1
  # the copy source inside the container must be that same mount destination
  [[ "$RUN_OUT" == *"cp"*"$dest"* ]]
}

case_stage_script_survives_remote_shell_reserialization() {
  local d command parsed
  d="$(new_case cred-stage-roundtrip)"
  printf 'cred' >"$d/codex/auth.json"
  launch "$d" --host box --id c --runtime codex --proc-status "$d/proc-status" --print-plan -- codex exec probe
  [ "$RUN_RC" = 0 ] || return 1
  command="${RUN_OUT#*agent-sandbox plan: remote-cmd}"
  parsed="$(bash -c 'eval "set -- $1"; printf "<%s>\n" "$@"' roundtrip "$command")" || return 1
  [[ "$parsed" == *'</bin/sh>'*'<-c>'* ]] || return 1
  [[ "$parsed" == *'<sh>'*'<codex>'*'<exec>'*'<probe>'* ]] || return 1
  [[ "$parsed" == *'dirname "$agent_cred_target"'* ]] || return 1
  [[ "$parsed" != *'<HOME%/>'* && "$parsed" != *'<HOME->'* && "$parsed" != *'<agent_cred_target%/*>'* ]]
}

case_credential_identity_is_the_realpath() {
  local d; d="$(new_case cred-identity)"
  mkdir -p "$d/accounts/acct-b"
  printf 'cred' >"$d/accounts/acct-b/auth.json"
  ln -s "$d/accounts/acct-b/auth.json" "$d/codex/auth.json"
  launch "$d" --host box --id c --runtime codex --proc-status "$d/proc-status" --print-plan
  [ "$RUN_RC" = 0 ] && [[ "$RUN_OUT" == *"credential identity: $d/accounts/acct-b/auth.json"* ]]
}

case_missing_credential_exits_10() {
  local d; d="$(new_case cred-missing)"
  launch "$d" --host box --id c --runtime codex --proc-status "$d/proc-status"
  [ "$RUN_RC" = 10 ] && ! stub_called "$d" ssh && ! stub_called "$d" rsync
}

case_credential_bytes_never_leak() {
  local d; d="$(new_case cred-leak)"
  printf 'TOP_SECRET_BYTES' >"$d/codex/auth.json"
  launch "$d" --host box --id c --runtime codex --proc-status "$d/proc-status" -- agent go
  [ "$RUN_RC" = 0 ] || return 1
  ! grep -q TOP_SECRET_BYTES <<<"$RUN_OUT" && ! grep -q TOP_SECRET_BYTES "$d/stub.log"
}

case_agent_runs_as_child_not_exec() {
  local d; d="$(new_case cred-child)"
  printf 'cred' >"$d/codex/auth.json"
  launch "$d" --host box --id c --runtime codex --proc-status "$d/proc-status" --print-plan -- agent go
  [ "$RUN_RC" = 0 ] && ! grep -qE '(^| )exec ' <<<"$RUN_OUT"
}

case_credential_staging_dir_is_removed() {
  local d; d="$(new_case cred-cleanup)"
  printf 'cred' >"$d/codex/auth.json"
  launch "$d" --host box --id c --runtime codex --proc-status "$d/proc-status" -- agent go
  [ "$RUN_RC" = 0 ] || return 1
  grep -qE '^ssh .*rm.*creds' "$d/stub.log"
}

case_credential_staging_failure_aborts() {
  local d; d="$(new_case cred-rsync-fail)"
  printf 'cred' >"$d/codex/auth.json"
  launch "$d" RSYNC_MODE=fail --host box --id c --runtime codex --proc-status "$d/proc-status" -- agent go
  [ "$RUN_RC" = 10 ]
}

case_git_mounts_same_path_rw() {
  local d; d="$(new_case git-rw)"
  launch "$d" --host box --id c --git-common /srv/git-common --git-dir /srv/gitdir \
    --proc-status "$d/proc-status" --print-plan
  [ "$RUN_RC" = 0 ] &&
    [[ "$RUN_OUT" == *"/srv/git-common:/srv/git-common:rw"* ]] &&
    [[ "$RUN_OUT" == *"/srv/gitdir:/srv/gitdir:rw"* ]]
}

case_git_relative_path_exits_14() {
  local d; d="$(new_case git-relative)"
  launch "$d" --host box --id c --git-common relative/path --proc-status "$d/proc-status"
  [ "$RUN_RC" = 14 ] && ! stub_called "$d" ssh
}

case_git_colon_path_exits_14() {
  local d; d="$(new_case git-colon)"
  launch "$d" --host box --id c --git-dir "/srv/with:colon" --proc-status "$d/proc-status"
  [ "$RUN_RC" = 14 ] && ! stub_called "$d" ssh
}

case_workspace_relative_exits_14() {
  local d; d="$(new_case workspace-relative)"
  launch "$d" --host box --id c --workspace rel/ws --proc-status "$d/proc-status"
  [ "$RUN_RC" = 14 ]
}

case_ssh_failure_has_no_local_fallback() {
  local d; d="$(new_case ssh-fail)"
  launch "$d" SSH_MODE=fail SSH_RC=77 --host box --id c --proc-status "$d/proc-status" -- agent go
  [ "$RUN_RC" = 77 ]
}

case_id_required() {
  local d; d="$(new_case usage-id)"
  launch "$d" --host box --proc-status "$d/proc-status"
  [ "$RUN_RC" = 2 ]
}

case_unknown_flag() {
  local d; d="$(new_case usage-flag)"
  launch "$d" --host box --id c --proc-status "$d/proc-status" --nope
  [ "$RUN_RC" = 2 ]
}

run() { # description function
  local out
  if out="$("$2" 2>&1)"; then ok "$1"; else bad "$1" "${out:-assertion failed}"; fi
}

# Sourcing the file exposes the cases individually without running the suite.
if [[ "${BASH_SOURCE[0]}" != "$0" ]]; then return 0; fi

run 'NoNewPrivs 0 does not re-exec'                     case_no_reexec_when_nonewprivs_zero
run 'NoNewPrivs 1 re-execs with the pinned systemd-run flags' case_reexec_when_nonewprivs_one
run 'a second NoNewPrivs 1 pass refuses to recurse (13)' case_recursion_guard
run 'a re-exec unit cgroup refuses to recurse (13)'      case_recursion_guard_by_cgroup
run 'systemd-run failing before the payload maps to 12'  case_systemd_run_failure_is_12
run 'a payload that records no status maps to 12'        case_payload_without_status_is_12
run '--print-plan performs no ssh and no rsync'          case_print_plan_makes_no_remote_call
run 'credential mount source expands on the node'        case_credential_mount_source_expands_on_the_node
run 'a payload exit of 7 is returned as 7, not 12'       case_payload_exit_is_passed_through
run 'preflight fails closed on linger (11)'              case_preflight_linger
run 'preflight fails closed on XDG_RUNTIME_DIR (11)'     case_preflight_runtime_dir
run 'preflight fails closed on a missing bus (11)'       case_preflight_missing_bus
run 'preflight is skipped when NoNewPrivs is 0'          case_preflight_skipped_without_nonewprivs
run 'an unreadable proc status takes the contained path' case_unreadable_proc_status_fails_closed
run 'credential is mounted ro under /sandbox-secrets'    case_credential_ro_mount
run 'stage script copies from the mount destination'     case_stage_script_reads_the_mount_destination
run 'stage script survives remote shell reserialization' case_stage_script_survives_remote_shell_reserialization
run 'credential identity is the resolved realpath'       case_credential_identity_is_the_realpath
run 'a missing credential exits 10 before any remote io' case_missing_credential_exits_10
run 'credential bytes never reach output or argv'        case_credential_bytes_never_leak
run 'the agent runs as a child, never exec'              case_agent_runs_as_child_not_exec
run 'the node-side credential dir is removed'            case_credential_staging_dir_is_removed
run 'a failed credential upload aborts with 10'          case_credential_staging_failure_aborts
run 'git mirrors mount rw at the same absolute path'     case_git_mounts_same_path_rw
run 'a relative git path exits 14 before ssh'            case_git_relative_path_exits_14
run 'a git path containing a colon exits 14 before ssh'  case_git_colon_path_exits_14
run 'a relative workspace exits 14'                      case_workspace_relative_exits_14
run 'an ssh failure never falls back to the host'        case_ssh_failure_has_no_local_fallback
run '--id is still required'                             case_id_required
run 'an unknown flag still exits 2'                      case_unknown_flag

printf '\nSummary: PASS=%d FAIL=%d\n' "$PASS" "$FAIL"
if ((FAIL > 0)); then
  printf 'Failures:\n'
  for f in "${FAIL_LIST[@]}"; do printf ' - %s\n' "$f"; done
  exit 1
fi