add auth to ws requests

jupyterlab_nvdashboard/apps/utils.py

@@ -1,10 +1,15 @@
 from tornado.websocket import WebSocketHandler
+from jupyter_server.base.handlers import JupyterHandler
 import tornado
 import json
 
 
-class CustomWebSocketHandler(WebSocketHandler):
+class CustomWebSocketHandler(JupyterHandler, WebSocketHandler):
  def open(self):
+ if not self.current_user:
+ self.write_message(json.dumps({"error": "Unauthorized access"}))
+ self.close()
+ return
  self.write_message(json.dumps({"status": "connected"}))
  self.set_nodelay(True)
  # Start a periodic callback to send data every 50ms

jupyterlab_nvdashboard/tests/test_cpu_handlers.py

@@ -20,9 +20,13 @@ def handler_args():
  with patch("tornado.web.Application") as mock_application, patch(
  "tornado.httputil.HTTPServerRequest"
  ) as mock_request:
+ # Mock the settings to return appropriate values
+ mock_settings = {
+ "base_url": "/",
+ }
+ mock_application.settings = mock_settings
  yield mock_application, mock_request
 
-
 def test_cpu_resource_handler(mock_handler, handler_args):
  handler = CPUResourceWebSocketHandler(*handler_args)
  handler.send_data()

jupyterlab_nvdashboard/tests/test_gpu_handlers.py

@@ -26,6 +26,11 @@ def handler_args():
  with patch("tornado.web.Application") as mock_application, patch(
  "tornado.httputil.HTTPServerRequest"
  ) as mock_request:
+ # Mock the settings to return appropriate values
+ mock_settings = {
+ "base_url": "/",
+ }
+ mock_application.settings = mock_settings
  yield mock_application, mock_request