bpo-28724: Add methods send_fds and recv_fds to the socket module (GH-12889)

nanjekyejoannah · vstinner · vstinner · commit 8d120f75fb8c · 2019-09-11T19:12:21.000+02:00
The socket module now has the socket.send_fds() and socket.recv.fds() functions.
Contributed by Joannah Nanjekye, Shinya Okano (original patch)
and Victor Stinner.

Co-Authored-By: Victor Stinner &lt;vstinner@redhat.com&gt;
diff --git a/Doc/library/socket.rst b/Doc/library/socket.rst
@@ -1584,6 +1584,29 @@ to sockets.
 
    .. versionadded:: 3.6
 
+.. method:: socket.send_fds(sock, buffers, fds[, flags[, address]])
+
+   Send the list of file descriptors *fds* over an :const:`AF_UNIX` socket.
+   The *fds* parameter is a sequence of file descriptors.
+   Consult :meth:`sendmsg` for the documentation of these parameters.
+
+   .. availability:: Unix supporting :meth:`~socket.sendmsg` and :const:`SCM_RIGHTS` mechanism.
+
+   .. versionadded:: 3.9
+
+.. method:: socket.recv_fds(sock, bufsize, maxfds[, flags])
+
+   Receive up to *maxfds* file descriptors. Return ``(msg, list(fds), flags, addr)``. Consult
+   :meth:`recvmsg` for the documentation of these parameters.
+
+   .. availability:: Unix supporting :meth:`~socket.recvmsg` and :const:`SCM_RIGHTS` mechanism.
+
+   .. versionadded:: 3.9
+
+   .. note::
+
+      Any truncated integers at the end of the list of file descriptors.
+
 .. method:: socket.sendfile(file, offset=0, count=None)
 
    Send a file until EOF is reached by using high-performance
diff --git a/Doc/whatsnew/3.8.rst b/Doc/whatsnew/3.8.rst
diff --git a/Lib/socket.py b/Lib/socket.py
@@ -12,6 +12,8 @@
 socket() -- create a new socket object
 socketpair() -- create a pair of new socket objects [*]
 fromfd() -- create a socket object from an open file descriptor [*]
+send_fds() -- Send file descriptor to the socket.
+recv_fds() -- Recieve file descriptors from the socket.
 fromshare() -- create a socket object from data received from socket.share() [*]
 gethostname() -- return the current hostname
 gethostbyname() -- map a hostname to its IP number
@@ -542,6 +544,40 @@ def fromfd(fd, family, type, proto=0):
     nfd = dup(fd)
     return socket(family, type, proto, nfd)
 
+if hasattr(_socket.socket, "sendmsg"):
+    import array
+
+    def send_fds(sock, buffers, fds, flags=0, address=None):
+        """ send_fds(sock, buffers, fds[, flags[, address]]) -> integer
+
+        Send the list of file descriptors fds over an AF_UNIX socket.
+        """
+        return sock.sendmsg(buffers, [(_socket.SOL_SOCKET,
+            _socket.SCM_RIGHTS, array.array("i", fds))])
+    __all__.append("send_fds")
+
+if hasattr(_socket.socket, "recvmsg"):
+    import array
+
+    def recv_fds(sock, bufsize, maxfds, flags=0):
+        """ recv_fds(sock, bufsize, maxfds[, flags]) -> (data, list of file
+        descriptors, msg_flags, address)
+
+        Receive up to maxfds file descriptors returning the message
+        data and a list containing the descriptors.
+        """
+        # Array of ints
+        fds = array.array("i")
+        msg, ancdata, flags, addr = sock.recvmsg(bufsize,
+            _socket.CMSG_LEN(maxfds * fds.itemsize))
+        for cmsg_level, cmsg_type, cmsg_data in ancdata:
+            if (cmsg_level == _socket.SOL_SOCKET and cmsg_type == _socket.SCM_RIGHTS):
+                fds.frombytes(cmsg_data[:
+                        len(cmsg_data) - (len(cmsg_data) % fds.itemsize)])
+
+        return msg, list(fds), flags, addr
+    __all__.append("recv_fds")
+
 if hasattr(_socket.socket, "share"):
     def fromshare(info):
         """ fromshare(info) -> socket object
diff --git a/Lib/test/test_socket.py b/Lib/test/test_socket.py
@@ -6436,11 +6436,53 @@ def test_dual_stack_client_v6(self):
             self.echo_server(sock)
             self.echo_client(("::1", port), socket.AF_INET6)
 
+@requireAttrs(socket, "send_fds")
+@requireAttrs(socket, "recv_fds")
+@requireAttrs(socket, "AF_UNIX")
+class SendRecvFdsTests(unittest.TestCase):
+    def testSendAndRecvFds(self):
+        def close_pipes(pipes):
+            for fd1, fd2 in pipes:
+                os.close(fd1)
+                os.close(fd2)
+
+        def close_fds(fds):
+            for fd in fds:
+                os.close(fd)
+
+        # send 10 file descriptors
+        pipes = [os.pipe() for _ in range(10)]
+        self.addCleanup(close_pipes, pipes)
+        fds = [rfd for rfd, wfd in pipes]
+
+        # use a UNIX socket pair to exchange file descriptors locally
+        sock1, sock2 = socket.socketpair(socket.AF_UNIX, socket.SOCK_STREAM)
+        with sock1, sock2:
+            socket.send_fds(sock1, [MSG], fds)
+            # request more data and file descriptors than expected
+            msg, fds2, flags, addr = socket.recv_fds(sock2, len(MSG) * 2, len(fds) * 2)
+            self.addCleanup(close_fds, fds2)
+
+        self.assertEqual(msg, MSG)
+        self.assertEqual(len(fds2), len(fds))
+        self.assertEqual(flags, 0)
+        # don't test addr
+
+        # test that file descriptors are connected
+        for index, fds in enumerate(pipes):
+            rfd, wfd = fds
+            os.write(wfd, str(index).encode())
+
+        for index, rfd in enumerate(fds2):
+            data = os.read(rfd, 100)
+            self.assertEqual(data,  str(index).encode())
+
 
 def test_main():
     tests = [GeneralModuleTests, BasicTCPTest, TCPCloserTest, TCPTimeoutTest,
              TestExceptions, BufferIOTest, BasicTCPTest2, BasicUDPTest,
-             UDPTimeoutTest, CreateServerTest, CreateServerFunctionalTest]
+             UDPTimeoutTest, CreateServerTest, CreateServerFunctionalTest,
+             SendRecvFdsTests]
 
     tests.extend([
         NonBlockingTCPTests,
@@ -6513,5 +6555,6 @@ def test_main():
     support.run_unittest(*tests)
     support.threading_cleanup(*thread_info)
 
+
 if __name__ == "__main__":
     test_main()
diff --git a/Misc/NEWS.d/next/Library/2019-05-30-15-51-42.bpo-28724.34TrS8.rst b/Misc/NEWS.d/next/Library/2019-05-30-15-51-42.bpo-28724.34TrS8.rst
@@ -0,0 +1,2 @@
+The socket module now has the :func:`socket.send_fds` and :func:`socket.recv.fds` methods.
+Contributed by Joannah Nanjekye, Shinya Okano and Victor Stinner.

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	+The socket module now has the :func:`socket.send_fds` and :func:`socket.recv.fds` methods.
	`2`	`+Contributed by Joannah Nanjekye, Shinya Okano and Victor Stinner.`