#!/usr/bin/env bash
# Run a command so that WRITES to /tmp land in project ./tmp instead of the real
# /tmp -- even code that hardcodes the literal "/tmp". 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 -> ./tmp, real /tmp readable, real /tmp untouched
#   2. bwrap overlay  : same, with the upper dir under TMPDIR when ./tmp is read-only
#   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 can distinguish OUR ./tmp from a project's own ./tmp and
  # only ever prune dirs we created.
  : > "$candidate_tmp/.tmpjail" 2>/dev/null || return 1
  tmp="$candidate_tmp"
  work="$candidate_work"
}

# 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 "$PWD" ||
  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