Repository: tornadoweb/tornado
Pull request #1917 (2016-12-20 20:12:14)

Allow custom websocket upgrade response headers
A new method, upgrade_response_headers, is introduced to WebSocketHandler that a user can override to return custom headers in the websocket upgrade response.


Patch:
diff --git a/tornado/websocket.py b/tornado/websocket.py
index d30aead8..1fafed77 100644
--- a/tornado/websocket.py
+++ b/tornado/websocket.py
@@ -65,6 +65,10 @@ class WebSocketHandler(tornado.web.RequestHandler):
     override `open` and `on_close` to handle opened and closed
     connections.
 
+    Custom upgrade response headers can be sent by overriding
+    `~tornado.web.RequestHandler.set_default_headers` or
+    `~tornado.web.RequestHandler.prepare`.
+
     See http://dev.w3.org/html5/websockets/ for details on the
     JavaScript interface.  The protocol is specified at
     http://tools.ietf.org/html/rfc6455.
@@ -176,18 +180,13 @@ class WebSocketHandler(tornado.web.RequestHandler):
             gen_log.debug(log_msg)
             return
 
-        self.stream = self.request.connection.detach()
-        self.stream.set_close_callback(self.on_connection_close)
-
         self.ws_connection = self.get_websocket_protocol()
         if self.ws_connection:
             self.ws_connection.accept_connection()
         else:
-            if not self.stream.closed():
-                self.stream.write(tornado.escape.utf8(
-                    "HTTP/1.1 426 Upgrade Required\r\n"
-                    "Sec-WebSocket-Version: 7, 8, 13\r\n\r\n"))
-                self.stream.close()
+            self.set_status(426, "Upgrade Required")
+            self.set_header("Sec-WebSocket-Version", "7, 8, 13")
+            self.finish()
 
     stream = None
 
@@ -429,18 +428,17 @@ class WebSocketHandler(tornado.web.RequestHandler):
             return WebSocketProtocol13(
                 self, compression_options=self.get_compression_options())
 
+    def _attach_stream(self):
+        self.stream = self.request.connection.detach()
+        self.stream.set_close_callback(self.on_connection_close)
+        # disable non-WS methods
+        for method in ["write", "redirect", "set_header", "set_cookie",
+                       "set_status", "flush", "finish"]:
+            setattr(self, method, _raise_not_supported_for_websockets)
 
-def _wrap_method(method):
-    def _disallow_for_websocket(self, *args, **kwargs):
-        if self.stream is None:
-            method(self, *args, **kwargs)
-        else:
-            raise RuntimeError("Method not supported for Web Sockets")
-    return _disallow_for_websocket
-for method in ["write", "redirect", "set_header", "set_cookie",
-               "set_status", "flush", "finish"]:
-    setattr(WebSocketHandler, method,
-            _wrap_method(getattr(WebSocketHandler, method)))
+
+def _raise_not_supported_for_websockets(*args, **kwargs):
+    raise RuntimeError("Method not supported for Web Sockets")
 
 
 class WebSocketProtocol(object):
@@ -625,8 +623,7 @@ class WebSocketProtocol13(WebSocketProtocol):
             selected = self.handler.select_subprotocol(subprotocols)
             if selected:
                 assert selected in subprotocols
-                subprotocol_header = ("Sec-WebSocket-Protocol: %s\r\n"
-                                      % selected)
+                self.handler.set_header("Sec-WebSocket-Protocol", selected)
 
         extension_header = ''
         extensions = self._parse_extensions_header(self.request.headers)
@@ -641,22 +638,20 @@ class WebSocketProtocol13(WebSocketProtocol):
                     # Don't echo an offered client_max_window_bits
                     # parameter with no value.
                     del ext[1]['client_max_window_bits']
-                extension_header = ('Sec-WebSocket-Extensions: %s\r\n' %
-                                    httputil._encode_header(
-                                        'permessage-deflate', ext[1]))
+                self.handler.set_header("Sec-WebSocket-Extensions",
+                                        httputil._encode_header(
+                                            'permessage-deflate', ext[1]))
                 break
 
