spawn_wrapper: Make pre_exec function picklable

Local functions are unpicklable, which triggered this error
with the multiprocessing "spawn" start method:

AttributeError: Can't pickle local object 'spawn_wrapper.__call__.<locals>._pre_exec'

Bug: https://bugs.gentoo.org/924273
Signed-off-by: Zac Medico <zmedico@gentoo.org>

2 files changed, 26 insertions, 17 deletions

diff --git a/lib/portage/_selinux.py b/lib/portage/_selinux.py
index bf6ad2489..5ae1b4e71 100644
--- a/lib/portage/_selinux.py
+++ b/lib/portage/_selinux.py
@@ -1,4 +1,4 @@
-# Copyright 1999-2020 Gentoo Authors
+# Copyright 1999-2024 Gentoo Authors
 # Distributed under the terms of the GNU General Public License v2
 
 # Don't use the unicode-wrapped os and shutil modules here since
@@ -6,6 +6,7 @@
 import os
 import shutil
 import warnings
+from functools import partial
 
 try:
     import selinux
@@ -134,14 +135,12 @@ class spawn_wrapper:
 
     def __call__(self, *args, **kwargs):
         if self._con is not None:
-            pre_exec = kwargs.get("pre_exec")
-
-            def _pre_exec():
-                if pre_exec is not None:
-                    pre_exec()
-                setexec(self._con)
-
-            kwargs["pre_exec"] = _pre_exec
+            pre_exec = partial(setexec, self._con)
+            kwargs["pre_exec"] = (
+                portage.process._chain_pre_exec_fns(pre_exec, kwargs["pre_exec"])
+                if kwargs.get("pre_exec")
+                else pre_exec
+            )
 
         return self._spawn_func(*args, **kwargs)
 
diff --git a/lib/portage/process.py b/lib/portage/process.py
index a33e7b474..6cad250e3 100644
--- a/lib/portage/process.py
+++ b/lib/portage/process.py
@@ -18,7 +18,7 @@ import os as _os
 import warnings
 
 from dataclasses import dataclass
-from functools import lru_cache
+from functools import lru_cache, partial
 from typing import Any, Optional, Callable, Union
 
 from portage import os
@@ -1383,18 +1383,28 @@ def _start_fork(
     return pid
 
 
-class _setup_pipes_after_fork:
-    def __init__(self, target, fd_pipes):
+class _chain_pre_exec_fns:
+    """
+    Wraps a target function to call pre_exec functions just before
+    the original target function.
+    """
+
+    def __init__(self, target, *args):
         self._target = target
-        self._fd_pipes = fd_pipes
+        self._pre_exec_fns = args
 
     def __call__(self, *args, **kwargs):
-        for fd in set(self._fd_pipes.values()):
-            os.set_inheritable(fd, True)
-        _setup_pipes(self._fd_pipes, close_fds=False, inheritable=True)
+        for pre_exec in self._pre_exec_fns:
+            pre_exec()
         return self._target(*args, **kwargs)
 
 
+def _setup_pipes_after_fork(fd_pipes):
+    for fd in set(fd_pipes.values()):
+        os.set_inheritable(fd, True)
+    _setup_pipes(fd_pipes, close_fds=False, inheritable=True)
+
+
 def _start_proc(
     target: Callable[..., None],
     args: Optional[tuple[Any, ...]] = (),
@@ -1419,7 +1429,7 @@ def _start_proc(
     # which ForkProcess does not handle because its target
     # function does not necessarily exec.
     if fd_pipes and multiprocessing.get_start_method() == "fork":
-        target = _setup_pipes_after_fork(target, fd_pipes)
+        target = _chain_pre_exec_fns(target, partial(_setup_pipes_after_fork, fd_pipes))
         fd_pipes = None
 
     proc = ForkProcess(