support/unix_connect: mimic XCB behavior

benoit-pierre · benoit-pierre · commit 3da78d203713 · 2017-12-02T20:49:21.000+01:00
Fallback to TCP when implicit Unix socket connection fails.
diff --git a/Xlib/protocol/display.py b/Xlib/protocol/display.py
@@ -82,12 +82,12 @@ class Display(object):
     event_classes = event.event_class.copy()
 
     def __init__(self, display = None):
-        name, host, displayno, screenno = connect.get_display(display)
+        name, protocol, host, displayno, screenno = connect.get_display(display)
 
         self.display_name = name
         self.default_screen = screenno
 
-        self.socket = connect.get_socket(name, host, displayno)
+        self.socket = connect.get_socket(name, protocol, host, displayno)
 
         auth_name, auth_data = connect.get_auth(self.socket,
                                                 name, host, displayno)
diff --git a/Xlib/support/connect.py b/Xlib/support/connect.py
@@ -56,34 +56,35 @@ def _relative_import(modname):
 
 
 def get_display(display):
-    """dname, host, dno, screen = get_display(display)
+    """dname, protocol, host, dno, screen = get_display(display)
 
     Parse DISPLAY into its components.  If DISPLAY is None, use
     the default display.  The return values are:
 
-      DNAME  -- the full display name (string)
-      HOST   -- the host name (string, possibly empty)
-      DNO    -- display number (integer)
-      SCREEN -- default screen number (integer)
+      DNAME    -- the full display name (string)
+      PROTOCOL -- the protocol to use (None if automatic)
+      HOST     -- the host name (string, possibly empty)
+      DNO      -- display number (integer)
+      SCREEN   -- default screen number (integer)
     """
 
     modname = _display_mods.get(platform, _default_display_mod)
     mod = _relative_import(modname)
     return mod.get_display(display)
 
 
-def get_socket(dname, host, dno):
-    """socket = get_socket(dname, host, dno)
+def get_socket(dname, protocol, host, dno):
+    """socket = get_socket(dname, protocol, host, dno)
 
-    Connect to the display specified by DNAME, HOST and DNO, which
+    Connect to the display specified by DNAME, PROTOCOL, HOST and DNO, which
     are the corresponding values from a previous call to get_display().
 
     Return SOCKET, a new socket object connected to the X server.
     """
 
     modname = _socket_mods.get(platform, _default_socket_mod)
     mod = _relative_import(modname)
-    return mod.get_socket(dname, host, dno)
+    return mod.get_socket(dname, protocol, host, dno)
 
 
 def get_auth(sock, dname, host, dno):
diff --git a/Xlib/support/unix_connect.py b/Xlib/support/unix_connect.py
@@ -62,50 +62,57 @@ def get_display(display):
 
     name = display
     protocol, host, dno, screen = m.group('proto', 'host', 'dno', 'screen')
-    if protocol == 'tcp':
+    if protocol == 'tcp' and not host:
         # Host is mandatory when protocol is TCP.
-        if not host:
-            raise error.DisplayNameError(display)
-    elif protocol == 'unix':
-        # Clear host to force Unix socket connection.
-        host = ''
-    else:
-        # Special case: `unix:0.0` is equivalent to `:0.0`.
-        if host == 'unix':
-            host = ''
+        raise error.DisplayNameError(display)
     dno = int(dno)
     if screen:
         screen = int(screen)
     else:
         screen = 0
 
-    return name, host, dno, screen
+    return name, protocol, host, dno, screen
+
 
+def _get_tcp_socket(host, dno):
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.connect((host, 6000 + dno))
+    return s
+
+def _get_unix_socket(address):
+    s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    s.connect(address)
+    return s
 
-def get_socket(dname, host, dno):
+def get_socket(dname, protocol, host, dno):
+    assert protocol in (None, 'tcp', 'unix')
     try:
-        # Darwin funky socket
-        if (uname[0] == 'Darwin') and host and host.startswith('/private/tmp/'):
-            s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-            s.connect(dname)
+        # Darwin funky socket.
+        if uname[0] == 'Darwin' and host and host.startswith('/private/tmp/'):
+            s = _get_unix_socket(dname)
 
-        # If hostname (or IP) is provided, use TCP socket
-        elif host:
-            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-            s.connect((host, 6000 + dno))
+        # TCP socket, note the special case: `unix:0.0` is equivalent to `:0.0`.
+        elif (not protocol or protocol != 'unix') and host and host != 'unix':
+            s = _get_tcp_socket(host, dno)
 
-        # Else use Unix socket
+        # Unix socket.
         else:
             address = '/tmp/.X11-unix/X%d' % dno
             if not os.path.exists(address):
                 # Use abstract address.
                 address = '\0' + address
-            s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-            s.connect(address)
+            try:
+                s = _get_unix_socket(address)
+            except socket.error:
+                if not protocol and not host:
+                    # If no protocol/host was specified, fallback to TCP.
+                    s = _get_tcp_socket(host, dno)
+                else:
+                    raise
     except socket.error as val:
         raise error.DisplayConnectionError(dname, str(val))
 
