From 3f519be04ff6175f2e76791ed0dd58650d9e891b Mon Sep 17 00:00:00 2001 From: Laela Zorana Date: Thu, 17 Sep 2026 02:39:20 -0500 Subject: [PATCH] Hide benchmark arguments from the imported kernel The benchmark call arguments were reachable from an imported kernel. The compiled do_bench runs the kernel on a second thread while _do_bench_impl is still on the stack of the main thread, so sys._current_frames hands the kernel that frame, with the result pipe and the expected output generator bound as locals. One frame further down BaseProcess.run holds self, and self._args is the tuple the child was started with. Issue #3. _do_bench_impl now hands the runner duplicated descriptors of its own and closes the three connection objects, clears _args and _kwargs on the child Process object, deletes the argument names from its own frame, and passes the packed call by popping it from a list. The tuple lives on the interpreter stack and on the C++ side while the kernel runs, and every frame is clear of it. os.dup returns descriptors with FD_CLOEXEC set, so an exec from user code leaves the pipes behind. The traceback pipe stays on the frame because the except handler needs it. test/test_hidden_arguments.py runs on CPU with the extension stubbed. It calls _do_bench_impl through a stand in for the child Process object and, in one test, in a real spawned child. From a worker thread it walks every thread through sys._current_frames and looks for the five names, and for the generator, the arguments dict and any Connection object by identity, in frame locals, inside tuples and lists, and inside _args. Before the change it finds all of them, after the change it finds none of them, and the result pipe still writes. --- python/pygpubench/__init__.py | 29 +++- test/test_hidden_arguments.py | 302 ++++++++++++++++++++++++++++++++++ 2 files changed, 326 insertions(+), 5 deletions(-) create mode 100644 test/test_hidden_arguments.py diff --git a/python/pygpubench/__init__.py b/python/pygpubench/__init__.py index 46abda1..2c968d5 100644 --- a/python/pygpubench/__init__.py +++ b/python/pygpubench/__init__.py @@ -53,10 +53,18 @@ def _do_bench_impl(out_fd: "multiprocessing.connection.Connection", in_fd: "mult try: with DeterministicContext(): - _pygpubench.do_bench( - out_fd.fileno(), - in_fd.fileno(), - supervisor_sock.fileno(), + # The runner only needs descriptors. Duplicate them so no python + # object owns what the runner uses, then close the originals. + result_fd = os.dup(out_fd.fileno()) + input_fd = os.dup(in_fd.fileno()) + supervisor_fd = os.dup(supervisor_sock.fileno()) + out_fd.close() + in_fd.close() + supervisor_sock.close() + benchmark_call = [( + result_fd, + input_fd, + supervisor_fd, qualname, test_generator, test_args, @@ -66,7 +74,18 @@ def _do_bench_impl(out_fd: "multiprocessing.connection.Connection", in_fd: "mult landlock, mseal, allow_root, - ) + )] + # The kernel is imported on a worker thread while this frame is live, + # so anything bound here, or in the Process object one frame down, is + # readable through sys._current_frames. Clear both, and pop the call + # so this frame holds no reference to it while the kernel runs. See issue #3. + current = mp.current_process() + if getattr(current, "_args", None): + current._args = () + current._kwargs = {} + del out_fd, in_fd, supervisor_sock, test_generator, test_args + del result_fd, input_fd, supervisor_fd, current + _pygpubench.do_bench(*benchmark_call.pop()) except BaseException: if tb_conn is not None: tb_conn.send(traceback.format_exc()) diff --git a/test/test_hidden_arguments.py b/test/test_hidden_arguments.py new file mode 100644 index 0000000..6e4880e --- /dev/null +++ b/test/test_hidden_arguments.py @@ -0,0 +1,302 @@ +"""Regression test for issue #3, hide the benchmark arguments from the python stack. + +The compiled ``do_bench`` imports and runs the user kernel on a second thread, +``run_thread`` in ``csrc/binding.cpp``, while ``_do_bench_impl`` is still on the +stack of the main thread. ``sys._current_frames()`` returns that frame, and one +frame further down ``BaseProcess.run`` holds ``self``, whose ``_args`` is the +tuple the child was started with. Before the fix both held the result pipe and +the expected output generator. + +The tests stub the compiled extension with a probe that runs the frame walk on +a worker thread, the way the kernel would. They run on CPU and need no GPU. One +test does it in a real spawned child, the others call ``_do_bench_impl`` +through a stand in for the child Process object so the ``run`` frame and +``_args`` are on the stack the same way. +""" + +import json +import multiprocessing as mp +import multiprocessing.connection +import multiprocessing.process +import os +import sys +import threading +import types +import unittest + +# Import the package from source when the test runs from a checkout, and +# register a stub for the compiled extension before the package is imported. +_TEST_DIR = os.path.dirname(os.path.abspath(__file__)) +_SRC_DIR = os.path.join(os.path.dirname(_TEST_DIR), "python") +if os.path.isdir(_SRC_DIR) and _SRC_DIR not in sys.path: + sys.path.insert(0, _SRC_DIR) + +_stub = types.ModuleType("pygpubench._pygpubench") + + +def _unset_do_bench(*args, **kwargs): + raise AssertionError("each test installs its own do_bench probe") + + +_stub.do_bench = _unset_do_bench +sys.modules["pygpubench._pygpubench"] = _stub + +# utils.py imports AESGCM at import time. It is a real dependency of the package, +# but nothing here touches the crypto path, so provide a placeholder when it is +# missing. An environment with cryptography installed uses it as normal. +try: + import cryptography # noqa: F401 +except ImportError: + _aead = types.ModuleType("cryptography.hazmat.primitives.ciphers.aead") + _aead.AESGCM = object + for _name in ( + "cryptography", + "cryptography.hazmat", + "cryptography.hazmat.primitives", + "cryptography.hazmat.primitives.ciphers", + "cryptography.hazmat.primitives.ciphers.aead", + ): + sys.modules.setdefault(_name, types.ModuleType(_name)) + sys.modules["cryptography.hazmat.primitives.ciphers.aead"] = _aead + +import pygpubench + + +class _NoOpDeterministicContext: + """Stand in for DeterministicContext so the test does not need torch.""" + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, tb): + return False + + +# The argument names an imported kernel would look for on the harness frame. +_SECRET_NAMES = ("out_fd", "in_fd", "supervisor_sock", "test_generator", "test_args") + + +def _a_generator(**kwargs): + return (), () + + +def _is_connection(value): + return isinstance(value, multiprocessing.connection.Connection) + + +def _walk_all_threads(secret_ids, skip_code): + """Walk every frame of every thread through sys._current_frames. + + Returns the secret names found in frame locals, and the places where a + secret object was found by identity, as a local, inside a tuple or list, or + inside the _args of a Process like object. Frames running skip_code model + the C++ side, which holds the arguments and is not a python frame. Module + frames are skipped, their locals are the module globals, which is where the + generator used by the spawn test has to live so it can be pickled. + """ + names = set() + places = [] + for frame in sys._current_frames().values(): + while frame is not None: + if frame.f_code is not skip_code and frame.f_locals is not frame.f_globals: + for name, value in frame.f_locals.items(): + if name in _SECRET_NAMES: + names.add(name) + if id(value) in secret_ids or _is_connection(value): + places.append((frame.f_code.co_name, name)) + if isinstance(value, (tuple, list)): + for item in value: + if id(item) in secret_ids or _is_connection(item): + places.append((frame.f_code.co_name, name + "[]")) + if hasattr(value, "_target") and hasattr(value, "_args"): + for item in value._args: + if id(item) in secret_ids or _is_connection(item): + places.append((frame.f_code.co_name, name + "._args")) + frame = frame.f_back + return names, places + + +class _Probe: + """Models the compiled do_bench. Runs the walk on a worker thread, records + what that thread can still reach, and writes to the result descriptor. With + report set it also writes its findings as json, for the spawned child.""" + + def __init__(self, report=False): + self.report = report + self.names = None + self.places = None + self.wrote = None + self.inheritable = None + self.args_cleared = None + + def __call__(self, *c_args): + result_fd = c_args[0] + secret_ids = {id(c_args[4]), id(c_args[5])} + self.inheritable = os.get_inheritable(result_fd) + self.args_cleared = getattr(mp.current_process(), "_args", ()) == () + + def user_code(): + self.names, self.places = _walk_all_threads(secret_ids, _Probe.__call__.__code__) + self.wrote = os.write(result_fd, b"probe-wrote") + + worker = threading.Thread(target=user_code) + worker.start() + worker.join() + if self.report: + findings = { + "names": sorted(self.names), + "places": self.places, + "wrote": self.wrote, + "inheritable": self.inheritable, + "args_cleared": self.args_cleared, + } + os.write(result_fd, json.dumps(findings).encode()) + + +class _FakeProcess: + """The child Process object as the spawned child sees it. run() binds self on + the frame below the target, the way BaseProcess.run does.""" + + def __init__(self, target, args): + self._target = target + self._args = args + self._kwargs = {} + + def run(self): + saved = multiprocessing.process._current_process + multiprocessing.process._current_process = self + try: + self._target(*self._args, **self._kwargs) + finally: + multiprocessing.process._current_process = saved + + +def _vulnerable_impl(out_fd, in_fd, supervisor_sock, qualname, test_generator, test_args, + stream=None, discard=True, nvtx=False, tb_conn=None, landlock=True, + mseal=True, allow_root=False): + """The shape of _do_bench_impl before the fix, kept as a control so the test + fails if the probe stops seeing a leak.""" + pygpubench._pygpubench.do_bench( + out_fd.fileno(), in_fd.fileno(), supervisor_sock.fileno(), qualname, + test_generator, test_args, stream, discard, nvtx, landlock, mseal, allow_root, + ) + + +def _spawn_entry(*packed): + """Target of the real spawned child. Installs the reporting probe and calls + the real _do_bench_impl. The call is handed over by popping it, so this frame + keeps no reference to the arguments while the probe walks the stack.""" + pygpubench.DeterministicContext = _NoOpDeterministicContext + pygpubench._pygpubench.do_bench = _Probe(report=True) + call = [packed] + del packed + pygpubench._do_bench_impl(*call.pop()) + + +class HiddenArgumentsTest(unittest.TestCase): + def setUp(self): + self._saved_ctx = pygpubench.DeterministicContext + pygpubench.DeterministicContext = _NoOpDeterministicContext + self._saved_do_bench = pygpubench._pygpubench.do_bench + # Names in this frame stay clear of the secret names, or the walk would + # find them here and report a false leak. + self._read_end, self._write_end = mp.Pipe(duplex=False) + self._sig_r, self._sig_w = mp.Pipe(duplex=False) + self._sup_a, self._sup_b = mp.Pipe(duplex=False) + self._generator = _a_generator + self._arguments = {"size": 256} + + def tearDown(self): + pygpubench.DeterministicContext = self._saved_ctx + pygpubench._pygpubench.do_bench = self._saved_do_bench + for conn in (self._read_end, self._write_end, self._sig_r, + self._sig_w, self._sup_a, self._sup_b): + try: + conn.close() + except OSError: + pass + + def _packed_args(self): + # (out_fd, in_fd, supervisor_sock, qualname, test_generator, test_args, + # stream, discard, nvtx, tb_conn, landlock, mseal, allow_root) + return ( + self._write_end, self._sig_r, self._sup_b, "some_module.kernel", + self._generator, self._arguments, 0, True, False, None, True, True, False, + ) + + def _probe(self): + probe = _Probe() + pygpubench._pygpubench.do_bench = probe + return probe + + def test_control_leaks_arguments(self): + # The probe must see the leak in the old shape, otherwise the fix proves nothing. + probe = self._probe() + _FakeProcess(_vulnerable_impl, self._packed_args()).run() + self.assertEqual(probe.wrote, len(b"probe-wrote")) + self.assertEqual(probe.names, set(_SECRET_NAMES)) + found = {place for _, place in probe.places} + self.assertIn("test_generator", found) + self.assertIn("out_fd", found) + self.assertIn("self._args", found) + + def test_fix_hides_arguments_from_frame_walk(self): + probe = self._probe() + _FakeProcess(pygpubench._do_bench_impl, self._packed_args()).run() + self.assertEqual(probe.names, set(), "a secret name is still bound on a reachable frame") + self.assertEqual(probe.places, [], "a secret object is still reachable from a frame") + self.assertTrue(probe.args_cleared) + + def test_result_pipe_still_writable(self): + probe = self._probe() + _FakeProcess(pygpubench._do_bench_impl, self._packed_args()).run() + self.assertEqual(probe.wrote, len(b"probe-wrote")) + self.assertTrue(self._read_end.poll(2), "the probe should have written to the pipe") + self.assertEqual(os.read(self._read_end.fileno(), 64), b"probe-wrote") + # The Connection objects are closed, the runner wrote through its own descriptor. + self.assertTrue(self._write_end.closed) + self.assertTrue(self._sig_r.closed) + self.assertTrue(self._sup_b.closed) + + def test_runner_descriptors_are_close_on_exec(self): + # Under spawn the pipes arrive inheritable. Model that, then check the + # descriptor the runner receives has FD_CLOEXEC set. + os.set_inheritable(self._write_end.fileno(), True) + os.set_inheritable(self._sig_r.fileno(), True) + probe = self._probe() + _FakeProcess(pygpubench._do_bench_impl, self._packed_args()).run() + self.assertFalse(probe.inheritable) + + def test_real_spawned_child(self): + ctx = mp.get_context("spawn") + read_end, write_end = ctx.Pipe(duplex=False) + sig_r, sig_w = ctx.Pipe(duplex=False) + sup_a, sup_b = ctx.Pipe(duplex=False) + child = ctx.Process(target=_spawn_entry, + args=(write_end, sig_r, sup_b, "some_module.kernel", _a_generator, + {"size": 256}, 0, True, False, None, True, True, False)) + child.start() + write_end.close() + sig_r.close() + sup_b.close() + chunks = b"" + while True: + chunk = os.read(read_end.fileno(), 65536) + if chunk == b"": + break + chunks += chunk + child.join(60) + for conn in (read_end, sig_w, sup_a): + conn.close() + self.assertEqual(child.exitcode, 0) + self.assertTrue(chunks.startswith(b"probe-wrote"), chunks) + findings = json.loads(chunks[len(b"probe-wrote"):]) + self.assertEqual(findings["names"], []) + self.assertEqual(findings["places"], []) + self.assertEqual(findings["wrote"], len(b"probe-wrote")) + self.assertFalse(findings["inheritable"]) + self.assertTrue(findings["args_cleared"]) + + +if __name__ == "__main__": + unittest.main(verbosity=2)