-        if self.stream.closed():
-            self._abort()
-            return
-        self.stream.write(tornado.escape.utf8(
-            "HTTP/1.1 101 Switching Protocols\r\n"
-            "Upgrade: websocket\r\n"
-            "Connection: Upgrade\r\n"
-            "Sec-WebSocket-Accept: %s\r\n"
-            "%s%s"
-            "\r\n" % (self._challenge_response(),
-                      subprotocol_header, extension_header)))
+        self.handler.clear_header("Content-Type")
+        self.handler.set_status(101)
+        self.handler.set_header("Upgrade", "websocket")
+        self.handler.set_header("Connection", "Upgrade")
+        self.handler.set_header("Sec-WebSocket-Accept", self._challenge_response())
+        self.handler.finish()
+
+        self.handler._attach_stream()
+        self.stream = self.handler.stream
 
         self.start_pinging()
         self._run_callback(self.handler.open, *self.handler.open_args,


Test patch:
diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py
index 2de22113..0875e4b9 100644
--- a/tornado/test/websocket_test.py
+++ b/tornado/test/websocket_test.py
@@ -1,5 +1,6 @@
 from __future__ import absolute_import, division, print_function
 
+import functools
 import sys
 import traceback
 
@@ -58,16 +59,36 @@ class ErrorInOnMessageHandler(TestWebSocketHandler):
 
 class HeaderHandler(TestWebSocketHandler):
     def open(self):
-        try:
-            # In a websocket context, many RequestHandler methods
-            # raise RuntimeErrors.
-            self.set_status(503)
-            raise Exception("did not get expected exception")
-        except RuntimeError:
-            pass
+        methods_to_test = [
+            functools.partial(self.write, 'This should not work'),
+            functools.partial(self.redirect, 'http://localhost/elsewhere'),
+            functools.partial(self.set_header, 'X-Test', ''),
+            functools.partial(self.set_cookie, 'Chocolate', 'Chip'),
+            functools.partial(self.set_status, 503),
+            self.flush,
+            self.finish,
+        ]
+        for method in methods_to_test:
+            try:
+                # In a websocket context, many RequestHandler methods
+                # raise RuntimeErrors.
+                method()
+                raise Exception("did not get expected exception")
+            except RuntimeError:
+                pass
         self.write_message(self.request.headers.get('X-Test', ''))
 
 
+class HeaderEchoHandler(TestWebSocketHandler):
+    def set_default_headers(self):
+        self.set_header("X-Extra-Response-Header", "Extra-Response-Value")
+
+    def prepare(self):
+        for k, v in self.request.headers.get_all():
+            if k.lower().startswith('x-test'):
+                self.set_header(k, v)
+
+
 class NonWebSocketHandler(RequestHandler):
     def get(self):
         self.write('ok')
@@ -135,6 +156,8 @@ class WebSocketTest(WebSocketBaseTestCase):
             ('/echo', EchoHandler, dict(close_future=self.close_future)),
             ('/non_ws', NonWebSocketHandler),
             ('/header', HeaderHandler, dict(close_future=self.close_future)),
+            ('/header_echo', HeaderEchoHandler,
+             dict(close_future=self.close_future)),
             ('/close_reason', CloseReasonHandler,
              dict(close_future=self.close_future)),
             ('/error_in_on_message', ErrorInOnMessageHandler,
@@ -152,6 +175,13 @@ class WebSocketTest(WebSocketBaseTestCase):
         response = self.fetch('/echo')
         self.assertEqual(response.code, 400)
 
+    def test_bad_websocket_version(self):
+        response = self.fetch('/echo',
+                              headers={'Connection': 'Upgrade',
+                                       'Upgrade': 'WebSocket',
+                                       'Sec-WebSocket-Version': '12'})
+        self.assertEqual(response.code, 426)
+
     @gen_test
     def test_websocket_gen(self):
         ws = yield self.ws_connect('/echo')
@@ -240,6 +270,18 @@ class WebSocketTest(WebSocketBaseTestCase):
         self.assertEqual(response, 'hello')
         yield self.close(ws)
 
+    @gen_test
+    def test_websocket_header_echo(self):
+        # Ensure that headers can be returned in the response.
+        # Specifically, that arbitrary headers passed through websocket_connect
+        # can be returned.
+        ws = yield websocket_connect(
+            HTTPRequest('ws://127.0.0.1:%d/header_echo' % self.get_http_port(),
+                        headers={'X-Test-Hello': 'hello'}))
+        self.assertEqual(ws.headers.get('X-Test-Hello'), 'hello')
+        self.assertEqual(ws.headers.get('X-Extra-Response-Header'), 'Extra-Response-Value')
+        yield self.close(ws)
+
     @gen_test
     def test_server_close_reason(self):
         ws = yield self.ws_connect('/close_reason')