session token test

relay.py

@@ -26,10 +26,14 @@ active_connections = {}
 # Dictionary to store undelivered messages
 message_store = {}
 
-async def register_client(websocket: WebSocket, username: str):
-    """Register a new client."""
+# Store session tokens (for simplicity, in-memory storage)
+session_tokens = {}
+
+async def register_client(websocket: WebSocket, username: str, token: str):
+    """Register a new client with a session token."""
     active_connections[username] = websocket
-    print(f"DEBUG: {username} connected.")
+    session_tokens[username] = token  # Store the session token
+    print(f"DEBUG: {username} connected with token.")
 
     # Deliver undelivered messages if any
     if username in message_store:
@@ -40,15 +44,17 @@ async def register_client(websocket: WebSocket, username: str):
         print(f"DEBUG: Cleared stored messages for {username}")
 
 async def unregister_client(username: str):
-    """Unregister a client."""
+    """Unregister a client and delete their session token."""
     if username in active_connections:
         del active_connections[username]
+        del session_tokens[username]  # Delete session token on disconnect
         print(f"DEBUG: {username} disconnected.")
 
 @app.websocket("/ws")
 async def relay_server(websocket: WebSocket):
     """Relay server handling WebSocket connections."""
     username = None
+    token = None
     try:
         await websocket.accept()
         print("DEBUG: WebSocket connection accepted.")
@@ -69,9 +75,17 @@ async def relay_server(websocket: WebSocket):
                 if response.status_code == 200:
                     print(f"DEBUG: Registration successful for {username}")
                     await websocket.send_text(json.dumps({"status": "success", "message": "Registration successful"}))
-                    # Register the client after successful signup
-                    await register_client(websocket, username)
-                    return  # Exit the websocket function after successful signup
+                    
+                    # Create session token upon registration
+                    token_response = await client.post(f"{SERVICES_URL}/login", json={"username": username, "password": password})
+                    if token_response.status_code == 200:
+                        token = token_response.json().get("token")
+                        await websocket.send_text(json.dumps({"status": "success", "token": token}))
+                        # Register the client after successful signup and token retrieval
+                        await register_client(websocket, username, token)
+                        return
+                    else:
+                        await websocket.send_text(json.dumps({"status": "error", "message": "Token retrieval failed"}))
                 else:
                     error_message = response.json().get("detail", "Registration failed")
                     print(f"DEBUG: Registration failed for {username}: {error_message}")
@@ -85,8 +99,9 @@ async def relay_server(websocket: WebSocket):
                 response = await client.post(f"{SERVICES_URL}/login", json={"username": username, "password": password})
                 if response.status_code == 200:
                     print(f"DEBUG: Authentication successful for {username}")
-                    await websocket.send_text(json.dumps({"status": "success", "message": "Authenticated"}))
-                    await register_client(websocket, username)
+                    token = response.json().get("token")
+                    await websocket.send_text(json.dumps({"status": "success", "message": "Authenticated", "token": token}))
+                    await register_client(websocket, username, token)
                 else:
                     error_message = response.json().get("detail", "Invalid credentials")
                     print(f"DEBUG: Authentication failed for {username}: {error_message}")
@@ -106,6 +121,13 @@ async def relay_server(websocket: WebSocket):
                 msg_content = msg_data.get("message")
                 timestamp = datetime.now().isoformat()
 
+                # Check session token for message validation
+                if recipient in session_tokens:
+                    received_token = msg_data.get("token")
+                    if received_token != session_tokens[recipient]:
+                        await websocket.send_text(json.dumps({"status": "error", "message": "Invalid session token"}))
+                        continue
+
                 # Create message object
                 message_obj = {
                     "from": username,