Add COW fork: Sandbox(policy, init, work) + sb.fork(env={})

Signed-off-by: Cong Wang <cwang@multikernel.io>

Author: congwang-mk

examples/fork.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+"""Demo: prove COW memory sharing works.
+
+1. init() allocates a large buffer and stamps it with a unique ID
+2. work() reads the buffer — proves data survived fork (not re-init'd)
+3. Parent checks /proc/pid/smaps — proves pages are physically shared
+
+Usage:
+    python3 examples/try_clone.py
+"""
+
+import os
+import sys
+import time
+import ctypes
+from sandlock import Sandbox, Policy
+
+# Unique token set by init(), never re-set
+_TOKEN = None
+_BUF_ADDR = None
+_BUF_SIZE = 10 * 1024 * 1024  # 10 MB
+
+
+def init():
+    """Allocate 10 MB, fill with a unique token. Runs once."""
+    global _TOKEN, _BUF_ADDR
+
+    _TOKEN = os.getpid()  # unique per-process — proves it's from init's PID
+
+    # Allocate 10 MB via mmap (anonymous, private)
+    libc = ctypes.CDLL("libc.so.6")
+    libc.mmap.restype = ctypes.c_void_p
+    libc.mmap.argtypes = [
+        ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int,
+        ctypes.c_int, ctypes.c_int, ctypes.c_long,
+    ]
+    addr = libc.mmap(None, _BUF_SIZE, 0x3, 0x22, -1, 0)  # RW, PRIVATE|ANON
+    _BUF_ADDR = addr
+
+    # Fill with pattern: each 4KB page starts with the token
+    buf = (ctypes.c_char * _BUF_SIZE).from_address(addr)
+    import struct
+    for offset in range(0, _BUF_SIZE, 4096):
+        struct.pack_into("<Q", buf, offset, _TOKEN)
+
+
+def work():
+    """Read the buffer. Proves COW — data is from init()'s process."""
+    import struct
+
+    seed = int(os.environ.get("SEED", "0"))
+    my_pid = os.getpid()
+
+    # Read token from first page
+    buf = (ctypes.c_char * 8).from_address(_BUF_ADDR)
+    token = struct.unpack_from("<Q", buf, 0)[0]
+
+    # Verify ALL pages have the same token (no corruption)
+    full_buf = (ctypes.c_char * _BUF_SIZE).from_address(_BUF_ADDR)
+    all_match = True
+    for offset in range(0, _BUF_SIZE, 4096):
+        page_token = struct.unpack_from("<Q", full_buf, offset)[0]
+        if page_token != token:
+            all_match = False
+            break
+
+    with open(f"/tmp/sandlock_cow_{seed}", "w") as f:
+        f.write(f"clone_pid={my_pid} init_pid={token} "
+                f"pages_ok={all_match} buf_addr=0x{_BUF_ADDR:x}")
+
+
+def get_shared_pages(pid):
+    """Read Shared_Clean + Shared_Dirty from /proc/pid/smaps."""
+    shared = 0
+    private = 0
+    try:
+        with open(f"/proc/{pid}/smaps") as f:
+            for line in f:
+                if line.startswith("Shared_Clean:") or line.startswith("Shared_Dirty:"):
+                    shared += int(line.split()[1])
+                elif line.startswith("Private_Clean:") or line.startswith("Private_Dirty:"):
+                    private += int(line.split()[1])
+    except OSError:
+        pass
+    return shared, private
+
+
+def main():
+    policy = Policy(
+        fs_writable=["/tmp"],
+        fs_readable=[sys.prefix, "/usr", "/lib", "/etc", "/proc", "/dev"],
+    )
+
+    print("=== COW Clone Proof ===\n", flush=True)
+
+    with Sandbox(policy, init, work) as sb:
+        # Fork 3 clones, keep them alive briefly to check smaps
+        clones = []
+        for seed in range(3):
+            clone = sb.fork(env={"SEED": str(seed)})
+            clones.append((seed, clone))
+
+        # Give clones time to run
+        time.sleep(1)
+
+        # Check shared pages for each clone
+        for seed, clone in clones:
+            pid = clone._clone_pid
+            if pid:
+                shared_kb, private_kb = get_shared_pages(pid)
+                print(f"  Clone {seed} (PID {pid}): "
+                      f"shared={shared_kb} KB, private={private_kb} KB",
+                      flush=True)
+
+        # Wait for all clones
+        for seed, clone in clones:
+            clone.wait(timeout=5)
+
+    # Read results
+    print(flush=True)
+    for seed in range(3):
+        f = f"/tmp/sandlock_cow_{seed}"
+        if os.path.exists(f):
+            print(f"  Clone {seed}: {open(f).read()}", flush=True)
+            os.unlink(f)
+        else:
+            print(f"  Clone {seed}: no output", flush=True)
+
+    print(f"\n  10 MB buffer allocated once in init().", flush=True)
+    print(f"  If shared > 0 KB, pages are physically shared (COW).", flush=True)
+    print(f"  If init_pid matches across clones, data survived fork.", flush=True)
+    print("\nDone.", flush=True)
+
+
+if __name__ == "__main__":
+    main()

