Repository: aaugustin/websockets
Pull request #528 (2018-12-13 00:50:36)

Handle redirects in client when connecting
Per https://tools.ietf.org/html/rfc6455.html#section-4.2.2 the server may redirect the client during the handshake. This allows the client to handle redirects properly instead of raising an InvalidStatusCode error.


Patch:
diff --git a/docs/changelog.rst b/docs/changelog.rst
index 3ec3544..a76e121 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -20,6 +20,9 @@ Also:
   :meth:`~protocol.WebSocketCommonProtocol.pong` support bytes-like types
   :class:`bytearray` and :class:`memoryview` in addition to :class:`bytes`.
 
+* :func:`~client.connect()` handles redirects from the server during the
+  handshake.
+
 7.0
 ...
 
diff --git a/src/websockets/client.py b/src/websockets/client.py
index 2de160e..7b0421a 100644
--- a/src/websockets/client.py
+++ b/src/websockets/client.py
@@ -13,6 +13,7 @@ from .exceptions import (
     InvalidMessage,
     InvalidStatusCode,
     NegotiationError,
+    RedirectHandshake,
 )
 from .extensions.permessage_deflate import ClientPerMessageDeflateFactory
 from .handshake import build_request, check_response
@@ -289,8 +290,11 @@ class WebSocketClientProtocol(WebSocketCommonProtocol):
         self.write_http_request(wsuri.resource_name, request_headers)
 
         status_code, response_headers = yield from self.read_http_response()
-
-        if status_code != 101:
+        if status_code in (301, 302, 303, 307, 308):
+            if 'Location' not in response_headers:
+                raise InvalidMessage('Redirect response missing Location')
+            raise RedirectHandshake(parse_uri(response_headers['Location']))
+        elif status_code != 101:
             raise InvalidStatusCode(status_code)
 
         check_response(response_headers, key)
