[kimi] Add WebSocket authentication for Matrix connections (#682) (#744)

src/config.py

@@ -155,6 +155,12 @@ class Settings(BaseSettings):
     # Example: "http://100.124.176.28:8080" or "https://alexanderwhitestone.com"
     matrix_frontend_url: str = ""  # Empty = disabled
 
+    # WebSocket authentication token for Matrix connections.
+    # When set, clients must provide this token via ?token= query param
+    # or in the first message as {"type": "auth", "token": "..."}.
+    # Empty/unset = auth disabled (dev mode).
+    matrix_ws_token: str = ""
+
     # Trusted hosts for the Host header check (TrustedHostMiddleware).
     # Set TRUSTED_HOSTS as a comma-separated list. Wildcards supported (e.g. "*.ts.net").
     # Defaults include localhost + Tailscale MagicDNS. Add your Tailscale IP if needed.

src/dashboard/routes/world.py

@@ -415,6 +415,50 @@ async def _heartbeat(websocket: WebSocket) -> None:
         logger.debug("Heartbeat stopped — connection gone")
 
 
+async def _authenticate_ws(websocket: WebSocket) -> bool:
+    """Authenticate WebSocket connection using matrix_ws_token.
+
+    Checks for token in query param ?token= first. If no query param,
+    accepts the connection and waits for first message with
+    {"type": "auth", "token": "..."}.
+
+    Returns True if authenticated (or if auth is disabled).
+    Returns False and closes connection with code 4001 if invalid.
+    """
+    token_setting = settings.matrix_ws_token
+
+    # Auth disabled in dev mode (empty/unset token)
+    if not token_setting:
+        return True
+
+    # Check query param first (can validate before accept)
+    query_token = websocket.query_params.get("token", "")
+    if query_token:
+        if query_token == token_setting:
+            return True
+        # Invalid token in query param - we need to accept to close properly
+        await websocket.accept()
+        await websocket.close(code=4001, reason="Invalid token")
+        return False
+
+    # No query token - accept and wait for auth message
+    await websocket.accept()
+
+    # Wait for auth message as first message
+    try:
+        raw = await websocket.receive_text()
+        data = json.loads(raw)
+        if data.get("type") == "auth" and data.get("token") == token_setting:
+            return True
+        # Invalid auth message
+        await websocket.close(code=4001, reason="Invalid token")
+        return False
+    except (json.JSONDecodeError, TypeError):
+        # Non-JSON first message without valid token
+        await websocket.close(code=4001, reason="Authentication required")
+        return False
+
+
 @router.websocket("/ws")
 async def world_ws(websocket: WebSocket) -> None:
     """Accept a Workshop client and keep it alive for state broadcasts.
@@ -423,8 +467,28 @@ async def world_ws(websocket: WebSocket) -> None:
     client never starts from a blank slate.  Incoming frames are parsed
     as JSON — ``visitor_message`` triggers a bark response.  A background
     heartbeat ping runs every 15 s to detect dead connections early.
+
+    Authentication:
+        - If matrix_ws_token is configured, clients must provide it via
+          ?token= query param or in the first message as
+          {"type": "auth", "token": "..."}.
+        - Invalid token results in close code 4001.
+        - Valid token receives a connection_ack message.
     """
-    await websocket.accept()
+    # Authenticate (may accept connection internally)
+    is_authed = await _authenticate_ws(websocket)
+    if not is_authed:
+        logger.info("World WS connection rejected — invalid token")
+        return
+
+    # Auth passed - accept if not already accepted
+    if websocket.client_state.name != "CONNECTED":
+        await websocket.accept()
+
+    # Send connection_ack if auth was required
+    if settings.matrix_ws_token:
+        await websocket.send_text(json.dumps({"type": "connection_ack"}))
+
     _ws_clients.append(websocket)
     logger.info("World WS connected — %d clients", len(_ws_clients))