Merge pull request #19 from DevanshuNEU/fix/issue-6-websocket-auth

Fix/issue 6 websocket auth

Author: DevanshuNEU

.gitignore

@@ -11,8 +11,11 @@ dist/
 # Environment
 .env
 .env.local
+.env.dev
+.env.prod
 backend/.env
 backend/.env.local
+frontend/.env.local
 mcp-server/.env
 mcp-server/.env.local
 
@@ -30,3 +33,4 @@ backend/repos/
 
 # MCP Server
 mcp-server/__pycache__/
+*.code-workspace

Makefile

@@ -5,76 +5,93 @@ help:
 	@echo "CodeIntel - Development Commands"
 	@echo ""
 	@echo "Local Development:"
-	@echo "  make dev          - Start all services with hot reload"
-	@echo "  make prod         - Start production-like environment"
-	@echo "  make build        - Build Docker images"
+	@echo "  make dev          - Start local dev (uses .env.dev)"
+	@echo "  make dev-prod     - Test prod config locally (uses .env.prod)"
 	@echo "  make stop         - Stop all services"
 	@echo "  make clean        - Stop and remove all containers/volumes"
 	@echo "  make logs         - View all logs"
+	@echo "  make health       - Check service health"
 	@echo ""
 	@echo "Testing:"
 	@echo "  make test         - Run backend tests"
-	@echo "  make test-watch   - Run tests in watch mode"
+	@echo "  make test-ws      - Run WebSocket auth tests only"
 	@echo "  make coverage     - Run tests with coverage report"
 	@echo ""
 	@echo "Deployment:"
 	@echo "  make deploy-backend  - Deploy backend to Railway"
 	@echo "  make deploy-frontend - Deploy frontend to Vercel"
-	@echo "  make deploy-all      - Deploy both backend and frontend"
 
-# Development with hot reload
+# ============================================
+# LOCAL DEVELOPMENT
+# ============================================
+
+# Development with .env.dev
 dev:
-	docker compose -f docker-compose.dev.yml up -d
+	@echo "🚀 Starting LOCAL DEV environment..."
+	@cp .env.dev .env
+	docker compose up -d --build
 	@echo ""
 	@echo "✅ Development environment started!"
 	@echo "   Backend:  http://localhost:8000"
-	@echo "   Docs:     http://localhost:8000/docs"
+	@echo "   API Docs: http://localhost:8000/docs"
 	@echo "   Frontend: http://localhost:3000"
 	@echo "   Redis:    localhost:6379"
 	@echo ""
 	@echo "View logs: make logs"
 
-# Production-like environment
-prod:
-	docker compose up -d
+# Test production config locally (uses .env.prod)
+dev-prod:
+	@echo "🚀 Starting LOCAL environment with PROD config..."
+	@cp .env.prod .env
+	docker compose up -d --build
 	@echo ""
-	@echo "✅ Production environment started!"
+	@echo "✅ Prod-config environment started!"
 	@echo "   Backend:  http://localhost:8000"
 	@echo "   Frontend: http://localhost:3000"
 
-# Build images
-build:
-	docker compose build
-
 # Stop services
 stop:
 	docker compose down
-	docker compose -f docker-compose.dev.yml down
+	@echo "✅ Services stopped"
 
 # Clean everything (including volumes)
 clean:
-	docker compose down -v
-	docker compose -f docker-compose.dev.yml down -v
+	docker compose down -v --remove-orphans
 	@echo "✅ Cleaned all containers and volumes"
 
 # View logs
 logs:
 	docker compose logs -f
 
-# Run backend tests
+# Logs for specific service
+logs-backend:
+	docker compose logs -f backend
+
+logs-frontend:
+	docker compose logs -f frontend
+
+# ============================================
+# TESTING
+# ============================================
+
+# Run all backend tests
 test:
-	cd backend && python -m pytest tests/ -v
+	cd backend && python3 -m pytest tests/ -v --no-cov
 
-# Run tests in watch mode
-test-watch:
-	cd backend && python -m pytest tests/ -v --looponfail
+# Run WebSocket auth tests only
+test-ws:
+	cd backend && python3 -m pytest tests/test_websocket_auth.py -v --no-cov
 
 # Run tests with coverage
 coverage:
-	cd backend && python -m pytest tests/ --cov=. --cov-report=html --cov-report=term
+	cd backend && python3 -m pytest tests/ --cov=. --cov-report=html --cov-report=term
 	@echo ""
 	@echo "Coverage report: backend/htmlcov/index.html"
 
+# ============================================
+# DEPLOYMENT
+# ============================================
+
 # Deploy backend to Railway
 deploy-backend:
 	@echo "🚀 Deploying backend to Railway..."
@@ -91,15 +108,16 @@ deploy-frontend:
 deploy-all: deploy-backend deploy-frontend
 	@echo "✅ All services deployed!"
 
-# Quick restart of backend (for dev)
-restart-backend:
-	docker compose restart backend
-	@echo "✅ Backend restarted"
+# ============================================
+# UTILITIES
+# ============================================
 
-# Quick restart of frontend (for dev)
-restart-frontend:
-	docker compose restart frontend
-	@echo "✅ Frontend restarted"
+# Check service health
+health:
+	@echo "Checking services..."
+	@curl -s http://localhost:8000/health | python3 -m json.tool 2>/dev/null || echo "❌ Backend not responding"
+	@curl -s -o /dev/null -w "" http://localhost:3000 && echo "✅ Frontend is up" || echo "❌ Frontend not responding"
+	@docker compose exec -T redis redis-cli ping 2>/dev/null | grep -q PONG && echo "✅ Redis is up" || echo "❌ Redis not responding"
 
 # Shell into backend container
 shell-backend:
@@ -109,27 +127,12 @@ shell-backend:
 shell-redis:
 	docker compose exec redis redis-cli
 
-# View Redis stats
-redis-stats:
-	docker compose exec redis redis-cli INFO
-
-# Check service health
-health:
-	@echo "Checking services..."
-	@curl -s http://localhost:8000/health | python -m json.tool || echo "❌ Backend not responding"
-	@curl -s http://localhost:3000 > /dev/null && echo "✅ Frontend is up" || echo "❌ Frontend not responding"
-	@docker compose exec redis redis-cli ping > /dev/null && echo "✅ Redis is up" || echo "❌ Redis not responding"
-
-# Install dependencies (local dev without Docker)
-install-backend:
-	cd backend && pip install -r requirements.txt
-
-install-frontend:
-	cd frontend && npm install
-
-# Run locally without Docker
-run-backend-local:
-	cd backend && uvicorn main:app --reload --port 8000
+# Quick rebuild backend only
+rebuild-backend:
+	docker compose up -d --build backend
+	@echo "✅ Backend rebuilt and restarted"
 
-run-frontend-local:
-	cd frontend && npm run dev
+# Quick rebuild frontend only  
+rebuild-frontend:
+	docker compose up -d --build frontend
+	@echo "✅ Frontend rebuilt and restarted"

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()