Merge pull request #180 from xhochy/proxy-ws-subprotocols

Proxy Websocket subprotocols

Author: ryanlovett

jupyter_server_proxy/handlers.py

@@ -44,6 +44,7 @@ def __init__(self, *args, **kwargs):
         self.proxy_base = ''
         self.absolute_url = kwargs.pop('absolute_url', False)
         self.host_whitelist = kwargs.pop('host_whitelist', ['localhost', '127.0.0.1'])
+        self.subprotocols = None
         super().__init__(*args, **kwargs)
 
     # Support all the methods that tornado does by default except for GET which
@@ -284,7 +285,8 @@ async def start_websocket_connection():
             self._record_activity()
             request = httpclient.HTTPRequest(url=client_uri, headers=headers)
             self.ws = await pingable_ws_connect(request=request,
-                on_message_callback=message_cb, on_ping_callback=ping_cb)
+                on_message_callback=message_cb, on_ping_callback=ping_cb,
+                subprotocols=self.subprotocols)
             ws_connected.set_result(True)
             self._record_activity()
             self.log.info('Websocket connection established to {}'.format(client_uri))
@@ -316,6 +318,7 @@ 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.info('Client sent subprotocols: {}'.format(subprotocols))
             return subprotocols[0]

jupyter_server_proxy/websocket.py

@@ -36,8 +36,8 @@ def on_ping(self, data):
             self._on_ping_callback(data)
 
 
-def pingable_ws_connect(request=None, on_message_callback=None,
-                        on_ping_callback=None):
+def pingable_ws_connect(request=None,on_message_callback=None,
+                        on_ping_callback=None, subprotocols=None):
     """
     A variation on websocket_connect that returns a PingableWSClientConnection
     with on_ping_callback.
@@ -60,7 +60,8 @@ def pingable_ws_connect(request=None, on_message_callback=None,
             compression_options={},
             on_message_callback=on_message_callback,
             on_ping_callback=on_ping_callback,
-            max_message_size=getattr(websocket, '_default_max_message_size', 10 * 1024 * 1024))
+            max_message_size=getattr(websocket, '_default_max_message_size', 10 * 1024 * 1024),
+            subprotocols=subprotocols)
 
     return conn.connect_future
 

tests/resources/jupyter_server_config.py

@@ -24,6 +24,9 @@ def mappathf(path):
         'command': ['python3', './tests/resources/httpinfo.py', '{port}'],
         'mappath': mappathf,
     },
+    'python-websocket' : {
+        'command': ['python3', './tests/resources/websocket.py', '--port={port}'],
+    }
 }
 
 import sys

tests/resources/websocket.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python
+#
+# Based on the chat demo from https://github.com/tornadoweb/tornado/blob/d6819307ee050bbd8ec5deb623e9150ce2220ef9/demos/websocket/chatdemo.py#L1
+# Original License:
+#
+# Copyright 2009 Facebook
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+
+import logging
+import json
+import tornado.escape
+import tornado.ioloop
+import tornado.options
+import tornado.web
+import tornado.websocket
+import os.path
+import uuid
+
+from tornado.options import define, options
+
+define("port", default=8888, help="run on the given port", type=int)
+
+
+class Application(tornado.web.Application):
+    def __init__(self):
+        handlers = [
+            (r"/", MainHandler),
+            (r"/echosocket", EchoWebSocket),
+            (r"/subprotocolsocket", SubprotocolWebSocket),
+        ]
+        settings = dict(
+            cookie_secret="__RANDOM_VALUE__",
+            template_path=os.path.join(os.path.dirname(__file__), "templates"),
+            static_path=os.path.join(os.path.dirname(__file__), "static"),
+            xsrf_cookies=True,
+        )
+        super(Application, self).__init__(handlers, **settings)
+
+
+class MainHandler(tornado.web.RequestHandler):
+    def get(self):
+        self.write("Hello, world!")
+
+
+class EchoWebSocket(tornado.websocket.WebSocketHandler):
+    def on_message(self, message):
+        self.write_message(message)
+
+
+class SubprotocolWebSocket(tornado.websocket.WebSocketHandler):
+    def __init__(self, *args, **kwargs):
+        self._subprotocols = None
+        super().__init__(*args, **kwargs)
+
+    def select_subprotocol(self, subprotocols):
+        self._subprotocols = subprotocols
+        return None
+
+    def on_message(self, message):
+        self.write_message(json.dumps(self._subprotocols))
+
+
+def main():
+    tornado.options.parse_command_line()
+    app = Application()
+    app.listen(options.port)
+    tornado.ioloop.IOLoop.current().start()
+
+
+if __name__ == "__main__":
+    main()

tests/test_proxies.py

@@ -1,6 +1,9 @@
+import asyncio
+import json
 import os
 from http.client import HTTPConnection
 import pytest
+from tornado.websocket import websocket_connect
 
 PORT = os.getenv('TEST_PORT', 8888)
 TOKEN = os.getenv('JUPYTER_TOKEN', 'secret')
@@ -99,3 +102,35 @@ def test_server_proxy_mappath_callable(requestpath, expected):
 def test_server_proxy_remote():
     r = request_get(PORT, '/newproxy', TOKEN, host='127.0.0.1')
     assert r.code == 200
+
+
+@pytest.fixture(scope="module")
+def event_loop():
+    loop = asyncio.get_event_loop()
+    yield loop
+    loop.close()
+
+
+async def _websocket_echo():
+    url = "ws://localhost:{}/python-websocket/echosocket".format(PORT)
+    conn = await websocket_connect(url)
+    expected_msg = "Hello, world!"
+    await conn.write_message(expected_msg)
+    msg = await conn.read_message()
+    assert msg == expected_msg
+
+def test_server_proxy_websocket(event_loop):
+    event_loop.run_until_complete(_websocket_echo())
+
+
+async def _websocket_subprotocols():
+    url = "ws://localhost:{}/python-websocket/subprotocolsocket".format(PORT)
+    conn = await websocket_connect(url, subprotocols=["protocol_1", "protocol_2"])
+    await conn.write_message("Hello, world!")
+    msg = await conn.read_message()
+    assert json.loads(msg) == ["protocol_1", "protocol_2"]
+
+
+def test_server_proxy_websocket_subprotocols(event_loop):
+    event_loop.run_until_complete(_websocket_subprotocols())
+