diff --git a/src/api/middleware.py b/src/api/middleware.py index c20092984..960842d06 100644 --- a/src/api/middleware.py +++ b/src/api/middleware.py @@ -2,7 +2,8 @@ import time import logging -from typing import Callable +from typing import Callable, Optional + from starlette.middleware.base import BaseHTTPMiddleware from starlette.requests import Request from starlette.responses import Response @@ -19,6 +20,33 @@ async def dispatch(self, request: Request, call_next: Callable) -> Response: return await call_next(request) +class CacheControlMiddleware(BaseHTTPMiddleware): + """Set Cache-Control: no-store on authenticated JSON responses. + + Prevents sensitive orchestration data from being cached by proxies or browsers. + Applied after the response is generated so it cannot be skipped on any code path. + """ + + async def dispatch(self, request: Request, call_next: Callable) -> Response: + response = await call_next(request) + + # Only apply to authenticated requests + auth_header = request.headers.get("Authorization", "") + has_session = request.headers.get("Cookie", "").startswith("session=") + if not auth_header.startswith("Bearer ") and not has_session: + return response + + # Only apply to JSON responses (application/json or any variant) + content_type = response.headers.get("content-type", "") + if "json" not in content_type.lower(): + return response + + # Set cache-control to prevent caching of sensitive data + response.headers["Cache-Control"] = "no-store" + response.headers["Pragma"] = "no-cache" + return response + + class RateLimitMiddleware(BaseHTTPMiddleware): def __init__(self, app, max_requests: int = 100, window: int = 60): super().__init__(app) @@ -177,3 +205,4 @@ async def dispatch(self, request: Request, call_next: Callable) -> Response: # 2026-03-27T12:58:53 update # 2026-05-12T17:19:36 update + diff --git a/src/api/server.py b/src/api/server.py index c6ae4bf2d..dfb16f9f6 100644 --- a/src/api/server.py +++ b/src/api/server.py @@ -8,7 +8,7 @@ from fastapi.middleware.trustedhost import TrustedHostMiddleware from .routes import router -from .middleware import AuthMiddleware, RateLimitMiddleware, LoggingMiddleware +from .middleware import AuthMiddleware, CacheControlMiddleware, RateLimitMiddleware, LoggingMiddleware def create_app(config: Dict = None) -> FastAPI: @@ -31,6 +31,7 @@ def create_app(config: Dict = None) -> FastAPI: app.add_middleware(TrustedHostMiddleware, allowed_hosts=os.getenv("TRUSTED_HOSTS", "*").split(",")) app.add_middleware(AuthMiddleware) + app.add_middleware(CacheControlMiddleware) app.add_middleware(RateLimitMiddleware) app.add_middleware(LoggingMiddleware) @@ -157,3 +158,4 @@ async def health(): # 2026-04-28T08:38:14 update # 2026-05-19T18:09:43 update + diff --git a/tests/test_cache_control.py b/tests/test_cache_control.py new file mode 100644 index 000000000..24cd91757 --- /dev/null +++ b/tests/test_cache_control.py @@ -0,0 +1,118 @@ +"""Tests for CacheControlMiddleware — verifies Cache-Control headers on authenticated JSON responses.""" + +import pytest +from unittest.mock import AsyncMock, MagicMock + +from starlette.requests import Request +from starlette.responses import Response, JSONResponse + +from src.api.middleware import CacheControlMiddleware + + +@pytest.mark.asyncio +async def test_authenticated_json_response_gets_no_store(): + """Authenticated JSON responses should have Cache-Control: no-store.""" + app = AsyncMock() + app.return_value = JSONResponse({"data": "secret"}) + + scope = { + "type": "http", + "method": "GET", + "path": "/api/v2/agents", + "headers": [ + (b"authorization", b"Bearer some-token"), + (b"content-type", b"application/json"), + ], + } + request = Request(scope) + middleware = CacheControlMiddleware(lambda: None) + + response = await middleware.dispatch(request, app) + assert response.headers.get("Cache-Control") == "no-store" + assert response.headers.get("Pragma") == "no-cache" + + +@pytest.mark.asyncio +async def test_unauthenticated_request_unchanged(): + """Unauthenticated requests should not have Cache-Control added.""" + app = AsyncMock() + app.return_value = JSONResponse({"public": "data"}) + + scope = { + "type": "http", + "method": "GET", + "path": "/health", + "headers": [(b"content-type", b"application/json")], + } + request = Request(scope) + middleware = CacheControlMiddleware(lambda: None) + + response = await middleware.dispatch(request, app) + assert "cache-control" not in response.headers + + +@pytest.mark.asyncio +async def test_non_json_response_unchanged(): + """Non-JSON authenticated responses should not have Cache-Control added.""" + app = AsyncMock() + app.return_value = Response(content="", media_type="text/html") + + scope = { + "type": "http", + "method": "GET", + "path": "/api/v2/agents", + "headers": [ + (b"authorization", b"Bearer token"), + (b"content-type", b"text/html"), + ], + } + request = Request(scope) + middleware = CacheControlMiddleware(lambda: None) + + response = await middleware.dispatch(request, app) + assert "cache-control" not in response.headers + + +@pytest.mark.asyncio +async def test_session_cookie_triggers_no_store(): + """Requests with session cookies should also get Cache-Control: no-store.""" + app = AsyncMock() + app.return_value = JSONResponse({"data": "session-data"}) + + scope = { + "type": "http", + "method": "GET", + "path": "/api/v2/agents", + "headers": [ + (b"cookie", b"session=abc123"), + (b"content-type", b"application/json"), + ], + } + request = Request(scope) + middleware = CacheControlMiddleware(lambda: None) + + response = await middleware.dispatch(request, app) + assert response.headers.get("Cache-Control") == "no-store" + + +@pytest.mark.asyncio +async def test_error_path_gets_no_store(): + """Error responses from authenticated requests should also be protected.""" + app = AsyncMock() + app.return_value = JSONResponse({"error": "not found"}, status_code=404) + + scope = { + "type": "http", + "method": "GET", + "path": "/api/v2/agents/unknown", + "headers": [ + (b"authorization", b"Bearer token"), + (b"content-type", b"application/json"), + ], + } + request = Request(scope) + middleware = CacheControlMiddleware(lambda: None) + + response = await middleware.dispatch(request, app) + assert response.headers.get("Cache-Control") == "no-store" + assert response.status_code == 404