feat: add per-endpoint rate limiting middleware

examples/19_rate_limit.py

@@ -0,0 +1,24 @@
+from sqlalchemy import create_engine
+from sqlalchemy.pool import StaticPool
+
+from lightapi import EndpointRateLimitMiddleware, HttpMethod, LightAPI, RestEndpoint
+from lightapi.fields import Field
+
+
+class BookEndpoint(RestEndpoint, HttpMethod.GET):
+    title: str = Field(min_length=1)
+    author: str = Field(min_length=1)
+
+    class Meta:
+        rate_limit = {"request": 3, "window": 60}
+
+
+if __name__ == "__main__":
+    engine = create_engine(
+        "sqlite:///:memory:",
+        connect_args={"check_same_thread": False},
+        poolclass=StaticPool,
+    )
+    app = LightAPI(engine=engine, middlewares=[EndpointRateLimitMiddleware])
+    app.register({"/books": BookEndpoint})
+    app.run(host="0.0.0.0", port=8000, reload=True)

lightapi/__init__.py

@@ -25,6 +25,7 @@
 from lightapi.lightapi import LightApi
 from lightapi.methods import HttpMethod
 from lightapi.rate_limiter import RateLimiter
+from lightapi.rate_limit_middleware import EndpointRateLimitMiddleware
 from lightapi.rest import RestEndpoint
 from lightapi.schema import SchemaFactory
 from lightapi.session import get_async_session, get_sync_session  # noqa: E402
@@ -51,6 +52,7 @@
     "IsAdminUser",
     # Rate Limiter
     "RateLimiter",
+    "EndpointRateLimitMiddleware",
     # Filters
     "FieldFilter",
     "SearchFilter",

lightapi/rate_limit_middleware.py

@@ -0,0 +1,114 @@
+"""Middleware for per-endpoint rate limiting."""
+
+import time
+
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response
+
+from lightapi.core import Middleware
+
+
+class RateLimitEntry:
+    """Stores request count and window expiration for a client/endpoint pair."""
+
+    def __init__(self, reset_time: float) -> None:
+        self.count = 0
+        self.reset_time = reset_time
+
+    def is_expired(self, now: float) -> bool:
+        return now >= self.reset_time
+
+
+class EndpointRateLimitMiddleware(Middleware):
+    """
+    Rate-limit requests per endpoint using the endpoint's `Meta.rate_limit`.
+
+    Example on a `RestEndpoint` subclass::
+
+        class BookEndpoint(RestEndpoint):
+            title: str = Field(min_length=1)
+
+            class Meta:
+                rate_limit = {"requests": 5, "window": 60}
+
+    The window is measured in seconds. The client is identified by IP
+    (or `X-Forwarded-For` when present). Responses include these headers::
+
+        X-RateLimit-Limit
+        X-RateLimit-Remaining
+        X-RateLimit-Reset
+
+    When the limit is exceeded the middleware short-circuits with HTTP 429.
+    """
+
+    _storage: dict[tuple[str, str], RateLimitEntry] = {}
+
+    @classmethod
+    def reset(cls) -> None:
+        """Clear all rate-limit counters (useful for tests)."""
+        cls._storage.clear()
+
+    def _get_client_id(self, request: Request) -> str:
+        forwarded = request.headers.get("X-Forwarded-For")
+        if forwarded:
+            return forwarded.split(",")[0].strip()
+        return request.client.host if request.client else "unknown"
+
+    def _get_endpoint_class(self, request: Request) -> type | None:
+        handler = request.scope.get("endpoint")
+        if handler is None:
+            return None
+        return getattr(handler, "__endpoint_cls__", None)
+
+    def _get_limit_config(self, request: Request) -> dict[str, int] | None:
+        endpoint_cls = self._get_endpoint_class(request)
+        if endpoint_cls is None:
+            return None
+        meta = getattr(endpoint_cls, "Meta", None)
+        if meta is None:
+            return None
+        return getattr(meta, "rate_limit", None)
+
+    def process(self, request: Request, response: Response | None) -> Response | None:
+        limit_config = self._get_limit_config(request)
+        if limit_config is None:
+            return response
+
+        now = time.monotonic()
+        max_requests = int(limit_config["requests"])
+        window_seconds = int(limit_config["window"])
+
+        client_id = self._get_client_id(request)
+        key = (client_id, request.url.path)
+
+        entry = self._storage.get(key)
+        if entry is None or entry.is_expired(now):
+            reset_time = now + window_seconds
+            entry = RateLimitEntry(reset_time)
+            self._storage[key] = entry
+
+        if response is None:
+            if entry.count >= max_requests:
+                headers = {
+                    "X-RateLimit-Limit": str(max_requests),
+                    "X-RateLimit-Remaining": "0",
+                    "X-RateLimit-Reset": str(int(entry.reset_time - now)),
+                }
+                return JSONResponse(
+                    {"detail": "Rate limit exceeded. Try again later."},
+                    status_code=429,
+                    headers=headers,
+                )
+
+            entry.count += 1
+            request.state._rate_limit_entry = entry
+            request.state._rate_limit_max = max_requests
+            return None
+
+        entry = getattr(request.state, "_rate_limit_entry", entry)
+        max_requests = getattr(request.state, "_rate_limit_max", max_requests)
+        remaining = max(0, max_requests - entry.count)
+        response.headers["X-RateLimit-Limit"] = str(max_requests)
+        response.headers["X-RateLimit-Remaining"] = str(remaining)
+        response.headers["X-RateLimit-Reset"] = str(int(entry.reset_time - now))
+        return response

