Fix : CSRF error

Author: Jones-peter

Tests/test_csrf_json.py

@@ -16,35 +16,51 @@
 
 async def run_csrf_test():
     """
-    Tests that CSRF protection works correctly for JSON endpoints.
+    Tests that CSRF protection works correctly for various request types.
     """
-    print("--- Starting CSRF JSON Test ---")
+    print("--- Starting CSRF Logic Test ---")
     async with httpx.AsyncClient(base_url=BASE_URL) as client:
         try:
-            # 1. Make a GET request to a page to get a CSRF token
+            # 1. Make a GET request to a page to get a CSRF token from the cookie
             print("Step 1: Getting CSRF token from homepage...")
             get_response = await client.get("/")
             get_response.raise_for_status()
             assert "csrf_token" in client.cookies, "CSRF token not found in cookie"
             csrf_token = client.cookies["csrf_token"]
-            print("   [PASS] CSRF token received.")
+            print(f"   [PASS] CSRF token received: {csrf_token[:10]}...")
 
-            # 2. Send a POST request to the JSON endpoint WITHOUT a CSRF token
+            # 2. Test POST without any CSRF token (should fail)
             print("\nStep 2: Testing POST to /api/test without CSRF token (expecting 403)...")
-            payload = {"message": "hello"}
-            fail_response = await client.post("/api/test", json=payload)
+            fail_response = await client.post("/api/test", json={"message": "hello"})
             assert fail_response.status_code == 403, f"Expected status 403, but got {fail_response.status_code}"
             assert "CSRF token missing or invalid" in fail_response.text
             print("   [PASS] Request was correctly forbidden.")
 
-            # 3. Send a POST request to the JSON endpoint WITH the correct CSRF token
-            print("\nStep 3: Testing POST to /api/test with CSRF token (expecting 200)...")
+            # 3. Test POST with CSRF token in JSON body (should pass)
+            print("\nStep 3: Testing POST to /api/test with CSRF token in JSON body (expecting 200)...")
             payload_with_token = {"message": "hello", "csrf_token": csrf_token}
-            success_response = await client.post("/api/test", json=payload_with_token)
-            assert success_response.status_code == 200, f"Expected status 200, but got {success_response.status_code}"
-            response_json = success_response.json()
-            assert response_json["message"] == "hello"
-            print("   [PASS] Request was successful.")
+            success_response_body = await client.post("/api/test", json=payload_with_token)
+            assert success_response_body.status_code == 200, f"Expected status 200, but got {success_response_body.status_code}"
+            assert success_response_body.json()["message"] == "hello"
+            print("   [PASS] Request with token in body was successful.")
+
+            # 4. Test POST with CSRF token in header (should pass)
+            print("\nStep 4: Testing POST to /api/test with CSRF token in header (expecting 200)...")
+            headers = {"X-CSRF-Token": csrf_token}
+            success_response_header = await client.post("/api/test", json={"message": "world"}, headers=headers)
+            assert success_response_header.status_code == 200, f"Expected status 200, but got {success_response_header.status_code}"
+            assert success_response_header.json()["message"] == "world"
+            print("   [PASS] Request with token in header was successful.")
+
+            # 5. Test empty-body POST with CSRF token in header (should pass validation, then redirect)
+            print("\nStep 5: Testing empty-body POST to /logout with CSRF token in header (expecting 302)...")
+            # Note: The /logout endpoint redirects after success, so we expect a 302
+            # We disable auto-redirects to verify the 302 status directly
+            empty_body_response = await client.post("/logout", headers=headers, follow_redirects=False)
+            
+            # If we got a 403, the CSRF check failed. If we got a 302, it passed!
+            assert empty_body_response.status_code == 302, f"Expected status 302 (Redirect), but got {empty_body_response.status_code}. (403 means CSRF failed)"
+            print("   [PASS] Empty-body request passed CSRF check and redirected.")
 
         except Exception as e:
             print(f"\n--- TEST FAILED ---")
@@ -53,7 +69,7 @@ async def run_csrf_test():
             traceback.print_exc()
             return False
 
-    print("\n--- TEST PASSED ---")
+    print("\n--- ALL CSRF TESTS PASSED ---")
     return True
 
 
@@ -89,13 +105,15 @@ def main():
         print("\nStopping test server...")
         server_process.terminate()
         # Get remaining output