@@ -358,6 +362,8 @@ class Connect:
 
     """
 
+    MAX_REDIRECTS_ALLOWED = 10
+
     def __init__(
         self,
         uri,
@@ -394,8 +400,8 @@ class Connect:
         if create_protocol is None:
             create_protocol = klass
 
-        wsuri = parse_uri(uri)
-        if wsuri.secure:
+        self._wsuri = parse_uri(uri)
+        if self._wsuri.secure:
             kwds.setdefault('ssl', True)
         elif kwds.get('ssl') is not None:
             raise ValueError(
@@ -416,53 +422,86 @@ class Connect:
         elif compression is not None:
             raise ValueError("Unsupported compression: {}".format(compression))
 
-        factory = lambda: create_protocol(
-            host=wsuri.host,
-            port=wsuri.port,
-            secure=wsuri.secure,
-            ping_interval=ping_interval,
-            ping_timeout=ping_timeout,
-            close_timeout=close_timeout,
-            max_size=max_size,
-            max_queue=max_queue,
-            read_limit=read_limit,
-            write_limit=write_limit,
-            loop=loop,
-            legacy_recv=legacy_recv,
-            origin=origin,
-            extensions=extensions,
-            subprotocols=subprotocols,
-            extra_headers=extra_headers,
+        self._create_protocol = create_protocol
+        self._ping_interval = ping_interval
+        self._ping_timeout = ping_timeout
+        self._close_timeout = close_timeout
+        self._max_size = max_size
+        self._max_queue = max_queue
+        self._read_limit = read_limit
+        self._write_limit = write_limit
+        self._loop = loop
+        self._legacy_recv = legacy_recv
+        self._klass = klass
+        self._timeout = timeout
+        self._compression = compression
+        self._origin = origin
+        self._extensions = extensions
+        self._subprotocols = subprotocols
+        self._extra_headers = extra_headers
+        self._kwds = kwds
+
+    def _creating_connection(self):
+        if self._wsuri.secure:
+            self._kwds.setdefault('ssl', True)
+
+        factory = lambda: self._create_protocol(
+            host=self._wsuri.host,
+            port=self._wsuri.port,
+            secure=self._wsuri.secure,
+            ping_interval=self._ping_interval,
+            ping_timeout=self._ping_timeout,
+            close_timeout=self._close_timeout,
+            max_size=self._max_size,
+            max_queue=self._max_queue,
+            read_limit=self._read_limit,
+            write_limit=self._write_limit,
+            loop=self._loop,
+            legacy_recv=self._legacy_recv,
+            origin=self._origin,
+            extensions=self._extensions,
+            subprotocols=self._subprotocols,
+            extra_headers=self._extra_headers,
         )
 
-        if kwds.get('sock') is None:
-            host, port = wsuri.host, wsuri.port
+        if self._kwds.get('sock') is None:
+            host, port = self._wsuri.host, self._wsuri.port
         else:
             # If sock is given, host and port mustn't be specified.
             host, port = None, None
 
-        self._wsuri = wsuri
-        self._origin = origin
+        self._wsuri = self._wsuri
+        self._origin = self._origin
 
         # This is a coroutine object.
-        self._creating_connection = loop.create_connection(factory, host, port, **kwds)
+        return self._loop.create_connection(factory, host, port, **self._kwds)
 
     @asyncio.coroutine
     def __iter__(self):  # pragma: no cover
-        transport, protocol = yield from self._creating_connection
-
-        try:
-            yield from protocol.handshake(
-                self._wsuri,
-                origin=self._origin,
-                available_extensions=protocol.available_extensions,
-                available_subprotocols=protocol.available_subprotocols,
-                extra_headers=protocol.extra_headers,
-            )
-        except Exception:
-            protocol.fail_connection()
-            yield from protocol.wait_closed()
-            raise
+        for redirects in range(self.MAX_REDIRECTS_ALLOWED):
+            transport, protocol = yield from self._creating_connection()
+
+            try:
+                try:
+                    yield from protocol.handshake(
+                        self._wsuri,
+                        origin=self._origin,
+                        available_extensions=protocol.available_extensions,
+                        available_subprotocols=protocol.available_subprotocols,
+                        extra_headers=protocol.extra_headers,
+                    )
+                    break  # redirection chain ended
+                except Exception:
+                    protocol.fail_connection()
+                    yield from protocol.wait_closed()
+                    raise
+            except RedirectHandshake as e:
+                if self._wsuri.secure and not e.wsuri.secure:
+                    raise InvalidHandshake('Redirect dropped TLS')
+                self._wsuri = e.wsuri
+                continue  # redirection chain continues
+        else:
+            raise InvalidHandshake('Maximum redirects exceeded')
 
         self.ws_client = protocol
         return protocol
diff --git a/src/websockets/compatibility.py b/src/websockets/compatibility.py
index b6506b7..369c63e 100644
--- a/src/websockets/compatibility.py
+++ b/src/websockets/compatibility.py
@@ -24,6 +24,11 @@ try:  # pragma: no cover
     UPGRADE_REQUIRED = http.HTTPStatus.UPGRADE_REQUIRED
     INTERNAL_SERVER_ERROR = http.HTTPStatus.INTERNAL_SERVER_ERROR
     SERVICE_UNAVAILABLE = http.HTTPStatus.SERVICE_UNAVAILABLE
+    MOVED_PERMANENTLY = http.HTTPStatus.MOVED_PERMANENTLY
+    FOUND = http.HTTPStatus.FOUND
+    SEE_OTHER = http.HTTPStatus.SEE_OTHER
+    TEMPORARY_REDIRECT = http.HTTPStatus.TEMPORARY_REDIRECT
+    PERMANENT_REDIRECT = http.HTTPStatus.PERMANENT_REDIRECT
 except AttributeError:  # pragma: no cover
     # Python < 3.5
     class SWITCHING_PROTOCOLS:
@@ -57,3 +62,23 @@ except AttributeError:  # pragma: no cover
     class SERVICE_UNAVAILABLE:
         value = 503
         phrase = "Service Unavailable"
+
+    class MOVED_PERMANENTLY:
+        value = 301
+        phrase = "Moved Permanently"
+
+    class FOUND:
+        value = 302
+        phrase = "Found"
+
+    class SEE_OTHER:
+        value = 303
+        phrase = "See Other"
+
+    class TEMPORARY_REDIRECT:
+        value = 307
+        phrase = "Temporary Redirect"
+
+    class PERMANENT_REDIRECT:
+        value = 308
+        phrase = "Permanent Redirect"
diff --git a/src/websockets/exceptions.py b/src/websockets/exceptions.py
index b34a2c0..39fa093 100644
--- a/src/websockets/exceptions.py
+++ b/src/websockets/exceptions.py
@@ -43,6 +43,16 @@ class AbortHandshake(InvalidHandshake):
         super().__init__(message)
 
 
+class RedirectHandshake(InvalidHandshake):
+    """
+    Exception raised when a handshake gets redirected.
+
+    """
+
+    def __init__(self, wsuri):
+        self.wsuri = wsuri
+
+
 class InvalidMessage(InvalidHandshake):
     """
     Exception raised when the HTTP message in a handshake request is malformed.
