Skip to content

Commit

Permalink
Improve websocket subprotocol request to backend server
Browse files Browse the repository at this point in the history
  • Loading branch information
consideRatio committed Feb 21, 2024
1 parent 141dbf0 commit fdaed94
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 7 deletions.
30 changes: 23 additions & 7 deletions jupyter_server_proxy/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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):
Expand Down
6 changes: 6 additions & 0 deletions tests/resources/websocket.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit fdaed94

Please sign in to comment.