Fix constant-time auth token validation

Author: Mateusz

src/core/security/middleware.py

@@ -6,6 +6,7 @@
 import json
 import logging
 import math
+import secrets
 import time
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
@@ -28,6 +29,29 @@
 logger = logging.getLogger(__name__)
 
 
+def _constant_time_equals(candidate: str | None, expected: str) -> bool:
+    """Compare secrets without leaking matching-prefix timing."""
+
+    return bool(
+        candidate
+        and isinstance(expected, str)
+        and secrets.compare_digest(candidate, expected)
+    )
+
+
+def _constant_time_member(candidate: str | None, expected_values: set[str]) -> bool:
+    """Check membership using constant-time comparisons for each expected value."""
+
+    if not candidate:
+        return False
+
+    matched = False
+    for expected in expected_values:
+        if isinstance(expected, str) and secrets.compare_digest(candidate, expected):
+            matched = True
+    return matched
+
+
 @dataclass
 class _BruteForceRecord:
     """Track failed attempts and blocking metadata for a client IP."""
@@ -219,7 +243,7 @@ async def dispatch(
         all_valid_keys: set[str] = self.valid_keys | app_state_keys
 
         method = request.method
-        if not api_key or api_key not in all_valid_keys:
+        if not _constant_time_member(api_key, all_valid_keys):
             logger.warning(
                 "Invalid or missing API key for %s %s from client %s",
                 method,
@@ -525,7 +549,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
             f"API Key authentication is enabled key_count={len(all_valid_keys)}"
         )
         method = scope.get("method", "UNKNOWN")
-        if not api_key or api_key not in all_valid_keys:
+        if not _constant_time_member(api_key, all_valid_keys):
             logger.warning(
                 "Invalid or missing API key for %s %s from client %s",
                 method,
@@ -659,7 +683,7 @@ async def dispatch(
         method = request.method
 
         # Validate the token
-        if not token or token != self.valid_token:
+        if not _constant_time_equals(token, self.valid_token):
             logger.warning(
                 "Invalid or missing auth token for %s %s from client %s",
                 method,
@@ -763,7 +787,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         method = scope.get("method", "UNKNOWN")
 
         # Validate the token
-        if not token or token != self.valid_token:
+        if not _constant_time_equals(token, self.valid_token):
             logger.warning(
                 "Invalid or missing auth token for %s %s from client %s",
                 method,

tests/unit/core/test_authentication_di.py

@@ -5,9 +5,10 @@
 refactored to use proper dependency injection instead of direct app.state access.
 """
 
-import json
-import os
-from unittest.mock import AsyncMock, MagicMock, patch
+import json
+import os
+from typing import Any, cast
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from fastapi import FastAPI, Request, Response
@@ -134,8 +135,8 @@ async def test_valid_query_key(self, api_key_middleware, mock_request):
         call_next.assert_called_once_with(mock_request)
         assert response == "next_response"
 
-    async def test_invalid_key(self, api_key_middleware, mock_request):
-        """Test that an invalid API key is rejected."""
+    async def test_invalid_key(self, api_key_middleware, mock_request):
+        """Test that an invalid API key is rejected."""
         # Setup
         mock_request.headers = {"Authorization": "Bearer invalid-key"}
         call_next = AsyncMock(return_value="next_response")
@@ -146,9 +147,31 @@ async def test_invalid_key(self, api_key_middleware, mock_request):
         # Verify
         call_next.assert_not_called()
         assert response.status_code == 401
-        assert (
-            response.body == f'{{"detail":"{HTTP_401_UNAUTHORIZED_MESSAGE}"}}'.encode()
-        )
+        assert (
+            response.body == f'{{"detail":"{HTTP_401_UNAUTHORIZED_MESSAGE}"}}'.encode()
+        )
+
+    async def test_valid_bearer_key_uses_constant_time_comparison(
+        self, api_key_middleware, mock_request, monkeypatch
+    ):
+        """API key validation should avoid short-circuit string membership."""
+
+        compared_values: list[tuple[str, str]] = []
+
+        def compare_digest(candidate: str, expected: str) -> bool:
+            compared_values.append((candidate, expected))
+            return candidate == expected
+
+        monkeypatch.setattr(
+            "src.core.security.middleware.secrets.compare_digest", compare_digest
+        )
+        mock_request.headers = {"Authorization": "Bearer test-key"}
+        call_next = AsyncMock(return_value="next_response")
+
+        response = await api_key_middleware.dispatch(mock_request, call_next)
+
+        assert response == "next_response"
+        assert ("test-key", "test-key") in compared_values
 
     async def test_missing_key(self, api_key_middleware, mock_request):
         """Test that a missing API key is rejected."""
@@ -278,7 +301,7 @@ async def _attempt(header_value: str) -> Response:
         blocked = await _attempt("bad-3")
         assert blocked.status_code == 429
         assert blocked.headers.get("Retry-After") == "10"
-        payload = json.loads(blocked.body.decode())
+        payload = json.loads(bytes(blocked.body).decode())
         assert payload["retry_after_seconds"] == 10
 
         # After the wait expires, another invalid attempt is allowed
@@ -289,7 +312,7 @@ async def _attempt(header_value: str) -> Response:
         blocked_again = await _attempt("bad-5")
         assert blocked_again.status_code == 429
         assert blocked_again.headers.get("Retry-After") == "20"
-        payload = json.loads(blocked_again.body.decode())
+        payload = json.loads(bytes(blocked_again.body).decode())
         assert payload["retry_after_seconds"] == 20
 
         # Provide a valid key to reset the tracker
@@ -340,7 +363,7 @@ def advance(self, seconds: float) -> None:
 
         disable_flag = {"value": False}
 
-        def get_setting(key: str, default=None):
+        def get_setting(key: str, default: Any = None) -> Any:
             if key == "disable_auth":
                 return disable_flag["value"]
             if key == "client_api_key":
@@ -355,7 +378,7 @@ def get_setting(key: str, default=None):
 
         app_state_service = MagicMock(spec=IApplicationState)
         app_state_service.get_setting.side_effect = get_setting
-        middleware.app_state_service = app_state_service
+        cast(Any, middleware).app_state_service = app_state_service
         mock_request.app.state.service_provider.get_service.return_value = (
             app_state_service
         )
@@ -405,8 +428,8 @@ async def test_valid_token(self, auth_token_middleware, mock_request):
         call_next.assert_called_once_with(mock_request)
         assert response == "next_response"
 
-    async def test_invalid_token(self, auth_token_middleware, mock_request):
-        """Test that an invalid auth token is rejected."""
+    async def test_invalid_token(self, auth_token_middleware, mock_request):
+        """Test that an invalid auth token is rejected."""
         # Setup
         mock_request.headers = {"X-Auth-Token": "invalid-token"}
         call_next = AsyncMock(return_value="next_response")
@@ -417,9 +440,31 @@ async def test_invalid_token(self, auth_token_middleware, mock_request):
         # Verify
         call_next.assert_not_called()
         assert response.status_code == 401
-        assert (
-            response.body == f'{{"detail":"{HTTP_401_UNAUTHORIZED_MESSAGE}"}}'.encode()
-        )
+        assert (
+            response.body == f'{{"detail":"{HTTP_401_UNAUTHORIZED_MESSAGE}"}}'.encode()
+        )
+
+    async def test_valid_token_uses_constant_time_comparison(
+        self, auth_token_middleware, mock_request, monkeypatch
+    ):
+        """Auth token validation should avoid short-circuit string equality."""
+
+        compared_values: list[tuple[str, str]] = []
+
+        def compare_digest(candidate: str, expected: str) -> bool:
+            compared_values.append((candidate, expected))
+            return candidate == expected
+
+        monkeypatch.setattr(
+            "src.core.security.middleware.secrets.compare_digest", compare_digest
+        )
+        mock_request.headers = {"X-Auth-Token": "test-token"}
+        call_next = AsyncMock(return_value="next_response")
+
+        response = await auth_token_middleware.dispatch(mock_request, call_next)
+
+        assert response == "next_response"
+        assert compared_values == [("test-token", "test-token")]
 
     async def test_missing_token(self, auth_token_middleware, mock_request):
         """Test that a missing auth token is rejected."""