Rate limiting com Redis para um SaaS freemium: a arquitetura por trás de 3 tiers

def get_rate_limit_key(identifier: str) -> str:
    today = datetime.utcnow().strftime("%Y-%m-%d")
    return f"rl:{identifier}:{today}"

async def check_rate_limit(redis, identifier: str, limit: int):
    key = get_rate_limit_key(identifier)

    current = await redis.incr(key)

    if current == 1:
        # Primeira request do dia — definir TTL de 24h
        await redis.expire(key, 86400)

    if current > limit:
        return False, current  # Rate limited

    return True, current

if current == 1:
    await redis.expire(key, 86400)
elif current == 2:
    # Safeguard: verificar se o TTL foi setado
    ttl = await redis.ttl(key)
    if ttl == -1:  # Sem TTL
        await redis.expire(key, 86400)

def get_identifier(request: Request) -> tuple[str, str]:
    # 1. Tentar extrair user do JWT (cookie ou header)
    token = extract_token(request)
    if token:
        user_id = decode_jwt(token).get("sub")
        if user_id:
            return f"user:{user_id}", "authenticated"

    # 2. Fallback: IP do cliente
    ip = (
        request.headers.get("X-Forwarded-For", "").split(",")[0].strip()
        or request.client.host
    )
    return ip, "anonymous"

async def get_effective_plan(redis, db, user_id: int) -> str:
    cache_key = f"uplan:{user_id}"

    cached = await redis.get(cache_key)
    if cached:
        return cached.decode()

    user = db.query(User).filter(User.id == user_id).first()
    plan = user.plan if user else "free"

    # Cache por 5 minutos
    await redis.setex(cache_key, 300, plan)

    return plan

COUNTED_PREFIXES = [
    "/api/cnpj/",
    "/api/search",
    "/api/busca-avancada",
    "/api/leads",
    "/api/socios/busca",
    "/api/pessoa/",
]

RATE_LIMITS = {
    "anon": 50,
    "free": 200,
    "pro": 5000,
}

class RateLimitMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        path = request.url.path

        # Só conta requests que consomem dados
        if not any(path.startswith(p) for p in COUNTED_PREFIXES):
            return await call_next(request)

        identifier, auth_type = get_identifier(request)

        if auth_type == "authenticated":
            plan = await get_effective_plan(redis, db, user_id)
        else:
            plan = "anon"

        limit = RATE_LIMITS[plan]
        allowed, current = await check_rate_limit(redis, identifier, limit)

        if not allowed:
            return JSONResponse(
                status_code=429,
                content={"detail": "Limite de consultas atingido"},
                headers=rate_limit_headers(limit, 0, plan),
            )

        response = await call_next(request)

        # Adicionar headers informativos
        response.headers.update(
            rate_limit_headers(limit, limit - current, plan)
        )

        return response

def rate_limit_headers(limit, remaining, plan):
    return {
        "X-RateLimit-Limit": str(limit),
        "X-RateLimit-Remaining": str(max(0, remaining)),
        "X-RateLimit-Plan": plan,
    }

// Frontend: lê os headers e avisa o usuário
function trackRateLimit(res: Response) {
  const remaining = res.headers.get("x-ratelimit-remaining");
  const limit = res.headers.get("x-ratelimit-limit");

  if (remaining && parseInt(remaining) <= 5) {
    // Mostra banner: "Você tem apenas X consultas restantes hoje"
    showRateLimitBanner(parseInt(limit), parseInt(remaining));
  }
}

async def check_auth_rate_limit(redis, ip: str):
    minute = datetime.utcnow().strftime("%Y%m%d%H%M")
    key = f"rl:auth:{ip}:{minute}"

    current = await redis.incr(key)
    if current == 1:
        await redis.expire(key, 120)  # 2 minutos de janela

    return current <= 10  # Max 10 tentativas por minuto

# Não conta:
# - Assets estáticos (/_next/*, /fonts/*)
# - Health checks (/api/health)
# - Auth endpoints (/api/auth/*)  ← tem rate limit próprio
# - Municipios/CNAEs search (autocomplete de filtros)

# Conta:
# - Consulta de CNPJ (/api/cnpj/{cnpj})
# - Busca textual (/api/search)
# - Busca avançada (/api/busca-avancada)
# - Leads (/api/leads)

async function apiFetch(url: string, init?: RequestInit) {
  const res = await fetch(url, init);

  if (res.status === 429 && typeof window !== "undefined") {
    window.location.href = "/limite";
  }

  return res;
}