align / code /docs /LEAN_SANDBOX.md
viai957's picture
Add files using upload-large-folder tool
d909bd8 verified
|
Raw History Blame Contribute Delete
16.6 kB

Lean REPL OS sandbox (H3)

Status: align/verify/lean/sandbox.py + tests/test_lean_sandbox.py exist and pass; nothing imports the sandbox yet. Turning it on is the ~45-line patch at the end (not applied: the P2 baseline run imports repl_pool.py, and its result must stay comparable to the post-SFT run).

The gap this closes

strict.py and the #print axioms audit decide whether a proof is accepted. They do not limit what Lean does while checking it. #eval, run_cmd, initialize, elab/macro run arbitrary code as the RL job's user, on a box that holds checkpoints, hf_cache/token, ~/.ssh and a neighbouring training run (same uid: kill(2), ptrace(2), /proc/<pid>/environ all work). Before this, containment was prlimit + a memory watchdog: resource limits, not capability limits. The tests demonstrate each of these succeeding in the real REPL outside the sandbox and failing inside (test_real_lean_elaborator_io_is_contained: read a secret, write outside scratch, signal another process).

Design (what is and is not a read-only /)

Probed on this host (kernel 6.1, max_user_namespaces=764532): unshare -Urnm works; remounting the inherited / read-only works only per mount and the view would still show every secret. Unprivileged pivot_root into a fresh tmpfs works, so the sandbox is an allow-list root instead: the child sees only

path mode
/usr (+ /bin /lib /lib64 /sbin symlinks), /etc/{ld.so.cache,ld.so.conf,nsswitch.conf,passwd,group,localtime} ro bind
ro_paths = <lean>/elan, <lean>/repl, <lean>/workspace ro bind
/dev/{null,zero,full,random,urandom} ro tmpfs + device binds
/proc fresh procfs of the new pid ns
/tmp private tmpfs, 64 MB, nosuid nodev noexec
scratch_dir (per worker) the only writable host directory
root / tmpfs remounted ro, old root detached (umount2(MNT_DETACH))

