From 8048b96f7e6db9166899573074f2b6ae7d86dcd7 Mon Sep 17 00:00:00 2001
From: Miguel Grinberg <miguel.grinberg@gmail.com>
Date: Tue, 8 Aug 2017 11:21:52 -0700
Subject: [PATCH] Weboscket subprotocol negotiation

Fixes #874
---
 sanic/app.py       |  7 +++++--
 sanic/websocket.py | 14 +++++++++++++-
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/sanic/app.py b/sanic/app.py
index f0ccad865c..954929f358 100644
--- a/sanic/app.py
+++ b/sanic/app.py
@@ -214,9 +214,12 @@ def add_route(self, handler, uri, methods=frozenset({'GET'}), host=None,
         return handler
 
     # Decorator
-    def websocket(self, uri, host=None, strict_slashes=False):
+    def websocket(self, uri, subprotocols=None, host=None,
+                  strict_slashes=False):
         """Decorate a function to be registered as a websocket route
         :param uri: path of the URL
+        :param subprotocols: optional list of strings with the supported
+                             subprotocols
         :param host:
         :return: decorated function
         """
@@ -236,7 +239,7 @@ async def websocket_handler(request, *args, **kwargs):
                     # On Python3.5 the Transport classes in asyncio do not
                     # have a get_protocol() method as in uvloop
                     protocol = request.transport._protocol
-                ws = await protocol.websocket_handshake(request)
+                ws = await protocol.websocket_handshake(request, subprotocols)
 
                 # schedule the application handler
                 # its future is kept in self.websocket_tasks in case it
diff --git a/sanic/websocket.py b/sanic/websocket.py
index 94320a5e29..e8e9922fac 100644
--- a/sanic/websocket.py
+++ b/sanic/websocket.py
@@ -41,7 +41,7 @@ def write_response(self, response):
         else:
             super().write_response(response)
 
-    async def websocket_handshake(self, request):
+    async def websocket_handshake(self, request, subprotocols=None):
         # let the websockets package do the handshake with the client
         headers = []
 
@@ -57,6 +57,17 @@ def set_header(k, v):
         except InvalidHandshake:
             raise InvalidUsage('Invalid websocket request')
 
+        subprotocol = None
+        if subprotocols and 'Sec-Websocket-Protocol' in request.headers:
+            # select a subprotocol
+            client_subprotocols = [p.strip() for p in request.headers[
+                'Sec-Websocket-Protocol'].split(',')]
+            for p in client_subprotocols:
+                if p in subprotocols:
+                    subprotocol = p
+                    set_header('Sec-Websocket-Protocol', subprotocol)
+                    break
+
         # write the 101 response back to the client
         rv = b'HTTP/1.1 101 Switching Protocols\r\n'
         for k, v in headers:
@@ -69,5 +80,6 @@ def set_header(k, v):
             max_size=self.websocket_max_size,
             max_queue=self.websocket_max_queue
         )
+        self.websocket.subprotocol = subprotocol
         self.websocket.connection_made(request.transport)
         return self.websocket