Minor edits per @.picnixz code review comments.

python · gpshead · Nov 20, 2024 · Nov 10, 2022 · Nov 11, 2022 · Nov 11, 2022
commit a53c01fc0f5cecbd006e5a37526fb9357386cfa1
@@ -355,12 +355,11 @@ json
 multiprocessing
 ---------------
 
-* :mod:`multiprocessing`'s ``"forkserver"`` start method gains authentication
-  on its control sockets so that it isn't solely reliant on filesystem
-  permissions to control what other processes can cause the fork server to
-  spawn workers and run code.
-  This improves the security story behind :gh:`97514`.
-  (Contributed by Gregory P. Smith.)
+* :mod:`multiprocessing`'s ``"forkserver"`` start method now authenticates
+  its control socket to avoid solely relying on filesystem permissions
+  to restrict what other processes could cause the forkserver to spawn workers
+  and run code.
+  (Contributed by Gregory P. Smith for :gh:`97514`.)
 
 
 operator

@@ -26,7 +26,7 @@
 
 MAXFDS_TO_SEND = 256
 SIGNED_STRUCT = struct.Struct('q')     # large enough for pid_t
-_authkey_len = 32  # <= PIPEBUF so it fits a single write to an empty pipe.
+_AUTHKEY_LEN = 32  # <= PIPEBUF so it fits a single write to an empty pipe.
 
 #
 # Forkserver class
@@ -87,6 +87,7 @@ def connect_to_new_process(self, fds):
         process data.
         '''
         self.ensure_running()
+        assert self._forkserver_authkey
         if len(fds) + 4 >= MAXFDS_TO_SEND:
             raise ValueError('too many fds')
         with socket.socket(socket.AF_UNIX) as client:
@@ -97,7 +98,6 @@ def connect_to_new_process(self, fds):
                       resource_tracker.getfd()]
             allfds += fds
             try:
-                assert self._forkserver_authkey
                 client.setblocking(True)
                 wrapped_client = connection.Connection(client.fileno())
                 # The other side of this exchange happens in the child as
@@ -183,7 +183,7 @@ def ensure_running(self):
                 # Authenticate our control socket to prevent access from
                 # processes we have not shared this key with.
                 try:
-                    self._forkserver_authkey = os.urandom(_authkey_len)
+                    self._forkserver_authkey = os.urandom(_AUTHKEY_LEN)
                     os.write(authkey_w, self._forkserver_authkey)
                 finally:
                     os.close(authkey_w)
@@ -199,9 +199,11 @@ def main(listener_fd, alive_r, preload, main_path=None, sys_path=None,
          *, authkey_r=None):
     """Run forkserver."""
     if authkey_r is not None:
-        authkey = os.read(authkey_r, _authkey_len)
-        assert len(authkey) == _authkey_len, f'{len(authkey)} < {_authkey_len}'
-        os.close(authkey_r)
+        try:
+            authkey = os.read(authkey_r, _AUTHKEY_LEN)
+            assert len(authkey) == _AUTHKEY_LEN, f'{len(authkey)} < {_AUTHKEY_LEN}'
+        finally:
+            os.close(authkey_r)
     else:
         authkey = b''
 

@@ -895,10 +895,6 @@ def test_forkserver_sigkill(self):
         if os.name != 'nt':
             self.check_forkserver_death(signal.SIGKILL)
 
-    @staticmethod
-    def _exit_process():
-        sys.exit(0)
-
     def test_forkserver_auth_is_enabled(self):
         if self.TYPE == "threads":
             self.skipTest(f"test not appropriate for {self.TYPE}")
@@ -920,7 +916,7 @@ def test_forkserver_auth_is_enabled(self):
         client.close()
 
         # That worked, now launch a quick process.
-        proc = self.Process(target=self._exit_process)
+        proc = self.Process(target=sys.exit)
         proc.start()
         proc.join()
         self.assertEqual(proc.exitcode, 0)
@@ -939,14 +935,14 @@ def test_forkserver_without_auth_fails(self):
                 forkserver, '_forkserver_authkey', None):
             # With an incorrect authkey we should get an auth rejection
             # rather than the above protocol error.
-            forkserver._forkserver_authkey = b'T'*authkey_len
-            proc = self.Process(target=self._exit_process)
+            forkserver._forkserver_authkey = b'T' * authkey_len
+            proc = self.Process(target=sys.exit)
             with self.assertRaises(multiprocessing.AuthenticationError):
                 proc.start()
             del proc
 
         # authkey restored, launching processes should work again.
-        proc = self.Process(target=self._exit_process)
+        proc = self.Process(target=sys.exit)
         proc.start()
         proc.join()