-        stdout, stderr = server_process.communicate(timeout=5)
-        
-        print("\n--- Server Output ---")
-        print("STDOUT:")
-        print(stdout)
-        print("\nSTDERR:")
-        print(stderr)
+        try:
+            stdout, stderr = server_process.communicate(timeout=5)
+            print("\n--- Server Output ---")
+            print("STDOUT:")
+            print(stdout)
+            print("\nSTDERR:")
+            print(stderr)
+        except subprocess.TimeoutExpired:
+            print("Server did not terminate gracefully.")
 
         if not test_passed:
             print("\nExiting with status 1 due to test failure.")

jsweb/middleware.py

@@ -2,6 +2,7 @@
 import logging
 from .static import serve_static
 from .response import Forbidden
+import json
 
 logger = logging.getLogger(__name__)
 
@@ -30,10 +31,11 @@ class CSRFMiddleware(Middleware):
     """
     Middleware to protect against Cross-Site Request Forgery (CSRF) attacks.
 
-    This middleware checks for a valid CSRF token in POST, PUT, PATCH, and DELETE
-    requests. It supports both form-based and JSON-based requests.
+    This middleware enforces CSRF protection for all state-changing HTTP methods
+    (POST, PUT, PATCH, DELETE). It requires a valid CSRF token to be present
+    in the request, either in the 'X-CSRF-Token' header or in the request body
+    (JSON or Form Data).
     """
-
     async def __call__(self, scope, receive, send):
         """
         Validates the CSRF token for state-changing HTTP methods.
@@ -53,34 +55,48 @@ async def __call__(self, scope, receive, send):
         req = scope['jsweb.request']
 
         if req.method in ("POST", "PUT", "PATCH", "DELETE"):
-            token = await self._get_token_from_request(req)
             cookie_token = req.cookies.get("csrf_token")
-
-            if not token or not cookie_token or not secrets.compare_digest(token, cookie_token):
-                logger.error("CSRF VALIDATION FAILED. Tokens do not match or are missing.")
+            submitted_token = None
+
+            # 1. Check header first (Best practice for AJAX/APIs)
+            submitted_token = req.headers.get("x-csrf-token")
+
+            # 2. If no header token, check the body based on content type
+            if not submitted_token:
+                content_type = req.headers.get("content-type", "")
+                
+                if "application/json" in content_type:
+                    try:
+                        # Request.json() safely returns {} for empty/invalid bodies
+                        data = await req.json()
+                        submitted_token = data.get("csrf_token")
+                    except Exception:
+                        # If JSON parsing fails, we treat it as no token found
+                        pass
+                
+                elif "application/x-www-form-urlencoded" in content_type or "multipart/form-data" in content_type:
+                    try:
+                        # Request.form() safely returns {} for empty/non-form bodies
+                        form = await req.form()
+                        submitted_token = form.get("csrf_token")
+                    except Exception:
+                        # If form parsing fails, we treat it as no token found
+                        pass
+
+            # 3. Perform the validation
+            # Both the cookie token and the submitted token MUST be present and match.
+            if not cookie_token or not submitted_token or not secrets.compare_digest(submitted_token, cookie_token):
+                logger.warning(
+                    f"CSRF validation failed for {req.method} {req.path}. "
+                    f"Cookie set: {'Yes' if cookie_token else 'No'}, "
+                    f"Token submitted: {'Yes' if submitted_token else 'No'}."
+                )
                 response = Forbidden("CSRF token missing or invalid.")
                 await response(scope, receive, send)
                 return
 
         await self.app(scope, receive, send)
 
-    async def _get_token_from_request(self, req):
-        """
-        Extracts the CSRF token from the request, handling both JSON and form data.
-        """
-        content_type = req.headers.get("content-type", "")
-        if "application/json" in content_type:
-            try:
-                data = await req.json()
-                return data.get("csrf_token")
-            except Exception:
-                # In case of malformed JSON, treat as if no token was sent
-                return None
-        else:
-            # Fallback for form data
-            form = await req.form()
-            return form.get("csrf_token")
-
 class StaticFilesMiddleware(Middleware):
     """
     Middleware for serving static files.

jsweb/request.py

@@ -34,6 +34,7 @@ def __init__(self, scope, receive, app):
         self.path = self.scope.get("path", "/")
         self.query = self._parse_query(self.scope.get("query_string", b"").decode())
         self.headers = self._parse_headers(self.scope.get("headers", []))
+        self.content_type = self.headers.get("content-type", "")
         self.cookies = self._parse_cookies(self.headers)
         self.user = None
 
@@ -101,8 +102,7 @@ async def json(self):
             dict: The parsed JSON data.
         """
         if self._json is None:
-            content_type = self.headers.get("content-type", "")
-            if "application/json" in content_type:
+            if "application/json" in self.content_type:
                 try:
                     body_bytes = await self.body()
                     self._json = json.loads(body_bytes) if body_bytes else {}