#!/usr/bin/env bash
# confine.sh — the single point where a resource ceiling is applied.
#
# usage: confine.sh <build|agent> COMMAND [ARG...]
#
# Fail-closed: every path either confines the command or refuses to run it.
# There is no branch that execs COMMAND unbounded.
#
#   1. cgroup (preferred)  — systemd-run --user --scope: MemorySwapMax, CPUQuota,
#                            TasksMax. Bounds the whole descendant tree
#                            regardless of PATH, absolute paths, or argv shape.
#   2. rlimit (floor)      — prlimit RLIMIT_AS, when the user manager is unreachable
#                            (e.g. from inside an agent sandbox with no session bus).
#                            Per-process and coarse, but bounds a single runaway.
#   3. refuse              — neither available: exit 78, command does not run.
#
# Tunables (per class, env):
#   BUILD_PER_JOB_CPU/BUILD_TASKS_MAX/BUILD_AS_MAX/BUILD_SWAP_MAX
#   AGENT_CPU/AGENT_TASKS_MAX/AGENT_AS_MAX/AGENT_SWAP_MAX
#
# *_AS_MAX bounds virtual address space, never resident memory: V8 reserves >20G of VA for
# Wasm cages (cursor-agent dies at 20G with "Cannot allocate Wasm memory"; 32G+ works).
# There is deliberately no MemoryMax/MemoryHigh here — see build.slice.d/90-ceiling.conf.
set -euo pipefail

if (( $# < 2 )); then
  printf 'usage: %s <build|agent> COMMAND [ARG...]\n' "$0" >&2
  exit 64
fi

class=$1
shift

case "$class" in
  build)
    slice=build.slice
    cpu_quota="${BUILD_PER_JOB_CPU:-200}%"
    tasks_max="${BUILD_TASKS_MAX:-512}"
    as_max="${BUILD_AS_MAX:-64G}"
    swap_max="${BUILD_SWAP_MAX:-0}"
    oom_policy=stop
    ;;
  agent)
    slice=agent.slice
    # Hard CPU and swap ceilings starved interactive tool shells and treated zram's normal
    # steady state as a failure. TasksMax plus pressure-based oomd remain the backstop.
    cpu_quota=""
    tasks_max="${AGENT_TASKS_MAX:-1024}"
    as_max="${AGENT_AS_MAX:-64G}"
    swap_max="infinity"
    oom_policy=continue
    ;;
  *)
    printf 'confine: invalid class: %s\n' "$class" >&2
    exit 64
    ;;
esac

