Add mprotect-based parent memory protection after branching

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

Author: congwang-mk

src/branching/exceptions.py

@@ -72,3 +72,9 @@ class MemoryBranchError(BranchingError):
     """Memory branching operation failed."""
 
     pass
+
+
+class MemoryProtectError(MemoryBranchError):
+    """mprotect() failed on a registered memory region."""
+
+    pass

src/branching/memory/__init__.py

@@ -1,3 +1,25 @@
-from .base import MemoryBackend, StubMemoryBackend
+from .base import MemoryBackend, StubMemoryBackend, MprotectMemoryBackend
+from ._mprotect import (
+    MemoryRegion,
+    mprotect,
+    protect_regions,
+    restore_regions,
+    PROT_NONE,
+    PROT_READ,
+    PROT_WRITE,
+    PROT_EXEC,
+)
 
-__all__ = ["MemoryBackend", "StubMemoryBackend"]
+__all__ = [
+    "MemoryBackend",
+    "StubMemoryBackend",
+    "MprotectMemoryBackend",
+    "MemoryRegion",
+    "mprotect",
+    "protect_regions",
+    "restore_regions",
+    "PROT_NONE",
+    "PROT_READ",
+    "PROT_WRITE",
+    "PROT_EXEC",
+]

src/branching/memory/_mprotect.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+"""ctypes bindings for mprotect(2).
+
+Provides memory protection manipulation for enforcing read-only
+invariants on parent memory regions after branching.
+"""
+
+import ctypes
+import ctypes.util
+import os
+from dataclasses import dataclass, field
+
+from ..exceptions import MemoryProtectError
+
+_libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)
+
+# Protection flags for mprotect(2)
+PROT_NONE = 0x0
+PROT_READ = 0x1
+PROT_WRITE = 0x2
+PROT_EXEC = 0x4
+
+
+def mprotect(addr: int, size: int, prot: int) -> None:
+    """Call mprotect(2) to change memory protection.
+
+    Args:
+        addr: Page-aligned start address of the region.
+        size: Size of the region in bytes.
+        prot: New protection flags (PROT_READ, PROT_WRITE, etc.).
+
+    Raises:
+        MemoryProtectError: If the syscall fails.
+    """
+    ret = _libc.mprotect(ctypes.c_void_p(addr), ctypes.c_size_t(size), ctypes.c_int(prot))
+    if ret != 0:
+        err = ctypes.get_errno()
+        raise MemoryProtectError(
+            f"mprotect(0x{addr:x}, {size}, 0x{prot:x}): {os.strerror(err)}"
+        )
+
+
+@dataclass
+class MemoryRegion:
+    """A memory region tracked for protection changes.
+
+    Attributes:
+        addr: Page-aligned start address.
+        size: Size of the region in bytes.
+        original_prot: Protection flags before we modified them.
+    """
+    addr: int
+    size: int
+    original_prot: int = field(default=PROT_READ | PROT_WRITE)
+
+
+def protect_regions(regions: list[MemoryRegion], prot: int = PROT_READ) -> None:
+    """Apply memory protection to a list of regions.
+
+    Each region's ``original_prot`` is assumed to already be set before
+    calling this function.
+
+    Args:
+        regions: Regions to protect.
+        prot: Protection flags to apply.
+
+    Raises:
+        MemoryProtectError: If any mprotect call fails.
+    """
+    for region in regions:
+        mprotect(region.addr, region.size, prot)
+
+
+def restore_regions(regions: list[MemoryRegion]) -> None:
+    """Restore each region to its original protection.
+
+    Args:
+        regions: Regions whose original_prot should be restored.
+
+    Raises:
+        MemoryProtectError: If any mprotect call fails.
+    """
+    for region in regions:
+        mprotect(region.addr, region.size, region.original_prot)

src/branching/memory/base.py

