diff --git a/document-parser/infra/rate_limiter.py b/document-parser/infra/rate_limiter.py new file mode 100644 index 0000000..bb82ee7 --- /dev/null +++ b/document-parser/infra/rate_limiter.py @@ -0,0 +1,89 @@ +"""Lightweight in-memory rate limiter middleware for FastAPI. + +Uses a sliding-window counter per client IP. No external dependency +required — suitable for single-process deployments with SQLite. + +For multi-process or distributed setups, replace with a Redis-backed +solution (e.g. slowapi). +""" + +from __future__ import annotations + +import logging +import time +from collections import defaultdict +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint +from starlette.responses import JSONResponse, Response + +if TYPE_CHECKING: + from starlette.requests import Request + +logger = logging.getLogger(__name__) + + +@dataclass +class _ClientBucket: + """Sliding window of request timestamps for a single client.""" + + timestamps: list[float] = field(default_factory=list) + + def count_recent(self, window: float, now: float) -> int: + """Remove expired entries and return the count of recent requests.""" + cutoff = now - window + self.timestamps = [t for t in self.timestamps if t > cutoff] + return len(self.timestamps) + + def add(self, now: float) -> None: + self.timestamps.append(now) + + +class RateLimiterMiddleware(BaseHTTPMiddleware): + """Per-IP rate limiter using in-memory sliding windows. + + Args: + app: The ASGI application. + requests_per_window: Max requests allowed per window. + window_seconds: Size of the sliding window in seconds. + exclude_paths: Paths exempt from rate limiting (e.g. health checks). + """ + + def __init__( + self, + app, + *, + requests_per_window: int = 60, + window_seconds: float = 60.0, + exclude_paths: tuple[str, ...] = ("/api/health",), + ): + super().__init__(app) + self._max_requests = requests_per_window + self._window = window_seconds + self._exclude = exclude_paths + self._buckets: dict[str, _ClientBucket] = defaultdict(_ClientBucket) + + async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: + if request.url.path in self._exclude: + return await call_next(request) + + client_ip = request.client.host if request.client else "unknown" + now = time.monotonic() + + bucket = self._buckets[client_ip] + recent = bucket.count_recent(self._window, now) + + if recent >= self._max_requests: + retry_after = int(self._window) + logger.warning( + "Rate limit exceeded for %s (%d/%d)", client_ip, recent, self._max_requests + ) + return JSONResponse( + status_code=429, + content={"detail": "Too many requests"}, + headers={"Retry-After": str(retry_after)}, + ) + + bucket.add(now) + return await call_next(request) diff --git a/document-parser/main.py b/document-parser/main.py index 91955f6..654c88e 100644 --- a/document-parser/main.py +++ b/document-parser/main.py @@ -20,6 +20,7 @@ from fastapi.middleware.cors import CORSMiddleware from api.analyses import router as analyses_router from api.documents import router as documents_router +from infra.rate_limiter import RateLimiterMiddleware from infra.settings import settings from persistence.database import get_connection, init_db from services.analysis_service import AnalysisService @@ -94,6 +95,7 @@ app.add_middleware( allow_methods=["GET", "POST", "DELETE", "OPTIONS"], allow_headers=["Content-Type", "Authorization"], ) +app.add_middleware(RateLimiterMiddleware, requests_per_window=100, window_seconds=60) app.include_router(documents_router) app.include_router(analyses_router) diff --git a/document-parser/tests/test_rate_limiter.py b/document-parser/tests/test_rate_limiter.py new file mode 100644 index 0000000..82c7fc5 --- /dev/null +++ b/document-parser/tests/test_rate_limiter.py @@ -0,0 +1,93 @@ +"""Tests for the in-memory rate limiter middleware.""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from infra.rate_limiter import RateLimiterMiddleware, _ClientBucket + + +class TestClientBucket: + def test_count_recent_filters_old_entries(self): + bucket = _ClientBucket(timestamps=[1.0, 2.0, 3.0, 10.0]) + count = bucket.count_recent(window=5.0, now=12.0) + assert count == 1 # only 10.0 is within [7.0, 12.0] + + def test_count_recent_keeps_all_when_within_window(self): + bucket = _ClientBucket(timestamps=[10.0, 11.0, 12.0]) + count = bucket.count_recent(window=60.0, now=15.0) + assert count == 3 + + def test_add(self): + bucket = _ClientBucket() + bucket.add(1.0) + bucket.add(2.0) + assert len(bucket.timestamps) == 2 + + +@pytest.fixture +def limited_app(): + """FastAPI app with a very low rate limit for testing.""" + app = FastAPI() + app.add_middleware( + RateLimiterMiddleware, + requests_per_window=3, + window_seconds=60, + exclude_paths=("/health",), + ) + + @app.get("/test") + def test_endpoint(): + return {"ok": True} + + @app.get("/health") + def health(): + return {"status": "ok"} + + return app + + +@pytest.fixture +def client(limited_app): + return TestClient(limited_app) + + +class TestRateLimiterMiddleware: + def test_allows_requests_under_limit(self, client): + for _ in range(3): + resp = client.get("/test") + assert resp.status_code == 200 + + def test_blocks_requests_over_limit(self, client): + for _ in range(3): + client.get("/test") + + resp = client.get("/test") + assert resp.status_code == 429 + assert resp.json()["detail"] == "Too many requests" + assert "Retry-After" in resp.headers + + def test_health_excluded_from_limit(self, client): + # Exhaust the limit + for _ in range(3): + client.get("/test") + + # Health should still work + resp = client.get("/health") + assert resp.status_code == 200 + + def test_window_resets(self, client): + """After the window expires, requests should be allowed again.""" + for _ in range(3): + client.get("/test") + + assert client.get("/test").status_code == 429 + + # Simulate time passing beyond the window + with patch("time.monotonic", return_value=1e12): + resp = client.get("/test") + assert resp.status_code == 200