Add security middleware

README.md

@@ -104,6 +104,12 @@ curl -N -X POST http://localhost:8000/chat/stream \
      -d '{"user":"demo","session":"default","prompt":"Hello"}'
 ```
 
+### Security
+
+Set one or more API keys in the ``API_KEYS`` environment variable. Requests must
+include the ``X-API-Key`` header when keys are configured. A simple rate limit is
+also enforced per key or client IP, configurable via ``RATE_LIMIT``.
+
 ## Command Line Interface
 
 Run the interactive CLI on any platform:

api_app/__init__.py

@@ -11,7 +11,12 @@ from pathlib import Path
 from typing import List
 import shutil
 
-from src.config import UPLOAD_DIR
+from src.config import UPLOAD_DIR, CORS_ORIGINS, RATE_LIMIT
+from src.security import (
+    APIKeyAuthMiddleware,
+    RateLimiterMiddleware,
+    SecurityHeadersMiddleware,
+)
 
 from src.team import TeamChatSession
 from src.log import get_logger
@@ -50,9 +55,13 @@ def _vm_host_path(user: str, vm_path: str) -> Path:
 def create_app() -> FastAPI:
     app = FastAPI(title="LLM Backend API")
 
+    app.add_middleware(APIKeyAuthMiddleware)
+    app.add_middleware(RateLimiterMiddleware, rate_limit=RATE_LIMIT)
+    app.add_middleware(SecurityHeadersMiddleware)
+
     app.add_middleware(
         CORSMiddleware,
-        allow_origins=["*"],
+        allow_origins=CORS_ORIGINS,
         allow_credentials=True,
         allow_methods=["*"],
         allow_headers=["*"],

src/config.py

@@ -59,3 +59,12 @@ Continue using tools until you have gathered everything the senior agent needs.
 Then send a brief, accurate summary so the senior agent can craft the final response.  
 Remember: you never speak to the user directly; all communication flows through the senior agent.
 """.strip()
+
+# Security settings
+API_KEYS: Final[list[str]] = (
+    os.getenv("API_KEYS", "").split(",") if os.getenv("API_KEYS") else []
+)
+RATE_LIMIT: Final[int] = int(os.getenv("RATE_LIMIT", "60"))
+CORS_ORIGINS: Final[list[str]] = (
+    os.getenv("CORS_ORIGINS", "*").split(",") if os.getenv("CORS_ORIGINS") else ["*"]
+)

src/security.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from fastapi import Request, HTTPException
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import Response
+from typing import Callable, Awaitable, MutableMapping
+from time import monotonic
+import asyncio
+
+from .config import API_KEYS, RATE_LIMIT
+
+
+class APIKeyAuthMiddleware(BaseHTTPMiddleware):
+    """Require a valid API key via the ``X-API-Key`` header."""
+
+    def __init__(self, app):
+        super().__init__(app)
+        self._keys = {k.strip() for k in API_KEYS if k.strip()}
+
+    async def dispatch(
+        self, request: Request, call_next: Callable[[Request], Awaitable[Response]]
+    ) -> Response:
+        if self._keys:
+            key = request.headers.get("X-API-Key")
+            if key not in self._keys:
+                raise HTTPException(status_code=401, detail="Invalid API key")
+        return await call_next(request)
+
+
+class RateLimiterMiddleware(BaseHTTPMiddleware):
+    """Simple in-memory rate limiter per client."""
+
+    def __init__(self, app, rate_limit: int = RATE_LIMIT) -> None:
+        super().__init__(app)
+        self.rate_limit = rate_limit
+        self._requests: MutableMapping[str, list[float]] = {}
+        self._lock = asyncio.Lock()
+
+    async def dispatch(
+        self, request: Request, call_next: Callable[[Request], Awaitable[Response]]
+    ) -> Response:
+        identifier = request.headers.get("X-API-Key") or request.client.host
+        now = monotonic()
+        async with self._lock:
+            timestamps = self._requests.setdefault(identifier, [])
+            while timestamps and now - timestamps[0] > 60:
+                timestamps.pop(0)
+            if len(timestamps) >= self.rate_limit:
+                raise HTTPException(status_code=429, detail="Rate limit exceeded")
+            timestamps.append(now)
+        return await call_next(request)
+
+
+class SecurityHeadersMiddleware(BaseHTTPMiddleware):
+    """Add common security-related HTTP headers."""
+
+    async def dispatch(
+        self, request: Request, call_next: Callable[[Request], Awaitable[Response]]
+    ) -> Response:
+        response = await call_next(request)
+        headers = response.headers
+        headers.setdefault("X-Frame-Options", "DENY")
+        headers.setdefault("X-Content-Type-Options", "nosniff")
+        headers.setdefault("Referrer-Policy", "same-origin")
+        headers.setdefault("Permissions-Policy", "geolocation=()")
+        headers.setdefault(
+            "Strict-Transport-Security", "max-age=63072000; includeSubDomains"
+        )
+        return response