websocket redirect ability

tornado/test/websocket_test.py

@@ -104,6 +104,14 @@ def get(self):
         self.write("ok")
 
 
+class RedirectHandler(RequestHandler):
+    def initialize(self, target_location):
+        self._target_location = target_location
+
+    def get(self):
+        self.redirect(self._target_location)
+
+
 class CloseReasonHandler(TestWebSocketHandler):
     def open(self):
         self.on_close_called = False
@@ -208,6 +216,12 @@ def get_app(self):
         return Application(
             [
                 ("/echo", EchoHandler, dict(close_future=self.close_future)),
+                ("/redirect", RedirectHandler, dict(target_location="/echo")),
+                (
+                    "/double_redirect",
+                    RedirectHandler,
+                    dict(target_location="/redirect"),
+                ),
                 ("/non_ws", NonWebSocketHandler),
                 ("/header", HeaderHandler, dict(close_future=self.close_future)),
                 (
@@ -355,6 +369,20 @@ def test_websocket_http_success(self):
         with self.assertRaises(WebSocketError):
             yield self.ws_connect("/non_ws")
 
+    @gen_test
+    def test_websocket_redirect(self):
+        conn = yield self.ws_connect("/redirect")
+        yield conn.write_message("hello redirect")
+        msg = yield conn.read_message()
+        self.assertEqual("hello redirect", msg)
+
+    @gen_test
+    def test_websocket_double_redirect(self):
+        conn = yield self.ws_connect("/double_redirect")
+        yield conn.write_message("hello redirect")
+        msg = yield conn.read_message()
+        self.assertEqual("hello redirect", msg)
+
     @gen_test
     def test_websocket_network_fail(self):
         sock, port = bind_unused_port()

tornado/websocket.py

@@ -18,16 +18,21 @@
 
 import abc
 import base64
+import copy
 import hashlib
 import os
 import sys
 import struct
 import tornado.escape
 import tornado.web
-from urllib.parse import urlparse
+from urllib.parse import urlparse, urljoin
 import zlib
 
-from tornado.concurrent import Future, future_set_result_unless_cancelled
+from tornado.concurrent import (
+    Future,
+    future_set_result_unless_cancelled,
+    future_set_exception_unless_cancelled,
+)
 from tornado.escape import utf8, native_str, to_unicode
 from tornado import gen, httpclient, httputil
 from tornado.ioloop import IOLoop, PeriodicCallback
@@ -143,6 +148,23 @@ class WebSocketClosedError(WebSocketError):
     pass
 
 
+class WebSocketRedirect(WebSocketError):
+    """Raised when code is 302/303 in received headers
+       while doing websocket handshake
+    """
+
+    def __init__(self, target_location: str) -> None:
+        self.target_location = target_location
+
+
+class RequestRedirectError(WebSocketError):
+    """Raised when request's max_redirects <= 0
+       or follow_redirect is False
+    """
+
+    pass
+
+
 class _DecompressTooLargeError(Exception):
     pass
 
@@ -1376,6 +1398,7 @@ def __init__(
         self.ping_interval = ping_interval
         self.ping_timeout = ping_timeout
         self.max_message_size = max_message_size
+        self._redirect_location = None  # type: Optional[str]
 
         scheme, sep, rest = request.url.partition(":")
         scheme = {"ws": "http", "wss": "https"}[scheme]
@@ -1430,7 +1453,14 @@ def close(self, code: int = None, reason: str = None) -> None:
 
     def on_connection_close(self) -> None:
         if not self.connect_future.done():
-            self.connect_future.set_exception(StreamClosedError())
+            if self._redirect_location is not None:
+                # websocket redirect, raise exception for reconnect
+                # to redirected url
+                self.connect_future.set_exception(
+                    WebSocketRedirect(self._redirect_location)
+                )
+            else:
+                self.connect_future.set_exception(StreamClosedError())
         self._on_message(None)
         self.tcp_client.close()
         super(WebSocketClientConnection, self).on_connection_close()
@@ -1450,6 +1480,14 @@ async def headers_received(
         headers: httputil.HTTPHeaders,
     ) -> None:
         assert isinstance(start_line, httputil.ResponseStartLine)
+        self.code = start_line.code
+        self.reason = start_line.reason
+        self.headers = headers
+
+        if self._should_follow_redirect():
+            self._redirect_location = headers["Location"]
+            return
+
         if start_line.code != 101:
             await super(WebSocketClientConnection, self).headers_received(
                 start_line, headers
@@ -1476,6 +1514,20 @@ async def headers_received(
 
         future_set_result_unless_cancelled(self.connect_future, self)
 
+    def finish(self) -> None:
+        if self._should_follow_redirect():
+            if self._redirect_location is not None:
+                self.connect_future.set_exception(
+                    WebSocketRedirect(self._redirect_location)
+                )
+            else:
+                raise ValueError("redirect location is None")
+            # close connection
+            self.on_connection_close()
+            self.connection.close()
+        else:
+            super().finish()
+
     def write_message(
         self, message: Union[str, bytes], binary: bool = False
     ) -> "Future[None]":
@@ -1634,6 +1686,103 @@ def websocket_connect(
         httpclient.HTTPRequest,
         httpclient._RequestProxy(request, httpclient.HTTPRequest._DEFAULTS),
     )
+
+    def do_connect(
+        request: httpclient.HTTPRequest
+    ) -> "Future[WebSocketClientConnection]":
+        conn_future = _websocket_connect(
+            request,
+            callback,
+            connect_timeout,
+            on_message_callback,
+            compression_options,
+            ping_interval,
+            ping_timeout,
+            max_message_size,
+            subprotocols,
+        )
+        return conn_future
+
+    def wrap_conn_future_callback(
+        conn_future: "Future[WebSocketClientConnection]"
+    ) -> None:
+        try:
+            conn = conn_future.result()
+        except WebSocketRedirect as e:
+            try:
+                new_request = redirect_request(request, e.target_location)
+            except RequestRedirectError as e:
+                future_set_exception_unless_cancelled(wrapped_conn_future, e)
+            else:
+                conn_future = do_connect(new_request)
+                IOLoop.current().add_future(conn_future, wrap_conn_future_callback)
+        except Exception as e:
+            future_set_exception_unless_cancelled(wrapped_conn_future, e)
+        else:
+            future_set_result_unless_cancelled(wrapped_conn_future, conn)
+
+    wrapped_conn_future = Future()  # type: Future[WebSocketClientConnection]
+    conn_future = do_connect(request)
+    IOLoop.current().add_future(conn_future, wrap_conn_future_callback)
+
+    return wrapped_conn_future
+
+
+def redirect_request(
+    request: httpclient.HTTPRequest, target_location: str
+) -> httpclient.HTTPRequest:
+    request_proxy = cast(httpclient._RequestProxy, request)
+
+    if not request_proxy.follow_redirects:
+        raise RequestRedirectError("follow_redirects is not True")
+
+    if request_proxy.max_redirects is None or request_proxy.max_redirects <= 0:
+        raise RequestRedirectError("max_redirects occurred, no more redirect")
+
+    original_request = getattr(request, "original_request", request)
+    new_request = copy.copy(request_proxy.request)
+    new_request.url = urljoin(request_proxy.request.url, target_location)
+
+    new_request.max_redirects = request_proxy.max_redirects - 1
+    del new_request.headers["Host"]
+    # http://www.w3.org/Protocols/rfc2616/rfc2616-sec10.html#sec10.3.4
+    # Client SHOULD make a GET request after a 303.
+    # According to the spec, 302 should be followed by the same
+    # method as the original request, but in practice browsers
+    # treat 302 the same as 303, and many servers use 302 for
+    # compatibility with pre-HTTP/1.1 user agents which don't
+    # understand the 303 status.
+    new_request.method = "GET"
+    new_request.body = cast(bytes, None)
+    for h in [
+        "Content-Length",
+        "Content-Type",
+        "Content-Encoding",
+        "Transfer-Encoding",
+    ]:
+        try:
+            del new_request.headers[h]
+        except KeyError:
+            pass
+    new_request.original_request = original_request  # type: ignore
+    new_request = cast(
+        httpclient.HTTPRequest,
+        httpclient._RequestProxy(new_request, httpclient.HTTPRequest._DEFAULTS),
+    )
+    return new_request
+
+
+def _websocket_connect(
+    request: httpclient.HTTPRequest,
+    callback: Callable[["Future[WebSocketClientConnection]"], None] = None,
+    connect_timeout: float = None,
+    on_message_callback: Callable[[Union[None, str, bytes]], None] = None,
+    compression_options: Dict[str, Any] = None,
+    ping_interval: float = None,
+    ping_timeout: float = None,
+    max_message_size: int = _default_max_message_size,
+    subprotocols: List[str] = None,
+) -> "Future[WebSocketClientConnection]":
     conn = WebSocketClientConnection(
         request,
         on_message_callback=on_message_callback,