@@ -1,14 +1,21 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Memory branching stub for future kernel support.
+"""Memory branching backends.
 
-This module defines the MemoryBackend ABC and a stub implementation.
-Real memory branching requires kernel support via the proposed branch()
-syscall with BR_MEMORY flag.
+This module defines the MemoryBackend ABC, a stub implementation, and
+an mprotect-based implementation that enforces read-only parent memory
+after branching.
 """
 
 from abc import ABC, abstractmethod
 from typing import Any
 
+from ._mprotect import (
+    MemoryRegion,
+    PROT_READ,
+    PROT_WRITE,
+    mprotect,
+)
+
 
 class MemoryBackend(ABC):
     """Abstract base class for memory branching.
@@ -71,3 +78,44 @@ def commit(self, handle: Any) -> None:
         raise NotImplementedError(
             "Memory branching requires kernel support via the branch() syscall."
         )
+
+
+class MprotectMemoryBackend(MemoryBackend):
+    """Memory backend using mprotect(2) to enforce read-only parent state.
+
+    After ``snapshot()``, the region is marked read-only so the parent
+    cannot mutate branched state. ``restore()`` and ``commit()`` both
+    re-enable write access (the parent regains control after branching).
+    """
+
+    def snapshot(self, addr: int, size: int) -> MemoryRegion:
+        """Mark a memory region as read-only and return a handle.
+
+        Args:
+            addr: Page-aligned start address of the region.
+            size: Size of the region in bytes.
+
+        Returns:
+            A MemoryRegion handle that can be passed to restore/commit.
+        """
+        region = MemoryRegion(addr=addr, size=size, original_prot=PROT_READ | PROT_WRITE)
+        mprotect(addr, size, PROT_READ)
+        return region
+
+    def restore(self, handle: Any) -> None:
+        """Restore write access to a previously snapshotted region.
+
+        Args:
+            handle: MemoryRegion from snapshot().
+        """
+        region: MemoryRegion = handle
+        mprotect(region.addr, region.size, region.original_prot)
+
+    def commit(self, handle: Any) -> None:
+        """Commit the snapshot — parent regains write access.
+
+        Args:
+            handle: MemoryRegion from snapshot().
+        """
+        region: MemoryRegion = handle
+        mprotect(region.addr, region.size, region.original_prot)

src/branching/process/context.py

@@ -16,7 +16,13 @@
 from pathlib import Path
 from typing import Callable, Iterator, Optional, Sequence, TYPE_CHECKING
 
-from ..exceptions import ForkError, ProcessBranchError
+from ..exceptions import ForkError, MemoryProtectError, ProcessBranchError
+from ..memory._mprotect import (
+    MemoryRegion,
+    PROT_READ,
+    PROT_WRITE,
+    mprotect as _mprotect,
+)
 from . import _cgroup
 from ._namespace import setup_user_ns, bind_mount
 
@@ -43,6 +49,7 @@ def __init__(
         close_fds: bool = False,
         limits: ResourceLimits | None = None,
         parent_cgroup: Path | None = None,
+        protected_regions: Sequence[tuple[int, int]] | None = None,
     ):
         """
         Args:
