Improve websocket subprotocol request to backend server

Author: consideRatio

jupyter_server_proxy/handlers.py

@@ -116,7 +116,6 @@ def __init__(self, *args, **kwargs):
             "rewrite_response",
             tuple(),
         )
-        self.subprotocols = None
         super().__init__(*args, **kwargs)
 
     # Support/use jupyter_server config arguments allow_origin and allow_origin_pat
@@ -493,7 +492,7 @@ async def start_websocket_connection():
                 request=request,
                 on_message_callback=message_cb,
                 on_ping_callback=ping_cb,
-                subprotocols=self.subprotocols,
+                subprotocols=[self.selected_subprotocol],
                 resolver=resolver,
             )
             self._record_activity()
@@ -531,12 +530,29 @@ def check_xsrf_cookie(self):
         """
 
     def select_subprotocol(self, subprotocols):
-        """Select a single Sec-WebSocket-Protocol during handshake."""
-        self.subprotocols = subprotocols
-        if isinstance(subprotocols, list) and subprotocols:
-            self.log.debug(f"Client sent subprotocols: {subprotocols}")
+        """
+        Select a single Sec-WebSocket-Protocol during handshake.
+
+        Note that this subprotocol selection should really be delegated to the
+        server we proxy to, but we don't! For this to happen, we would need to
+        delay accepting the handshake with the client until we have successfully
+        handshaked with the server.
+
+        Overrides `tornado.websocket.WebSocketHandler.select_subprotocol` that
+        includes an informative docstring:
+        https://github.com/tornadoweb/tornado/blob/v6.4.0/tornado/websocket.py#L337-L360.
+        """
+        if subprotocols:
+            # Tornado 5.0 doesn't pass an empty list, but a list with a an empty
+            # string element.
+            if subprotocols[0] == "":
+                return None
+            self.log.debug(
+                f"Client sent subprotocols: {subprotocols}, selecting the first"
+            )
+            # TODO: warn if we select one out of multiple!
             return subprotocols[0]
-        return super().select_subprotocol(subprotocols)
+        return None
 
 
 class LocalProxyHandler(ProxyHandler):

tests/resources/websocket.py

@@ -54,16 +54,22 @@ def get(self):
 
 
 class EchoWebSocket(tornado.websocket.WebSocketHandler):
+    """Echoes back received messages."""
+
     def on_message(self, message):
         self.write_message(message)
 
 
 class HeadersWebSocket(tornado.websocket.WebSocketHandler):
+    """Echoes back incoming request headers."""
+
     def on_message(self, message):
         self.write_message(json.dumps(dict(self.request.headers)))
 
 
 class SubprotocolWebSocket(tornado.websocket.WebSocketHandler):
+    """Echoes back incoming requested subprotocols."""
+
     def __init__(self, *args, **kwargs):
         self._subprotocols = None
         super().__init__(*args, **kwargs)