Preserve sandbox capabilities through snapshot/recycle/restore

pyisolate/checkpoint.py

@@ -84,6 +84,7 @@ def restore(blob: bytes, key: bytes) -> Sandbox:
         mem_bytes=state.get("mem_bytes"),
         allowed_imports=state.get("allowed_imports"),
         numa_node=state.get("numa_node"),
+        capabilities=state.get("capabilities"),
     )
 
 

pyisolate/runtime/thread.py

@@ -29,13 +29,20 @@
 from typing import Any, Callable, Iterable, Optional
 
 from .. import errors
+from ..capabilities import (
+    ClockCapability,
+    FilesystemCapability,
+    NetworkCapability,
+    RandomCapability,
+    SecretCapability,
+    SubprocessCapability,
+)
 from .protocol import (
     AttachCgroupRequest,
     CallRequest,
     ExecRequest,
     StopRequest,
 )
-from ..capabilities import ClockCapability
 from ..numa import bind_current_thread
 from ..observability.trace import Tracer
 
@@ -44,6 +51,83 @@
 _ORIG_OPEN = builtins.open
 _ORIG_SOCKET_CONNECT = socket.socket.connect
 _ORIG_THREAD_START = threading.Thread.start
+_CAPABILITY_MARKER = "__pyisolate_capability__"
+
+
+def _serialize_capability(capability: Any) -> Any:
+    if isinstance(capability, FilesystemCapability):
+        return {
+            _CAPABILITY_MARKER: "filesystem",
+            "roots": [str(root) for root in capability.roots],
+        }
+    if isinstance(capability, NetworkCapability):
+        return {
+            _CAPABILITY_MARKER: "network",
+            "destinations": sorted(capability.destinations),
+        }
+    if isinstance(capability, SecretCapability):
+        return {
+            _CAPABILITY_MARKER: "secrets",
+            "values": {
+                key: value.hex() for key, value in sorted(capability.values.items())
+            },
+        }
+    if isinstance(capability, SubprocessCapability):
+        return {
+            _CAPABILITY_MARKER: "subprocess",
+            "allowed_commands": sorted(capability.allowed_commands),
+            "allow_shell": capability.allow_shell,
+        }
+    if isinstance(capability, ClockCapability):
+        return {_CAPABILITY_MARKER: "clock"}
+    if isinstance(capability, RandomCapability):
+        return {_CAPABILITY_MARKER: "random"}
+    return capability
+
+
+def _deserialize_capability(capability: Any) -> Any:
+    if not isinstance(capability, dict):
+        return capability
+    kind = capability.get(_CAPABILITY_MARKER)
+    if kind == "filesystem":
+        roots = capability.get("roots", [])
+        return FilesystemCapability.from_paths(*roots)
+    if kind == "network":
+        destinations = capability.get("destinations", [])
+        return NetworkCapability.from_destinations(*destinations)
+    if kind == "secrets":
+        encoded_values = capability.get("values", {})
+        decoded_values = {
+            key: bytes.fromhex(value) for key, value in encoded_values.items()
+        }
+        return SecretCapability(values=decoded_values)
+    if kind == "subprocess":
+        commands = capability.get("allowed_commands", [])
+        allow_shell = bool(capability.get("allow_shell", False))
+        return SubprocessCapability.from_commands(*commands, allow_shell=allow_shell)
+    if kind == "clock":
+        return ClockCapability()
+    if kind == "random":
+        return RandomCapability()
+    return capability
+
+
+def serialize_capabilities(capabilities: Optional[dict[str, Any]]) -> dict[str, Any]:
+    if not capabilities:
+        return {}
+    return {
+        name: _serialize_capability(capability)
+        for name, capability in sorted(capabilities.items())
+    }
+
+
+def deserialize_capabilities(capabilities: Optional[dict[str, Any]]) -> dict[str, Any]:
+    if not capabilities:
+        return {}
+    return {
+        name: _deserialize_capability(capability)
+        for name, capability in capabilities.items()
+    }
 
 
 def _blocked_open(file, *args, **kwargs):
