Check authentication when websockets open

Author: yuvipanda

jupyter_server_proxy/handlers.py

@@ -130,6 +130,39 @@ def check_origin(self, origin=None):
     async def open(self, port, proxied_path):
         raise NotImplementedError("Subclasses of ProxyHandler should implement open")
 
+    async def prepare(self, *args, **kwargs):
+        """
+        Enforce authentication on *all* requests.
+
+        This method is called *before* any other method for all requests.
+        See https://www.tornadoweb.org/en/stable/web.html#tornado.web.RequestHandler.prepare.
+        """
+        # Due to https://github.com/jupyter-server/jupyter_server/issues/1012,
+        # we can not decorate `prepare` with `@web.authenticated`.
+        # `super().prepare`, which calls `JupyterHandler.prepare`, *must* be called
+        # before `@web.authenticated` can work. Since `@web.authenticated` is a decorator
+        # that relies on the decorated method to get access to request information, we can
+        # not call it directly. Instead, we create an empty lambda that takes a request_handler,
+        # decorate that with web.authenticated, and call the decorated function.
+        # super().prepare became async with jupyter_server v2
+        _prepared = super().prepare(*args, **kwargs)
+        if _prepared is not None:
+            await _prepared
+
+        # If this is a GET request that wants to be upgraded to a websocket, users not
+        # already authenticated gets a straightforward 403. Everything else is dealt
+        # with by `web.authenticated`, which does a 302 to the appropriate login url.
+        # Websockets are purely API calls made by JS rather than a direct user facing page,
+        # so redirects do not make sense for them.
+        if (
+            self.request.method == "GET"
+            and self.request.headers.get("Upgrade", "").lower() == "websocket"
+        ):
+            if not self.current_user:
+                raise web.HTTPError(403)
+        else:
+            web.authenticated(lambda request_handler: None)(self)
+
     async def http_get(self, host, port, proxy_path=""):
         """Our non-websocket GET."""
         raise NotImplementedError(
@@ -280,7 +313,6 @@ def _check_host_allowlist(self, host):
         else:
             return host in self.host_allowlist
 
-    @web.authenticated
     async def proxy(self, host, port, proxied_path):
         """
         This serverextension handles:
@@ -682,7 +714,6 @@ def _realize_rendered_template(self, attribute):
             attribute = call_with_asked_args(attribute, self.process_args)
         return self._render_template(attribute)
 
-    @web.authenticated
     async def proxy(self, port, path):
         if not path.startswith("/"):
             path = "/" + path
@@ -866,7 +897,6 @@ async def ensure_process(self):
                     del self.state["proc"]
                     raise
 
-    @web.authenticated
     async def proxy(self, port, path):
         await self.ensure_process()
         return await ensure_async(super().proxy(port, path))

tests/test_proxies.py

@@ -7,6 +7,7 @@
 from urllib.parse import quote
 
 import pytest
+from tornado.httpclient import HTTPClientError
 from tornado.websocket import websocket_connect
 
 # use ipv4 for CI, etc.
@@ -334,8 +335,8 @@ def test_server_content_encoding_header(
 async def test_server_proxy_websocket_messages(
     a_server_port_and_token: Tuple[int, str]
 ) -> None:
-    PORT = a_server_port_and_token[0]
-    url = f"ws://{LOCALHOST}:{PORT}/python-websocket/echosocket"
+    PORT, TOKEN = a_server_port_and_token
+    url = f"ws://{LOCALHOST}:{PORT}/python-websocket/echosocket?token={TOKEN}"
     conn = await websocket_connect(url)
     expected_msg = "Hello, world!"
     await conn.write_message(expected_msg)
@@ -344,8 +345,8 @@ async def test_server_proxy_websocket_messages(
 
 
 async def test_server_proxy_websocket_headers(a_server_port_and_token: Tuple[int, str]):
-    PORT = a_server_port_and_token[0]
-    url = f"ws://{LOCALHOST}:{PORT}/python-websocket/headerssocket"
+    PORT, TOKEN = a_server_port_and_token
+    url = f"ws://{LOCALHOST}:{PORT}/python-websocket/headerssocket?token={TOKEN}"
     conn = await websocket_connect(url)
     await conn.write_message("Hello")
     msg = await conn.read_message()
@@ -394,7 +395,7 @@ async def test_server_proxy_websocket_subprotocols(
     proxy_responded,
 ):
     PORT, TOKEN = a_server_port_and_token
-    url = f"ws://{LOCALHOST}:{PORT}/python-websocket/subprotocolsocket"
+    url = f"ws://{LOCALHOST}:{PORT}/python-websocket/subprotocolsocket?token={TOKEN}"
     conn = await websocket_connect(url, subprotocols=client_requested)
     await conn.write_message("Hello, world!")
 
@@ -418,6 +419,17 @@ async def test_server_proxy_websocket_subprotocols(
         assert "Sec-Websocket-Protocol" in conn.headers
 
 
+async def test_websocket_no_auth_failure(
+    a_server_port_and_token: Tuple[int, str]
+) -> None:
+    PORT = a_server_port_and_token[0]
+    # Intentionally do not pass an appropriate token, which should cause a 403
+    url = f"ws://{LOCALHOST}:{PORT}/python-websocket/headerssocket"
+
+    with pytest.raises(HTTPClientError, match=r".*HTTP 403: Forbidden.*"):
+        await websocket_connect(url)
+
+
 @pytest.mark.parametrize(
     "proxy_path, status",
     [