feat(ratelimit): implement dual-key in-memory sliding window rate limiter

app/core/dependencies.py

@@ -54,3 +54,12 @@ def get_memory_service() -> 'MemoryService':
     from services.memory_service import MemoryService
     return MemoryService(get_llm_service())
 
+
+from core.rate_limiter import DualKeyRateLimitChecker
+
+rate_limit_ask = DualKeyRateLimitChecker("ask", default_limit=20, env_var_name="RATE_LIMIT_ASK_PER_MIN")
+rate_limit_upload = DualKeyRateLimitChecker("upload", default_limit=5, env_var_name="RATE_LIMIT_UPLOAD_PER_MIN")
+rate_limit_quiz = DualKeyRateLimitChecker("quiz", default_limit=5, env_var_name="RATE_LIMIT_QUIZ_PER_MIN")
+rate_limit_tools = DualKeyRateLimitChecker("tools", default_limit=20, env_var_name="RATE_LIMIT_TOOLS_PER_MIN")
+
+

app/core/rate_limiter.py

@@ -0,0 +1,142 @@
+"""Rate Limiter Module for Atlas API.
+
+Provides a thread-safe, in-memory sliding-window rate limiter.
+Dual-keyed on client_id AND client IP to prevent:
+1. Limit bypass via rotating client UUIDs (guarded by IP bucket).
+2. Shared bucket starvation behind proxies (guarded by client_id bucket).
+
+Behind Hugging Face Spaces / reverse proxies, request.client.host points to the
+internal proxy infrastructure. The originating client IP is extracted from the
+first hop of the X-Forwarded-For header.
+"""
+
+import os
+import threading
+import time
+from typing import Dict, List, Tuple, Optional
+from fastapi import Request, HTTPException, status
+from utils.logger import app_logger
+
+
+def get_client_ip(request: Request) -> str:
+    """
+    Extracts the originating client IP address.
+    Behind Hugging Face / reverse proxies, request.client.host is the internal proxy IP.
+    The first hop in X-Forwarded-For contains the true client IP.
+    """
+    xff = request.headers.get("x-forwarded-for")
+    if xff:
+        # First hop is the original client IP before traversing upstream proxies
+        return xff.split(",")[0].strip()
+    if request.client and request.client.host:
+        return request.client.host
+    return "127.0.0.1"
+
+
+class SlidingWindowRateLimiter:
+    """Thread-safe sliding-window rate limiter with sub-second precision."""
+
+    def __init__(self, window_seconds: float = 60.0):
+        self.window_seconds = window_seconds
+        self._lock = threading.Lock()
+        self._hits: Dict[str, List[float]] = {}
+
+    def check_and_record(
+        self,
+        key: str,
+        limit: int,
+    ) -> Tuple[bool, int]:
+        """
+        Checks if key is within limit. If allowed, records the hit.
+        Returns (is_allowed, retry_after_seconds).
+        """
+        now = time.monotonic()
+        cutoff = now - self.window_seconds
+
+        with self._lock:
+            timestamps = self._hits.get(key, [])
+            # Evict timestamps outside sliding window
+            timestamps = [t for t in timestamps if t > cutoff]
+
+            if len(timestamps) >= limit:
+                # Oldest timestamp in window determines remaining cooldown
+                oldest_hit = timestamps[0]
+                retry_after = max(1, int(self.window_seconds - (now - oldest_hit) + 0.99))
+                self._hits[key] = timestamps
+                return False, retry_after
+
+            timestamps.append(now)
+            self._hits[key] = timestamps
+            return True, 0
+
+    def reset(self) -> None:
+        """Clear all recorded hits (useful for test isolation)."""
+        with self._lock:
+            self._hits.clear()
+
+
+# Global rate limiter instance
+_global_limiter = SlidingWindowRateLimiter(window_seconds=60.0)
+
+
+def get_rate_limiter() -> SlidingWindowRateLimiter:
+    return _global_limiter
+
+
+class DualKeyRateLimitChecker:
+    """
+    FastAPI dependency callable that checks rate limits against both
+    client_id and originating IP address.
+    """
+
+    def __init__(self, action_name: str, default_limit: int, env_var_name: str):
+        self.action_name = action_name
+        self.default_limit = default_limit
+        self.env_var_name = env_var_name
+
+    def _get_limit(self) -> int:
+        raw = os.getenv(self.env_var_name)
+        if raw:
+            try:
+                val = int(raw.strip())
+                if val > 0:
+                    return val
+            except ValueError:
+                pass
+        return self.default_limit
+
+    async def __call__(self, request: Request) -> None:
+        client_id = request.headers.get("x-client-id", "").strip().lower()
+        ip_addr = get_client_ip(request)
+        limit = self._get_limit()
+
+        limiter = get_rate_limiter()
+
+        # Check client_id bucket if present
+        if client_id:
+            client_key = f"client:{self.action_name}:{client_id}"
+            allowed, retry_after = limiter.check_and_record(client_key, limit)
+            if not allowed:
+                app_logger.warning(
+                    f"[RateLimit] Client rate limit exceeded for {client_id} on {self.action_name} "
+                    f"({limit}/min). IP: {ip_addr}"
+                )
+                raise HTTPException(
+                    status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+                    detail=f"Rate limit exceeded for {self.action_name}. Maximum {limit} requests per minute.",
+                    headers={"Retry-After": str(retry_after)},
+                )
+
+        # Check IP bucket
+        ip_key = f"ip:{self.action_name}:{ip_addr}"
+        allowed_ip, retry_after_ip = limiter.check_and_record(ip_key, limit)
+        if not allowed_ip:
+            app_logger.warning(
+                f"[RateLimit] IP rate limit exceeded for {ip_addr} on {self.action_name} "
+                f"({limit}/min). Client ID: {client_id}"
+            )
+            raise HTTPException(
+                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+                detail=f"Rate limit exceeded for {self.action_name}. Maximum {limit} requests per minute.",
+                headers={"Retry-After": str(retry_after_ip)},
+            )

