Stop multiple authentications on single websocket

Author: stopnoanime

oioioi/notifications/server/server.py

@@ -17,11 +17,14 @@ def __init__(self, port: int, amqp_url: str, auth_url: str) -> None:
         self.logger = logging.getLogger('oioioi')
 
         self.app.on_start(self.on_start)
-        self.app.ws("/", {
-            "upgrade": self.on_ws_upgrade,
-            "message": self.on_ws_message,
-            "close": self.on_ws_close,
-        })
+        self.app.ws(
+            "/",
+            {
+                "upgrade": self.on_ws_upgrade,
+                "message": self.on_ws_message,
+                "close": self.on_ws_close,
+            },
+        )
 
     def run(self) -> None:
         """Start the notification server."""
@@ -34,9 +37,9 @@ async def on_start(self) -> None:
         await self.queue.connect()
 
     def on_ws_upgrade(self, res, req, socket_context):
-        """ 
+        """
         Taken from socketify's documentation.
-        This method allows for storing extra data inside the websocket object. 
+        This method allows for storing extra data inside the websocket object.
         """
 
         key = req.get_header("sec-websocket-key")
@@ -47,7 +50,9 @@ def on_ws_upgrade(self, res, req, socket_context):
 
         res.upgrade(key, protocol, extensions, socket_context, user_data)
 
-    async def on_ws_message(self, ws: WebSocket, msg: Union[bytes, str], opcode: OpCode) -> None:
+    async def on_ws_message(
+        self, ws: WebSocket, msg: Union[bytes, str], opcode: OpCode
+    ) -> None:
         """Handle incoming WebSocket messages."""
         try:
             data = json.loads(msg)
@@ -61,15 +66,17 @@ async def on_ws_message(self, ws: WebSocket, msg: Union[bytes, str], opcode: OpC
         except Exception as e:
             self.logger.error(f"Error processing message: {str(e)}")
 
-    async def on_ws_close(self, ws: WebSocket, code: int, msg: Union[bytes, str]) -> None:
+    async def on_ws_close(
+        self, ws: WebSocket, code: int, msg: Union[bytes, str]
+    ) -> None:
         """Handle WebSocket connection closure."""
         try:
             user_id = ws.get_user_data()["user_id"]
 
             # If there are no more active connections for this user, unsubscribe from the RabbitMQ queue
             if user_id and self.app.num_subscribers(user_id) == 0:
                 await self.queue.unsubscribe(user_id)
-                
+
             self.logger.debug(f"WebSocket closed for user {user_id}")
 
         except Exception as e:
@@ -85,6 +92,10 @@ def on_rabbit_message(self, user_name: str, msg: str) -> None:
 
     async def on_ws_auth_message(self, ws: WebSocket, session_id: str) -> None:
         try:
+            current_user_id = ws.get_user_data()["user_id"]
+            if current_user_id:
+                raise RuntimeError("Socket is already authenticated.")
+
             user_id = await self.auth.authenticate(session_id)
 
             ws.subscribe(user_id)
@@ -96,6 +107,8 @@ async def on_ws_auth_message(self, ws: WebSocket, session_id: str) -> None:
 
         except Exception as e:
             self.logger.error(
-                f"Authentication error for session {session_id}: {str(e)}")
-            ws.send({"type": "SOCKET_AUTH_RESULT",
-                    "status": "ERR_AUTH_FAILED"}, OpCode.TEXT)
+                f"Authentication error for session {session_id}: {str(e)}"
+            )
+            ws.send(
+                {"type": "SOCKET_AUTH_RESULT", "status": "ERR_AUTH_FAILED"}, OpCode.TEXT
+            )

oioioi/notifications/server/tests.py

@@ -199,6 +199,7 @@ async def test_on_ws_message_auth_success(self, mock_auth, mock_queue, mock_app)
         mock_auth.return_value.authenticate.return_value = "user_id"
 
         mock_ws = MagicMock()
+        mock_ws.get_user_data.return_value = {"user_id": None}
         message = json.dumps({"type": "SOCKET_AUTH", "session_id": "session_id"})
         await server.on_ws_message(mock_ws, message, OpCode.TEXT)
 
@@ -217,6 +218,7 @@ async def test_on_ws_message_auth_failure(self, mock_auth, mock_queue, mock_app)
         )
 
         mock_ws = MagicMock()
+        mock_ws.get_user_data.return_value = {"user_id": None}
         message = json.dumps({"type": "SOCKET_AUTH", "session_id": "session_id"})
         await server.on_ws_message(mock_ws, message, OpCode.TEXT)
 
@@ -226,6 +228,24 @@ async def test_on_ws_message_auth_failure(self, mock_auth, mock_queue, mock_app)
         mock_ws.send.assert_called_once_with(
             {"type": "SOCKET_AUTH_RESULT", "status": "ERR_AUTH_FAILED"}, OpCode.TEXT
         )
+        
+    async def test_on_ws_message_auth_multiple_times(self, mock_auth, mock_queue, mock_app):
+        server = Server(self.port, self.amqp_url, self.auth_url)
+
+        mock_ws = MagicMock()
+        # Set up the mock websocket to return user data indicating it's already authenticated
+        mock_ws.get_user_data.return_value = {"user_id": "user_id"}
+        
+        message = json.dumps({"type": "SOCKET_AUTH", "session_id": "session_id"})
+        await server.on_ws_message(mock_ws, message, OpCode.TEXT)
+
+        # Verify authentication was not attempted again
+        mock_auth.return_value.authenticate.assert_not_called()
+        mock_queue.return_value.subscribe.assert_not_called()
+        mock_ws.subscribe.assert_not_called()
+        mock_ws.send.assert_called_once_with(
+            {"type": "SOCKET_AUTH_RESULT", "status": "ERR_AUTH_FAILED"}, OpCode.TEXT
+        )
 
     async def test_on_ws_close(self, mock_auth, mock_queue, mock_app):
         server = Server(self.port, self.amqp_url, self.auth_url)