# Same class already applied to this tree via a REAL cgroup scope: re-scoping would only
# re-derive the same ceiling. A DIFFERENT class still scopes — a build started inside an
# agent session belongs in build.slice under the build ceiling, not the agent's.
#
# The marker carries the mechanism (":cgroup" vs ":rlimit"), not just the class. An ancestor
# that fell back to rlimit_floor (no systemd-run reachable at that moment) is NOT actually
# capped on TasksMax, so a bare "$class" flag would make every descendant trust that weak
# floor forever and skip cgroup scoping permanently.
#
# The marker is still only a hint: `systemd-run --scope` forks the caller, so a nested
# scope created elsewhere (app.slice, no ceiling) inherits the variable while leaving the
# ceiling behind. The cgroup is the fact — both must agree before a launch skips scoping.
ceiling_in_force() {
  local leaf max
  leaf="$(cut -d: -f3 /proc/self/cgroup 2>/dev/null | tail -1)"
  [[ "$leaf" == */"$slice"/*.scope ]] || return 1
  max="$(cat "/sys/fs/cgroup${leaf}/pids.max" 2>/dev/null)" || return 1
  [[ -n "$max" && "$max" != max ]]
}

if [[ "${CONFINE_ACTIVE:-}" == "$class:cgroup" ]] && ceiling_in_force; then
  exec "$@"
fi

to_bytes() {
  local v="${1^^}" n
  case "$v" in
    *K) n=$(( ${v%K} * 1024 )) ;;
    *M) n=$(( ${v%M} * 1024 * 1024 )) ;;
    *G) n=$(( ${v%G} * 1024 * 1024 * 1024 )) ;;
    *) n="$v" ;;
  esac
  printf '%s' "$n"
}

# An unprivileged process can lower a hard rlimit but never raise it. Nested inside an
# outer confinement the requested ceiling may exceed what is already in force, so clamp
# instead of failing with EPERM.
as_bytes() {
  local bytes hard
  bytes="$(to_bytes "$as_max")"
  hard="$(ulimit -H -v)"
  if [[ "$hard" != unlimited ]] && (( bytes > hard * 1024 )); then
    bytes=$(( hard * 1024 ))
  fi
  printf '%s' "$bytes"
}

rlimit_floor() {
  local bytes
  bytes="$(as_bytes)"
  if command -v prlimit >/dev/null 2>&1; then
    exec prlimit --as="$bytes" -- "$@"
  fi
  # ulimit -v takes KiB. The subshell lowers its own hard limit before exec, and an
  # unprivileged process can never raise a hard limit back.
  if ulimit -v "$(( bytes / 1024 ))" 2>/dev/null; then
    exec "$@"
  fi
  printf 'confine: no cgroup and no rlimit mechanism available; refusing to run unbounded: %s\n' "$1" >&2
  exit 78
}

if [[ "${CONFINE_NO_SYSTEMD:-0}" == 1 ]] || ! command -v systemd-run >/dev/null 2>&1; then
  export CONFINE_ACTIVE="$class:rlimit"
  rlimit_floor "$@"
fi

# systemd-run is on PATH inside sandboxes that cannot reach the user manager, so its
# presence proves nothing; ask the manager itself. Creates no unit, unlike a probe run.
if ! command -v systemctl >/dev/null 2>&1 \
   || ! systemctl --user show-environment >/dev/null 2>&1; then
  export CONFINE_ACTIVE="$class:rlimit"
  rlimit_floor "$@"
fi

# systemd derives a scope name from the caller's PID; nested inside an existing scope
# that name collides ("already loaded or has a fragment file"), so name it explicitly.
unit="confine-${class}-$$-${RANDOM}"

properties=(
  # Do not inherit a stale parent ceiling into a new agent/build scope.
  -p MemoryHigh=infinity
  -p MemoryMax=infinity
  -p MemorySwapMax="$swap_max"
  -p TasksMax="$tasks_max"
  -p OOMPolicy="$oom_policy"
)
[[ -n "$cpu_quota" ]] && properties+=(-p CPUQuota="$cpu_quota")

if [[ -n "${CONFINE_CPU_WEIGHT:-}" ]]; then
  properties+=(-p CPUWeight="$CONFINE_CPU_WEIGHT")
fi
# io.weight is absent unless the io controller is delegated to the user manager.
if [[ -n "${CONFINE_IO_WEIGHT:-}" ]]; then
  user_cgroup="$(systemctl --user show --property=ControlGroup --value 2>/dev/null || true)"
  [[ -n "$user_cgroup" && -e "/sys/fs/cgroup${user_cgroup}/io.weight" ]] \
    && properties+=(-p IOWeight="$CONFINE_IO_WEIGHT")
fi

environment=(-E "PATH=$PATH" -E "CONFINE_ACTIVE=$class:cgroup")
for name in BUILD_SLOT_HELD CPU_GUARD_ACTIVE BUILD_SLOT_DIR BUILD_SLOTS \
  BUILD_SLOT_FIXED BUILD_SLOT_FLOOR BUILD_SLOT_TIMEOUT BUILD_SLOT_STAT \
  AGENT_BUILD_SCOPE_ACTIVE LOCAL_GATE_ACTIVE NODE_OPTIONS; do
  if [[ -v "$name" ]]; then
    environment+=(-E "$name=${!name}")
  fi
done

# prlimit inside the scope as well: the cgroup bounds the tree's total, the rlimit
# bounds any single process, and the cheaper one trips first on a single runaway.
prlimit_prefix=()
if command -v prlimit >/dev/null 2>&1; then
  prlimit_prefix=(prlimit --as="$(as_bytes)" --)
fi

exec systemd-run --user --scope --quiet --collect \
  --unit="$unit" \
  --slice="$slice" \
  --same-dir \
  "${environment[@]}" \
  "${properties[@]}" \
  -- "${prlimit_prefix[@]}" "$@"