From a728d39fce8d6f943190421341678a99487fa2e2 Mon Sep 17 00:00:00 2001
From: Cizz22 <cisatraa@gmail.com>
Date: Thu, 12 Mar 2026 12:12:06 +0700
Subject: [PATCH] feat: Add CSRF protection middleware to the application.

---
 src/csrf_protect.py | 83 +++++++++++++++++++++++++++++++++++++++++++++
 src/main.py         |  5 +++
 2 files changed, 88 insertions(+)
 create mode 100644 src/csrf_protect.py

diff --git a/src/csrf_protect.py b/src/csrf_protect.py
new file mode 100644
index 0000000..728b231
--- /dev/null
+++ b/src/csrf_protect.py
@@ -0,0 +1,83 @@
+
+import secrets
+import logging
+from typing import Optional, List
+from fastapi import Request, HTTPException
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import Response
+
+log = logging.getLogger(__name__)
+
+class CSRFProtectMiddleware(BaseHTTPMiddleware):
+    """
+    CSRF Protection Middleware using Double Submit Cookie pattern.
+    
+    This middleware:
+    1. Ensures a 'csrf_token' cookie exists for every request.
+    2. For mutating requests (POST, PUT, PATCH, DELETE):
+       - Requires an 'X-CSRF-Token' header.
+       - Validates that the header value matches the 'csrf_token' cookie.
+    """
+    
+    def __init__(
+        self, 
+        app, 
+        cookie_name: str = "csrf_token",
+        header_name: str = "X-CSRF-Token",
+        excluded_paths: Optional[List[str]] = None,
+        safe_methods: List[str] = ("GET", "HEAD", "OPTIONS", "TRACE")
+    ):
+        super().__init__(app)
+        self.cookie_name = cookie_name
+        self.header_name = header_name
+        self.excluded_paths = excluded_paths or []
+        self.safe_methods = safe_methods
+
+    async def dispatch(self, request: Request, call_next):
+        # 1. Skip validation for "safe" methods (GET, HEAD, etc.) or excluded paths
+        is_safe_method = request.method in self.safe_methods
+        is_excluded = any(request.url.path.startswith(path) for path in self.excluded_paths)
+        
+        if not is_safe_method and not is_excluded:
+            # 2. Extract token from cookie and header
+            csrf_cookie = request.cookies.get(self.cookie_name)
+            csrf_header = request.headers.get(self.header_name.lower()) # Starlette lowercase headers
+            
+            # 3. Validate
+            if not csrf_cookie:
+                log.warning(f"CSRF validation failed: Missing cookie '{self.cookie_name}'")
+                raise HTTPException(
+                    status_code=403, 
+                    detail="CSRF validation failed: Missing session cookie."
+                )
+                
+            if not csrf_header:
+                log.warning(f"CSRF validation failed: Missing header '{self.header_name}'")
+                raise HTTPException(
+                    status_code=403, 
+                    detail=f"CSRF validation failed: Missing header '{self.header_name}'."
+                )
+                
+            if not secrets.compare_digest(csrf_cookie, csrf_header):
+                log.warning(f"CSRF validation failed: Token mismatch for path {request.url.path}")
+                raise HTTPException(
+                    status_code=403, 
+                    detail="CSRF validation failed: Token mismatch."
+                )
+
+        # 4. Proceed with the request
+        response: Response = await call_next(request)
+
+        # 5. Ensure the cookie is set if missing (e.g., on first GET request)
+        if not request.cookies.get(self.cookie_name):
+            new_token = secrets.token_urlsafe(32)
+            response.set_cookie(
+                key=self.cookie_name,
+                value=new_token,
+                httponly=False,  # Set to False so frontend JS can read it for Double Submit
+                samesite="lax",
+                secure=True,     # Recommended for production with HTTPS
+                path="/"
+            )
+            
+        return response
diff --git a/src/main.py b/src/main.py
index 0f1bbf0..8eeb0cf 100644
--- a/src/main.py
+++ b/src/main.py
@@ -29,6 +29,7 @@ from src.enums import ResponseStatus
 from src.exceptions import handle_exception
 from src.logging import configure_logging
 from src.middleware import RequestValidationMiddleware
+from src.csrf_protect import CSRFProtectMiddleware
 from src.rate_limiter import limiter
 from fastapi.exceptions import RequestValidationError
 from starlette.exceptions import HTTPException as StarletteHTTPException
@@ -124,6 +125,10 @@ security_headers_middleware(app)
 
 
 app.add_middleware(RequestValidationMiddleware)
+app.add_middleware(
+    CSRFProtectMiddleware,
+    excluded_paths=["/healthcheck"]
+)