-    # Make sure that the connection isn't inherited in child processes
+    # Make sure that the connection isn't inherited in child processes.
     fcntl.fcntl(s.fileno(), F_SETFD, FD_CLOEXEC)
 
     return s
diff --git a/Xlib/support/vms_connect.py b/Xlib/support/vms_connect.py
@@ -32,7 +32,7 @@ def get_display(display):
     # check DECW$DISPLAY instead, but that has to wait
 
     if display is None:
-        return ':0.0', 'localhost', 0, 0
+        return ':0.0', None, 'localhost', 0, 0
 
     m = display_re.match(display)
     if not m:
@@ -52,10 +52,10 @@ def get_display(display):
     else:
         screen = 0
 
-    return name, host, dno, screen
+    return name, None, host, dno, screen
 
 
-def get_socket(dname, host, dno):
+def get_socket(dname, protocol, host, dno):
     try:
         # Always use TCP/IP sockets.  Later it would be nice to
         # be able to use DECNET och LOCAL connections.
diff --git a/dev-requirements.txt b/dev-requirements.txt
@@ -1,4 +1,5 @@
 coverage
 codecov
+mock
 nose
 setuptools-scm
diff --git a/test/test_unix_connect.py b/test/test_unix_connect.py
@@ -1,59 +1,158 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*
 
+from functools import partial
+import socket
 import sys
 import unittest
 
+from mock import patch
+
 from Xlib.support import unix_connect
-from Xlib.error import DisplayNameError
+from Xlib.error import DisplayConnectionError, DisplayNameError
 
 
 @unittest.skipUnless(sys.platform.startswith('linux'), 'Linux specific tests')
 class TestUnixConnect(unittest.TestCase):
 
