Added support for subprotocols.

diff --git a/tornado/websocket.py b/tornado/websocket.py
index 3a17f93e1058efc1de6a0cf02f80a2d9c3ec50cf..a69a5e8a7ef076bdb34919b8ca1ba7368e58879c 100644 (file)
--- a/tornado/websocket.py
+++ b/tornado/websocket.py
@@ -6,8 +6,8 @@ communication between the browser and server.
 .. warning::
 
    The WebSocket protocol is still in development.  This module currently
-   implements the "hixie-76" and "hybi-10" versions of the protocol.  
-   See this `browser compatibility table 
+   implements the "hixie-76" and "hybi-10" versions of the protocol.
+   See this `browser compatibility table
    <http://en.wikipedia.org/wiki/WebSockets#Browser_support>`_ on Wikipedia.
 """
 # Author: Jacob Kristhammar, 2010
@@ -83,13 +83,13 @@ class WebSocketHandler(tornado.web.RequestHandler):
             self.request.headers.get("Sec-WebSocket-Version") == "7"):
             self.ws_connection = WebSocketProtocol8(self)
             self.ws_connection.accept_connection()
-            
+
         elif self.request.headers.get("Sec-WebSocket-Version"):
             self.stream.write(tornado.escape.utf8(
                 "HTTP/1.1 426 Upgrade Required\r\n"
                 "Sec-WebSocket-Version: 8\r\n\r\n"))
             self.stream.close()
-            
+
         else:
             self.ws_connection = WebSocketProtocol76(self)
             self.ws_connection.accept_connection()
@@ -98,6 +98,10 @@ class WebSocketHandler(tornado.web.RequestHandler):
         """Sends the given message to the client of this Web Socket."""
         self.ws_connection.write_message(message)
 
+    def validate_subprotocol(self, subprotocols):
+        """Invoked when a new WebSocket requests specific subprotocols."""
+        return None
+
     def open(self, *args, **kwargs):
         """Invoked when a new WebSocket is opened."""
         pass
@@ -199,6 +203,18 @@ class WebSocketProtocol76(WebSocketProtocol):
             logging.debug("Malformed WebSocket request received")
             self._abort()
             return
+
+        subprotocols = self.request.headers.get("Sec-WebSocket-Protocol", None)
+        if subprotocols:
+            subprotocol = self.handler.validate_subprotocol(subprotocols)
+            if not subprotocol:
+                logging.debug("Subprotocol rejected by handler.")
+                self._abort()
+                return
+            subprotocol = "Sec-WebSocket-Protocol: %s\r\n" % subprotocol
+        else:
+            subprotocol = ''
+
         scheme = "wss" if self.request.protocol == "https" else "ws"
         # Write the initial headers before attempting to read the challenge.
         # This is necessary when using proxies (such as HAProxy), which
@@ -210,12 +226,15 @@ class WebSocketProtocol76(WebSocketProtocol):
             "Connection: Upgrade\r\n"
             "Server: TornadoServer/%(version)s\r\n"
             "Sec-WebSocket-Origin: %(origin)s\r\n"
-            "Sec-WebSocket-Location: %(scheme)s://%(host)s%(uri)s\r\n\r\n" % (dict(
+            "Sec-WebSocket-Location: %(scheme)s://%(host)s%(uri)s\r\n"
+            "%(subprotocol)s"
+            "\r\n" % (dict(
                     version=tornado.version,
                     origin=self.request.headers["Origin"],
                     scheme=scheme,
                     host=self.request.host,
-                    uri=self.request.uri))))
+                    uri=self.request.uri,
+                    subprotocol=subprotocol))))
         self.stream.read_bytes(8, self._handle_challenge)
 
     def challenge_response(self, challenge):
@@ -350,7 +369,7 @@ class WebSocketProtocol8(WebSocketProtocol):
             logging.debug("Malformed WebSocket request received")
             self._abort()
             return
-    
+
     def _handle_websocket_headers(self):
         """Verifies all invariant- and required headers
 
@@ -374,11 +393,24 @@ class WebSocketProtocol8(WebSocketProtocol):
         return tornado.escape.native_str(base64.b64encode(sha1.digest()))
 
     def _accept_connection(self):
+        subprotocols = self.request.headers.get("Sec-WebSocket-Protocol", None)
+        if subprotocols:
+            subprotocol = self.handler.validate_subprotocol(subprotocols)
+            if not subprotocol:
+                logging.debug("Subprotocol rejected by handler.")
+                self._abort()
+                return
+            subprotocol = "Sec-WebSocket-Protocol: %s\r\n" % subprotocol
+        else:
+            subprotocol = ''
+
         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\r\n" % self._challenge_response()))
+            "Sec-WebSocket-Accept: %s\r\n"
+            "%s"
+            "\r\n" % (self._challenge_response(), subprotocol)))
 
         self.async_callback(self.handler.open)(*self.handler.open_args, **self.handler.open_kwargs)
         self._receive_frame()
@@ -434,7 +466,7 @@ class WebSocketProtocol8(WebSocketProtocol):
     def _on_frame_length_16(self, data):
         self._frame_length = struct.unpack("!H", data)[0];
         self.stream.read_bytes(4, self._on_masking_key);
-        
+
     def _on_frame_length_64(self, data):
         self._frame_length = struct.unpack("!Q", data)[0];
         self.stream.read_bytes(4, self._on_masking_key);
@@ -466,11 +498,11 @@ class WebSocketProtocol8(WebSocketProtocol):
 
         if not self.client_terminated:
             self._receive_frame()
-        
+
 
     def _handle_message(self, opcode, data):
         if self.client_terminated: return
-        
+
         if opcode == 0x1:
             # UTF-8 data
             self.async_callback(self.handler.on_message)(data.decode("utf-8", "replace"))
@@ -491,10 +523,9 @@ class WebSocketProtocol8(WebSocketProtocol):
             pass
         else:
             self._abort()
-        
+
     def close(self):
         """Closes the WebSocket connection."""
         self._write_frame(True, 0x8, b(""))
         self._started_closing_handshake = True
         self._waiting = tornado.ioloop.IOLoop.instance().add_timeout(time.time() + 5, self._abort)
-