diff --git a/src/websockets/py35/client.py b/src/websockets/py35/client.py
index a016ba4..bd90284 100644
--- a/src/websockets/py35/client.py
+++ b/src/websockets/py35/client.py
@@ -1,3 +1,6 @@
+from ..exceptions import InvalidHandshake, RedirectHandshake
+
+
 async def __aenter__(self):
     return await self
 
@@ -9,20 +12,30 @@ async def __aexit__(self, exc_type, exc_value, traceback):
 async def __await_impl__(self):
     # Duplicated with __iter__ because Python 3.7 requires an async function
     # (as explained in __await__ below) which Python 3.4 doesn't support.
-    transport, protocol = await self._creating_connection
-
-    try:
-        await protocol.handshake(
-            self._wsuri,
-            origin=self._origin,
-            available_extensions=protocol.available_extensions,
-            available_subprotocols=protocol.available_subprotocols,
-            extra_headers=protocol.extra_headers,
-        )
-    except Exception:
-        protocol.fail_connection()
-        await protocol.wait_closed()
-        raise
+    for redirects in range(self.MAX_REDIRECTS_ALLOWED):
+        transport, protocol = await self._creating_connection()
+
+        try:
+            try:
+                await protocol.handshake(
+                    self._wsuri,
+                    origin=self._origin,
+                    available_extensions=protocol.available_extensions,
+                    available_subprotocols=protocol.available_subprotocols,
+                    extra_headers=protocol.extra_headers,
+                )
+                break  # redirection chain ended
+            except Exception:
+                protocol.fail_connection()
+                await protocol.wait_closed()
+                raise
+        except RedirectHandshake as e:
+            if self._wsuri.secure and not e.wsuri.secure:
+                raise InvalidHandshake('Redirect dropped TLS')
+            self._wsuri = e.wsuri
+            continue  # redirection chain continues
+    else:
+        raise InvalidHandshake('Maximum redirects exceeded')
 
     self.ws_client = protocol
     return protocol


Test patch:
diff --git a/tests/test_client_server.py b/tests/test_client_server.py
index 73866ff..eade7e0 100644
--- a/tests/test_client_server.py
+++ b/tests/test_client_server.py
@@ -15,10 +15,20 @@ import urllib.request
 import warnings
 
 from websockets.client import *
-from websockets.compatibility import FORBIDDEN, OK, UNAUTHORIZED
+from websockets.compatibility import (
+    FORBIDDEN,
+    FOUND,
+    MOVED_PERMANENTLY,
+    OK,
+    PERMANENT_REDIRECT,
+    SEE_OTHER,
+    TEMPORARY_REDIRECT,
+    UNAUTHORIZED,
+)
 from websockets.exceptions import (
     ConnectionClosed,
     InvalidHandshake,
+    InvalidMessage,
     InvalidStatusCode,
     NegotiationError,
 )
@@ -79,6 +89,17 @@ def temp_test_server(test, **kwds):
         test.stop_server()
 
 
+@contextlib.contextmanager
+def temp_test_redirecting_server(
+    test, status, include_location=True, force_insecure=False
+):
+    test.start_redirecting_server(status, include_location, force_insecure)
+    try:
+        yield
+    finally:
+        test.stop_redirecting_server()
+
+
 @contextlib.contextmanager
 def temp_test_client(test, *args, **kwds):
     test.start_client(*args, **kwds)
@@ -227,6 +248,8 @@ class ClientServerTests(unittest.TestCase):
     def setUp(self):
         self.loop = asyncio.new_event_loop()
         asyncio.set_event_loop(self.loop)
+        self.server = None
+        self.redirecting_server = None
 
     def tearDown(self):
         self.loop.close()
@@ -237,6 +260,10 @@ class ClientServerTests(unittest.TestCase):
         self.loop.call_soon(self.loop.stop)
         self.loop.run_forever()
 