-     def test_get_display(self):
-         # Valid cases.
-         for display, expected in (
-             # Implicit Unix socket connections.
-             (':0.1', ('', 0, 1)),
-             (':4', ('', 4, 0)),
-             # Implicit TCP connections.
-             ('foo:1.2', ('foo', 1, 2)),
-             ('bar:5', ('bar', 5, 0)),
-             # Explicit Unix socket connections.
-             ('unix/foo:4.3', ('', 4, 3)),
-             ('unix/:66', ('', 66, 0)),
-             # Explicit TCP connections.
-             ('tcp/foo:11.1', ('foo', 11, 1)),
-             ('tcp/bar:66.6', ('bar', 66, 6)),
-             ('tcp/unix:54.3', ('unix', 54, 3)),
-             # Special case: `unix:0.0` is equivalent to `:0.0`.
-             ('unix:99.5', ('', 99, 5)),
-             ('unix:42', ('', 42, 0)),
-         ):
-             result = unix_connect.get_display(display)
-             self.assertEqual(result, (display,) + expected)
-         # Invalid cases.
-         for display in (
-             # No display number.
-             '',
-             ':',
-             'foo',
-             'bar:',
-             # Bad screen number.
-             ':48.',
-             ':47.f',
-             # Bad hostname.
-             u'fòó:0',
-             u'tcp/bàr:1',
-             u'unix/fòóbàr:2',
-             # Bad protocol.
-             'udp/foo:0'
-             # With explicit TCP connections, hostname must be set.
-             'tcp/:0',
-         ):
-             with self.assertRaises(DisplayNameError):
-                 unix_connect.get_display(display)
+    def test_get_display(self):
+        # Valid cases.
+        for display, expected in (
+            # Implicit Unix socket connections.
+            (':0.1', (None, '', 0, 1)),
+            (':4', (None, '', 4, 0)),
+            # Implicit TCP connections.
+            ('foo:1.2', (None, 'foo', 1, 2)),
+            ('bar:5', (None, 'bar', 5, 0)),
+            # Explicit Unix socket connections.
+            ('unix/foo:4.3', ('unix', 'foo', 4, 3)),
+            ('unix/:66', ('unix', '', 66, 0)),
+            # Explicit TCP connections.
+            ('tcp/foo:11.1', ('tcp', 'foo', 11, 1)),
+            ('tcp/bar:66.6', ('tcp', 'bar', 66, 6)),
+            ('tcp/unix:54.3', ('tcp', 'unix', 54, 3)),
+            # Special case: `unix:0.0` is equivalent to `:0.0`.
+            ('unix:99.5', (None, 'unix', 99, 5)),
+            ('unix:42', (None, 'unix', 42, 0)),
+        ):
+            result = unix_connect.get_display(display)
+            self.assertEqual(result, (display,) + expected)
+        # Invalid cases.
+        for display in (
+            # No display number.
+            '',
+            ':',
+            'foo',
+            'bar:',
+            # Bad screen number.
+            ':48.',
+            ':47.f',
+            # Bad hostname.
+            u'fòó:0',
+            u'tcp/bàr:1',
+            u'unix/fòóbàr:2',
+            # Bad protocol.
+            'udp/foo:0'
+            # With explicit TCP connections, hostname must be set.
+            'tcp/:0',
+        ):
+            with self.assertRaises(DisplayNameError):
+                unix_connect.get_display(display)
+
+    def test_get_socket(self):
+        class FakeSocket(object):
+            def fileno(self):
+                return 42
+        calls = []
+        def _get_socket(socket_type, raises, *params):
+            calls.append(('_get_%s_socket' % socket_type,) + params)
+            if raises:
+                raise socket.error()
+            return FakeSocket()
+        def path_exists(returns, path):
+            calls.append(('os.path.exists', path))
+            return returns
+        def fcntl(*args):
+            calls.append(('fcntl',) + args)
+        for params, allow_unix, unix_addr_exists, allow_tcp, expect_connection_error, expected_calls in (
+            # Successful explicit TCP socket connection.
+            (('tcp/host:6', None, 'host', 6), False, False, True, False, [
+                ('_get_tcp_socket', 'host', 6),
+            ]),
+            # Failed explicit TCP socket connection.
+            (('tcp/host:6', None, 'host', 6), False, False, False, True, [
+                ('_get_tcp_socket', 'host', 6),
+            ]),
+            # Successful implicit TCP socket connection.
+            (('host:5', None, 'host', 5), False, False, True, False, [
+                ('_get_tcp_socket', 'host', 5),
+            ]),
+            # Failed implicit TCP socket connection.
+            (('host:5', None, 'host', 5), False, False, False, True, [
+                ('_get_tcp_socket', 'host', 5),
+            ]),
+            # Successful explicit Unix socket connection.
+            (('unix/name:0', 'unix', 'name', 0), True, True, False, False, [
+                ('os.path.exists', '/tmp/.X11-unix/X0'),
+                ('_get_unix_socket', '/tmp/.X11-unix/X0'),
+            ]),
+            # Failed explicit Unix socket connection.
+            (('unix/name:0', 'unix', 'name', 0), False, True, False, True, [
+                ('os.path.exists', '/tmp/.X11-unix/X0'),
+                ('_get_unix_socket', '/tmp/.X11-unix/X0'),
+            ]),
+            # Successful explicit Unix socket connection, variant.
+            (('unix:0', None, 'unix', 0), True, True, False, False, [
+                ('os.path.exists', '/tmp/.X11-unix/X0'),
+                ('_get_unix_socket', '/tmp/.X11-unix/X0'),
+            ]),
+            # Failed explicit Unix socket connection, variant.
+            (('unix:0', None, 'unix', 0), False, True, False, True, [
+                ('os.path.exists', '/tmp/.X11-unix/X0'),
+                ('_get_unix_socket', '/tmp/.X11-unix/X0'),
+            ]),
+            # Successful implicit Unix socket connection.
+            ((':4', None, '', 4), True, True, False, False, [
+                ('os.path.exists', '/tmp/.X11-unix/X4'),
+                ('_get_unix_socket', '/tmp/.X11-unix/X4'),
+            ]),
+            # Successful implicit Unix socket connection, abstract address.
+            ((':3', None, '', 3), True, False, False, False, [
+                ('os.path.exists', '/tmp/.X11-unix/X3'),
+                ('_get_unix_socket', '\0/tmp/.X11-unix/X3'),
+            ]),
+            # Failed implicit Unix socket connection, successful fallback on TCP.
+            ((':2', None, '', 2), False, False, True, False, [
+                ('os.path.exists', '/tmp/.X11-unix/X2'),
+                ('_get_unix_socket', '\0/tmp/.X11-unix/X2'),
+                ('_get_tcp_socket', '', 2),
+            ]),
+            # Failed implicit Unix socket connection, failed fallback on TCP.
+            ((':1', None, '', 1), False, False, False, True, [
+                ('os.path.exists', '/tmp/.X11-unix/X1'),
+                ('_get_unix_socket', '\0/tmp/.X11-unix/X1'),
+                ('_get_tcp_socket', '', 1),
+            ]),
+        ):
+            with \
+                    patch('Xlib.support.unix_connect._get_unix_socket',
+                       partial(_get_socket, 'unix', not allow_unix)), \
+                    patch('Xlib.support.unix_connect._get_tcp_socket',
+                          partial(_get_socket, 'tcp', not allow_tcp)), \
+                    patch('os.path.exists',
+                          partial(path_exists, unix_addr_exists)), \
+                    patch('fcntl.fcntl', fcntl):
+                del calls[:]
+                if expect_connection_error:
+                    with self.assertRaises(DisplayConnectionError):
+                        unix_connect.get_socket(*params)
+                else:
+                    s = unix_connect.get_socket(*params)
+                    self.assertIsInstance(s, FakeSocket)
+                    expected_calls.append(('fcntl', 42,
+                                           unix_connect.F_SETFD,
+                                           unix_connect.FD_CLOEXEC))
+                self.assertEqual(calls, expected_calls)
 
 
 if __name__ == '__main__':

-Original file line number
+Diff line change
@@ @@ -1,4 +1,5 @@ @@
 coverage
 codecov
 +mock
 nose
 setuptools-scm