#!/usr/bin/env bash
# Run a command so that WRITES to /tmp land in a per-tree jail dir under the user
# cache instead of the real /tmp -- even code that hardcodes the literal "/tmp".
# The jail is keyed to the enclosing git work tree, so every invocation from
# anywhere inside one checkout sees the SAME /tmp regardless of cwd. Uses an
# overlay so that
# EXISTING /tmp entries (X11 display socket, Postgres sockets, askpass sockets,
# ...) stay READABLE inside the jail -- only new writes are diverted.
#
# Layering, best -> fallback:
#   1. bwrap overlay  : /tmp writes -> cache jail, real /tmp readable and untouched
#   2. bwrap overlay  : same, with the upper dir under TMPDIR when the cache is unusable
#   3. bwrap bind     : /tmp fully replaced by the upper dir (hides real /tmp sockets)
#   4. TMPDIR env only: redirect honest tempfile APIs; hardcoded /tmp NOT blocked
set -euo pipefail

tmp=""
work=""

# overlayfs needs upper and work on the same writable filesystem, and neither may
# sit inside the lower dir (/tmp).
try_base() {
  local base="$1"
  local candidate_tmp="$base/tmp"
  local candidate_work="$base/.tmpjail-work"
  case "$base" in /tmp | /tmp/*) return 1 ;; esac
  mkdir -p "$candidate_tmp" "$candidate_work" 2>/dev/null || return 1
  [[ -w "$candidate_tmp" && -w "$candidate_work" ]] || return 1
  # Marker so tmpjail-gc only ever prunes dirs we created.
  : > "$candidate_tmp/.tmpjail" 2>/dev/null || return 1
  tmp="$candidate_tmp"
  work="$candidate_work"
}

# One jail per work tree, not per cwd: a `cd` inside a checkout must not swap /tmp
# out from under files an earlier invocation wrote there.
jail_key() {
  local root
  root="$(/usr/bin/git rev-parse --show-toplevel 2>/dev/null)" || root=""
  [[ -n "$root" ]] || root="$PWD"
  printf '%s' "$root" | tr -c 'A-Za-z0-9' '_'
}

# Already jailed (nested launch): the child inherits our /tmp overlay; don't
# stack another bwrap (nested userns can fail and it buys nothing).
if [[ -n "${TMPJAIL_ACTIVE:-}" ]]; then
  exec "$@"
fi

try_base "${XDG_CACHE_HOME:-$HOME/.cache}/agent-tmp/tmpjail/$(jail_key)" ||
  try_base "${TMPDIR:-}/tmpjail-$$" ||
  try_base "${XDG_CACHE_HOME:-$HOME/.cache}/tmpjail-$$" || {
    echo "tmpjail: no writable jail directory -- running unjailed" >&2
    exec "$@"
  }

# The jail's userns maps only uid 1000, so every root-owned file reads as `nobody` inside it and
# ssh rejects /etc/ssh/ssh_config.d/* on its owner check — exit 255 for every host, which kills all
# remote offload. An empty tmpfs there is owned by the jailed uid, so ssh finds nothing to reject.
if command -v bwrap >/dev/null 2>&1; then
  if bwrap --help 2>&1 | grep -q -- '--overlay\b'; then
    exec bwrap \
      --dev-bind / / \
      --overlay-src /tmp --overlay "$tmp" "$work" /tmp \
      --tmpfs /etc/ssh/ssh_config.d \
      --setenv TMPDIR /tmp \
      --setenv TMPJAIL_ACTIVE 1 \
      --die-with-parent \
      -- "$@"
  fi
  # Overlay unsupported: fall back to a full bind (hides real /tmp contents).
  exec bwrap \
    --dev-bind / / \
    --bind "$tmp" /tmp \
    --tmpfs /etc/ssh/ssh_config.d \
    --setenv TMPDIR /tmp \
    --setenv TMPJAIL_ACTIVE 1 \
    --die-with-parent \
    -- "$@"
else
  echo "tmpjail: bwrap not found -- redirecting via TMPDIR only (hardcoded /tmp NOT blocked)" >&2
  exec env TMPDIR="$tmp" TMP="$tmp" TEMP="$tmp" TMPJAIL_ACTIVE=1 "$@"
fi