src/sandlock/_checkpoint.py

@@ -407,6 +407,100 @@ def start_child_listener(
     _CheckpointListener(control_fd, save_fn).start()
 
 
+# --- Clone-ready loop (runs on main thread after init_fn returns) ---
+
+TRIGGER_FORK = b"\x03"
+
+
+def clone_ready_loop(control_fd: int, work_fn: "Callable") -> None:
+    """Main-thread loop: wait for fork commands, fork and run work_fn.
+
+    After init_fn returns, the main thread enters this loop.  It blocks
+    on ``os.read()`` (GIL released), so no CPU is wasted.  When the
+    parent sends TRIGGER_FORK, the main thread calls ``os.fork()`` —
+    giving the clone a full COW copy of all memory (including
+    everything init_fn set up in globals/heap).
+
+    Args:
+        control_fd: Child's end of the control socket.
+        work_fn: Function to run in each clone.
+    """
+    import sys
+
+    while True:
+        try:
+            trigger = os.read(control_fd, 1)
+        except OSError:
+            break
+        if not trigger:
+            break
+
+        if trigger == TRIGGER_FORK:
+            try:
+                env_json = _recv_bytes(control_fd)
+            except (EOFError, OSError):
+                break
+
+            env = json.loads(env_json) if env_json else {}
+
+            try:
+                sys.stdout.flush()
+                sys.stderr.flush()
+            except Exception:
+                pass
+
+            try:
+                pid = os.fork()
+            except OSError:
+                os.write(control_fd, struct.pack(">I", 0))
+                continue
+
+            if pid == 0:
+                # === Clone child ===
+                try:
+                    os.close(control_fd)
+                    os.setpgid(0, 0)
+                    os.environ.update(env)
+                    work_fn()
+                except SystemExit as e:
+                    os._exit(e.code if isinstance(e.code, int) else 1)
+                except BaseException:
+                    os._exit(1)
+                os._exit(0)
+            else:
+                # === Template: send clone PID back ===
+                os.write(control_fd, struct.pack(">I", pid))
+
+
+def request_fork(
+    control_fd: int,
+    env: dict[str, str] | None = None,
+) -> int:
+    """Send a fork command to the template's clone-ready loop.
+
+    Args:
+        control_fd: Parent's end of the control socket.
+        env: Environment variable overrides for the clone.
+
+    Returns:
+        PID of the clone child.
+
+    Raises:
+        RuntimeError: If fork failed in the child.
+    """
+    env_json = json.dumps(env or {}).encode()
+    os.write(control_fd, TRIGGER_FORK)
+    _send_bytes(control_fd, env_json)
+
+    raw = os.read(control_fd, 4)
+    if len(raw) < 4:
+        raise RuntimeError("Fork: no response from child")
+    pid = struct.unpack(">I", raw)[0]
+    if pid == 0:
+        raise RuntimeError("Fork: fork() failed in child")
+    return pid
+
+
 # --- Parent side ---
 
 def request_app_state(control_fd: int) -> bytes:

src/sandlock/_context.py

@@ -197,11 +197,13 @@ def __init__(
         save_fn: Callable[[], bytes] | None = None,
         overlay_branch: "object | None" = None,
         cow_branch: "object | None" = None,
+        clone_loop_fn: "Callable[[int], None] | None" = None,
     ):
         self._target = target
         self._policy = policy
         self._sandbox_id = sandbox_id
         self._save_fn = save_fn
+        self._clone_loop_fn = clone_loop_fn
         self._overlay_branch = overlay_branch or getattr(policy, '_overlay_branch', None)
         self._cow_branch = cow_branch or getattr(policy, '_cow_branch', None)
         self._pid: Optional[int] = None
@@ -650,10 +652,12 @@ def __enter__(self) -> "SandboxContext":
                         if self._policy.strict:
                             raise
 
-                # 5. Start checkpoint listener thread (if save_fn provided)
+                # 5. Start checkpoint/clone listener thread (if save_fn provided)
                 #    Must happen BEFORE seccomp — seccomp blocks clone3
                 #    which Python's threading module uses.
-                if self._save_fn is not None:
+                if self._clone_loop_fn is not None:
+                    pass  # Keep ctrl_child_fd open for clone_ready_loop
+                elif self._save_fn is not None:
                     try:
                         start_child_listener(ctrl_child_fd, self._save_fn)
                     except RuntimeError:
@@ -752,8 +756,11 @@ def __enter__(self) -> "SandboxContext":
                     n = self._policy.max_open_files
                     resource.setrlimit(resource.RLIMIT_NOFILE, (n, n))
 
-                # 10. Run target
-                self._target()
+                # 10. Run target (or clone-ready loop)
+                if self._clone_loop_fn is not None:
+                    self._clone_loop_fn(ctrl_child_fd)
+                else:
+                    self._target()
                 os._exit(0)
             except SystemExit as e:
                 os._exit(e.code if isinstance(e.code, int) else 1)