#!/usr/bin/env bash
# Asserts the containment chain as it is actually configured on this workstation,
# not as the repo describes it. Every check here failed silently at least once:
# a launcher resolved past the shim, a ceiling was shadowed by a same-named user
# drop-in, and oomd held a stale view that left the account unmonitored.
set -uo pipefail

AGENT_CEILING=/etc/systemd/user/agent.slice.d/90-ceiling.conf
BUILD_CEILING=/etc/systemd/user/build.slice.d/90-ceiling.conf

# Skipping on the ceilings themselves would let a machine that lost them report green,
# so the gate is the launcher: wherever agents can start, the ceilings must be present.
if [[ ! -r "$HOME/.claude/bin/_agent-build-scope" ]]; then
  echo "SKIP containment invariants (no agent install: _agent-build-scope absent)"
  exit 0
fi

pass=0
fail=0
skip=0

check() { # $1=name $2=actual $3=expected
  if [[ "$2" == "$3" ]]; then
    pass=$((pass + 1))
    echo "PASS $1"
  else
    fail=$((fail + 1))
    echo "FAIL $1 (want=$3 got=$2)"
  fi
}

BIN="$HOME/.claude/bin"

for name in claude codex cursor-agent; do
  target="$(basename "$(readlink -f "$BIN/$name" 2>/dev/null)" 2>/dev/null)"
  check "entry-point-$name-is-jailed" "$target" "_tmpjail-shim.sh"
done

scope_invariants() { # $1=name-prefix $2=file
  local body
  body="$(<"$2")"
  case "$body" in
    *"confine.sh"*) check "$1-confines" yes yes ;;
    *) check "$1-confines" no yes ;;
  esac
  case "$body" in
    *"exec \"\$@\""*) check "$1-has-no-raw-exec" no yes ;;
    *) check "$1-has-no-raw-exec" yes yes ;;
  esac
}

scope_invariants agent-scope "$BIN/_agent-build-scope"

# The installed copy passing says nothing about the file a land would install next: a branch
# carrying the fail-open variant merges green while the box still holds the good one.
SCOPE_SRC="$(cd "$(dirname "${BASH_SOURCE[0]}")/../bin" && pwd)/_agent-build-scope"
if [[ -r "$SCOPE_SRC" ]]; then
  scope_invariants agent-scope-src "$SCOPE_SRC"
else
  check "agent-scope-src-present" no yes
fi

# confine.sh's reentry guard must only skip re-scoping for a descendant whose ancestor
# actually holds a cgroup scope. An ancestor that only reached the rlimit floor (no
# systemd-run reachable at that moment) is NOT capped on TasksMax; trusting a bare class
# marker there is exactly how session-364.scope and a run-p*.scope both ran an agent tree
# with no TasksMax ceiling at all after one transient systemd-run failure.
# Baseline: this test process may itself already run inside an outer confinement (its own
# ulimit -v already clamped by an ancestor), so "untouched by confine.sh" is whatever a bare
# child sees with confine.sh out of the picture entirely -- not a hardcoded "unlimited".
CONFINE_BASELINE_ULIMIT="$(bash -c 'ulimit -v')"

confine_reentry_check() { # $1=name $2=confine.sh path $3=ancestor CONFINE_ACTIVE $4=expect-skip(yes|no)
  local name="$1" src="$2" ancestor="$3" expect_skip="$4" out
  # AGENT_AS_MAX pinned strictly below the ambient baseline: an unambiguous drop proves
  # rlimit_floor actually ran, rather than "happened to already equal the default".
  # The marker alone no longer earns a skip — the guard also reads the live cgroup — so the
  # skip case has to run under a real agent.slice scope rather than claim one in the env.
  if [[ "$expect_skip" == yes ]]; then
    out="$(systemd-run --user --scope --quiet --collect \
      --unit="confine-reentry-test-$$" --slice=agent.slice -p TasksMax=2048 -p MemoryMax=8G -- \
      env CONFINE_ACTIVE="$ancestor" CONFINE_NO_SYSTEMD=1 AGENT_AS_MAX=1G \
      bash "$src" agent bash -c 'ulimit -v' 2>/dev/null | tail -1)"
  else
    out="$(CONFINE_ACTIVE="$ancestor" CONFINE_NO_SYSTEMD=1 AGENT_AS_MAX=1G \
      bash "$src" agent bash -c 'ulimit -v' 2>/dev/null)"
  fi
  if [[ "$expect_skip" == yes ]]; then
    # A real cgroup ancestor legitimately short-circuits before rlimit_floor ever runs,
    # so the child sees exactly the ambient ulimit -- confine.sh touched nothing.
    check "$name" "$out" "$CONFINE_BASELINE_ULIMIT"
  else
    # Must NOT be a no-op skip: rlimit_floor must have run and clamped to AGENT_AS_MAX.
    check "$name" "$out" 1048576
  fi
}
CONFINE_INSTALLED="$HOME/.claude/lib/confine.sh"
CONFINE_SRC="$(cd "$(dirname "${BASH_SOURCE[0]}")/../lib" && pwd)/confine.sh"

