utils/request_limits.py

import time
from typing import Dict, Tuple
from starlette.requests import Request  # type: ignore
from starlette.responses import Response  # type: ignore
from starlette.types import ASGIApp, Receive, Scope, Send  # type: ignore


class SimpleRateLimiter:
    """Naive in-memory rate limiter (per IP). Suitable for single-instance deployments."""

    def __init__(self, app: ASGIApp, max_requests: int = 100, window_seconds: int = 60):
        self.app = app
        self.max_requests = max_requests
        self.window_seconds = window_seconds
        self.ip_buckets: Dict[str, Tuple[int, float]] = {}

    async def __call__(self, scope: Scope, receive: Receive, send: Send):
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        request = Request(scope, receive=receive)
        client_ip = request.client.host if request.client else "unknown"
        now = time.time()
        count, start = self.ip_buckets.get(client_ip, (0, now))

        if now - start > self.window_seconds:
            count, start = 0, now

        count += 1
        self.ip_buckets[client_ip] = (count, start)

        if count > self.max_requests:
            response = Response("Rate limit exceeded", status_code=429)
            await response(scope, receive, send)
            return

        await self.app(scope, receive, send)


class MaxUploadSizeLimiter:
    """Limit maximum request body size for uploads."""

    def __init__(self, app: ASGIApp, max_bytes: int = 10 * 1024 * 1024):
        self.app = app
        self.max_bytes = max_bytes

    async def __call__(self, scope: Scope, receive: Receive, send: Send):
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        received = 0

        async def limited_receive():
            nonlocal received
            message = await receive()
            if message.get("type") == "http.request":
                body = message.get("body", b"")
                received += len(body)
                if received > self.max_bytes:
                    response = Response("Payload too large", status_code=413)
                    await response(scope, receive, send)
                    return {"type": "http.disconnect"}
            return message

        await self.app(scope, limited_receive, send)