+    @property
+    def server_context(self):
+        return None
+
     def start_server(self, **kwds):
         # Disable compression by default in tests.
         kwds.setdefault('compression', None)
@@ -245,13 +272,35 @@ class ClientServerTests(unittest.TestCase):
         start_server = serve(handler, 'localhost', 0, **kwds)
         self.server = self.loop.run_until_complete(start_server)
 
+    def start_redirecting_server(
+        self, status, include_location=True, force_insecure=False
+    ):
+        def _process_request(path, headers):
+            server_uri = get_server_uri(self.server, self.secure, path)
+            if force_insecure:
+                server_uri = server_uri.replace('wss:', 'ws:')
+            headers = {'Location': server_uri} if include_location else []
+            return status, headers, b""
+
+        start_server = serve(
+            handler,
+            'localhost',
+            0,
+            compression=None,
+            ping_interval=None,
+            process_request=_process_request,
+            ssl=self.server_context,
+        )
+        self.redirecting_server = self.loop.run_until_complete(start_server)
+
     def start_client(self, resource_name='/', user_info=None, **kwds):
         # Disable compression by default in tests.
         kwds.setdefault('compression', None)
         # Disable pings by default in tests.
         kwds.setdefault('ping_interval', None)
         secure = kwds.get('ssl') is not None
-        server_uri = get_server_uri(self.server, secure, resource_name, user_info)
+        server = self.redirecting_server if self.redirecting_server else self.server
+        server_uri = get_server_uri(server, secure, resource_name, user_info)
         start_client = connect(server_uri, **kwds)
         self.client = self.loop.run_until_complete(start_client)
 
@@ -272,6 +321,17 @@ class ClientServerTests(unittest.TestCase):
         except asyncio.TimeoutError:  # pragma: no cover
             self.fail("Server failed to stop")
 
+    def stop_redirecting_server(self):
+        self.redirecting_server.close()
+        try:
+            self.loop.run_until_complete(
+                asyncio.wait_for(self.redirecting_server.wait_closed(), timeout=1)
+            )
+        except asyncio.TimeoutError:  # pragma: no cover
+            self.fail("Redirecting server failed to stop")
+        finally:
+            self.redirecting_server = None
+
     @contextlib.contextmanager
     def temp_server(self, **kwds):
         with temp_test_server(self, **kwds):
@@ -289,6 +349,36 @@ class ClientServerTests(unittest.TestCase):
         reply = self.loop.run_until_complete(self.client.recv())
         self.assertEqual(reply, "Hello!")
 
+    @with_server()
+    def test_redirect(self):
+        redirect_statuses = [
+            MOVED_PERMANENTLY,
+            FOUND,
+            SEE_OTHER,
+            TEMPORARY_REDIRECT,
+            PERMANENT_REDIRECT,
+        ]
+        for status in redirect_statuses:
+            with temp_test_redirecting_server(self, status):
+                with temp_test_client(self):
+                    self.loop.run_until_complete(self.client.send("Hello!"))
+                    reply = self.loop.run_until_complete(self.client.recv())
+                    self.assertEqual(reply, "Hello!")
+
+    def test_infinite_redirect(self):
+        with temp_test_redirecting_server(self, FOUND):
+            self.server = self.redirecting_server
+            with self.assertRaises(InvalidHandshake):
+                with temp_test_client(self):
+                    self.fail('Did not raise')  # pragma: no cover
+
+    @with_server()
+    def test_redirect_missing_location(self):
+        with temp_test_redirecting_server(self, FOUND, include_location=False):
+            with self.assertRaises(InvalidMessage):
+                with temp_test_client(self):
+                    self.fail('Did not raise')  # pragma: no cover
+
     def test_explicit_event_loop(self):
         with self.temp_server(loop=self.loop):
             with self.temp_client(loop=self.loop):
@@ -1070,6 +1160,13 @@ class SSLClientServerTests(ClientServerTests):
             # raised only when awaiting.
             self.loop.run_until_complete(client)  # pragma: no cover
 
+    @with_server()
+    def test_redirect_insecure(self):
+        with temp_test_redirecting_server(self, FOUND, force_insecure=True):
+            with self.assertRaises(InvalidHandshake):
+                with temp_test_client(self):
+                    self.fail('Did not raise')  # pragma: no cover
+
 
 class ClientServerOriginTests(unittest.TestCase):
     def setUp(self):