@@ -356,7 +440,7 @@ def __init__(
         self._cgroup_path = cgroup_path
         self._trace_enabled = False
         self._syscall_log: list[str] = []
-        self._capabilities = dict(capabilities or {})
+        self._capabilities = deserialize_capabilities(capabilities)
         self._quarantine_reason: str | None = None
         self.termination_reason: str | None = None
         self._open_files = 0
@@ -377,7 +461,7 @@ def snapshot(self) -> dict:
             if self.allowed_imports is not None
             else None,
             "numa_node": self.numa_node,
-            "capabilities": sorted(self._capabilities),
+            "capabilities": serialize_capabilities(self._capabilities),
             "wall_time_ms": self.wall_time_ms,
             "open_files_max": self.open_files_max,
             "network_ops_max": self.network_ops_max,
@@ -547,7 +631,7 @@ def reset(
         self._syscall_log = []
         self._start_time = None
         self._cgroup_path = cgroup_path
-        self._capabilities = dict(capabilities or {})
+        self._capabilities = deserialize_capabilities(capabilities)
         self.termination_reason = None
         self._open_files = 0
         self._network_ops = 0

pyisolate/supervisor.py

@@ -414,6 +414,7 @@ def recycle(self, name: str) -> Sandbox:
             child_work_max=snap["child_work_max"],
             allowed_imports=snap["allowed_imports"],
             numa_node=snap["numa_node"],
+            capabilities=snap["capabilities"],
         )
 
     def _cleanup(self) -> None:

tests/test_checkpoint.py

@@ -50,6 +50,7 @@ def open_ring_buffer(self):
 
 import pyisolate as iso
 import pyisolate.policy as policy
+from pyisolate.capabilities import FilesystemCapability
 
 
 def _make_blob(payload, key):
@@ -133,6 +134,29 @@ def test_checkpoint_restores_imports_and_numa():
         pass
 
 
+def test_checkpoint_restores_capabilities(tmp_path):
+    key = os.urandom(32)
+    allowed = tmp_path / "allowed"
+    allowed.mkdir()
+    target = allowed / "ok.txt"
+    target.write_text("ok", encoding="utf-8")
+
+    sb = iso.spawn(
+        "cap-cp",
+        capabilities={"filesystem": FilesystemCapability.from_paths(str(allowed))},
+    )
+    try:
+        blob = iso.checkpoint(sb, key)
+        sb2 = iso.restore(blob, key)
+        try:
+            sb2.exec(f"post(open({str(target)!r}).read())")
+            assert sb2.recv(timeout=0.5) == "ok"
+        finally:
+            sb2.close()
+    finally:
+        pass
+
+
 @pytest.mark.parametrize(
     "payload, message",
     [

tests/test_supervisor.py

@@ -178,6 +178,30 @@ def test_quarantine_and_recycle():
         sup.shutdown()
 
 
+def test_recycle_preserves_capabilities(tmp_path):
+    allowed = tmp_path / "allowed"
+    allowed.mkdir()
+    target = allowed / "ok.txt"
+    target.write_text("ok")
+
+    sup = iso.Supervisor()
+    try:
+        sb = sup.spawn(
+            "recycle-cap",
+            capabilities={
+                "filesystem": iso.FilesystemCapability.from_paths(str(allowed)),
+            },
+        )
+        sb.exec(f"post(open({str(target)!r}).read())")
+        assert sb.recv(timeout=0.5) == "ok"
+
+        recycled = sb.recycle()
+        recycled.exec(f"post(open({str(target)!r}).read())")
+        assert recycled.recv(timeout=0.5) == "ok"
+    finally:
+        sup.shutdown()
+
+
 def test_sandbox_termination_reason_passthrough():
     sup = iso.Supervisor()
     try: