Ticket #4173: progrium-websocket-76.patch

File progrium-websocket-76.patch, 8.1 KB (added by progrium, 3 years ago)

simpler implementation of 76 spec (works with older twisted as well), but also general refactoring

  • twisted/web/websocket.py

     
    22# Copyright (c) 2009 Twisted Matrix Laboratories. 
    33# See LICENSE for details. 
    44 
     5# Based on http://twistedmatrix.com/trac/browser/branches/websocket-4173-2/twisted/web/websocket.py 
     6 
    57""" 
    68WebSocket server protocol. 
    79 
    … …  
    1113@since: 10.1 
    1214""" 
    1315 
    14  
    1516from twisted.web.http import datetimeToString 
    1617from twisted.web.server import Request, Site, version, unquote 
     18import struct 
     19import re 
     20import hashlib 
    1721 
    1822 
    19  
    2023class WebSocketRequest(Request): 
    2124    """ 
    2225    A general purpose L{Request} supporting connection upgrade for WebSocket. 
    2326    """ 
     27    handlerFactory = None 
    2428 
     29    def isWebSocket(self): 
     30        return self.requestHeaders.getRawHeaders("Upgrade") == ["WebSocket"] and \ 
     31            self.requestHeaders.getRawHeaders("Connection") == ["Upgrade"] 
     32 
    2533    def process(self): 
    26         if (self.requestHeaders.getRawHeaders("Upgrade") == ["WebSocket"] and 
    27             self.requestHeaders.getRawHeaders("Connection") == ["Upgrade"]): 
     34        if self.isWebSocket(): 
    2835            return self.processWebSocket() 
    2936        else: 
    3037            return Request.process(self) 
    … …  
    3744        # get site from channel 
    3845        self.site = self.channel.site 
    3946 
     47        # set an empty handler attribute 
     48        self.handler = None 
     49 
    4050        # set various default headers 
    4151        self.setHeader("server", version) 
    4252        self.setHeader("date", datetimeToString()) 
    … …  
    4757        self.renderWebSocket() 
    4858 
    4959 
    50     def _checkClientHandshake(self): 
    51         """ 
    52         Verify client handshake, closing the connection in case of problem. 
     60    def _handshake75(self): 
     61        origin  = self.requestHeaders.getRawHeaders("Origin",   [None])[0] 
     62        host    = self.requestHeaders.getRawHeaders("Host",     [None])[0] 
     63        if not origin or not host: 
     64            return 
     65         
     66        protocol = self.requestHeaders.getRawHeaders("WebSocket-Protocol", [None])[0] 
     67        if protocol and protocol not in self.site.supportedProtocols: 
     68            return 
     69             
     70        if self.isSecure(): 
     71            scheme = "wss" 
     72        else: 
     73            scheme = "ws" 
     74        location = "%s://%s%s" % (scheme, host, self.uri) 
     75        handshake = [ 
     76            "HTTP/1.1 101 Web Socket Protocol Handshake", 
     77            "Upgrade: WebSocket", 
     78            "Connection: Upgrade", 
     79            "WebSocket-Origin: %s" % origin, 
     80            "WebSocket-Location: %s" % location, 
     81            ] 
     82        if protocol is not None: 
     83            handshake.append("WebSocket-Protocol: %s" % protocol) 
     84                 
     85        return handshake 
     86     
     87    def _handshake76(self): 
     88        origin  = self.requestHeaders.getRawHeaders("Origin",   [None])[0] 
     89        host    = self.requestHeaders.getRawHeaders("Host",     [None])[0] 
     90        if not origin or not host: 
     91            return None, None 
     92         
     93        protocol = self.requestHeaders.getRawHeaders("Sec-WebSocket-Protocol", [None])[0] 
     94        if protocol and protocol not in self.site.supportedProtocols: 
     95            return None, None 
    5396 
    54         @return: C{None} if a problem was detected, or a tuple of I{Origin} 
    55             header, I{Host} header, I{WebSocket-Protocol} header, and 
    56             C{WebSocketHandler} instance. The I{WebSocket-Protocol} header will 
    57             be C{None} if not specified by the client. 
    58         """ 
    59         def finish(): 
    60             self.channel.transport.loseConnection() 
    61         if self.queued: 
    62             return finish() 
    63         originHeaders = self.requestHeaders.getRawHeaders("Origin", []) 
    64         if len(originHeaders) != 1: 
    65             return finish() 
    66         hostHeaders = self.requestHeaders.getRawHeaders("Host", []) 
    67         if len(hostHeaders) != 1: 
    68             return finish() 
    69  
    70         handlerFactory = self.site.handlers.get(self.uri) 
    71         if not handlerFactory: 
    72             return finish() 
    73         transport = WebSocketTransport(self) 
    74         handler = handlerFactory(transport) 
    75         transport._attachHandler(handler) 
    76  
    77         protocolHeaders = self.requestHeaders.getRawHeaders( 
    78             "WebSocket-Protocol", []) 
    79         if len(protocolHeaders) not in (0,  1): 
    80             return finish() 
    81         if protocolHeaders: 
    82             if protocolHeaders[0] not in self.site.supportedProtocols: 
    83                 return finish() 
    84             protocolHeader = protocolHeaders[0] 
     97        if self.isSecure(): 
     98            scheme = "wss" 
    8599        else: 
    86             protocolHeader = None 
    87         return originHeaders[0], hostHeaders[0], protocolHeader, handler 
     100            scheme = "ws" 
     101        location = "%s://%s%s" % (scheme, host, self.uri) 
     102        handshake = [ 
     103            "HTTP/1.1 101 Web Socket Protocol Handshake", 
     104            "Upgrade: WebSocket", 
     105            "Connection: Upgrade", 
     106            "Sec-WebSocket-Origin: %s" % origin, 
     107            "Sec-WebSocket-Location: %s" % location, 
     108            ] 
     109        if protocol is not None: 
     110            handshake.append("Sec-WebSocket-Protocol: %s" % protocol) 
     111         
     112        self.channel.setRawMode()  
     113         
     114        # Refer to 5.2 4-9 of the draft 76 
     115        key1 = self.requestHeaders.getRawHeaders('Sec-WebSocket-Key1', [None])[0] 
     116        key2 = self.requestHeaders.getRawHeaders('Sec-WebSocket-Key2', [None])[0] 
     117        key3 = self.content.getvalue() 
     118         
     119        def extract_nums(s): return int(''.join(re.findall(r'[0-9]', s))) 
     120        def count_spaces(s): return len(re.findall(r' ', s)) 
     121        part1 = extract_nums(key1) / count_spaces(key1) 
     122        part2 = extract_nums(key2) / count_spaces(key2) 
     123        challenge = hashlib.md5(struct.pack('>ii8s', part1, part2, key3)).digest() 
     124         
     125        return handshake, challenge 
    88126 
     127    def gotLength(self, length): 
     128        spec76 = self.requestHeaders.getRawHeaders("Sec-WebSocket-Key1", [None])[0] 
     129        if self.isWebSocket() and spec76: 
     130            self.channel.headerReceived("content-length: 8") 
     131        Request.gotLength(self, length) 
    89132 
    90133    def renderWebSocket(self): 
    91134        """ 
    … …  
    95138        connection will be closed. Otherwise, the response to the handshake is 
    96139        sent and a C{WebSocketHandler} is created to handle the request. 
    97140        """ 
    98         check =  self._checkClientHandshake() 
    99         if check is None: 
     141        if self.queued: 
     142            self.channel.transport.loseConnection() 
    100143            return 
    101         originHeader, hostHeader, protocolHeader, handler = check 
    102         self.startedWriting = True 
    103         handshake = [ 
    104             "HTTP/1.1 101 Web Socket Protocol Handshake", 
    105             "Upgrade: WebSocket", 
    106             "Connection: Upgrade"] 
    107         handshake.append("WebSocket-Origin: %s" % (originHeader)) 
    108         if self.isSecure(): 
    109             scheme = "wss" 
     144         
     145        if self.requestHeaders.getRawHeaders("Sec-WebSocket-Key1", [None])[0]: 
     146            handshake, challenge_response = self._handshake76() 
    110147        else: 
    111             scheme = "ws" 
    112         handshake.append( 
    113             "WebSocket-Location: %s://%s%s" % ( 
    114             scheme, hostHeader, self.uri)) 
     148            handshake = self._handshake75() 
     149            challenge_response = None 
     150         
     151        if not handshake: 
     152            self.channel.transport.loseConnection() 
     153            return 
    115154 
    116         if protocolHeader is not None: 
    117             handshake.append("WebSocket-Protocol: %s" % protocolHeader) 
     155        handlerFactory = self.site.handlers.get(self.uri) or self.handlerFactory 
     156        if not handlerFactory: 
     157            return self.channel.transport.loseConnection() 
     158        transport = WebSocketTransport(self) 
     159        handler = handlerFactory(transport) 
     160        transport._attachHandler(handler) 
     161        self.handler = handler 
    118162 
     163        self.startedWriting = True 
     164         
    119165        for header in handshake: 
    120166            self.write("%s\r\n" % header) 
    121167 
    122168        self.write("\r\n") 
     169        if challenge_response: 
     170            self.write(challenge_response) 
     171         
    123172        self.channel.setRawMode() 
    124173        # XXX we probably don't want to set _transferDecoder 
    125174        self.channel._transferDecoder = WebSocketFrameDecoder(