Update Tornado

libs/tornado/websocket.py

@@ -33,7 +33,7 @@ import tornado.web
 
 from tornado.concurrent import TracebackFuture
 from tornado.escape import utf8, native_str
-from tornado import httpclient
+from tornado import httpclient, httputil
 from tornado.ioloop import IOLoop
 from tornado.iostream import StreamClosedError
 from tornado.log import gen_log, app_log
@@ -52,6 +52,10 @@ class WebSocketError(Exception):
 
 
 class WebSocketClosedError(WebSocketError):
+    """Raised by operations on a closed connection.
+
+    .. versionadded:: 3.2
+    """
     pass
 
 
@@ -163,6 +167,12 @@ class WebSocketHandler(tornado.web.RequestHandler):
         encoded as json).  If the ``binary`` argument is false, the
         message will be sent as utf8; in binary mode any byte string
         is allowed.
+
+        If the connection is already closed, raises `WebSocketClosedError`.
+
+        .. versionchanged:: 3.2
+           `WebSocketClosedError` was added (previously a closed connection
+           would raise an `AttributeError`)
         """
         if self.ws_connection is None:
             raise WebSocketClosedError()
@@ -586,7 +596,7 @@ class WebSocketProtocol13(WebSocketProtocol):
             frame += struct.pack("!BQ", 127 | mask_bit, l)
         if self.mask_outgoing:
             mask = os.urandom(4)
-            data = mask + self._apply_mask(mask, data)
+            data = mask + _websocket_mask(mask, data)
         frame += data
         self.stream.write(frame)
 
@@ -671,21 +681,8 @@ class WebSocketProtocol13(WebSocketProtocol):
         except StreamClosedError:
             self._abort()
 
-    def _apply_mask(self, mask, data):
-        mask = array.array("B", mask)
-        unmasked = array.array("B", data)
-        for i in xrange(len(data)):
-            unmasked[i] = unmasked[i] ^ mask[i % 4]
-        if hasattr(unmasked, 'tobytes'):
-            # tostring was deprecated in py32.  It hasn't been removed,
-            # but since we turn on deprecation warnings in our tests
-            # we need to use the right one.
-            return unmasked.tobytes()
-        else:
-            return unmasked.tostring()
-
     def _on_masked_frame_data(self, data):
-        self._on_frame_data(self._apply_mask(self._frame_mask, data))
+        self._on_frame_data(_websocket_mask(self._frame_mask, data))
 
     def _on_frame_data(self, data):
         if self._frame_opcode_is_control:
@@ -771,7 +768,11 @@ class WebSocketProtocol13(WebSocketProtocol):
 
 
 class WebSocketClientConnection(simple_httpclient._HTTPConnection):
-    """WebSocket client connection."""
+    """WebSocket client connection.
+
+    This class should not be instantiated directly; use the
+    `websocket_connect` function instead.
+    """
     def __init__(self, io_loop, request):
         self.connect_future = TracebackFuture()
         self.read_future = None
@@ -793,9 +794,19 @@ class WebSocketClientConnection(simple_httpclient._HTTPConnection):
             io_loop, None, request, lambda: None, self._on_http_response,
             104857600, self.resolver)
 
+    def close(self):
+        """Closes the websocket connection.
+
+        .. versionadded:: 3.2
+        """
+        if self.protocol is not None:
+            self.protocol.close()
+            self.protocol = None
+
     def _on_close(self):
         self.on_message(None)
         self.resolver.close()
+        super(WebSocketClientConnection, self)._on_close()
 
     def _on_http_response(self, response):
         if not self.connect_future.done():
@@ -859,13 +870,54 @@ def websocket_connect(url, io_loop=None, callback=None, connect_timeout=None):
 
     Takes a url and returns a Future whose result is a
     `WebSocketClientConnection`.
+
+    .. versionchanged:: 3.2
+       Also accepts ``HTTPRequest`` objects in place of urls.
     """
     if io_loop is None:
         io_loop = IOLoop.current()
-    request = httpclient.HTTPRequest(url, connect_timeout=connect_timeout)
+    if isinstance(url, httpclient.HTTPRequest):
+        assert connect_timeout is None
+        request = url
+        # Copy and convert the headers dict/object (see comments in
+        # AsyncHTTPClient.fetch)
+        request.headers = httputil.HTTPHeaders(request.headers)
+    else:
+        request = httpclient.HTTPRequest(url, connect_timeout=connect_timeout)
     request = httpclient._RequestProxy(
         request, httpclient.HTTPRequest._DEFAULTS)
     conn = WebSocketClientConnection(io_loop, request)
     if callback is not None:
         io_loop.add_future(conn.connect_future, callback)
     return conn.connect_future
+
+def _websocket_mask_python(mask, data):
+    """Websocket masking function.
+
+    `mask` is a `bytes` object of length 4; `data` is a `bytes` object of any length.
+    Returns a `bytes` object of the same length as `data` with the mask applied
+    as specified in section 5.3 of RFC 6455.
+
+    This pure-python implementation may be replaced by an optimized version when available.
+    """
+    mask = array.array("B", mask)
+    unmasked = array.array("B", data)
+    for i in xrange(len(data)):
+        unmasked[i] = unmasked[i] ^ mask[i % 4]
+    if hasattr(unmasked, 'tobytes'):
+        # tostring was deprecated in py32.  It hasn't been removed,
+        # but since we turn on deprecation warnings in our tests
+        # we need to use the right one.
+        return unmasked.tobytes()
+    else:
+        return unmasked.tostring()
+
+if os.environ.get('TORNADO_NO_EXTENSION'):
+    # This environment variable exists to make it easier to do performance comparisons;
+    # it's not guaranteed to remain supported in the future.
+    _websocket_mask = _websocket_mask_python
+else:
+    try:
+        from tornado.speedups import websocket_mask as _websocket_mask
+    except ImportError:
+        _websocket_mask = _websocket_mask_python