Fix vDSO patching: patch before dispatch, detect exec, pre-import modules

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

Author: congwang-mk

src/sandlock/_context.py

@@ -475,6 +475,10 @@ def __enter__(self) -> "SandboxContext":
         # just sys.modules lookups.
         if use_notif:
             from ._notif import install_notif_filter, send_fd  # noqa: F811
+        if (self._notif_policy is not None
+                and self._notif_policy.time_start is not None):
+            import time as _time  # noqa: F811
+            from ._vdso import disable_vdso_local  # noqa: F811
         if self._save_fn is not None:
             from ._checkpoint import start_child_listener  # noqa: F811
         # User namespace is needed for privileged mode or overlayfs

src/sandlock/_notif.py

@@ -424,8 +424,7 @@ def __init__(
         # Deterministic time
         self._time_offset = None  # TimeOffset | None
         self._mono_offset_s: int = 0  # monotonic offset for vDSO stubs
-        self._vdso_patch_fd: int = -1   # pre-opened /proc/pid/mem
-        self._vdso_patch_writes: list[tuple[int, bytes]] = []  # (offset, stub)
+        self._vdso_patched_addr: int = 0  # vDSO base address we last patched
         self._virtual_btime: int = 0  # virtual boot time for /proc/stat
         if policy.time_start is not None:
             import time as _time
@@ -556,6 +555,12 @@ def _handle_one(self) -> None:
         if ret < 0:
             return  # ENOENT = child died, EINTR = interrupted
 
+        # Patch vDSO before dispatching (child is stopped in seccomp
+        # notification state, so /proc/pid/mem writes are reliable).
+        # Re-patch when the vDSO address changes (exec replaces the vDSO).
+        if self._time_offset is not None:
+            self._maybe_patch_vdso(notif.pid)
+
         try:
             self._dispatch(notif)
         except Exception:
@@ -566,53 +571,45 @@ def _handle_one(self) -> None:
             except Exception:
                 pass
 
-        # Post-dispatch: patch vDSO for new PIDs (after exec).
-        # Two-phase approach:
-        #   Phase 1 (this notification): pre-compute — open fd, parse
-        #     vDSO ELF, compute offsets.  This is slow but the child
-        #     will be unfrozen when dispatch responds above.
-        #   Phase 2 (next notification): minimal lseek+write using the
-        #     pre-computed fd and offsets.  Fast enough to land while
-        #     the child is briefly running after the previous response.
-        if self._vdso_patch_writes:
-            # Phase 2: fast write (child briefly running after respond)
-            fd = self._vdso_patch_fd
-            for off, stub in self._vdso_patch_writes:
-                os.lseek(fd, off, os.SEEK_SET)
-                os.write(fd, stub)
-            os.close(fd)
-            self._vdso_patch_fd = -1
-            self._vdso_patch_writes = []
-        elif self._time_offset is not None and self._vdso_patch_fd == -1:
-            # Phase 1: pre-compute (first notification from new PID)
-            pid = notif.pid
-            from ._vdso import _find_vdso, _parse_vdso_symbols, _get_stubs
-            info = _find_vdso(pid)
-            stubs = _get_stubs(self._mono_offset_s)
-            if info and stubs:
-                addr, size = info
-                try:
-                    fd = os.open(f"/proc/{pid}/mem", os.O_RDWR)
-                    os.lseek(fd, addr, os.SEEK_SET)
-                    data = os.read(fd, size)
-                    writes = []
-                    for name, off in _parse_vdso_symbols(data):
-                        stub = stubs.get(name)
-                        if stub:
-                            writes.append((addr + off, stub))
-                    if writes:
-                        self._vdso_patch_fd = fd
-                        self._vdso_patch_writes = writes
-                    else:
-                        os.close(fd)
-                except OSError:
-                    pass
-
     @property
     def tracked_pids(self) -> set[int]:
         """All PIDs known to belong to this sandbox."""
         return set(self._proc_pids)
 
+    def _maybe_patch_vdso(self, pid: int) -> None:
+        """Patch the child's vDSO to force real syscalls, if needed.
+
+        Called while the child is stopped in seccomp notification state,
+        so /proc/pid/mem writes land reliably before the child resumes.
+        Tracks the vDSO base address to detect exec (which replaces the
+        vDSO at a new address) and re-patch automatically.
+        """
+        from ._vdso import _find_vdso, _parse_vdso_symbols, _get_stubs
+        info = _find_vdso(pid)
+        if not info:
+            return
+        addr, size = info
+        if addr == self._vdso_patched_addr:
+            return  # already patched this vDSO
+        stubs = _get_stubs(self._mono_offset_s)
+        if not stubs:
+            return
+        try:
+            fd = os.open(f"/proc/{pid}/mem", os.O_RDWR)
+            try:
+                os.lseek(fd, addr, os.SEEK_SET)
+                data = os.read(fd, size)
+                for name, off in _parse_vdso_symbols(data):
+                    stub = stubs.get(name)
+                    if stub:
+                        os.lseek(fd, addr + off, os.SEEK_SET)
+                        os.write(fd, stub)
+                self._vdso_patched_addr = addr
+            finally:
+                os.close(fd)
+        except OSError:
+            pass
+
     def _dispatch(self, notif: SeccompNotif) -> None:
         """Route a notification to the appropriate handler."""
         # Lazily track every PID that makes an intercepted syscall

src/sandlock/_vdso.py

@@ -7,8 +7,8 @@
 so seccomp can intercept them for time virtualization.
 
 For the forked child (before exec): patches in-process via mprotect + write.
-For Sandbox.run (after exec): patches via /proc/pid/mem with retries,
-since writes only take effect when the child is not in seccomp-stop.
+For Sandbox.run (after exec): patches via /proc/pid/mem while the child
+is stopped in seccomp notification state, ensuring reliable delivery.
 """
 
 from __future__ import annotations

tests/test_integration.py

@@ -1323,7 +1323,7 @@ def test_cow_utimensat_committed(self, isolation):
 # --- Deterministic time ---
 
 class TestDeterministicTime:
-    @pytest.mark.xfail(reason="vDSO patching via /proc/pid/mem is racy for exec'd processes")
+
     def test_time_start_accepts_unix_timestamp(self):
         """time_start accepts a numeric Unix timestamp."""
         # 946684800 = 2000-01-01T00:00:00Z
@@ -1399,7 +1399,7 @@ def test_time_start_monotonic_near_zero_run(self):
         mono = float(result.stdout.strip())
         assert mono < 5.0, f"Monotonic too high: {mono}"
 
-    @pytest.mark.xfail(reason="vDSO patching via /proc/pid/mem is racy for exec'd processes")
+
     def test_time_syscall_shifted_run_ctypes(self):
         """time() syscall is shifted (used by uptime, w, etc.)."""
         policy = Policy(
@@ -1408,8 +1408,8 @@ def test_time_syscall_shifted_run_ctypes(self):
         )
         result = Sandbox(policy).run(
             ["python3", "-c",
-             "import ctypes, ctypes.util; "
-             "libc = ctypes.CDLL(ctypes.util.find_library('c')); "
+             "import ctypes; "
+             "libc = ctypes.CDLL(None); "
              "libc.time.restype = ctypes.c_long; "
              "print(libc.time(0))"]
         )
@@ -1448,12 +1448,12 @@ def test_proc_stat_btime_virtualized(self):
         btime = int(result.stdout.strip().split()[1])
         assert btime == 946684800, f"btime not virtualized: {btime}"
 
-    @pytest.mark.xfail(reason="vDSO patching via /proc/pid/mem is racy for exec'd processes")
+
     def test_timerfd_abstime_works(self):
         """timerfd_settime with TFD_TIMER_ABSTIME fires correctly."""
         code = (
-            "import ctypes, ctypes.util, struct, os, time\n"
-            "libc = ctypes.CDLL(ctypes.util.find_library('c'))\n"
+            "import ctypes, struct, os, time\n"
+            "libc = ctypes.CDLL(None)\n"
             "CLOCK_MONOTONIC = 1\n"
             "TFD_TIMER_ABSTIME = 1\n"
             "fd = libc.timerfd_create(CLOCK_MONOTONIC, 0)\n"