Ticket #4173: 4173-5.patch

File 4173-5.patch, 21.2 KB (added by MostAwesomeDude, 3 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")