Merge pull request #150 from flying-sheep/patch-1

Add websocket subprotocol support

Author: fafhrd91

CONTRIBUTORS.txt

@@ -22,3 +22,4 @@ Sebastian Hanula
 Simon Kennedy
 Vaibhav Sagar
 Vitaly Haritonsky
+Philipp A.

aiohttp/websocket.py

@@ -21,7 +21,7 @@
 
 WS_KEY = b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
 WS_HDRS = ('UPGRADE', 'CONNECTION',
-           'SEC-WEBSOCKET-VERSION', 'SEC-WEBSOCKET-KEY')
+           'SEC-WEBSOCKET-VERSION', 'SEC-WEBSOCKET-KEY', 'SEC-WEBSOCKET-PROTOCOL')
 
 Message = collections.namedtuple('Message', ['tp', 'data', 'extra'])
 
@@ -182,10 +182,14 @@ def close(self, code=1000, message=b''):
             opcode=OPCODE_CLOSE)
 
 
-def do_handshake(method, headers, transport):
+def do_handshake(method, headers, transport, protocols=()):
     """Prepare WebSocket handshake. It return http response code,
     response headers, websocket parser, websocket writer. It does not
-    perform any IO."""
+    perform any IO.
+
+    `protocols` is a sequence of known protocols. On successful handshake,
+    the returned response headers contain the first protocol in this list
+    which the server also knows."""
 
     # WebSocket accepts only GET
     if method.upper() != 'GET':
@@ -201,6 +205,21 @@ def do_handshake(method, headers, transport):
         raise errors.HttpBadRequest(
             'No CONNECTION upgrade hdr: {}'.format(
                 headers.get('CONNECTION')))
+    
+    # find common sub-protocol between client and server
+    protocol = None
+    if 'SEC-WEBSOCKET-PROTOCOL' in headers:
+        req_protocols = {str(proto.strip()) for proto in
+            headers['SEC-WEBSOCKET-PROTOCOL'].split(',')}
+        
+        for proto in protocols:
+            if proto in req_protocols:
+                protocol = proto
+                break
+        else:
+            raise errors.HttpBadRequest(
+                'Client protocols {!r} don’t overlap server-known ones {!r}'
+                .format(protocols, req_protocols))
 
     # check supported version
     version = headers.get('SEC-WEBSOCKET-VERSION')
@@ -218,12 +237,18 @@ def do_handshake(method, headers, transport):
         raise errors.HttpBadRequest(
             'Handshake error: {!r}'.format(key)) from None
 
-    # response code, headers, parser, writer
-    return (101,
-            (('UPGRADE', 'websocket'),
+    response_headers = [('UPGRADE', 'websocket'),
              ('CONNECTION', 'upgrade'),
              ('TRANSFER-ENCODING', 'chunked'),
              ('SEC-WEBSOCKET-ACCEPT', base64.b64encode(
-                 hashlib.sha1(key.encode() + WS_KEY).digest()).decode())),
+                 hashlib.sha1(key.encode() + WS_KEY).digest()).decode())]
+    
+    if protocol:
+        response_headers.append(('SEC-WEBSOCKET-PROTOCOL', protocol))
+
+    # response code, headers, parser, writer, protocol
+    return (101,
+            response_headers,
             WebSocketParser,
-            WebSocketWriter(transport))
+            WebSocketWriter(transport),
+            protocol)

tests/test_websocket.py

@@ -419,18 +419,54 @@ def test_protocol_key(self):
             websocket.do_handshake,
             self.message.method, self.message.headers, self.transport)
 
+    def gen_ws_headers(self, protocols=''):
+        key = base64.b64encode(os.urandom(16)).decode()
+        hdrs = [('UPGRADE', 'websocket'),
+                ('CONNECTION', 'upgrade'),
+                ('SEC-WEBSOCKET-VERSION', '13'),
+                ('SEC-WEBSOCKET-KEY', key)]
+        if protocols:
+            hdrs += [('SEC-WEBSOCKET-PROTOCOL', protocols)]
+        return hdrs, key
+
     def test_handshake(self):
-        sec_key = base64.b64encode(os.urandom(16)).decode()
+        hdrs, sec_key = self.gen_ws_headers()
 
-        self.headers.extend([('UPGRADE', 'websocket'),
-                             ('CONNECTION', 'upgrade'),
-                             ('SEC-WEBSOCKET-VERSION', '13'),
-                             ('SEC-WEBSOCKET-KEY', sec_key)])
-        status, headers, parser, writer = websocket.do_handshake(
+        self.headers.extend(hdrs)
+        status, headers, parser, writer, protocol = websocket.do_handshake(
             self.message.method, self.message.headers, self.transport)
         self.assertEqual(status, 101)
+        self.assertIsNone(protocol)
 
         key = base64.b64encode(
             hashlib.sha1(sec_key.encode() + websocket.WS_KEY).digest())
         headers = dict(headers)
         self.assertEqual(headers['SEC-WEBSOCKET-ACCEPT'], key.decode())
+
+    def test_handshake_protocol(self):
+        '''Tests if one protocol is returned by do_handshake'''
+        proto = 'chat'
+
+        self.headers.extend(self.gen_ws_headers(proto)[0])
+        _, resp_headers, _, _, protocol = websocket.do_handshake(
+            self.message.method, self.message.headers, self.transport,
+            protocols=[proto])
+
+        self.assertEqual(protocol, proto)
+
+        #also test if we reply with the protocol
+        resp_headers = dict(resp_headers)
+        self.assertEqual(resp_headers['SEC-WEBSOCKET-PROTOCOL'], proto)
+
+    def test_handshake_protocol_agreement(self):
+        '''Tests if the right protocol is selected given multiple'''
+        best_proto = 'chat'
+        wanted_protos = ['best', 'chat', 'worse_proto']
+        server_protos = 'worse_proto,chat'
+
+        self.headers.extend(self.gen_ws_headers(server_protos)[0])
+        _, resp_headers, _, _, protocol = websocket.do_handshake(
+            self.message.method, self.message.headers, self.transport,
+            protocols=wanted_protos)
+
+        self.assertEqual(protocol, best_proto)