Add authentication middleware and update auth controller for API key validation

src/app.py

@@ -6,6 +6,7 @@ from fastapi.middleware.cors import CORSMiddleware
 
 from src.controllers import api_router, websocket_router
 from src.config import logger, DatabaseConfig
+from src.middlewares import AuthenticationMiddleware
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -37,6 +38,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+app.add_middleware(AuthenticationMiddleware)
 
 @app.get("/")
 async def check_health():

src/controllers/__init__.py

@@ -1,4 +1,5 @@
-from fastapi import APIRouter
+from fastapi import APIRouter, Depends
+from fastapi.security import APIKeyHeader
 from ._auth_controller import AuthController
 from ._ai_voice_controller import AIVoiceController
 from ._ai_text_controller import AITextController
@@ -12,10 +13,12 @@ _ai_voice_controller = AIVoiceController()
 _ai_text_controller = AITextController()
 _pinecone_controller = PineconeController()
 
+API_KEY_HEADER = APIKeyHeader(name="Authorization", auto_error=False)
+
 websocket_router.include_router(_ai_voice_controller.websocket_router, prefix="/v1")
 websocket_router.include_router(_ai_text_controller.websocket_router, prefix="/v1")
 
-api_router.include_router(_auth_controller.api_router, prefix="/v1")
+api_router.include_router(_auth_controller.api_router, prefix="/v1", dependencies=[Depends(API_KEY_HEADER)])
 api_router.include_router(_ai_voice_controller.api_router, prefix="/v1")
 api_router.include_router(_ai_text_controller.api_router, prefix="/v1")
 api_router.include_router(_pinecone_controller.api_router, prefix="/v1")

src/controllers/_auth_controller.py

@@ -1,5 +1,4 @@
-from fastapi import APIRouter, HTTPException, Depends, status
-from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from fastapi import APIRouter, HTTPException, status, Request
 from src.services import AuthService  
 from src.schemas import (
     UserSignInSchema,
@@ -51,9 +50,9 @@ class AuthController:
              logger.error(f"Error during signin: {e}")
              raise HTTPException(status_code=500, detail="Error during signin")
     
-    async def _signout(self, token: HTTPAuthorizationCredentials = Depends(HTTPBearer())):
-        try: 
-            token = token.credentials.split("Bearer ")[1]
+    async def _signout(self, request: Request, ):
+        try:
+            token = request.headers.get("Authorization").split(" ")[1]
             await self._auth_service.sign_out(token)
             return status.HTTP_200_OK
         except HTTPException as e:

src/middlewares/__init__.py

@@ -0,0 +1,7 @@
+from ._authentication import AuthenticationMiddleware
+
+__all__ = [
+    "AuthenticationMiddleware",
+]
+__version__ = "0.1.0"
+__author__ = "Narinder Singh"

src/middlewares/_authentication.py

@@ -0,0 +1,52 @@
+import re
+from fastapi.responses import JSONResponse
+from fastapi import HTTPException, Request
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from src.utils import JWTUtil
+from src.services import SessionService
+
+
+EXEMPT_ROUTES = [
+    "/",
+    "/docs",
+    "/openapi.json",
+    "/api/v1/auth/login",
+    "/api/v1/auth/register",
+]
+class AuthenticationMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app):
+        super().__init__(app)
+        self._jwt = JWTUtil()
+        self._session_service = SessionService()
+
+    async def dispatch(self, request: Request, call_next):
+        if self._require_auth(request.url.path):
+            try:
+                user_payload =await self._validate_api_key(request.headers.get("Authorization"))
+                request.state.user = user_payload
+            except HTTPException as e:
+                return JSONResponse(status_code=e.status_code, content=e.detail)
+        return await call_next(request)
+
+    def _require_auth(self, path: str):
+        if str(path) in EXEMPT_ROUTES:
+            return False
+        return True
+
+    async def _validate_api_key(self, api_key: str):
+        if not api_key:
+            raise HTTPException(status_code=401, detail="No API key provided")
+        if not re.match(r"Bearer .+", api_key):
+            raise HTTPException(status_code=401, detail="Invalid API key")
+
+        token = api_key.split(" ")[1]
+
+        is_expired = await self._session_service.is_session_expired(token)
+        if is_expired:
+            raise HTTPException(status_code=401, detail="Token expired")
+
+        jwt_payload = self._jwt.validate_jwt(token)
+        if not jwt_payload:
+            raise HTTPException(status_code=401, detail="Invalid token")
+        return jwt_payload