Ticket #4173: 4173-5.patch

File 4173-5.patch, 21.2 KB (added by MostAwesomeDude, 5 years ago)

Yet another version of #4173.

  • twisted/web/websockets.py

     
     1# Copyright (c) 2011-2012 Oregon State University Open Source Lab
     2#
     3# Permission is hereby granted, free of charge, to any person obtaining a copy
     4# of this software and associated documentation files (the "Software"), to
     5# deal in the Software without restriction, including without limitation the
     6# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
     7# sell copies of the Software, and to permit persons to whom the Software is
     8# furnished to do so, subject to the following conditions:
     9#
     10#    The above copyright notice and this permission notice shall be included
     11#    in all copies or substantial portions of the Software.
     12#
     13#    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
     14#    OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
     15#    MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
     16#    NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
     17#    DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
     18#    OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
     19#    USE OR OTHER DEALINGS IN THE SOFTWARE.
     20
     21"""
     22The WebSockets protocol (RFC 6455), provided as a resource which wraps a
     23factory.
     24"""
     25
     26from base64 import b64encode, b64decode
     27from hashlib import sha1
     28from struct import pack, unpack
     29
     30from twisted.protocols.policies import ProtocolWrapper, WrappingFactory
     31from twisted.python import log
     32from twisted.web.error import NoResource
     33from twisted.web.resource import IResource
     34from twisted.web.server import NOT_DONE_YET
     35from zope.interface import implements
     36
     37class WSException(Exception):
     38    """
     39    Something stupid happened here.
     40
     41    If this class escapes txWS, then something stupid happened in multiple
     42    places.
     43    """
     44
     45# Control frame specifiers. Some versions of WS have control signals sent
     46# in-band. Adorable, right?
     47
     48NORMAL, CLOSE, PING, PONG = range(4)
     49
     50opcode_types = {
     51    0x0: NORMAL,
     52    0x1: NORMAL,
     53    0x2: NORMAL,
     54    0x8: CLOSE,
     55    0x9: PING,
     56    0xa: PONG,
     57}
     58
     59opcode_for_type = {
     60    NORMAL: 0x1,
     61    CLOSE: 0x8,
     62    PING: 0x9,
     63    PONG: 0xa,
     64}
     65
     66encoders = {
     67    "base64": b64encode,
     68}
     69
     70decoders = {
     71    "base64": b64decode,
     72}
     73
     74# Authentication for WS.
     75
     76def make_accept(key):
     77    """
     78    Create an "accept" response for a given key.
     79
     80    This dance is expected to somehow magically make WebSockets secure.
     81    """
     82
     83    guid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
     84
     85    return sha1("%s%s" % (key, guid)).digest().encode("base64").strip()
     86
     87# Frame helpers.
     88# Separated out to make unit testing a lot easier.
     89# Frames are bonghits in newer WS versions, so helpers are appreciated.
     90
     91def mask(buf, key):
     92    """
     93    Mask or unmask a buffer of bytes with a masking key.
     94
     95    The key must be exactly four bytes long.
     96    """
     97
     98    # This is super-secure, I promise~
     99    key = [ord(i) for i in key]
     100    buf = list(buf)
     101    for i, char in enumerate(buf):
     102        buf[i] = chr(ord(char) ^ key[i % 4])
     103    return "".join(buf)
     104
     105def make_hybi07_frame(buf, opcode=NORMAL):
     106    """
     107    Make a HyBi-07 frame.
     108
     109    This function always creates unmasked frames, and attempts to use the
     110    smallest possible lengths.
     111    """
     112
     113    if len(buf) > 0xffff:
     114        length = "\x7f%s" % pack(">Q", len(buf))
     115    elif len(buf) > 0x7d:
     116        length = "\x7e%s" % pack(">H", len(buf))
     117    else:
     118        length = chr(len(buf))
     119
     120    # Always make a normal packet.
     121    header = chr(0x80 | opcode_for_type[opcode])
     122    frame = "%s%s%s" % (header, length, buf)
     123    return frame
     124
     125def parse_hybi07_frames(buf):
     126    """
     127    Parse HyBi-07 frames in a highly compliant manner.
     128    """
     129
     130    start = 0
     131    frames = []
     132
     133    while True:
     134        # If there's not at least two bytes in the buffer, bail.
     135        if len(buf) - start < 2:
     136            break
     137
     138        # Grab the header. This single byte holds some flags nobody cares
     139        # about, and an opcode which nobody cares about.
     140        header = ord(buf[start])
     141        if header & 0x70:
     142            # At least one of the reserved flags is set. Pork chop sandwiches!
     143            raise WSException("Reserved flag in HyBi-07 frame (%d)" % header)
     144            frames.append(("", CLOSE))
     145            return frames, buf
     146
     147        # Get the opcode, and translate it to a local enum which we actually
     148        # care about.
     149        opcode = header & 0xf
     150        try:
     151            opcode = opcode_types[opcode]
     152        except KeyError:
     153            raise WSException("Unknown opcode %d in HyBi-07 frame" % opcode)
     154
     155        # Get the payload length and determine whether we need to look for an
     156        # extra length.
     157        length = ord(buf[start + 1])
     158        masked = length & 0x80
     159        length &= 0x7f
     160
     161        # The offset we're gonna be using to walk through the frame. We use
     162        # this because the offset is variable depending on the length and
     163        # mask.
     164        offset = 2
     165
     166        # Extra length fields.
     167        if length == 0x7e:
     168            if len(buf) - start < 4:
     169                break
     170
     171            length = buf[start + 2:start + 4]
     172            length = unpack(">H", length)[0]
     173            offset += 2
     174        elif length == 0x7f:
     175            if len(buf) - start < 10:
     176                break
     177
     178            # Protocol bug: The top bit of this long long *must* be cleared;
     179            # that is, it is expected to be interpreted as signed. That's
     180            # fucking stupid, if you don't mind me saying so, and so we're
     181            # interpreting it as unsigned anyway. If you wanna send exabytes
     182            # of data down the wire, then go ahead!
     183            length = buf[start + 2:start + 10]
     184            length = unpack(">Q", length)[0]
     185            offset += 8
     186
     187        if masked:
     188            if len(buf) - (start + offset) < 4:
     189                break
     190
     191            key = buf[start + offset:start + offset + 4]
     192            offset += 4
     193
     194        if len(buf) - (start + offset) < length:
     195            break
     196
     197        data = buf[start + offset:start + offset + length]
     198
     199        if masked:
     200            data = mask(data, key)
     201
     202        if opcode == CLOSE:
     203            if len(data) >= 2:
     204                # Gotta unpack the opcode and return usable data here.
     205                data = unpack(">H", data[:2])[0], data[2:]
     206            else:
     207                # No reason given; use generic data.
     208                data = 1000, "No reason given"
     209
     210        frames.append((opcode, data))
     211        start += offset + length
     212
     213    return frames, buf[start:]
     214
     215class WebSocketsProtocol(ProtocolWrapper):
     216    """
     217    Protocol which wraps another protocol to provide a WebSockets transport
     218    layer.
     219    """
     220
     221    buf = ""
     222    codec = None
     223
     224    def __init__(self, *args, **kwargs):
     225        ProtocolWrapper.__init__(self, *args, **kwargs)
     226        self.pending_frames = []
     227
     228    def connectionMade(self):
     229        ProtocolWrapper.connectionMade(self)
     230        log.msg("Opening connection with %s" % self.transport.getPeer())
     231
     232    def parseFrames(self):
     233        """
     234        Find frames in incoming data and pass them to the underlying protocol.
     235        """
     236
     237        try:
     238            frames, self.buf = parse_hybi07_frames(self.buf)
     239        except WSException, wse:
     240            # Couldn't parse all the frames, something went wrong, let's bail.
     241            log.err()
     242            self.loseConnection()
     243            return
     244
     245        for frame in frames:
     246            opcode, data = frame
     247            if opcode == NORMAL:
     248                # Business as usual. Decode the frame, if we have a decoder.
     249                if self.codec:
     250                    data = decoders[self.codec](data)
     251                # Pass the frame to the underlying protocol.
     252                ProtocolWrapper.dataReceived(self, data)
     253            elif opcode == CLOSE:
     254                # The other side wants us to close. I wonder why?
     255                reason, text = data
     256                log.msg("Closing connection: %r (%d)" % (text, reason))
     257
     258                # Close the connection.
     259                self.loseConnection()
     260                return
     261            elif opcode == PING:
     262                # 5.5.2 PINGs must be responded to with PONGs.
     263                # 5.5.3 PONGs must contain the data that was sent with the
     264                # provoking PING.
     265                self.transport.write(make_hybi07_packet(data, opcode=PONG))
     266
     267    def sendFrames(self):
     268        """
     269        Send all pending frames.
     270        """
     271
     272        for frame in self.pending_frames:
     273            # Encode the frame before sending it.
     274            if self.codec:
     275                frame = encoders[self.codec](frame)
     276            packet = make_hybi07_frame(frame)
     277            self.transport.write(packet)
     278        self.pending_frames = []
     279
     280    def dataReceived(self, data):
     281        self.buf += data
     282
     283        self.parseFrames()
     284
     285        # Kick any pending frames. This is needed because frames might have
     286        # started piling up early; we can get write()s from our protocol above
     287        # when they makeConnection() immediately, before our browser client
     288        # actually sends any data. In those cases, we need to manually kick
     289        # pending frames.
     290        if self.pending_frames:
     291            self.sendFrames()
     292
     293    def write(self, data):
     294        """
     295        Write to the transport.
     296
     297        This method will only be called by the underlying protocol.
     298        """
     299
     300        self.pending_frames.append(data)
     301        self.sendFrames()
     302
     303    def writeSequence(self, data):
     304        """
     305        Write a sequence of data to the transport.
     306
     307        This method will only be called by the underlying protocol.
     308        """
     309
     310        self.pending_frames.extend(data)
     311        self.sendFrames()
     312
     313    def loseConnection(self):
     314        """
     315        Close the connection.
     316
     317        This includes telling the other side we're closing the connection.
     318
     319        If the other side didn't signal that the connection is being closed,
     320        then we might not see their last message, but since their last message
     321        should, according to the spec, be a simple acknowledgement, it
     322        shouldn't be a problem.
     323        """
     324
     325        # Send a closing frame. It's only polite. (And might keep the browser
     326        # from hanging.)
     327        if not self.disconnecting:
     328            frame = make_hybi07_frame("", opcode=CLOSE)
     329            self.transport.write(frame)
     330
     331            ProtocolWrapper.loseConnection(self)
     332
     333class WebSocketsFactory(WrappingFactory):
     334    """
     335    Factory which wraps another factory to provide WebSockets frames for all
     336    of its protocols.
     337
     338    This factory does not provide the HTTP headers required to perform a
     339    WebSockets handshake; see C{WebSocketsResource}.
     340    """
     341
     342    protocol = WebSocketsProtocol
     343
     344class WebSocketsResource(object):
     345    """
     346    A resource for serving a protocol through WebSockets.
     347
     348    This class wraps a factory and connects it to WebSockets clients. Each
     349    connecting client will be connected to a new protocol of the factory.
     350
     351    Due to unresolved questions of logistics, this resource cannot have
     352    children.
     353    """
     354
     355    implements(IResource)
     356
     357    isLeaf = True
     358
     359    def __init__(self, factory):
     360        self._factory = WebSocketsFactory(factory)
     361
     362    def getChildWithDefault(self, name, request):
     363        return NoResource("No such child resource.")
     364
     365    def putChild(self, path, child):
     366        pass
     367
     368    def render(self, request):
     369        """
     370        Render a request.
     371
     372        We're not actually rendering a request. We are secretly going to
     373        handle a WebSockets connection instead.
     374        """
     375
     376        # If we fail at all, we're gonna fail with 400 and no response.
     377        # You might want to pop open the RFC and read along.
     378        failed = False
     379
     380        if request.method != "GET":
     381            # 4.2.1.1 GET is required.
     382            failed = True
     383
     384        upgrade = request.getHeader("Upgrade")
     385        if upgrade is None or "websocket" not in upgrade.lower():
     386            # 4.2.1.3 Upgrade: WebSocket is required.
     387            failed = True
     388
     389        connection = request.getHeader("Connection")
     390        if connection is None or "upgrade" not in connection.lower():
     391            # 4.2.1.4 Connection: Upgrade is required.
     392            failed = True
     393
     394        key = request.getHeader("Sec-WebSocket-Key")
     395        if key is None:
     396            # 4.2.1.5 The challenge key is required.
     397            failed = True
     398
     399        version = request.getHeader("Sec-WebSocket-Version")
     400        if version != "13":
     401            # 4.2.1.6 Only version 13 works.
     402            failed = True
     403            # 4.4 Forward-compatible version checking.
     404            request.setHeader("Sec-WebSocket-Version", "13")
     405
     406        # Check whether a codec is needed. WS calls this a "protocol" for
     407        # reasons I cannot fathom. The specification permits multiple,
     408        # comma-separated codecs to be listed, but this functionality isn't
     409        # used in the wild. (If that ever changes, we'll have already added
     410        # the requisite codecs here anyway.) The main reason why we check for
     411        # codecs at all is that older draft versions of WebSockets used base64
     412        # encoding to work around the inability to send \x00 bytes, and those
     413        # runtimes would request base64 encoding during the handshake. We
     414        # stand prepared to engage that behavior should any of those runtimes
     415        # start supporting RFC WebSockets.
     416        #
     417        # We probably should remove this altogether, but I'd rather leave it
     418        # because it will prove to be a useful reference if/when extensions
     419        # are added, and it *does* work as advertised.
     420        codec = request.getHeader("Sec-WebSocket-Protocol")
     421
     422        if codec:
     423            if codec not in encoders or codec not in decoders:
     424                log.msg("Codec %s is not implemented" % codec)
     425                failed = True
     426
     427        if failed:
     428            request.setResponseCode(400)
     429            return ""
     430
     431        # We are going to finish this handshake. We will return a valid status
     432        # code.
     433        # 4.2.2.5.1 101 Switching Protocols
     434        request.setResponseCode(101)
     435        # 4.2.2.5.2 Upgrade: websocket
     436        request.setHeader("Upgrade", "WebSocket")
     437        # 4.2.2.5.3 Connection: Upgrade
     438        request.setHeader("Connection", "Upgrade")
     439        # 4.2.2.5.4 Response to the key challenge
     440        request.setHeader("Sec-WebSocket-Accept", make_accept(key))
     441        # 4.2.2.5.5 Optional codec declaration
     442        if codec:
     443            request.setHeader("Sec-WebSocket-Protocol", codec)
     444
     445        # Create the protocol. This could fail, in which case we deliver an
     446        # error status. Status 502 was decreed by glyph; blame him.
     447        protocol = self._factory.buildProtocol(request.transport.getPeer())
     448        if not protocol:
     449            request.setResponseCode(502)
     450            return ""
     451        if codec:
     452            protocol.codec = codec
     453
     454        # Provoke request into flushing headers and finishing the handshake.
     455        request.write("")
     456
     457        # And now take matters into our own hands. We shall manage the
     458        # transport's lifecycle.
     459        transport, request.transport = request.transport, None
     460
     461        # Connect the transport to our factory, and make things go. We need to
     462        # do some stupid stuff here; see #3204, which could fix it.
     463        transport.protocol = protocol
     464        protocol.makeConnection(transport)
     465
     466        return NOT_DONE_YET
     467
     468__all__ = ("WebSocketsResource",)
  • twisted/web/test/test_websockets.py

     
     1from twisted.trial import unittest
     2
     3from twisted.web.websockets import (make_accept, mask, CLOSE, NORMAL, PING,
     4    PONG, parse_hybi07_frames)
     5
     6class TestKeys(unittest.TestCase):
     7
     8    def test_make_accept_rfc(self):
     9        """
     10        Test ``make_accept()`` using the keys listed in the RFC for HyBi-07
     11        through HyBi-10.
     12        """
     13
     14        key = "dGhlIHNhbXBsZSBub25jZQ=="
     15
     16        self.assertEqual(make_accept(key), "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=")
     17
     18    def test_make_accept_wikipedia(self):
     19        """
     20        Test ``make_accept()`` using the keys listed on Wikipedia.
     21        """
     22
     23        key = "x3JJHMbDL1EzLkh9GBhXDw=="
     24
     25        self.assertEqual(make_accept(key), "HSmrc0sMlYUkAGmm5OPpG2HaGWk=")
     26
     27class TestHyBi07Helpers(unittest.TestCase):
     28    """
     29    HyBi-07 is best understood as a large family of helper functions which
     30    work together, somewhat dysfunctionally, to produce a mediocre
     31    Thanksgiving every other year.
     32    """
     33
     34    def test_mask_noop(self):
     35        key = "\x00\x00\x00\x00"
     36        self.assertEqual(mask("Test", key), "Test")
     37
     38    def test_mask_noop_long(self):
     39        key = "\x00\x00\x00\x00"
     40        self.assertEqual(mask("LongTest", key), "LongTest")
     41
     42    def test_mask_noop_odd(self):
     43        """
     44        Masking works even when the data to be masked isn't a multiple of four
     45        in length.
     46        """
     47
     48        key = "\x00\x00\x00\x00"
     49        self.assertEqual(mask("LongestTest", key), "LongestTest")
     50
     51    def test_mask_hello(self):
     52        """
     53        From RFC 6455, 5.7.
     54        """
     55
     56        key = "\x37\xfa\x21\x3d"
     57        self.assertEqual(mask("Hello", key), "\x7f\x9f\x4d\x51\x58")
     58
     59    def test_parse_hybi07_unmasked_text(self):
     60        """
     61        From HyBi-10, 4.7.
     62        """
     63
     64        frame = "\x81\x05Hello"
     65        frames, buf = parse_hybi07_frames(frame)
     66        self.assertEqual(len(frames), 1)
     67        self.assertEqual(frames[0], (NORMAL, "Hello"))
     68        self.assertEqual(buf, "")
     69
     70    def test_parse_hybi07_masked_text(self):
     71        """
     72        From HyBi-10, 4.7.
     73        """
     74
     75        frame = "\x81\x857\xfa!=\x7f\x9fMQX"
     76        frames, buf = parse_hybi07_frames(frame)
     77        self.assertEqual(len(frames), 1)
     78        self.assertEqual(frames[0], (NORMAL, "Hello"))
     79        self.assertEqual(buf, "")
     80
     81    def test_parse_hybi07_unmasked_text_fragments(self):
     82        """
     83        We don't care about fragments. We are totally unfazed.
     84
     85        From HyBi-10, 4.7.
     86        """
     87
     88        frame = "\x01\x03Hel\x80\x02lo"
     89        frames, buf = parse_hybi07_frames(frame)
     90        self.assertEqual(len(frames), 2)
     91        self.assertEqual(frames[0], (NORMAL, "Hel"))
     92        self.assertEqual(frames[1], (NORMAL, "lo"))
     93        self.assertEqual(buf, "")
     94
     95    def test_parse_hybi07_ping(self):
     96        """
     97        From HyBi-10, 4.7.
     98        """
     99
     100        frame = "\x89\x05Hello"
     101        frames, buf = parse_hybi07_frames(frame)
     102        self.assertEqual(len(frames), 1)
     103        self.assertEqual(frames[0], (PING, "Hello"))
     104        self.assertEqual(buf, "")
     105
     106    def test_parse_hybi07_pong(self):
     107        """
     108        From HyBi-10, 4.7.
     109        """
     110
     111        frame = "\x8a\x05Hello"
     112        frames, buf = parse_hybi07_frames(frame)
     113        self.assertEqual(len(frames), 1)
     114        self.assertEqual(frames[0], (PONG, "Hello"))
     115        self.assertEqual(buf, "")
     116
     117    def test_parse_hybi07_close_empty(self):
     118        """
     119        A HyBi-07 close packet may have no body. In that case, it should use
     120        the generic error code 1000, and have no reason.
     121        """
     122
     123        frame = "\x88\x00"
     124        frames, buf = parse_hybi07_frames(frame)
     125        self.assertEqual(len(frames), 1)
     126        self.assertEqual(frames[0], (CLOSE, (1000, "No reason given")))
     127        self.assertEqual(buf, "")
     128
     129    def test_parse_hybi07_close_reason(self):
     130        """
     131        A HyBi-07 close packet must have its first two bytes be a numeric
     132        error code, and may optionally include trailing text explaining why
     133        the connection was closed.
     134        """
     135
     136        frame = "\x88\x0b\x03\xe8No reason"
     137        frames, buf = parse_hybi07_frames(frame)
     138        self.assertEqual(len(frames), 1)
     139        self.assertEqual(frames[0], (CLOSE, (1000, "No reason")))
     140        self.assertEqual(buf, "")
     141
     142    def test_parse_hybi07_partial_no_length(self):
     143        frame = "\x81"
     144        frames, buf = parse_hybi07_frames(frame)
     145        self.assertFalse(frames)
     146        self.assertEqual(buf, "\x81")
     147
     148    def test_parse_hybi07_partial_truncated_length_int(self):
     149        frame = "\x81\xfe"
     150        frames, buf = parse_hybi07_frames(frame)
     151        self.assertFalse(frames)
     152        self.assertEqual(buf, "\x81\xfe")
     153
     154    def test_parse_hybi07_partial_truncated_length_double(self):
     155        frame = "\x81\xff"
     156        frames, buf = parse_hybi07_frames(frame)
     157        self.assertFalse(frames)
     158        self.assertEqual(buf, "\x81\xff")
     159
     160    def test_parse_hybi07_partial_no_data(self):
     161        frame = "\x81\x05"
     162        frames, buf = parse_hybi07_frames(frame)
     163        self.assertFalse(frames)
     164        self.assertEqual(buf, "\x81\x05")
     165
     166    def test_parse_hybi07_partial_truncated_data(self):
     167        frame = "\x81\x05Hel"
     168        frames, buf = parse_hybi07_frames(frame)
     169        self.assertFalse(frames)
     170        self.assertEqual(buf, "\x81\x05Hel")