tests/test_rate_limit_middleware.py

@@ -0,0 +1,93 @@
+"""Test for EndpointRateLimitMiddleware"""
+
+import pytest
+from starlette.testclient import TestClient
+from sqlalchemy import create_engine
+from sqlalchemy.pool import StaticPool
+
+from lightapi import EndpointRateLimitMiddleware, HttpMethod, LightApi, RestEndpoint
+from lightapi.fields import Field
+
+
+def _make_engine():
+    return create_engine(
+        "sqlite:///:memory:",
+        connect_args={"check_same_thread": False},
+        poolclass=StaticPool,
+    )
+
+
+@pytest.fixture(autouse=True)
+def _reset_rate_limit_storage():
+    EndpointRateLimitMiddleware.reset()
+    yield
+    EndpointRateLimitMiddleware.reset()
+
+
+class LimitedEndpoint(RestEndpoint, HttpMethod.GET):
+    """Endpoint with a 3-requests-per-minute rate limit."""
+
+    title: str = Field(min_length=1)
+
+    class Meta:
+        rate_limit = {"requests": 3, "window": 60}
+
+
+class UnlimitedEndpoint(RestEndpoint, HttpMethod.GET):
+    """Endpoint without rate limiting."""
+
+    name: str = Field(min_length=1)
+
+
+@pytest.fixture
+def limited_client():
+    engine = _make_engine()
+    app = LightApi(engine=engine, middlewares=[EndpointRateLimitMiddleware])
+    app.register({"/books": LimitedEndpoint})
+    return TestClient(app.build_app())
+
+
+@pytest.fixture
+def mixed_client():
+    engine = _make_engine()
+    app = LightApi(engine=engine, middlewares=[EndpointRateLimitMiddleware])
+    app.register({"/books": LimitedEndpoint, "/authors": UnlimitedEndpoint})
+    return TestClient(app.build_app())
+
+
+class TestRateLimitHeaders:
+    def test_includes_rate_limit_headers(self, limited_client):
+        response = limited_client.get("/books")
+        assert response.status_code == 200
+        assert response.headers["X-RateLimit-Limit"] == "3"
+        assert response.headers["X-RateLimit-Remaining"] == "2"
+        assert int(response.headers["X-RateLimit-Reset"]) <= 60
+
+
+class TestRateLimitEnforcement:
+    def test_requests_within_limit_succeed(self, limited_client):
+        for _ in range(3):
+            response = limited_client.get("/books")
+            assert response.status_code == 200
+
+    def test_request_over_limit_returns_429(self, limited_client):
+        for _ in range(3):
+            limited_client.get("/books")
+
+        response = limited_client.get("/books")
+        assert response.status_code == 429
+        assert response.headers["X-RateLimit-Limit"] == "3"
+        assert response.headers["X-RateLimit-Remaining"] == "0"
+        assert int(response.headers["X-RateLimit-Reset"]) <= 60
+        assert response.json()["detail"] == "Rate limit exceeded. Try again later."
+
+
+class TestUnlimitedEndpoint:
+    def test_unlimited_endpoint_not_blocked(self, mixed_client):
+        for _ in range(3):
+            mixed_client.get("/books")
+        assert mixed_client.get("/books").status_code == 429
+
+        response = mixed_client.get("/authors")
+        assert response.status_code == 200
+        assert "X-RateLimit-Limit" not in response.headers