From 255d66deb75b1345648f5264ac9ba3009edeace7 Mon Sep 17 00:00:00 2001 From: Iago Homsi <–iagohomsi@gmail.com> Date: Mon, 21 Sep 2026 15:38:35 -0300 Subject: [PATCH] feature: add per-endpoint rate limiting middleware - Add EndpointRateLimitMiddleware to enforce rate limits per endpoint - Configure rate limits via endpoint Meta.rate_limit = {requests, window} - Add X-RateLimit-Limit, X-RateLimit-Remaining and X-RateLimit-Reset headers - Return 429 when limit exceeded - Add example (19_rate_limit.py) and tests --- examples/19_rate_limit.py | 24 ++++++ lightapi/__init__.py | 2 + lightapi/rate_limit_middleware.py | 114 ++++++++++++++++++++++++++++ tests/test_rate_limit_middleware.py | 93 +++++++++++++++++++++++ 4 files changed, 233 insertions(+) create mode 100644 examples/19_rate_limit.py create mode 100644 lightapi/rate_limit_middleware.py create mode 100644 tests/test_rate_limit_middleware.py diff --git a/examples/19_rate_limit.py b/examples/19_rate_limit.py new file mode 100644 index 0000000..885c984 --- /dev/null +++ b/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) diff --git a/lightapi/__init__.py b/lightapi/__init__.py index 3c17ff4..d697913 100644 --- a/lightapi/__init__.py +++ b/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", diff --git a/lightapi/rate_limit_middleware.py b/lightapi/rate_limit_middleware.py new file mode 100644 index 0000000..09beb9e --- /dev/null +++ b/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 diff --git a/tests/test_rate_limit_middleware.py b/tests/test_rate_limit_middleware.py new file mode 100644 index 0000000..014d352 --- /dev/null +++ b/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