Everything else (/home, /root, /apps/* incl. checkpoints and hf_cache, <lean>/tmp, <lean>/logs) does not exist. Paths leading to scratch exist as empty skeleton dirs.

Other properties: new user ns (uid/gid mapped 1:1, not fake root, so exec drops all caps: CapEff=0; setgroups denied), net ns (loopback down: even 127.0.0.1 is unreachable; abstract unix sockets of the host invisible), pid ns (neighbours invisible and unkillable; when the REPL, pid 1, dies the kernel kills every descendant incl. double-fork daemons), ipc/uts ns, PR_SET_NO_NEW_PRIVS, RLIMIT_CORE=0, own session/pgrp, minimal env (ENV_KEEP allow-list; HOME=scratch, TMPDIR=/tmp, PATH pruned to dirs that exist). nice/prlimit/stdbuf run inside the sandbox (sandbox -- nice prlimit stdbuf repl), so the pool's rlimits are unchanged and stdbuf's LD_PRELOAD still reaches the REPL (/usr is visible).

Fail closed: if any step of setup fails the command is not executed (exit 125/126, message in the worker's stderr file -> ReplCrash with tail). ReplPool.from_toolchain(sandbox=True) raises if namespaces are unavailable rather than falling back.

Why a supervisor, and what Popen.pid means now

ReplPool spawns from ThreadPoolExecutor threads, so preexec_fn is unsafe and unshare(CLONE_NEWUSER) needs a single-threaded caller. The launcher is a tiny python -I -S process: unshare, fork the REPL into the new pid ns, stay outside and poll its parent process. Popen.pid is the supervisor; the REPL's host pid is published in a pidfile outside the sandbox (inner_pid()), which the memory watchdog uses (patch below; verified: identical pss_anon readings). killpg(Popen.pid) still kills everything (tested). PR_SET_PDEATHSIG is not used: prctl(2) ties it to the spawning thread; on this kernel 6.1 it was measured not to fire for a non-leader thread, but it differs across kernels, so nothing depends on it.

What is NOT claimed

  • No seccomp filter. Namespaces are a capability boundary, not a proof; kernel bugs reachable from an unprivileged userns (the price of using one) are out of scope. The adversarial false-accept hunt is a separate task.
  • The toolchain (elan, repl, workspace) is shared read-only between workers. A proof cannot modify it, but can read all of it (Mathlib is public).
  • GPUs are not exposed. /proc of the new pid ns still shows the kernel-global files (/proc/meminfo, /proc/cpuinfo, /proc/sys, ...): host load and RAM are observable, not secret.
  • A proof can still burn CPU/RAM up to the pool's limits and write fsize_mb into scratch (the scratch dir is on /apps: size it/clean it per worker; .sandbox-root is an empty dir the setup leaves in it).
  • Tiny race: if the supervisor is SIGKILLed between fork and the child's prctl(PDEATHSIG), the REPL survives until its stdin closes.
  • Linux x86_64/aarch64 only (pivot_root syscall number); other arch -> setup fails closed.

Evidence

All numbers from this host while the P2 baseline held the GPUs (1-2 Lean workers, one import Mathlib header per worker, no heavy-job lock taken).

Property tests (tests/test_lean_sandbox.py, 31 pass, ~13 s, plus 2 slow real-Lean tests)

Every property: probe run unsandboxed (control, must succeed) then inside (must fail).

property control outside inside
connect to a live 127.0.0.1 listener connects ENETUNREACH (101)
connect to host abstract unix socket connects fails
write a host-writable file outside scratch / / / /usr / ro toolchain OK fails (EROFS/ENOENT), host file absent
write inside scratch OK OK, visible on host
read a host secret outside the allow-list; /home, /root, ~, /etc/shadow readable does not exist
exec a binary dropped in /tmp - EACCES (noexec); /tmp not the host's
rename/modify an allow-listed file - fails; content unchanged
mount() remount-rw / bind (setuid mount(8) too) - EPERM / non-zero
.., /proc/1/root, /proc/self/cwd/.. to reach old root; open fds - unreachable; <=4 fds
NoNewPrivs, CapEff/CapPrm, uid, core rlimit 0, host caps, uid, unlimited 1, 0, same non-zero uid, (0,0)
see/signal a live same-uid process, read its environ visible, kill 0 OK pid 1 only, ESRCH
double-fork + setsid daemon - gone when the REPL exits
env: HF_TOKEN-style var, LD_PRELOAD passes dropped; PATH pruned
own session / pgrp (no start_new_session) - sid == pgid == supervisor
parent process SIGKILLed - REPL and supervisor gone <5 s
spawning thread exits (pool pattern) - REPL survives
killpg(Popen.pid) - REPL dies
setup failure (missing scratch) - command NOT run, exit 125/126

Real REPL (slow; run outside pytest because the instructions forbid -m slow while the baseline runs): #eval that reads a host secret / writes outside scratch / kill -0s a neighbour succeeds without the sandbox and fails inside; 10 known-good + sorry_plain + eval_side_effect + exit_after_theorem verdicts are byte-identical inside and outside.

Mutation record

Each line removed/flipped in sandbox.py, suite re-run (the file was restored; diff clean):

mutant killed by
no CLONE_NEWNET test_network_connect_fails_inside_succeeds_outside
no PR_SET_NO_NEW_PRIVS test_no_new_privs_no_capabilities_same_uid
no CLONE_NEWPID (host /proc bound) test_other_processes_are_invisible_and_unsignalable
scratch bound read-only test_scratch_is_writable_and_shared_with_host
toolchain bound read-write test_allow_listed_path_is_readable_but_read_only
/home or a sibling dir bound in test_read_outside_allow_list_fails_inside_succeeds_outside
old root not detached test_no_old_root_reachable
parent-death poll disabled test_repl_dies_when_the_parent_process_dies
uid mapped to fake root test_cannot_remount_or_mount_over_the_view
no setsid test_own_session_and_process_group
/tmp exec allowed test_tmp_is_private_noexec_and_not_the_host_tmp
setup error falls through to exec test_setup_failure_never_runs_the_command
scratch existence unchecked (host dirs created) same test
env not filtered / PATH not pruned test_minimal_env_... / test_child_sees_only_the_spec_env_...
RLIMIT_CORE left test_no_new_privs_no_capabilities_same_uid
no pidfile test_pidfile_names_the_repl_host_pid
root tmpfs left writable test_write_outside_scratch_...
/tmp tmpfs without a size cap (defaults to 50% of RAM, box has no swap) test_tmp_tmpfs_is_size_capped (review)
closerange removed (fd passed by a caller crosses the exec) test_inherited_fds_do_not_cross_the_exec (review)
recursive bind (MS_BIND|MS_REC) of a ro path test_nested_writable_mount_under_a_ro_path_is_not_carried (review; before it, only caught by accident)
REPL's own PDEATHSIG removed (supervisor alone SIGKILLed) test_killing_only_the_supervisor_kills_the_repl (review)
no alias for a symlinked ro_paths / scratch_dir test_toolchain_reached_through_a_symlink_is_usable (review)
inner_pid trusts a stale pidfile test_inner_pid_ignores_a_stale_pidfile (review)
launcher broken wholesale (3 variants) test_launcher_works_whenever_the_kernel_allows_namespaces (before it existed the same mutants turned all 24 tests into SKIPS)

Survivors, honestly: (a) PR_SET_PDEATHSIG on the supervisor: kernel 6.1 does not deliver it on thread exit, so test_repl_survives_the_spawning_thread_exiting cannot fail here (it pins the contract on kernels that do). (b) Deleting the supervisor's explicit "kill inner on parent death": equivalent, the inner's own PDEATHSIG on the (single-threaded) supervisor covers it - kept as two layers. Not mutation-tested: MS_REC|MS_PRIVATE propagation (new mount ns already does not propagate upward here). Equivalent mutants noted by review: binding /etc files read-write (the host forbids the write anyway).

Review notes: a mount nested under a ro_paths entry makes setup REFUSE (kernel EINVAL on a non-recursive bind over a locked child mount; fail closed, exit 126) -- list such a mount as its own ro_paths entry. test_pidfile_names_the_repl_host_pid was flaky (the pidfile is written before the child's exec; the test now waits for it). The pidfile is published BEFORE the exec, so for a few ms inner_pid names the launcher's own child still running python.

Overhead (real Lean 4.9 / Mathlib d1d1e4b72c REPL)

Single worker, 10 known-good ceiling proofs x3 reps: cold import Mathlib header 3.18 s -> 3.33 s (+0.15 s, once per worker); per-proof median 20.8 -> 20.4 ms, mean 80.9 -> 79.8 ms (no measurable difference). Pool of 2 workers via ReplPool.from_toolchain(sandbox=True) using the patched module (patch applied to a COPY of the tree, not the repo), 10 proofs x5 over map: warm 3.33 -> 3.67 s (+0.33 s, once), median 23.9 -> 22.1 ms, mean 85.2 -> 81.5 ms; verdicts identical; memory_report() identical pss_anon (256/252 MB) through inner_pid; kill() + respawn works; no leftover sandbox.py processes. Existing tests/test_lean_verifier.py passes (198 passed with the sandbox tests) on the patched copy. Cost: one extra ~10 MB python per worker.

The patch (NOT applied)

Touches only align/verify/lean/repl_pool.py; imports the sandbox lazily, default sandbox=False is behaviour-identical. sandbox_supported() actually launches the helper once (a sysctl is not proof).

--- a/align/verify/lean/repl_pool.py	2026-10-13 04:43:15.450960951 +0530
+++ b/align/verify/lean/repl_pool.py	2026-10-13 04:43:15.460961181 +0530
@@ -54,7 +54,10 @@
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from typing import Any, Callable, Iterator, Optional, Sequence
+from typing import TYPE_CHECKING, Any, Callable, Iterator, Optional, Sequence
+
+if TYPE_CHECKING:  # the sandbox module is imported lazily, only when a pool opts in
+    from align.verify.lean.sandbox import SandboxSpec
 
 # --------------------------------------------------------------------------- errors
 
@@ -199,7 +202,10 @@
         limits: ReplLimits = ReplLimits(),
         stderr_path: Optional[str] = None,
         wrap: bool = True,
+        sandbox: Optional["SandboxSpec"] = None,
     ) -> None:
+        #: opt-in OS sandbox (align/verify/lean/sandbox.py); None = unchanged behaviour
+        self.sandbox = sandbox
         self.argv = list(argv)
         self.env = dict(env)
         self.cwd = cwd
@@ -216,6 +222,14 @@
     # -- lifecycle
 
     def _wrapped_argv(self) -> list[str]:
+        argv = self._limited_argv()
+        if self.sandbox is not None:  # sandbox OUTSIDE nice/prlimit/stdbuf: they run inside it
+            from align.verify.lean.sandbox import sandbox_argv
+
+            argv = sandbox_argv(self.sandbox, argv)
+        return argv
+
+    def _limited_argv(self) -> list[str]:
         if not self.wrap:
             return list(self.argv)
         lim = self.limits
@@ -237,6 +251,13 @@
             pre += [stdbuf, "-o0", "-e0"]
         return pre + self.argv
 
+    def _spawn_env(self) -> dict[str, str]:
+        if self.sandbox is None:
+            return self.env
+        from align.verify.lean.sandbox import helper_env
+
+        return helper_env()  # the launcher gets nothing; the REPL's env travels in sandbox.env
+
     def start(self) -> None:
         if self.alive:
             return
@@ -252,7 +273,7 @@
             stdin=subprocess.PIPE,
             stdout=subprocess.PIPE,
             stderr=err,
-            env=self.env,
+            env=self._spawn_env(),
             cwd=self.cwd,
             start_new_session=True,  # own process group => killpg reaps every descendant
             close_fds=True,
@@ -316,7 +337,12 @@
             return ""
 
     def memory(self) -> Optional[dict[str, float]]:
-        return process_memory(self.pid) if self.pid else None
+        pid = self.pid
+        if self.sandbox is not None:  # Popen.pid is the supervisor; the REPL is its child
+            from align.verify.lean.sandbox import inner_pid
+
+            pid = inner_pid(self.sandbox, self.pid)
+        return process_memory(pid) if pid else None
 
     # -- I/O
 
@@ -671,6 +697,7 @@
         max_requests: int = 1000,
         default_timeout_s: float = 120.0,
         preflight: bool = False,
+        sandbox: bool = False,
     ) -> "ReplPool":
         n = n_workers if n_workers is not None else default_workers()
         if preflight:
@@ -679,12 +706,29 @@
         os.makedirs(scratch, exist_ok=True)
 
         def launcher(idx: int) -> ReplProcess:
+            spec = None
+            if sandbox:
+                from align.verify.lean.sandbox import SandboxSpec, minimal_env, sandbox_supported
+
+                ok, why = sandbox_supported()
+                if not ok:  # fail closed: never silently fall back to no isolation
+                    raise RuntimeError(f"sandbox=True but namespaces are unavailable: {why}")
+                wdir = os.path.join(scratch, f"sandbox{idx}")
+                os.makedirs(wdir, exist_ok=True)
+                spec = SandboxSpec(
+                    scratch_dir=wdir,
+                    ro_paths=tuple(str(toolchain.root / d) for d in ("elan", "repl", "workspace")),
+                    cwd=toolchain.cwd,
+                    pidfile=os.path.join(scratch, f"worker{idx}.pid"),
+                    env=minimal_env(toolchain.child_env, wdir),
+                )
             return ReplProcess(
                 toolchain.argv(),
                 toolchain.child_env,
                 cwd=toolchain.cwd,
                 limits=limits,
                 stderr_path=os.path.join(scratch, f"worker{idx}.stderr"),
+                sandbox=spec,
             )
 
         return cls(launcher, n, max_requests, default_timeout_s)

To use: ReplPool.from_toolchain(tc, n_workers=..., scratch_dir=".../pool", sandbox=True); pass scratch_dir OUTSIDE the Lean root (the default <lean>/tmp/pool is fine because only elan, repl and workspace are mapped, not <lean>/tmp). Remember to enable it for BOTH the baseline-comparable pre-SFT and post-SFT runs, or neither: verdicts were identical in every case tested, but the comparability argument should not rest on that.