Add authentication to ws requests (#196)

* add auth to ws requests

* style check fix

Author: AjayThorve

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,6 +20,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
 
 

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