Ticket #4173: 4173-3.patch

twisted/web/websockets.py

@@ +1 @@
+# Copyright (c) 2011-2012 Oregon State University Open Source Lab
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+#    The above copyright notice and this permission notice shall be included
+#    in all copies or substantial portions of the Software.
+#
+#    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+#    OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+#    MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
+#    NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+#    DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+#    OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
+#    USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+"""
+The WebSockets protocol (RFC 6455), provided as a resource which wraps a
+protocol.
+"""
+
+from base64 import b64encode, b64decode
+from hashlib import sha1
+from struct import pack, unpack
+
+from twisted.protocols.policies import ProtocolWrapper, WrappingFactory
+from twisted.python import log
+from twisted.web.error import NoResource
+from twisted.web.resource import IResource
+from twisted.web.server import NOT_DONE_YET
+from zope.interface import implements
+
+class WSException(Exception):
+    """
+    Something stupid happened here.
+
+    If this class escapes txWS, then something stupid happened in multiple
+    places.
+    """
+
+# Control frame specifiers. Some versions of WS have control signals sent
+# in-band. Adorable, right?
+
+NORMAL, CLOSE, PING, PONG = range(4)
+
+opcode_types = {
+    0x0: NORMAL,
+    0x1: NORMAL,
+    0x2: NORMAL,
+    0x8: CLOSE,
+    0x9: PING,
+    0xa: PONG,
+}
+
+encoders = {
+    "base64": b64encode,
+}
+
+decoders = {
+    "base64": b64decode,
+}
+
+# Authentication for WS.
+
+def make_accept(key):
+    """
+    Create an "accept" response for a given key.
+
+    This dance is expected to somehow magically make WebSockets secure.
+    """
+
+    guid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
+
+    return sha1("%s%s" % (key, guid)).digest().encode("base64").strip()
+
+# Frame helpers.
+# Separated out to make unit testing a lot easier.
+# Frames are bonghits in newer WS versions, so helpers are appreciated.
+
+def mask(buf, key):
+    """
+    Mask or unmask a buffer of bytes with a masking key.
+
+    The key must be exactly four bytes long.
+    """
+
+    # This is super-secure, I promise~
+    key = [ord(i) for i in key]
+    buf = list(buf)
+    for i, char in enumerate(buf):
+        buf[i] = chr(ord(char) ^ key[i % 4])
+    return "".join(buf)
+
+def make_hybi07_frame(buf, opcode=0x1):
+    """
+    Make a HyBi-07 frame.
+
+    This function always creates unmasked frames, and attempts to use the
+    smallest possible lengths.
+    """
+
+    if len(buf) > 0xffff:
+        length = "\x7f%s" % pack(">Q", len(buf))
+    elif len(buf) > 0x7d:
+        length = "\x7e%s" % pack(">H", len(buf))
+    else:
+        length = chr(len(buf))
+
+    # Always make a normal packet.
+    header = chr(0x80 | opcode)
+    frame = "%s%s%s" % (header, length, buf)
+    return frame
+
+def parse_hybi07_frames(buf):
+    """
+    Parse HyBi-07 frames in a highly compliant manner.
+    """
+
+    start = 0
+    frames = []
+
+    while True:
+        # If there's not at least two bytes in the buffer, bail.
+        if len(buf) - start < 2:
+            break
+
+        # Grab the header. This single byte holds some flags nobody cares
+        # about, and an opcode which nobody cares about.
+        header = ord(buf[start])
+        if header & 0x70:
+            # At least one of the reserved flags is set. Pork chop sandwiches!
+            raise WSException("Reserved flag in HyBi-07 frame (%d)" % header)
+            frames.append(("", CLOSE))
+            return frames, buf
+
+        # Get the opcode, and translate it to a local enum which we actually
+        # care about.
+        opcode = header & 0xf
+        try:
+            opcode = opcode_types[opcode]
+        except KeyError:
+            raise WSException("Unknown opcode %d in HyBi-07 frame" % opcode)
+
+        # Get the payload length and determine whether we need to look for an
+        # extra length.
+        length = ord(buf[start + 1])
+        masked = length & 0x80
+        length &= 0x7f
+
+        # The offset we're gonna be using to walk through the frame. We use
+        # this because the offset is variable depending on the length and
+        # mask.
+        offset = 2
+
+        # Extra length fields.
+        if length == 0x7e:
+            if len(buf) - start < 4:
+                break
+
+            length = buf[start + 2:start + 4]
+            length = unpack(">H", length)[0]
+            offset += 2
+        elif length == 0x7f:
+            if len(buf) - start < 10:
+                break
+
+            # Protocol bug: The top bit of this long long *must* be cleared;
+            # that is, it is expected to be interpreted as signed. That's
+            # fucking stupid, if you don't mind me saying so, and so we're
+            # interpreting it as unsigned anyway. If you wanna send exabytes
+            # of data down the wire, then go ahead!
+            length = buf[start + 2:start + 10]
+            length = unpack(">Q", length)[0]
+            offset += 8
+
+        if masked:
+            if len(buf) - (start + offset) < 4:
+                break
+
+            key = buf[start + offset:start + offset + 4]
+            offset += 4
+
+        if len(buf) - (start + offset) < length:
+            break
+
+        data = buf[start + offset:start + offset + length]
+
+        if masked:
+            data = mask(data, key)
+
+        if opcode == CLOSE:
+            if len(data) >= 2:
+                # Gotta unpack the opcode and return usable data here.
+                data = unpack(">H", data[:2])[0], data[2:]
+            else:
+                # No reason given; use generic data.
+                data = 1000, "No reason given"
+
+        frames.append((opcode, data))
+        start += offset + length
+
+    return frames, buf[start:]
+
+class WebSocketsProtocol(ProtocolWrapper):
+    """
+    Protocol which wraps another protocol to provide a WebSockets transport
+    layer.
+    """
+
+    buf = ""
+    codec = None
+
+    def __init__(self, *args, **kwargs):
+        ProtocolWrapper.__init__(self, *args, **kwargs)
+        self.pending_frames = []
+
+    def parseFrames(self):
+        """
+        Find frames in incoming data and pass them to the underlying protocol.
+        """
+
+        try:
+            frames, self.buf = parse_hybi07_frames(self.buf)
+        except WSException, wse:
+            # Couldn't parse all the frames, something went wrong, let's bail.
+            self.close(wse.args[0])
+            return
+
+        for frame in frames:
+            opcode, data = frame
+            if opcode == NORMAL:
+                # Business as usual. Decode the frame, if we have a decoder.
+                if self.codec:
+                    data = decoders[self.codec](data)
+                # Pass the frame to the underlying protocol.
+                ProtocolWrapper.dataReceived(self, data)
+            elif opcode == CLOSE:
+                # The other side wants us to close. I wonder why?
+                reason, text = data
+                log.msg("Closing connection: %r (%d)" % (text, reason))
+
+                # Close the connection.
+                self.close()
+
+    def sendFrames(self):
+        """
+        Send all pending frames.
+        """
+
+        for frame in self.pending_frames:
+            # Encode the frame before sending it.
+            if self.codec:
+                frame = encoders[self.codec](frame)
+            packet = make_hybi07_frame(frame)
+            self.transport.write(packet)
+        self.pending_frames = []
+
+    def dataReceived(self, data):
+        self.buf += data
+
+        self.parseFrames()
+
+        # Kick any pending frames. This is needed because frames might have
+        # started piling up early; we can get write()s from our protocol above
+        # when they makeConnection() immediately, before our browser client
+        # actually sends any data. In those cases, we need to manually kick
+        # pending frames.
+        if self.pending_frames:
+            self.sendFrames()
+
+    def write(self, data):
+        """
+        Write to the transport.
+
+        This method will only be called by the underlying protocol.
+        """
+
+        self.pending_frames.append(data)
+        self.sendFrames()
+
+    def writeSequence(self, data):
+        """
+        Write a sequence of data to the transport.
+
+        This method will only be called by the underlying protocol.
+        """
+
+        self.pending_frames.extend(data)
+        self.sendFrames()
+
+    def close(self, reason=""):
+        """
+        Close the connection.
+
+        This includes telling the other side we're closing the connection.
+
+        If the other side didn't signal that the connection is being closed,
+        then we might not see their last message, but since their last message
+        should, according to the spec, be a simple acknowledgement, it
+        shouldn't be a problem.
+        """
+
+        # Send a closing frame. It's only polite. (And might keep the browser
+        # from hanging.)
+        frame = make_hybi07_frame(reason, opcode=0x8)
+        self.transport.write(frame)
+
+        self.loseConnection()
+
+class WebSocketsFactory(WrappingFactory):
+    """
+    Factory which wraps another factory to provide WebSockets frames for all
+    of its protocols.
+
+    This factory does not provide the HTTP headers required to perform a
+    WebSockets handshake; see C{WebSocketsResource}.
+    """
+
+    protocol = WebSocketsProtocol
+
+class WebSocketsResource(object):
+
+    implements(IResource)
+
+    isLeaf = True
+
+    def __init__(self, factory):
+        self._factory = WebSocketsFactory(factory)
+
+    def getChildWithDefault(self, name, request):
+        return NoResource("No such child resource.")
+
+    def putChild(self, path, child):
+        pass
+
+    def render(self, request):
+        """
+        Render a request.
+
+        We're not actually rendering a request. We are secretly going to
+        handle a WebSockets connection instead.
+        """
+
+        # If we fail at all, we're gonna fail with 400 and no response.
+        # You might want to pop open the RFC and read along.
+        failed = False
+
+        if request.method != "GET":
+            # 4.2.1.1 GET is required.
+            failed = True
+
+        upgrade = request.getHeader("Upgrade")
+        if upgrade is None or "websocket" not in upgrade.lower():
+            # 4.2.1.3 Upgrade: WebSocket is required.
+            failed = True
+
+        connection = request.getHeader("Connection")
+        if connection is None or "upgrade" not in connection.lower():
+            # 4.2.1.4 Connection: Upgrade is required.
+            failed = True
+
+        key = request.getHeader("Sec-WebSocket-Key")
+        if key is None:
+            # 4.2.1.5 The challenge key is required.
+            failed = True
+
+        version = request.getHeader("Sec-WebSocket-Version")
+        if version is None or version != "13":
+            # 4.2.1.6 Only version 13 works.
+            failed = True
+            # 4.4 Forward-compatible version checking.
+            request.setHeader("Sec-WebSocket-Version", "13")
+
+        # Stash host and origin for those browsers that care about it.
+        host = request.getHeader("Host")
+        origin = request.getHeader("Origin")
+
+        # Check whether a codec is needed. WS calls this a "protocol" for
+        # reasons I cannot fathom.
+        protocol = request.getHeader("Sec-WebSocket-Protocol")
+
+        if protocol:
+            if protocol not in encoders or protocol not in decoders:
+                log.msg("Protocol %s is not implemented" % protocol)
+                failed = True
+
+        if failed:
+            request.setResponseCode(400)
+            return ""
+
+        # We are going to finish this handshake. We will return a valid status
+        # code.
+        # 4.2.2.5.1 101 Switching Protocols
+        request.setResponseCode(101)
+        # 4.2.2.5.2 Upgrade: websocket
+        request.setHeader("Upgrade", "WebSocket")
+        # 4.2.2.5.3 Connection: Upgrade
+        request.setHeader("Connection", "Upgrade")
+        # 4.2.2.5.4 Response to the key challenge
+        request.setHeader("Sec-WebSocket-Accept", make_accept(key))
+
+        # Provoke request into flushing headers and finishing the handshake.
+        request.write("")
+
+        # And now take matters into our own hands. We shall manage the
+        # transport's lifecycle.
+        transport, request.transport = request.transport, None
+
+        # Connect the transport to our factory, and make things go. We need to
+        # do some stupid stuff here; see #3204, which could fix it.
+        protocol = self._factory.buildProtocol(transport.getPeer())
+        transport.protocol = protocol
+        protocol.makeConnection(transport)
+
+        return NOT_DONE_YET

twisted/web/test/test_websockets.py

@@ +1 @@
+from twisted.trial import unittest
+
+from twisted.web.websockets import (make_accept, mask, CLOSE, NORMAL, PING,
+    PONG, parse_hybi07_frames)
+
+class TestKeys(unittest.TestCase):
+
+    def test_make_accept_rfc(self):
+        """
+        Test ``make_accept()`` using the keys listed in the RFC for HyBi-07
+        through HyBi-10.
+        """
+
+        key = "dGhlIHNhbXBsZSBub25jZQ=="
+
+        self.assertEqual(make_accept(key), "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=")
+
+    def test_make_accept_wikipedia(self):
+        """
+        Test ``make_accept()`` using the keys listed on Wikipedia.
+        """
+
+        key = "x3JJHMbDL1EzLkh9GBhXDw=="
+
+        self.assertEqual(make_accept(key), "HSmrc0sMlYUkAGmm5OPpG2HaGWk=")
+
+class TestHyBi07Helpers(unittest.TestCase):
+    """
+    HyBi-07 is best understood as a large family of helper functions which
+    work together, somewhat dysfunctionally, to produce a mediocre
+    Thanksgiving every other year.
+    """
+
+    def test_mask_noop(self):
+        key = "\x00\x00\x00\x00"
+        self.assertEqual(mask("Test", key), "Test")
+
+    def test_mask_noop_long(self):
+        key = "\x00\x00\x00\x00"
+        self.assertEqual(mask("LongTest", key), "LongTest")
+
+    def test_parse_hybi07_unmasked_text(self):
+        """
+        From HyBi-10, 4.7.
+        """
+
+        frame = "\x81\x05Hello"
+        frames, buf = parse_hybi07_frames(frame)
+        self.assertEqual(len(frames), 1)
+        self.assertEqual(frames[0], (NORMAL, "Hello"))
+        self.assertEqual(buf, "")
+
+    def test_parse_hybi07_masked_text(self):
+        """
+        From HyBi-10, 4.7.
+        """
+
+        frame = "\x81\x857\xfa!=\x7f\x9fMQX"
+        frames, buf = parse_hybi07_frames(frame)
+        self.assertEqual(len(frames), 1)
+        self.assertEqual(frames[0], (NORMAL, "Hello"))
+        self.assertEqual(buf, "")
+
+    def test_parse_hybi07_unmasked_text_fragments(self):
+        """
+        We don't care about fragments. We are totally unfazed.
+
+        From HyBi-10, 4.7.
+        """
+
+        frame = "\x01\x03Hel\x80\x02lo"
+        frames, buf = parse_hybi07_frames(frame)
+        self.assertEqual(len(frames), 2)
+        self.assertEqual(frames[0], (NORMAL, "Hel"))
+        self.assertEqual(frames[1], (NORMAL, "lo"))
+        self.assertEqual(buf, "")
+
+    def test_parse_hybi07_ping(self):
+        """
+        From HyBi-10, 4.7.
+        """
+
+        frame = "\x89\x05Hello"
+        frames, buf = parse_hybi07_frames(frame)
+        self.assertEqual(len(frames), 1)
+        self.assertEqual(frames[0], (PING, "Hello"))
+        self.assertEqual(buf, "")
+
+    def test_parse_hybi07_pong(self):
+        """
+        From HyBi-10, 4.7.
+        """
+
+        frame = "\x8a\x05Hello"
+        frames, buf = parse_hybi07_frames(frame)
+        self.assertEqual(len(frames), 1)
+        self.assertEqual(frames[0], (PONG, "Hello"))
+        self.assertEqual(buf, "")
+
+    def test_parse_hybi07_close_empty(self):
+        """
+        A HyBi-07 close packet may have no body. In that case, it should use
+        the generic error code 1000, and have no reason.
+        """
+
+        frame = "\x88\x00"
+        frames, buf = parse_hybi07_frames(frame)
+        self.assertEqual(len(frames), 1)
+        self.assertEqual(frames[0], (CLOSE, (1000, "No reason given")))
+        self.assertEqual(buf, "")
+
+    def test_parse_hybi07_close_reason(self):
+        """
+        A HyBi-07 close packet must have its first two bytes be a numeric
+        error code, and may optionally include trailing text explaining why
+        the connection was closed.
+        """
+
+        frame = "\x88\x0b\x03\xe8No reason"
+        frames, buf = parse_hybi07_frames(frame)
+        self.assertEqual(len(frames), 1)
+        self.assertEqual(frames[0], (CLOSE, (1000, "No reason")))
+        self.assertEqual(buf, "")
+
+    def test_parse_hybi07_partial_no_length(self):
+        frame = "\x81"
+        frames, buf = parse_hybi07_frames(frame)
+        self.assertFalse(frames)
+        self.assertEqual(buf, "\x81")
+
+    def test_parse_hybi07_partial_truncated_length_int(self):
+        frame = "\x81\xfe"
+        frames, buf = parse_hybi07_frames(frame)
+        self.assertFalse(frames)
+        self.assertEqual(buf, "\x81\xfe")
+
+    def test_parse_hybi07_partial_truncated_length_double(self):
+        frame = "\x81\xff"
+        frames, buf = parse_hybi07_frames(frame)
+        self.assertFalse(frames)
+        self.assertEqual(buf, "\x81\xff")
+
+    def test_parse_hybi07_partial_no_data(self):
+        frame = "\x81\x05"
+        frames, buf = parse_hybi07_frames(frame)
+        self.assertFalse(frames)
+        self.assertEqual(buf, "\x81\x05")
+
+    def test_parse_hybi07_partial_truncated_data(self):
+        frame = "\x81\x05Hel"
+        frames, buf = parse_hybi07_frames(frame)
+        self.assertFalse(frames)
+        self.assertEqual(buf, "\x81\x05Hel")