@@ -56,16 +63,21 @@ def __init__(
             parent_cgroup: Optional parent cgroup directory.  When given,
                 the child's scope is created under this directory instead
                 of the process's own cgroup, enabling hierarchical nesting.
+            protected_regions: Optional list of (addr, size) tuples.
+                After fork, these regions are marked read-only in the
+                parent via mprotect(2) to enforce the branch invariant.
         """
         self._target = target
         self._workspace = workspace
         self._isolate = isolate
         self._close_fds = close_fds
         self._limits = limits
         self._parent_cgroup = parent_cgroup
+        self._protected_regions = protected_regions
         self._pid: Optional[int] = None
         self._exited = False
         self._cgroup_scope: Optional[Path] = None
+        self._memory_regions: list[MemoryRegion] = []
 
     @property
     def pid(self) -> int:
@@ -91,6 +103,9 @@ def alive(self) -> bool:
     def wait(self, timeout: Optional[float] = None) -> None:
         """Wait for the child process to exit.
 
+        On completion (success or failure), restores write access to any
+        protected memory regions so the parent can proceed.
+
         Args:
             timeout: Maximum seconds to wait (None = wait forever).
 
@@ -99,6 +114,7 @@ def wait(self, timeout: Optional[float] = None) -> None:
             TimeoutError: If the child doesn't exit within timeout.
         """
         exit_code = self._wait_raw(timeout)
+        self._restore_memory_regions()
         if exit_code != 0:
             raise ProcessBranchError(
                 f"Child {self._pid} exited with status {exit_code}"
@@ -132,12 +148,15 @@ def _wait_raw(self, timeout: Optional[float] = None) -> int:
     def abort(self, timeout: float = 5.0) -> None:
         """Abort the child and all its descendants.
 
-        Uses both cgroup kill (catches escapees) and killpg (POSIX standard).
-        Escalates SIGTERM -> SIGKILL after timeout.
+        Restores write access to protected memory regions, then kills the
+        child. Uses both cgroup kill (catches escapees) and killpg (POSIX
+        standard). Escalates SIGTERM -> SIGKILL after timeout.
         """
         if self._pid is None or self._exited:
             return
 
+        self._restore_memory_regions()
+
         # Cgroup kill — catches descendants that escaped the process group
         if self._cgroup_scope is not None:
             _cgroup.kill_scope(self._cgroup_scope)
@@ -172,6 +191,12 @@ def abort(self, timeout: float = 5.0) -> None:
         self._reap()
         self._exited = True
 
+    def _restore_memory_regions(self) -> None:
+        """Restore write access to all protected memory regions."""
+        for region in self._memory_regions:
+            _mprotect(region.addr, region.size, region.original_prot)
+        self._memory_regions.clear()
+
     def _reap(self) -> None:
         """Reap the child process (non-blocking, best-effort)."""
         try:
@@ -243,11 +268,26 @@ def __enter__(self) -> "BranchContext":
                 except OSError:
                     pass  # Best-effort
 
+            # Protect registered memory regions in the parent
+            if self._protected_regions:
+                for addr, size in self._protected_regions:
+                    region = MemoryRegion(
+                        addr=addr,
+                        size=size,
+                        original_prot=PROT_READ | PROT_WRITE,
+                    )
+                    _mprotect(addr, size, PROT_READ)
+                    self._memory_regions.append(region)
+
             return self
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
         self.abort()
 
+        # If abort() was a no-op (child already waited/exited), ensure
+        # regions are still restored.
+        self._restore_memory_regions()
+
         # Clean up private mount dir (best-effort)
         try:
             os.rmdir(self._private_dir)
@@ -266,6 +306,7 @@ def create(
         close_fds: bool = False,
         limits: ResourceLimits | None = None,
         parent_cgroup: Path | None = None,
+        protected_regions: Sequence[tuple[int, int]] | None = None,
     ) -> Iterator[list["BranchContext"]]:
         """Create N branch contexts, mirroring branch(BR_CREATE, n_branches=N).
 
@@ -278,6 +319,8 @@ def create(
             close_fds: BR_CLOSE_FDS — close inherited fds in children.
             limits: Optional resource limits applied via cgroup v2.
             parent_cgroup: Optional parent cgroup for hierarchical nesting.
+            protected_regions: Optional list of (addr, size) tuples to
+                mark read-only in the parent after each fork.
 
         Yields:
             List of entered BranchContext instances (already forked).
@@ -291,6 +334,7 @@ def create(
                 ctx = BranchContext(
                     target, workspace, isolate=isolate, close_fds=close_fds,
                     limits=limits, parent_cgroup=parent_cgroup,
+                    protected_regions=protected_regions,
                 )
                 ctx.__enter__()
                 contexts.append(ctx)