Keep proxying all requested subprotocols

When jupyter-server-proxy proxies websockets, its finalizes the
websocket handshake between client/proxy before it initiates the
proxy/server websocket handshake.

In 4.1.1 the thinking was that it was a better compromise to not forward
all subprotocol choices in the proxy/server handshake if we had
prematurely picked a single choice in the client/proxy handshake. This
turns out to have introduced a regression though, as at least bokeh
had been using secondary subprotocol choices to pass other information
such as base64 encoded JSON with keys like `session_id`, `session_expiry`
and `__bk__zlib_`.

This commit makes sure we keep passing all requested subprotocols, even
though we have prematurely picked a specific ahead of time - which is a
bug tracked in
#459.

Author: consideRatio

jupyter_server_proxy/handlers.py

@@ -116,6 +116,7 @@ def __init__(self, *args, **kwargs):
             "rewrite_response",
             tuple(),
         )
+        self._requested_subprotocols = None
         super().__init__(*args, **kwargs)
 
     # Support/use jupyter_server config arguments allow_origin and allow_origin_pat
@@ -520,27 +521,21 @@ async def start_websocket_connection():
             self.log.info(f"Trying to establish websocket connection to {client_uri}")
             self._record_activity()
             request = httpclient.HTTPRequest(url=client_uri, headers=headers)
-            subprotocols = (
-                [self.selected_subprotocol] if self.selected_subprotocol else None
-            )
             self.ws = await pingable_ws_connect(
                 request=request,
                 on_message_callback=message_cb,
                 on_ping_callback=ping_cb,
-                subprotocols=subprotocols,
+                subprotocols=self._requested_subprotocols,
                 resolver=resolver,
             )
             self._record_activity()
             self.log.info(f"Websocket connection established to {client_uri}")
-            if (
-                subprotocols
-                and self.ws.selected_subprotocol != self.selected_subprotocol
-            ):
+            if self.ws.selected_subprotocol != self.selected_subprotocol:
                 self.log.warn(
                     f"Websocket subprotocol between proxy/server ({self.ws.selected_subprotocol}) "
                     f"became different than for client/proxy ({self.selected_subprotocol}) "
                     "due to https://github.com/jupyterhub/jupyter-server-proxy/issues/459. "
-                    f"Requested subprotocols were {subprotocols}."
+                    f"Requested subprotocols were {self._requested_subprotocols}."
                 )
 
         # Wait for the WebSocket to be connected before resolving.
@@ -578,20 +573,27 @@ def select_subprotocol(self, 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. This issue is tracked via
-        https://github.com/jupyterhub/jupyter-server-proxy/issues/459.
-
         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.
         """
+        # Stash all requested subprotocols to be re-used as requested
+        # subprotocols in the proxy/server handshake to be performed later. At
+        # least bokeh has used additional subprotocols to pass credentials,
+        # making this a required workaround for now.
+        #
+        self._requested_subprotocols = subprotocols if subprotocols else None
+
         if subprotocols:
             self.log.debug(
                 f"Client sent subprotocols: {subprotocols}, selecting the first"
             )
+            # FIXME: Subprotocol selection should 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. This issue is
+            #        tracked in https://github.com/jupyterhub/jupyter-server-proxy/issues/459.
+            #
             return subprotocols[0]
         return None
 

tests/test_proxies.py

@@ -360,6 +360,7 @@ async def test_server_proxy_websocket_headers(a_server_port_and_token: Tuple[int
     [
         (None, None, None, None),
         (["first"], ["first"], "first", "first"),
+        (["first", "second"], ["first", "second"], "first", "first"),
         # IMPORTANT: The tests below verify current bugged behavior, and the
         #            commented out tests is what we want to succeed!
         #
@@ -369,14 +370,12 @@ async def test_server_proxy_websocket_headers(a_server_port_and_token: Tuple[int
         #            before the proxy/server handshake, and that makes it
         #            impossible. We currently instead just pick the first
         #            requested protocol no matter what what subprotocol the
-        #            server picks.
+        #            server picks and warn if there is a mismatch retroactively.
         #
-        # Bug 1 - server wasn't passed all subprotocols:
-        (["first", "second"], ["first"], "first", "first"),
-        # (["first", "second"], ["first", "second"], "first", "first"),
+        #            Tracked in https://github.com/jupyterhub/jupyter-server-proxy/issues/459.
         #
-        # Bug 2 - server_responded doesn't match proxy_responded:
-        (["first", "favored"], ["first"], "first", "first"),
+        # Bug - server_responded doesn't match proxy_responded:
+        (["first", "favored"], ["first", "favored"], "favored", "first"),
         # (["first", "favored"], ["first", "favored"], "favored", "favored"),
         (
             ["please_select_no_protocol"],