confine_reentry_suite() { # $1=label $2=confine.sh path
  confine_reentry_check "confine-reentry-cgroup-ancestor-skips-$1" "$2" "agent:cgroup" yes
  confine_reentry_check "confine-reentry-rlimit-ancestor-rescopes-$1" "$2" "agent:rlimit" no
  confine_reentry_check "confine-reentry-no-ancestor-rescopes-$1" "$2" "" no
}

# The installed leg is deploy-drift coverage: it is only meaningful where the installed path
# resolves to THIS repo's deploy artifact. A host that runs the suite without deploying
# overdeck (an offloaded build box) carries an unrelated confine.sh copy that no land can
# update, so asserting against it reds every branch that introduces a confine.sh change.
installed_deploy_artifact() {
  local real root
  real="$(readlink -f "$CONFINE_INSTALLED" 2>/dev/null)"
  [[ -n "$real" && -r "$real" ]] || return 1
  root="$(git -C "${real%/*}" rev-parse --show-toplevel 2>/dev/null)" || return 1
  [[ -n "$root" && "$real" == "$root/modules/workstation/claude/lib/confine.sh" ]]
}

if installed_deploy_artifact; then
  confine_reentry_suite installed "$CONFINE_INSTALLED"
else
  skip=$((skip + 1))
  echo "SKIP confine-reentry-installed ($CONFINE_INSTALLED is not this repo's deploy artifact)"
fi

if [[ -r "$CONFINE_SRC" ]]; then
  confine_reentry_suite src "$CONFINE_SRC"
else
  check "confine-present-src" no yes
fi

# A confine scope must carry its own non-memory control and no size ceiling.
check_confine_scope_invariant() { # $1=label $2=class $3=cpu quota $4=tasks $5=oom $6=require-nonzero-swap
  local label="$1" class="$2" expect_cpu="$3" expect_tasks="$4" expect_oom="$5" require_swap="$6"
  shift 6
  local out leaf memory_high memory_max memory_swap cpu_quota tasks_max oom_policy
  local cpu_percent expect_cpu_quota cpu_quota_normalized
  cpu_percent="${expect_cpu%\%}"
  if [[ "$cpu_percent" =~ ^[0-9]+$ ]]; then
    expect_cpu_quota="$(awk -v n="$cpu_percent" 'BEGIN { printf "%g", n / 100 }')"
  elif [[ "$expect_cpu" == "infinity" ]]; then
    expect_cpu_quota="infinity"
  else
    expect_cpu_quota=""
  fi
  out="$(env "$@" "$CONFINE_SRC" "$class" bash -c '
    leaf="$(cut -d: -f3 /proc/self/cgroup | tail -1)"
    unit="$(basename "$leaf")"
    printf "%s|%s|%s|%s|%s|%s|%s\n" \
      "$leaf" \
      "$(cat "/sys/fs/cgroup$leaf/memory.high")" \
      "$(cat "/sys/fs/cgroup$leaf/memory.max")" \
      "$(cat "/sys/fs/cgroup$leaf/memory.swap.max")" \
      "$(systemctl --user show "$unit" -p CPUQuotaPerSecUSec --value 2>/dev/null)" \
      "$(cat "/sys/fs/cgroup$leaf/pids.max")" \
      "$(systemctl --user show "$unit" -p OOMPolicy --value 2>/dev/null)"' 2>/dev/null)"
  IFS='|' read -r leaf memory_high memory_max memory_swap cpu_quota tasks_max oom_policy <<<"$out"
  [[ "$leaf" == */$class.slice/* ]] || { check_glob "$label-scope-path" "$leaf" "*/$class.slice/*"; return; }
  check "$label-no-memory-high" "$memory_high" max
  check "$label-no-memory-max" "$memory_max" max
  cpu_quota_normalized="${cpu_quota%s}"
  cpu_quota_normalized="$(awk -v n="$cpu_quota_normalized" 'BEGIN { sub(/\.?0+$/,"",n); if (n ~ /\.$/) sub(/\.$/,"",n); if (n == "") n="0"; print n }')"
  if [[ "$require_swap" == "positive" ]]; then
    if [[ "$memory_swap" =~ ^[1-9][0-9]*$ ]]; then
      pass=$((pass + 1)); echo "PASS $label-memory-swap"
    else
      fail=$((fail + 1)); echo "FAIL $label-memory-swap (want=a positive byte count got=$memory_swap)"
    fi
  elif [[ "$require_swap" == "max" ]]; then
    check "$label-memory-swap" "$memory_swap" max
  else
    check "$label-memory-swap" "$memory_swap" "0"
  fi
  check "$label-cpu-quota" "$cpu_quota_normalized" "$expect_cpu_quota"
  check "$label-tasks-max" "$tasks_max" "$expect_tasks"
  check "$label-oom-policy" "$oom_policy" "$expect_oom"
}

if systemctl --user show-environment >/dev/null 2>&1; then
  # Agent CPU and swap are deliberately unlimited: a hard CPU quota starved tool shells,
  # while zram makes swap occupancy a normal steady state. TasksMax remains the backstop.
  check_confine_scope_invariant "agent" agent "infinity" "1024" "continue" max \
    AGENT_CPU=600 AGENT_TASKS_MAX=1024 AGENT_AS_MAX=1G
  check_confine_scope_invariant "build" build "240%" "512" "stop" 0 \
    BUILD_PER_JOB_CPU=240 BUILD_TASKS_MAX=512 BUILD_AS_MAX=4G BUILD_SWAP_MAX=0
else
  echo "SKIP confine scope invariants (no user manager)"
fi

# A scope or service placed in agent.slice/build.slice without going through confine.sh
# carries no TasksMax and no CPUQuota, so one runaway consumes the whole slice. Memory is
# deliberately not part of this check any more: no confine scope carries a size ceiling.
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)"
if command -v git >/dev/null 2>&1 && [[ -d "$REPO_ROOT/.git" || -f "$REPO_ROOT/.git" ]]; then
  unceiled=""
  while IFS= read -r f; do
    [[ "$f" == *"/tests/"* || "$f" == *.test.sh ]] && continue
    case "$f" in *.sh|*.py|*.js|*.ts|*.service|*.conf) ;; *) continue ;; esac
    grep -q -E '^[^#]*(TasksMax|CPUQuota|confine\.sh)' "$REPO_ROOT/$f" || unceiled="$unceiled $f"
  done < <(git -C "$REPO_ROOT" grep -lE -- '(--slice=|^Slice=)(agent|build)\.slice' 2>/dev/null)
  check "in-repo-slice-placements-carry-a-ceiling" "${unceiled:-none}" none
else
  echo "SKIP in-repo slice placements (not a git checkout)"
fi

# Locks the invariant in the repo drop-ins themselves: re-adding a size ceiling to either
# file fails here, not silently on the next deploy.
for ceiling in \
  "$REPO_ROOT/modules/workstation/claude/systemd/user-root/agent.slice.d/90-ceiling.conf" \
  "$REPO_ROOT/modules/workstation/claude/systemd/user-root/build.slice.d/90-ceiling.conf"; do
  check "repo-dropin-declares-no-size-ceiling-$(basename "$(dirname "$ceiling")")" \
    "$( { grep -cE '^(MemoryMax|MemoryHigh)=' "$ceiling" 2>/dev/null; } || true )" 0
  check "repo-dropin-declares-managedoom-pressure-kill-$(basename "$(dirname "$ceiling")")" \
    "$(grep -E '^ManagedOOMMemoryPressure=' "$ceiling" 2>/dev/null | tail -1)" ManagedOOMMemoryPressure=kill
done

if [[ "${CI:-}" == true ]]; then
  skip=$((skip + 2))
  echo "SKIP installed slice ceilings (CI validates repository drop-ins and source-created scopes)"
elif systemctl --user show-environment >/dev/null 2>&1; then
  [[ -r "$AGENT_CEILING" ]] && check "agent-ceiling-installed" yes yes || check "agent-ceiling-installed" no yes
  [[ -r "$BUILD_CEILING" ]] && check "build-ceiling-installed" yes yes || check "build-ceiling-installed" no yes
else
  echo "SKIP installed slice ceilings (no user manager)"
fi

MEM_TOTAL_BYTES=$(awk '/^MemTotal:/ {print $2 * 1024}' /proc/meminfo)

declared() { # $1=drop-in $2=key — bytes, or pct:<n> for a host-relative ceiling
  local raw
  raw="$(grep -E "^$2=" "$1" 2>/dev/null | tail -1)"
  raw="${raw#*=}"
  [[ -n "$raw" ]] || { printf 'unset'; return; }
  case "$raw" in
    infinity|0) printf '%s' "$raw" ;;
    *%) printf 'pct:%s' "${raw%\%}" ;;
    *) numfmt --from=iec "$raw" 2>/dev/null || printf 'unparsed:%s' "$raw" ;;
  esac
}

check_ceiling() { # $1=name $2=effective $3=declared — percentages match within 2%
  local want=$3 target
  # A key the drop-in never declares reads back as the uncapped systemd default.
  if [[ "$want" == unset ]]; then
    check "$1" "$2" infinity
    return
  fi
  if [[ "$want" == pct:* ]]; then
    target=$(( MEM_TOTAL_BYTES * ${want#pct:} / 100 ))
    if [[ "$2" =~ ^[0-9]+$ ]] && (( $2 * 100 >= target * 98 && $2 * 100 <= target * 102 )); then
      pass=$((pass + 1)); echo "PASS $1"
    else
      fail=$((fail + 1)); echo "FAIL $1 (want~=$target got=$2)"
    fi
    return
  fi
  check "$1" "$2" "$want"
}

if [[ "${CI:-}" == true ]]; then
  skip=$((skip + 1))
  echo "SKIP effective slice ceilings (CI host state is not deployed from this checkout)"
elif systemctl --user show-environment >/dev/null 2>&1; then
  effective() { systemctl --user show "$1" -p "$2" --value 2>/dev/null; }
  # The drop-in is the intended ceiling and systemctl reports the effective one:
  # a shadowing user drop-in or a skipped daemon-reload shows up as a mismatch.
  for key in MemoryHigh MemoryMax MemorySwapMax; do
    # agent.slice is governed by sustained memory pressure, never by a number: agent work
    # is legitimately heavy and no byte figure separates a 35G build from a runaway. The
    # want is hardcoded rather than derived from the drop-in so that re-adding a ceiling
    # to that file fails here instead of silently redefining the invariant.
    check "agent-slice-${key,,}" "$(effective agent.slice "$key")" infinity
    check_ceiling "build-slice-${key,,}" "$(effective build.slice "$key")" "$(declared "$BUILD_CEILING" "$key")"
  done
else
  echo "SKIP slice ceilings (no user manager)"
fi

if [[ "${CI:-}" == true ]]; then
  skip=$((skip + 1))
  echo "SKIP oomd installed state (CI host is not deployed from this checkout)"
elif command -v oomctl >/dev/null 2>&1 && systemctl is-active systemd-oomd >/dev/null 2>&1; then
  # A slice already carrying processes has already pushed its ManagedOOM registration to
  # oomd -- reading oomctl costs nothing and proves the real, deployed agent.slice/
  # build.slice are watched, which is the actual doctrine this check guards. Only a
  # genuinely empty slice needs a probe unit parked in it, and an empty slice has nothing
  # running for a pressure squeeze to hit.
  #
  # Never `daemon-reload` here: that used to be how this check forced a re-check, but
  # forcing it while build.slice hosts this suite's own live process tree synchronizes
  # oomd's re-evaluation against a slice that may already be over its pressure limit --
  # that mid-run re-check is what killed two full factory gates. Starting a unit already
  # pushes ManagedOOM registration on its own; no reload is needed to observe it.
  slice_occupied() { # $1=slice
    local uid dir
    uid="$(id -u)"
    dir="/sys/fs/cgroup/user.slice/user-${uid}.slice/user@${uid}.service/$1"
    [[ -d "$dir" ]] && find "$dir" -name cgroup.procs -exec cat {} \; 2>/dev/null | grep -q .
  }
  watched="$(oomctl 2>/dev/null || true)"
  for slice in agent.slice build.slice; do
    if slice_occupied "$slice"; then
      case "$watched" in
        *"/$slice"*) check "oomd-watches-$slice" yes yes ;;
        *) check "oomd-watches-$slice" no yes ;;
      esac
    else
      unit="oomd-invariant-${slice%%.*}.service"
      systemd-run --user --quiet --collect --slice="$slice" --unit="$unit" sleep 20 2>/dev/null || true
      sleep 1
      out="$(oomctl 2>/dev/null || true)"
      systemctl --user stop "$unit" 2>/dev/null || true
      case "$out" in
        *"/$slice"*) check "oomd-watches-$slice" yes yes ;;
        *) check "oomd-watches-$slice" no yes ;;
      esac
    fi
  done
  # Pure read: the effective ManagedOOM keys on the live slices match doctrine
  # (pressure-only kill, never a size ceiling) without spawning or reloading anything.
  for slice in agent.slice build.slice; do
    check "effective-$slice-managedoom-pressure" \
      "$(systemctl --user show "$slice" -p ManagedOOMMemoryPressure --value 2>/dev/null)" kill
    # Reported as a ratio scaled to UINT32_MAX (2576980377 == 60%), not a "60%" string.
    limit_raw="$(systemctl --user show "$slice" -p ManagedOOMMemoryPressureLimit --value 2>/dev/null)"
    limit_pct="$(awk -v n="$limit_raw" 'BEGIN { printf "%.0f", n / 4294967295 * 100 }' 2>/dev/null)"
    check "effective-$slice-managedoom-pressure-limit" "${limit_pct}%" 60%
  done
else
  echo "SKIP oomd (not active)"
fi

echo "passed=$pass failed=$fail skipped=$skip"
(( fail == 0 ))