diff --git a/main.py b/main.py index 8142a94..9013c5b 100644 --- a/main.py +++ b/main.py @@ -4,6 +4,7 @@ Rodar local: uvicorn main:app --reload (ou: python main.py) """ import html +import json import os import platform import time @@ -12,9 +13,7 @@ from pathlib import Path import fastapi from fastapi import FastAPI, Request -from fastapi.exceptions import RequestValidationError from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response -from pydantic import BaseModel from links import FORMATO_CODIGO, ErroDeValidacao, criar_link, listar_links, registrar_clique @@ -95,16 +94,33 @@ def para_json(link: dict, request: Request) -> dict: } -class NovoLink(BaseModel): - url: str | None = None +CORPO_INVALIDO = 'Corpo inválido. Envie JSON no formato {"url": "https://..."}.' +# Um pedido de encurtar link cabe com folga em 10 KB; mais que isso é recusado +# sem ler o resto (o app é público: ninguém deve conseguir mandar megabytes). +TAMANHO_MAXIMO_CORPO = 10 * 1024 -@app.exception_handler(RequestValidationError) -async def corpo_invalido(request: Request, erro: RequestValidationError): - return JSONResponse( - status_code=400, - content={"erro": 'Corpo inválido. Envie JSON no formato {"url": "https://..."}.'}, - ) +class CorpoInvalido(Exception): + pass + + +async def ler_json_limitado(request: Request) -> dict: + """Lê o corpo em pedaços e para assim que passar do limite (não confia só no Content-Length).""" + declarado = request.headers.get("content-length") + if declarado is not None and (not declarado.isdigit() or int(declarado) > TAMANHO_MAXIMO_CORPO): + raise CorpoInvalido() + corpo = bytearray() + async for pedaco in request.stream(): + corpo += pedaco + if len(corpo) > TAMANHO_MAXIMO_CORPO: + raise CorpoInvalido() + try: + dados = json.loads(corpo) + except ValueError: + raise CorpoInvalido() from None + if not isinstance(dados, dict): + raise CorpoInvalido() + return dados @app.get("/", response_class=HTMLResponse, include_in_schema=False) @@ -120,8 +136,24 @@ def listar(request: Request): ) -@app.post("/api/links", status_code=201, summary="Encurta um link") -def criar(novo: NovoLink, request: Request): +@app.post( + "/api/links", + status_code=201, + summary="Encurta um link", + # O corpo é lido à mão (com limite de tamanho); isto só descreve o formato no /docs. + openapi_extra={ + "requestBody": { + "required": True, + "content": { + "application/json": { + "schema": {"type": "object", "properties": {"url": {"type": "string"}}, "required": ["url"]}, + "example": {"url": "https://exemplo.com.br/uma-pagina"}, + } + }, + } + }, +) +async def criar(request: Request): ip = request.headers.get("cf-connecting-ip") or (request.client.host if request.client else "?") if passou_do_limite(ip): return JSONResponse( @@ -129,7 +161,11 @@ def criar(novo: NovoLink, request: Request): content={"erro": "Muitos links em pouco tempo. Espere alguns minutos e tente de novo."}, ) try: - link = criar_link(novo.url) + dados = await ler_json_limitado(request) + except CorpoInvalido: + return JSONResponse(status_code=400, content={"erro": CORPO_INVALIDO}) + try: + link = criar_link(dados.get("url")) except ErroDeValidacao as erro: return JSONResponse(status_code=400, content={"erro": str(erro)}) return JSONResponse(status_code=201, content=para_json(link, request))