Ticket #4173: progrium-websocket-76.patch

File progrium-websocket-76.patch, 8.1 KB (added by progrium, 6 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(