From aecc406cd820a3bdb21316c1a37ffd1569db4fbc Mon Sep 17 00:00:00 2001
From: Samuel Grayson <sam@samgrayson.me>
Date: Wed, 14 Feb 2024 19:35:31 -0600
Subject: [PATCH] Add child_pid kwarg to parent_setup_fn

---
 benchexec/baseexecutor.py     | 18 ++++++++++++++++--
 benchexec/test_runexecutor.py | 31 +++++++++++++++++++++++++++++--
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/benchexec/baseexecutor.py b/benchexec/baseexecutor.py
index d81fdda27..dab59f397 100644
--- a/benchexec/baseexecutor.py
+++ b/benchexec/baseexecutor.py
@@ -88,7 +88,11 @@ def _start_execution(
             and the result of parent_cleanup_fn (do not use os.wait)
         """
 
+        from_child, to_parent = os.pipe()
+
         def pre_subprocess():
+            os.close(from_child)
+
             # Do some other setup the caller wants.
             child_setup_fn()
 
@@ -96,6 +100,9 @@ def pre_subprocess():
             pid = os.getpid()
             cgroups.add_task(pid)
 
+            os.write(to_parent, str(pid).encode())
+            os.close(to_parent)
+
         # Set HOME and TMPDIR to fresh directories.
         tmp_dir = os.path.join(temp_dir, "tmp")
         home_dir = os.path.join(temp_dir, "home")
@@ -108,8 +115,6 @@ def pre_subprocess():
         env["TEMP"] = tmp_dir
         logging.debug("Executing run with $HOME and $TMPDIR below %s.", temp_dir)
 
-        parent_setup = parent_setup_fn()
-
         p = subprocess.Popen(
             args,
             stdin=stdin,
@@ -118,8 +123,17 @@ def pre_subprocess():
             env=env,
             cwd=cwd,
             close_fds=True,
+            pass_fds=(to_parent,),
             preexec_fn=pre_subprocess,
         )
+        print("Continuing")
+
+        os.close(to_parent)
+
+        # read at most 10 bytes because this is enough for 32bit int
+        child_pid = int(os.read(from_child, 10))
+        os.close(from_child)
+        parent_setup = parent_setup_fn(child_pid=child_pid)
 
         def wait_and_get_result():
             exitcode, ru_child = self._wait_for_process(p.pid, args[0])
diff --git a/benchexec/test_runexecutor.py b/benchexec/test_runexecutor.py
index cf7dd54d5..44d27e643 100644
--- a/benchexec/test_runexecutor.py
+++ b/benchexec/test_runexecutor.py
@@ -887,6 +887,35 @@ def test_frozen_process(self):
             "run output misses command output and was not executed properly",
         )
 
+    def test_parent_fns(self):
+        if not os.path.exists("/bin/sh"):
+            self.skipTest("missing /bin/sh")
+        parent_setup_ran = False
+        parent_cleanup_ran = False
+
+        def parent_setup_fn(*, child_pid, **kwargs):
+            # I don't want to require psutil just for this
+            # I'll just read the procfs
+            assert os.path.exists(f"/proc/{child_pid}")
+            nonlocal parent_setup_ran
+            parent_setup_ran = True
+            return 12345
+
+        def parent_cleanup_fn(parent_setup, exit_code, path):
+            assert parent_setup == 12345
+            nonlocal parent_cleanup_ran
+            parent_cleanup_ran = True
+
+        self.execute_run(
+            "/bin/sh",
+            "-c",
+            "echo hi",
+            parent_setup_fn=parent_setup_fn,
+            parent_cleanup_fn=parent_cleanup_fn,
+        )
+        assert parent_setup_ran
+        assert parent_cleanup_ran
+
 
 class TestRunExecutorWithContainer(TestRunExecutor):
     def setUp(self, *args, **kwargs):
@@ -1101,7 +1130,6 @@ def test_parent_fns(self):
         def parent_setup_fn(*, grandchild_pid, child_pid, **kwargs):
             # I don't want to require psutil just for this
             # I'll just read the procfs
-            print("parent setup fn")
             assert os.path.exists(f"/proc/{grandchild_pid}")
             assert os.path.exists(f"/proc/{child_pid}")
             nonlocal parent_setup_ran
@@ -1109,7 +1137,6 @@ def parent_setup_fn(*, grandchild_pid, child_pid, **kwargs):
             return 12345
 
         def parent_cleanup_fn(parent_setup, exit_code, path):
-            print("parent cleanup fn")
             assert parent_setup == 12345
             nonlocal parent_cleanup_ran
             parent_cleanup_ran = True