Implement auth and tests for websockets

`WebsocketMixing` may be used with or without `JupyterHandler`;
- if used with it, we want to have custom auth implementation
  because redirecting to login page does not make sense for
  a websocket's GET request
- if these are used without `JupyterHandler `we want the auth
  rules to still apply, even though the `current_user`
  logic may differ slightly

Author: krassowski

jupyter_server/base/websocket.py

@@ -3,7 +3,7 @@
 from typing import Optional, no_type_check
 from urllib.parse import urlparse
 
-from tornado import ioloop
+from tornado import ioloop, web
 from tornado.iostream import IOStream
 
 # ping interval for keeping websockets alive (30 seconds)
@@ -82,6 +82,26 @@ def check_origin(self, origin: Optional[str] = None) -> bool:
     def clear_cookie(self, *args, **kwargs):
         """meaningless for websockets"""
 
+    @no_type_check
+    def _maybe_auth(self):
+        """Verify authentication if required"""
+        if not self.settings.get("allow_unauthenticated_access", False):
+            if not self.request.method:
+                raise web.HTTPError(403)
+            method = getattr(self, self.request.method.lower())
+            if not getattr(method, "__allow_unauthenticated", False):
+                # rather than re-using `web.authenticated` which also redirects
+                # to login page on GET, just raise 403 if user is not known
+                user = self.current_user
+                if user is None:
+                    self.log.warning("Couldn't authenticate WebSocket connection")
+                    raise web.HTTPError(403)
+
+    def prepare(self, *args, **kwargs):
+        """Handle a get request."""
+        self._maybe_auth()
+        return super().prepare(*args, **kwargs)
+
     @no_type_check
     def open(self, *args, **kwargs):
         """Open the websocket."""

tests/base/test_websocket.py

@@ -4,12 +4,15 @@
 from unittest.mock import MagicMock
 
 import pytest
+from tornado.httpclient import HTTPClientError
 from tornado.httpserver import HTTPRequest
 from tornado.httputil import HTTPHeaders
 from tornado.websocket import WebSocketClosedError, WebSocketHandler
 
+from jupyter_server.auth.decorator import allow_unauthenticated
 from jupyter_server.base.websocket import WebSocketMixin
 from jupyter_server.serverapp import ServerApp
+from jupyter_server.utils import url_path_join
 
 
 class MockHandler(WebSocketMixin, WebSocketHandler):
@@ -60,3 +63,58 @@ async def test_ping_client_timeout(mixin):
     mixin.send_ping()
     with pytest.raises(WebSocketClosedError):
         mixin.write_message("hello")
+
+
+class NoAuthRulesWebsocketHandler(MockHandler):
+    pass
+
+
+class PermissiveWebsocketHandler(MockHandler):
+    @allow_unauthenticated
+    def get(self, *args, **kwargs) -> None:
+        return super().get(*args, **kwargs)
+
+
+@pytest.mark.parametrize(
+    "jp_server_config", [{"ServerApp": {"allow_unauthenticated_access": True}}]
+)
+async def test_websocket_auth_permissive(jp_serverapp, jp_ws_fetch):
+    app: ServerApp = jp_serverapp
+    app.web_app.add_handlers(
+        ".*$",
+        [
+            (url_path_join(app.base_url, "no-rules"), NoAuthRulesWebsocketHandler),
+            (url_path_join(app.base_url, "permissive"), PermissiveWebsocketHandler),
+        ],
+    )
+
+    # should always permit access when `@allow_unauthenticated` is used
+    ws = await jp_ws_fetch("permissive", headers={"Authorization": ""})
+    ws.close()
+
+    # should allow access when no authentication rules are set up
+    ws = await jp_ws_fetch("no-rules", headers={"Authorization": ""})
+    ws.close()
+
+
+@pytest.mark.parametrize(
+    "jp_server_config", [{"ServerApp": {"allow_unauthenticated_access": False}}]
+)
+async def test_websocket_auth_required(jp_serverapp, jp_ws_fetch):
+    app: ServerApp = jp_serverapp
+    app.web_app.add_handlers(
+        ".*$",
+        [
+            (url_path_join(app.base_url, "no-rules"), NoAuthRulesWebsocketHandler),
+            (url_path_join(app.base_url, "permissive"), PermissiveWebsocketHandler),
+        ],
+    )
+
+    # should always permit access when `@allow_unauthenticated` is used
+    ws = await jp_ws_fetch("permissive", headers={"Authorization": ""})
+    ws.close()
+
+    # should forbid access when no authentication rules are set up
+    with pytest.raises(HTTPClientError) as exception:
+        ws = await jp_ws_fetch("no-rules", headers={"Authorization": ""})
+    assert exception.value.code == 403