From ddf4cc574395a04af7570b05e81809d50482423f Mon Sep 17 00:00:00 2001
From: Amr Gaber <amrtgaber@gmail.com>
Date: Wed, 27 May 2026 00:12:12 -0500
Subject: [PATCH] feat(rate-limit): add Retry-After header + structured 429
 body (#10)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When the rate-limit middleware refuses a request, the response now
carries enough information for the frontend to render a precise
countdown instead of a generic "try again later":

  - ``Retry-After`` header with integer seconds until the next request
    will succeed (per RFC 7231 §7.1.3).
  - Structured body shape
    ``{"detail": {"code, "message", "limit", "retry_after"}}``. The
    ``code`` is a stable ``rate_limited`` string so the frontend can
    branch on it; ``message`` is a human-readable string with the
    correct singular/plural for ``second`` vs ``seconds``; ``limit`` is
    the configured ratio (e.g. ``"5 per 1 minute"``) for diagnostic
    surfacing; ``retry_after`` mirrors the header so callers can use
    either source.

The reset timestamp comes from ``limits.strategies.MovingWindowRateLimiter.
get_window_stats(item, key)``, which after a refused hit reports when
the oldest in-window request rolls off. Clamped to ≥1s so the frontend
never gets a zero/negative countdown from edge timing.
---
 app/main.py                       |  31 ++++++++-
 tests/test_rate_limit_feedback.py | 108 ++++++++++++++++++++++++++++++
 2 files changed, 136 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_rate_limit_feedback.py

diff --git a/app/main.py b/app/main.py
index 08dfa9e..a72a382 100644
--- a/app/main.py
+++ b/app/main.py
@@ -243,19 +243,44 @@ async def delete_current_user(
 
 @app.middleware("http")
 async def rate_limit_auth(request: Request, call_next) -> Response:
-    """Apply rate limits to auth and enumeration-adjacent endpoints."""
+    """Apply rate limits to auth and enumeration-adjacent endpoints.
+
+    On rejection the response includes both a ``Retry-After`` header
+    (per RFC 7231 §7.1.3) and a structured body so the frontend can
+    render a precise countdown rather than a generic "try again later".
+    """
     rate_limit = _AUTH_RATE_LIMITS.get(f"{request.method} {request.url.path}")
     if rate_limit:
         key = get_remote_address(request)
         if not limiter._limiter.hit(rate_limit, key):
+            # window_stats: (reset_unix_ts, remaining). After a refused
+            # hit, ``remaining`` is 0 and ``reset_unix_ts`` is when the
+            # oldest in-window request rolls off — i.e., the earliest
+            # moment a fresh hit will succeed. Clamp to >=1 so the
+            # frontend never gets a zero/negative countdown.
+            reset_ts, _ = limiter._limiter.get_window_stats(rate_limit, key)
+            retry_after = max(1, int(reset_ts - time.time()))
+
             log_security_event(
                 SecurityEvent.RATE_LIMIT_HIT,
                 request=request,
-                detail=f"path={request.url.path}",
+                detail=f"path={request.url.path} retry_after={retry_after}",
             )
             return JSONResponse(
                 status_code=429,
-                content={"detail": "Rate limit exceeded"},
+                headers={"Retry-After": str(retry_after)},
+                content={
+                    "detail": {
+                        "code": "rate_limited",
+                        "message": (
+                            f"Too many requests on {request.method} "
+                            f"{request.url.path}. Try again in {retry_after} "
+                            f"second{'s' if retry_after != 1 else ''}."
+                        ),
+                        "limit": str(rate_limit),
+                        "retry_after": retry_after,
+                    }
+                },
             )
     return await call_next(request)
 
diff --git a/tests/test_rate_limit_feedback.py b/tests/test_rate_limit_feedback.py
new file mode 100644
index 0000000..ff830f8
--- /dev/null
+++ b/tests/test_rate_limit_feedback.py
@@ -0,0 +1,108 @@
+"""Tests for the rate-limit response shape (issue #10).
+
+The middleware that gates `/auth/jwt/login`, `/auth/register`, `/auth/refresh`,
+`/users/search`, and `/users/lookup` used to return a bare ``429`` with
+``{"detail": "Rate limit exceeded"}``. That's not enough for the frontend
+to render a countdown.
+
+After issue #10 the 429 response includes:
+
+- A ``Retry-After`` header (per RFC 7231 §7.1.3) with the integer
+  seconds until the next request will succeed.
+- A structured body ``{"detail": {"code", "message", "limit", "retry_after"}}``
+  the frontend can use to render a precise "try again in 42 seconds"
+  countdown.
+"""
+
+from __future__ import annotations
+
+import pytest
+from httpx import AsyncClient
+
+
+@pytest.fixture(autouse=True)
+def reset_rate_limiter():
+    """Per-test reset of the process-global rate limiter storage so test
+    order doesn't matter and the test population starts clean."""
+    from app.main import limiter
+
+    try:
+        limiter._limiter.storage.reset()
+    except Exception:
+        pass
+    yield
+    try:
+        limiter._limiter.storage.reset()
+    except Exception:
+        pass
+
+
+async def test_429_includes_retry_after_header_and_structured_body(
+    client: AsyncClient,
+) -> None:
+    """POST /auth/register is limited to 3/min. The 4th call in quick
+    succession must come back with the new shape."""
+    # First three are allowed (whether they succeed or fail at the
+    # registration layer is irrelevant — the rate limiter runs before
+    # the route handler).
+    for i in range(3):
+        await client.post(
+            "/auth/register",
+            json={"email": f"u{i}@example.com", "password": "correct-horse-battery"},
+        )
+
+    # The 4th should be rejected by the limiter.
+    resp = await client.post(
+        "/auth/register",
+        json={"email": "u3@example.com", "password": "correct-horse-battery"},
+    )
+    assert resp.status_code == 429, resp.text
+
+    # Header — must be numeric, must be >= 1.
+    retry_after_header = resp.headers.get("retry-after")
+    assert retry_after_header is not None, "Retry-After header missing"
+    retry_after_int = int(retry_after_header)
+    assert retry_after_int >= 1
+
+    # Structured body.
+    body = resp.json()
+    detail = body["detail"]
+    assert detail["code"] == "rate_limited"
+    assert detail["retry_after"] == retry_after_int, (
+        "header and body retry_after must agree so the frontend can use either"
+    )
+    assert detail["limit"] == "3 per 1 minute"
+    assert "POST /auth/register" in detail["message"]
+    assert str(retry_after_int) in detail["message"]
+
+
+async def test_429_message_singularizes_seconds(client: AsyncClient) -> None:
+    """Cosmetic but tested: the message renders 'second' vs 'seconds'
+    correctly. Cheap regression guard against ' seconds' creeping in
+    for ``retry_after == 1``."""
+    # Exhaust the limit so we get a 429 with a real retry_after value.
+    for i in range(3):
+        await client.post(
+            "/auth/register",
+            json={"email": f"v{i}@example.com", "password": "correct-horse-battery"},
+        )
+    resp = await client.post(
+        "/auth/register",
+        json={"email": "v3@example.com", "password": "correct-horse-battery"},
+    )
+    detail = resp.json()["detail"]
+    if detail["retry_after"] == 1:
+        assert "1 second." in detail["message"]
+        assert "1 seconds" not in detail["message"]
+    else:
+        assert f"{detail['retry_after']} seconds." in detail["message"]
+
+
+async def test_429_does_not_fire_on_non_rate_limited_endpoint(
+    client: AsyncClient,
+) -> None:
+    """Sanity: hammering ``/health`` (not in the rate-limit table) doesn't
+    produce 429s — confirms the middleware is path-scoped."""
+    for _ in range(20):
+        resp = await client.get("/health")
+        assert resp.status_code == 200