fix(security): add JWT authentication to WebSocket endpoint

- Add authenticate_websocket() helper for token validation via query param
- Reject connections without token (4001) or with invalid token (4001)
- Validate repo exists before accepting connection (4004)
- Add TODO for repo ownership validation (needs user_id column)
- Add unit tests for WebSocket authentication
- Configure pytest for async tests

Fixes #6

Author: DevanshuNEU

backend/main.py

@@ -180,19 +180,56 @@ async def add_repository(
         raise HTTPException(status_code=400, detail=str(e))
 
 
+async def authenticate_websocket(websocket: WebSocket) -> Optional[dict]:
+    """
+    Authenticate WebSocket connection via query parameter token.
+    
+    WebSockets can't use Authorization headers during handshake,
+    so we pass the JWT token as a query parameter instead.
+    
+    Returns:
+        User dict if authenticated, None otherwise (connection closed with error)
+    """
+    token = websocket.query_params.get("token")
+    if not token:
+        await websocket.close(code=4001, reason="Missing authentication token")
+        return None
+    
+    try:
+        from services.auth import get_auth_service
+        auth_service = get_auth_service()
+        return auth_service.verify_jwt(token)
+    except Exception:
+        await websocket.close(code=4001, reason="Invalid or expired token")
+        return None
+
+
 @app.websocket("/ws/index/{repo_id}")
 async def websocket_index(websocket: WebSocket, repo_id: str):
-    """Real-time indexing with progress updates"""
+    """
+    Real-time repository indexing with progress updates.
+    
+    Requires JWT token passed as query parameter: ?token=<jwt>
+    Sends progress updates via JSON messages during indexing.
+    """
+    # Authenticate before accepting connection
+    user = await authenticate_websocket(websocket)
+    if not user:
+        return
+    
+    # TODO: Add repo ownership validation once user_id column exists in repos table
+    # For now, any authenticated user can index any repo they know the ID of
+    
+    # Validate repo exists before accepting connection
+    repo = repo_manager.get_repo(repo_id)
+    if not repo:
+        await websocket.close(code=4004, reason="Repository not found")
+        return
+    
+    # Connection authenticated and repo valid - accept
     await websocket.accept()
 
     try:
-        # Get repo info
-        repo = repo_manager.get_repo(repo_id)
-        if not repo:
-            await websocket.send_json({"error": "Repository not found"})
-            await websocket.close()
-            return
-        
         repo_manager.update_status(repo_id, "indexing")
 
         # Index with progress callback

backend/pytest.ini

@@ -3,6 +3,7 @@ testpaths = tests
 python_files = test_*.py
 python_classes = Test*
 python_functions = test_*
+asyncio_mode = auto
 addopts = 
     -v
     --tb=short

backend/tests/test_websocket_auth.py

@@ -0,0 +1,89 @@
+"""
+WebSocket Authentication Tests
+Tests for issue #6: Secure WebSocket endpoints
+"""
+import pytest
+from unittest.mock import MagicMock, AsyncMock, patch
+
+
+class TestWebSocketAuthentication:
+    """Integration tests for WebSocket authentication via query parameter token"""
+    
+    def test_websocket_rejects_missing_token(self, client):
+        """WebSocket should reject connections without token (4001)"""
+        with pytest.raises(Exception):
+            with client.websocket_connect("/ws/index/test-repo-id"):
+                pass
+    
+    def test_websocket_rejects_invalid_token(self, client):
+        """WebSocket should reject connections with invalid token (4001)"""
+        with pytest.raises(Exception):
+            with client.websocket_connect("/ws/index/test-repo-id?token=invalid-token"):
+                pass
+    
+    def test_websocket_rejects_nonexistent_repo(self, client):
+        """WebSocket should reject if repo doesn't exist (4004)"""
+        with patch('main.authenticate_websocket') as mock_auth:
+            mock_auth.return_value = {"user_id": "test-user", "email": "test@example.com"}
+            
+            with pytest.raises(Exception):
+                with client.websocket_connect("/ws/index/nonexistent-repo?token=valid"):
+                    pass
+
+
+class TestAuthenticateWebsocketFunction:
+    """Unit tests for the authenticate_websocket helper"""
+    
+    @pytest.mark.asyncio
+    async def test_returns_none_without_token(self):
+        """Should return None and close connection if no token provided"""
+        from main import authenticate_websocket
+        
+        mock_ws = MagicMock()
+        mock_ws.query_params = {}
+        mock_ws.close = AsyncMock()
+        
+        result = await authenticate_websocket(mock_ws)
+        
+        assert result is None
+        mock_ws.close.assert_called_once_with(code=4001, reason="Missing authentication token")
+    
+    @pytest.mark.asyncio
+    async def test_returns_none_with_invalid_token(self):
+        """Should return None and close connection if token is invalid"""
+        from main import authenticate_websocket
+        
+        mock_ws = MagicMock()
+        mock_ws.query_params = {"token": "invalid-token"}
+        mock_ws.close = AsyncMock()
+        
+        with patch('services.auth.get_auth_service') as mock_get_service:
+            mock_service = MagicMock()
+            mock_service.verify_jwt.side_effect = Exception("Invalid token")
+            mock_get_service.return_value = mock_service
+            
+            result = await authenticate_websocket(mock_ws)
+        
+        assert result is None
+        mock_ws.close.assert_called_once_with(code=4001, reason="Invalid or expired token")
+    
+    @pytest.mark.asyncio
+    async def test_returns_user_with_valid_token(self):
+        """Should return user dict if token is valid"""
+        from main import authenticate_websocket
+        
+        mock_ws = MagicMock()
+        mock_ws.query_params = {"token": "valid-jwt-token"}
+        mock_ws.close = AsyncMock()
+        
+        expected_user = {"user_id": "user-123", "email": "test@example.com"}
+        
+        with patch('services.auth.get_auth_service') as mock_get_service:
+            mock_service = MagicMock()
+            mock_service.verify_jwt.return_value = expected_user
+            mock_get_service.return_value = mock_service
+            
+            result = await authenticate_websocket(mock_ws)
+        
+        assert result == expected_user
+        mock_ws.close.assert_not_called()