app/routers/chat.py

@@ -19,6 +19,7 @@ from core.dependencies import (
     get_storage,
     get_memory_service,
     get_client_id,
+    rate_limit_ask,
 )
 from core.rag_pipeline import VectorlessRAG
 from services.llm_service import LLMService
@@ -138,6 +139,7 @@ async def ask_question(
     rag_pipeline: VectorlessRAG = Depends(get_rag_pipeline),
     memory_service: MemoryService = Depends(get_memory_service),
     client_id: str = Depends(get_client_id),
+    _rl: None = Depends(rate_limit_ask),
 ):
     # Ensure chat session exists in persistent storage and is owned by client_id
     if storage.chat_exists(request.chat_id):

app/routers/documents.py

@@ -8,7 +8,13 @@ from fastapi import APIRouter, UploadFile, File, Form, HTTPException, Depends, R
 from services.document_service import DocumentService
 from services.llm_service import LLMService
 from core.storage import MemoryStorage
-from core.dependencies import get_document_service, get_llm_service, get_storage, get_client_id
+from core.dependencies import (
+    get_document_service,
+    get_llm_service,
+    get_storage,
+    get_client_id,
+    rate_limit_upload,
+)
 from schemas.schemas import UploadResponse, DocumentStatusResponse
 from utils.logger import app_logger
 from core.persistent_db import DATA_DIR
@@ -129,6 +135,7 @@ async def upload_pdf(
     llm_service: LLMService = Depends(get_llm_service),
     storage: MemoryStorage = Depends(get_storage),
     client_id: str = Depends(get_client_id),
+    _rl: None = Depends(rate_limit_upload),
 ):
     """
     Asynchronous PDF upload and indexing endpoint.

app/routers/studio.py

@@ -4,7 +4,14 @@ from typing import List, Dict, Any, Optional
 from fastapi import APIRouter, Depends, HTTPException, status
 from pydantic import BaseModel
 
-from core.dependencies import get_storage, get_document_service, get_llm_service, get_memory_service, get_client_id
+from core.dependencies import (
+    get_storage,
+    get_document_service,
+    get_llm_service,
+    get_memory_service,
+    get_client_id,
+    rate_limit_quiz,
+)
 from core.storage import MemoryStorage
 from services.document_service import DocumentService
 from services.llm_service import LLMService
@@ -170,6 +177,7 @@ async def generate_document_quiz(
     storage: MemoryStorage = Depends(get_storage),
     doc_service: DocumentService = Depends(get_document_service),
     llm_service: LLMService = Depends(get_llm_service),
+    _rl: None = Depends(rate_limit_quiz),
 ):
     """
     Generates 5 Multiple Choice Questions and 5 Study Flashcards for a document

app/routers/tools.py

@@ -12,7 +12,7 @@ from typing import Any, Dict, Optional
 from fastapi import APIRouter, HTTPException, Depends
 from pydantic import BaseModel, Field
 
