osutil: implement pure version of recvfds() for PyPy
authorYuya Nishihara <yuya@tcha.org>
Thu, 17 Dec 2015 23:53:09 +0900
changeset 27474 e517a89c24e1
parent 27473 34a37a2e03e6
child 27475 a2e2a8fa5fd1
osutil: implement pure version of recvfds() for PyPy This is less portable than the C version, but PyPy can't load CPython extensions. So for now, this will be used on PyPy. I've tested it on Linux amd64 and Mac OS X.
mercurial/pure/osutil.py
--- a/mercurial/pure/osutil.py	Thu Dec 17 23:41:46 2015 +0900
+++ b/mercurial/pure/osutil.py	Thu Dec 17 23:53:09 2015 +0900
@@ -7,8 +7,12 @@
 
 from __future__ import absolute_import
 
+import ctypes
+import ctypes.util
 import os
+import socket
 import stat as statmod
+import sys
 
 def _mode_to_kind(mode):
     if statmod.S_ISREG(mode):
@@ -59,8 +63,87 @@
 
 if os.name != 'nt':
     posixfile = open
+
+    _SCM_RIGHTS = 0x01
+    _socklen_t = ctypes.c_uint
+
+    if sys.platform == 'linux2':
+        # socket.h says "the type should be socklen_t but the definition of
+        # the kernel is incompatible with this."
+        _cmsg_len_t = ctypes.c_size_t
+        _msg_controllen_t = ctypes.c_size_t
+        _msg_iovlen_t = ctypes.c_size_t
+    else:
+        _cmsg_len_t = _socklen_t
+        _msg_controllen_t = _socklen_t
+        _msg_iovlen_t = ctypes.c_int
+
+    class _iovec(ctypes.Structure):
+        _fields_ = [
+            ('iov_base', ctypes.c_void_p),
+            ('iov_len', ctypes.c_size_t),
+        ]
+
+    class _msghdr(ctypes.Structure):
+        _fields_ = [
+            ('msg_name', ctypes.c_void_p),
+            ('msg_namelen', _socklen_t),
+            ('msg_iov', ctypes.POINTER(_iovec)),
+            ('msg_iovlen', _msg_iovlen_t),
+            ('msg_control', ctypes.c_void_p),
+            ('msg_controllen', _msg_controllen_t),
+            ('msg_flags', ctypes.c_int),
+        ]
+
+    class _cmsghdr(ctypes.Structure):
+        _fields_ = [
+            ('cmsg_len', _cmsg_len_t),
+            ('cmsg_level', ctypes.c_int),
+            ('cmsg_type', ctypes.c_int),
+            ('cmsg_data', ctypes.c_ubyte * 0),
+        ]
+
+    _libc = ctypes.CDLL(ctypes.util.find_library('c'), use_errno=True)
+    _recvmsg = _libc.recvmsg
+    _recvmsg.restype = ctypes.c_ssize_t
+    _recvmsg.argtypes = (ctypes.c_int, ctypes.POINTER(_msghdr), ctypes.c_int)
+
+    def _CMSG_FIRSTHDR(msgh):
+        if msgh.msg_controllen < ctypes.sizeof(_cmsghdr):
+            return
+        cmsgptr = ctypes.cast(msgh.msg_control, ctypes.POINTER(_cmsghdr))
+        return cmsgptr.contents
+
+    # The pure version is less portable than the native version because the
+    # handling of socket ancillary data heavily depends on C preprocessor.
+    # Also, some length fields are wrongly typed in Linux kernel.
+    def recvfds(sockfd):
+        """receive list of file descriptors via socket"""
+        dummy = (ctypes.c_ubyte * 1)()
+        iov = _iovec(ctypes.cast(dummy, ctypes.c_void_p), ctypes.sizeof(dummy))
+        cbuf = ctypes.create_string_buffer(256)
+        msgh = _msghdr(None, 0,
+                       ctypes.pointer(iov), 1,
+                       ctypes.cast(cbuf, ctypes.c_void_p), ctypes.sizeof(cbuf),
+                       0)
+        r = _recvmsg(sockfd, ctypes.byref(msgh), 0)
+        if r < 0:
+            e = ctypes.get_errno()
+            raise OSError(e, os.strerror(e))
+        # assumes that the first cmsg has fds because it isn't easy to write
+        # portable CMSG_NXTHDR() with ctypes.
+        cmsg = _CMSG_FIRSTHDR(msgh)
+        if not cmsg:
+            return []
+        if (cmsg.cmsg_level != socket.SOL_SOCKET or
+            cmsg.cmsg_type != _SCM_RIGHTS):
+            return []
+        rfds = ctypes.cast(cmsg.cmsg_data, ctypes.POINTER(ctypes.c_int))
+        rfdscount = ((cmsg.cmsg_len - _cmsghdr.cmsg_data.offset) /
+                     ctypes.sizeof(ctypes.c_int))
+        return [rfds[i] for i in xrange(rfdscount)]
+
 else:
-    import ctypes
     import msvcrt
 
     _kernel32 = ctypes.windll.kernel32