improve some error handling and add a test.

python · gpshead · Nov 20, 2024 · Nov 10, 2022 · Nov 11, 2022 · Nov 11, 2022
commit 72f3843fb9c0145f6fac7f1fea9520a4382387e5
@@ -8,6 +8,7 @@
 import threading
 import warnings
 
+from . import AuthenticationError
 from . import connection
 from . import process
 from .context import reduction
@@ -290,15 +291,22 @@ def sigchld_handler(*_unused):
                 if listener in rfds:
                     # Incoming fork request
                     with listener.accept()[0] as s:
-                        if authkey:
-                            wrapped_s = connection.Connection(s.fileno())
-                            try:
-                                connection.deliver_challenge(wrapped_s, authkey)
-                                connection.answer_challenge(wrapped_s, authkey)
-                            finally:
-                                wrapped_s._detach()
-                        # Receive fds from client
-                        fds = reduction.recvfds(s, MAXFDS_TO_SEND + 1)
+                        try:
+                            if authkey:
+                                wrapped_s = connection.Connection(s.fileno())
+                                try:
+                                    connection.deliver_challenge(
+                                            wrapped_s, authkey)
+                                    connection.answer_challenge(
+                                            wrapped_s, authkey)
+                                finally:
+                                    wrapped_s._detach()
+                            # Receive fds from client
+                            fds = reduction.recvfds(s, MAXFDS_TO_SEND + 1)
+                        except (EOFError, OSError, AuthenticationError):
+                            # broken pipe or failed authentication
+                            s.close()
+                            continue
                         if len(fds) > MAXFDS_TO_SEND:
                             raise RuntimeError(
                                 "Too many ({0:n}) fds to send".format(

@@ -139,15 +139,12 @@ def detach(self):
     __all__ += ['DupFd', 'sendfds', 'recvfds']
     import array
 
-    # On MacOSX we should acknowledge receipt of fds -- see Issue14669
-    ACKNOWLEDGE = sys.platform == 'darwin'
-
     def sendfds(sock, fds):
         '''Send an array of fds over an AF_UNIX socket.'''
         fds = array.array('i', fds)
         msg = bytes([len(fds) % 256])
         sock.sendmsg([msg], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, fds)])
-        if ACKNOWLEDGE and sock.recv(1) != b'A':
+        if sock.recv(1) != b'A':
             raise RuntimeError('did not receive acknowledgement of fd')
 
     def recvfds(sock, size):
@@ -158,8 +155,7 @@ def recvfds(sock, size):
         if not msg and not ancdata:
             raise EOFError
         try:
-            if ACKNOWLEDGE:
-                sock.send(b'A')
+            sock.send(b'A')  # Acknowledge
             if len(ancdata) != 1:
                 raise RuntimeError('received %d items of ancdata' %
                                    len(ancdata))

diff --git a/Lib/test/test_multiprocessing_forkserver.py b/Lib/test/test_multiprocessing_forkserver.py
@@ -1,6 +1,8 @@
 import unittest
+from unittest import mock
 import test._test_multiprocessing
 
+import os
 import sys
 from test import support
 
@@ -10,7 +12,56 @@
 if sys.platform == "win32":
     raise unittest.SkipTest("forkserver is not available on Windows")
 
+import multiprocessing
+import multiprocessing.connection
+import multiprocessing.forkserver
+
 test._test_multiprocessing.install_tests_in_module_dict(globals(), 'forkserver')
 
+
+class TestForkserverControlAuthentication(unittest.TestCase):
+    def setUp(self):
+        super().setUp()
+        self.context = multiprocessing.get_context("forkserver")
+        self.pool = self.context.Pool(processes=1, maxtasksperchild=4)
+        self.assertEqual(self.pool.apply(eval, ("2+2",)), 4)
+        self.forkserver = multiprocessing.forkserver._forkserver
+        self.addr = self.forkserver._forkserver_address
+        self.assertTrue(self.addr)
+        self.authkey = self.forkserver._forkserver_authkey
+        self.assertGreater(len(self.authkey), 15)
+        self.assertTrue(self.forkserver._forkserver_pid)
+
+    def tearDown(self):
+        self.pool.terminate()
+        self.pool.join()
+        super().tearDown()
+
+    def test_auth_works(self):
+        """FYI: An 'EOFError: ran out of input' from a worker is normal."""
+        # First, demonstrate that a raw auth handshake as Client makes
+        # does not raise.
+        client = multiprocessing.connection.Client(
+                self.addr, authkey=self.authkey)
+        client.close()
+
+        # Now use forkserver code to do the same thing and more.
+        status_r, data_w = self.forkserver.connect_to_new_process([])
+        # It is normal for this to trigger an EOFError on stderr from the
+        # process... it is expecting us to send over a pickle of a Process
+        # instance to tell it what to do.
+        # If the authentication handshake and subsequent file descriptor
+        # sending dance had failed, an exception would've been raised.
+        os.close(data_w)
+        os.close(status_r)
+
+    def test_no_auth_fails(self):
+        with mock.patch.object(self.forkserver, '_forkserver_authkey', None):
+            # With no authkey set, the connection this makes will fail to
+            # do the file descriptor transfer over the pipe.
+            with self.assertRaisesRegex(RuntimeError, 'not receive ack'):
+                status_r, data_w = self.forkserver.connect_to_new_process([])
+
+
 if __name__ == '__main__':
     unittest.main()