-from core.dependencies import get_client_id
+from core.dependencies import get_client_id, rate_limit_tools
 from mcp_server.web_search_mcp import execute_web_search
 from mcp_server.research_tools import (
     fetch_webpage_markdown,
@@ -21,7 +21,11 @@ from mcp_server.research_tools import (
     evaluate_math_expression_safe,
 )
 
-router = APIRouter(prefix="/tools", tags=["Tools & MCP"], dependencies=[Depends(get_client_id)])
+router = APIRouter(
+    prefix="/tools",
+    tags=["Tools & MCP"],
+    dependencies=[Depends(get_client_id), Depends(rate_limit_tools)],
+)
 
 
 class WebSearchRequest(BaseModel):

tests/test_rate_limiter.py

@@ -0,0 +1,110 @@
+import uuid
+import pytest
+from fastapi.testclient import TestClient
+
+from app import app
+from core.rate_limiter import get_rate_limiter, get_client_ip
+
+
+@pytest.fixture(autouse=True)
+def reset_limiter():
+    """Reset in-memory rate limiter buckets before and after each test."""
+    limiter = get_rate_limiter()
+    limiter.reset()
+    yield
+    limiter.reset()
+
+
+@pytest.fixture
+def client():
+    return TestClient(app)
+
+
+def test_first_hop_ip_extraction():
+    """Verifies that only the first hop of X-Forwarded-For is extracted (ignoring proxy hops)."""
+    from starlette.requests import Request
+
+    scope = {
+        "type": "http",
+        "headers": [
+            (b"x-forwarded-for", b"198.51.100.42, 10.0.0.1, 172.16.0.2"),
+        ],
+        "client": ("10.0.0.1", 54321),
+    }
+    req = Request(scope)
+    assert get_client_ip(req) == "198.51.100.42"
+
+
+def test_rate_limiter_blocks_after_limit_and_sets_retry_after(client, monkeypatch):
+    """Enforces rate limit and checks HTTP 429 and Retry-After header."""
+    monkeypatch.setenv("RATE_LIMIT_TOOLS_PER_MIN", "3")
+
+    client_id = str(uuid.uuid4())
+    headers = {
+        "X-Client-ID": client_id,
+        "X-Forwarded-For": "203.0.113.10",
+    }
+
+    # 3 allowed requests
+    for i in range(3):
+        r = client.post("/tools/calculate-math", headers=headers, json={"expression": f"{i} + 1"})
+        assert r.status_code == 200
+
+    # 4th request must be rejected with 429
+    r4 = client.post("/tools/calculate-math", headers=headers, json={"expression": "99 + 1"})
+    assert r4.status_code == 429
+    assert "Rate limit exceeded" in r4.json()["detail"]
+    assert "Retry-After" in r4.headers
+    retry_after = int(r4.headers["Retry-After"])
+    assert retry_after > 0
+
+
+def test_rotating_uuids_cannot_bypass_rate_limit(client, monkeypatch):
+    """
+    Rotating client_id UUIDs from the SAME IP must NOT bypass rate limiting.
+    The IP-level bucket must catch and throttle the attacker.
+    """
+    monkeypatch.setenv("RATE_LIMIT_TOOLS_PER_MIN", "3")
+    attacker_ip = "198.51.100.99"
+
+    # Make 3 requests, each with a different UUID but same IP
+    for i in range(3):
+        headers = {
+            "X-Client-ID": str(uuid.uuid4()),
+            "X-Forwarded-For": f"{attacker_ip}, 10.0.0.1",
+        }
+        r = client.post("/tools/calculate-math", headers=headers, json={"expression": f"{i} + 10"})
+        assert r.status_code == 200
+
+    # 4th request with yet another brand-new UUID from same IP must be rejected
+    headers_new_uuid = {
+        "X-Client-ID": str(uuid.uuid4()),
+        "X-Forwarded-For": f"{attacker_ip}, 10.0.0.1",
+    }
+    r4 = client.post("/tools/calculate-math", headers=headers_new_uuid, json={"expression": "100 + 1"})
+    assert r4.status_code == 429
+    assert "Rate limit exceeded" in r4.json()["detail"]
+
+
+def test_distinct_clients_on_distinct_ips_have_independent_limits(client, monkeypatch):
+    """Client A hitting limit does not block Client B on a different IP."""
+    monkeypatch.setenv("RATE_LIMIT_TOOLS_PER_MIN", "2")
+
+    headers_a = {
+        "X-Client-ID": str(uuid.uuid4()),
+        "X-Forwarded-For": "192.0.2.1",
+    }
+    headers_b = {
+        "X-Client-ID": str(uuid.uuid4()),
+        "X-Forwarded-For": "192.0.2.2",
+    }
+
+    # Exhaust Client A
+    client.post("/tools/calculate-math", headers=headers_a, json={"expression": "1 + 1"})
+    client.post("/tools/calculate-math", headers=headers_a, json={"expression": "2 + 2"})
+    r_a_blocked = client.post("/tools/calculate-math", headers=headers_a, json={"expression": "3 + 3"})
+    assert r_a_blocked.status_code == 429
+
+    # Client B should still succeed
+    r_b = client.post("/tools/calculate-math", headers=headers_b, json={"expression": "4 + 4"})
+    assert r_b.status_code == 200