feat: per-IP rate limit and write bans

Apply sliding-window 2/s limits and BannedIP checks on write routes.
Also force SVG attachment disposition and claim burn-after-read in DB
before streaming, deleting the file via BackgroundTask after the response.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
肉末
2026-07-14 12:52:45 +08:00
parent f2a1f3a1d6
commit 175a4ad0d1
6 changed files with 197 additions and 19 deletions
+46 -16
View File
@@ -1,13 +1,15 @@
from pathlib import Path
from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile
from fastapi.responses import FileResponse
from sqlalchemy.orm import Session
from starlette.background import BackgroundTask
from app import db as db_module
from app.db import get_db
from app.models import Item
from app.schemas import ItemCreate, ItemOut, TTLOption, WallResponse
from app.services import bans as bans_service
from app.services import items as items_service
from app.services import rate_limit as rate_limit_service
from app.services import storage as storage_service
router = APIRouter(prefix="/api/items", tags=["items"])
@@ -16,21 +18,47 @@ _INLINE_IMAGE_PREFIXES = ("image/",)
def _client_ip(request: Request) -> str:
"""IP used for rate limits, bans, and created_ip.
Prefer ``request.client.host`` (TestClient yields ``testclient``; direct
connections use the peer address). ``X-Forwarded-For`` is only a fallback
when client is missing — trusting XFF without a trusted reverse proxy
allows clients to spoof their IP for ban/rate-limit evasion.
"""
if request.client and request.client.host:
return request.client.host
forwarded = request.headers.get("x-forwarded-for")
if forwarded:
return forwarded.split(",")[0].strip()
if request.client and request.client.host:
return request.client.host
return ""
def _enforce_write_guards(request: Request, db: Session) -> str:
ip = _client_ip(request)
if bans_service.is_banned(db, ip):
raise HTTPException(status_code=403, detail="Forbidden")
if not rate_limit_service.allow(ip):
raise HTTPException(status_code=429, detail="Too Many Requests")
return ip
def _content_disposition_type(media_type: str, filename: str) -> str:
# SVG can carry script; never inline even when mime is image/svg+xml.
if media_type == "image/svg+xml" or filename.lower().endswith(".svg"):
return "attachment"
if media_type.startswith(_INLINE_IMAGE_PREFIXES):
return "inline"
return "attachment"
@router.post("", response_model=ItemOut)
def create_item(
payload: ItemCreate,
request: Request,
db: Session = Depends(get_db),
):
item = items_service.create_text_item(db, payload, _client_ip(request))
ip = _enforce_write_guards(request, db)
item = items_service.create_text_item(db, payload, ip)
return items_service.item_to_dict(item)
@@ -44,6 +72,7 @@ async def upload_item(
ttl: TTLOption = Form("24h"),
title: str | None = Form(None),
):
ip = _enforce_write_guards(request, db)
rel_path, size, mime, safe_name = await storage_service.save_upload(file)
item = items_service.create_file_item(
db,
@@ -54,7 +83,7 @@ async def upload_item(
is_public=is_public,
burn_after_read=burn_after_read,
ttl=ttl,
created_ip=_client_ip(request),
created_ip=ip,
title=title,
)
return items_service.item_to_dict(item)
@@ -75,14 +104,11 @@ def wall(
)
def _burn_file_after_response(item_id: int) -> None:
db = db_module.SessionLocal()
def _delete_file_path(path: str) -> None:
try:
item = db.get(Item, item_id)
if item is not None and item.burn_after_read and not item.burned:
items_service.burn_file_item(db, item)
finally:
db.close()
Path(path).unlink(missing_ok=True)
except OSError:
pass
@router.get("/{slug}/file")
@@ -96,15 +122,19 @@ def download_file(slug: str, db: Session = Depends(get_db)):
media_type = item.mime or "application/octet-stream"
filename = item.file_name or path.name
inline = media_type.startswith(_INLINE_IMAGE_PREFIXES)
disposition = _content_disposition_type(media_type, filename)
background = None
if item.burn_after_read:
background = BackgroundTask(_burn_file_after_response, item.id)
# Claim burn before streaming so concurrent fetches cannot re-read.
items_service.claim_burn(db, item)
background = BackgroundTask(_delete_file_path, str(path))
return FileResponse(
path,
media_type=media_type,
filename=filename,
content_disposition_type="inline" if inline else "attachment",
content_disposition_type=disposition,
background=background,
)
+14
View File
@@ -0,0 +1,14 @@
"""IP ban checks for write endpoints."""
from __future__ import annotations
from sqlalchemy import select
from sqlalchemy.orm import Session
from app.models import BannedIP
def is_banned(db: Session, ip: str) -> bool:
if not ip:
return False
return db.scalar(select(BannedIP).where(BannedIP.ip == ip)) is not None
+13 -3
View File
@@ -134,11 +134,15 @@ def get_file_item(db: Session, slug: str) -> Item | None:
return item
def burn_file_item(db: Session, item: Item) -> None:
"""Mark a file item burned and remove its bytes from disk."""
def claim_burn(db: Session, item: Item) -> None:
"""Mark burn consumed in DB before streaming bytes."""
item.burned = True
path = resolve_file_path(item)
db.commit()
def delete_item_file(item: Item) -> None:
"""Remove file bytes from disk (safe after response has finished streaming)."""
path = resolve_file_path(item)
if path is not None:
try:
path.unlink(missing_ok=True)
@@ -146,6 +150,12 @@ def burn_file_item(db: Session, item: Item) -> None:
pass
def burn_file_item(db: Session, item: Item) -> None:
"""Mark a file item burned and remove its bytes from disk."""
claim_burn(db, item)
delete_item_file(item)
def get_item(db: Session, slug: str) -> Item | None:
item = db.scalar(select(Item).where(Item.slug == slug))
if item is None or _is_unavailable(item):
+40
View File
@@ -0,0 +1,40 @@
"""In-memory sliding-window rate limit (per IP)."""
from __future__ import annotations
import threading
import time
from collections import defaultdict, deque
from app import config
_lock = threading.Lock()
_hits: dict[str, deque[float]] = defaultdict(deque)
def reset() -> None:
"""Clear all tracked hits (tests)."""
with _lock:
_hits.clear()
def allow(ip: str, *, now: float | None = None) -> bool:
"""Return True if ``ip`` may proceed; record the hit when allowed.
Allows when the number of hits in the last 1.0s is strictly less than
``settings.rate_limit_per_second``.
"""
if not ip:
ip = ""
limit = config.settings.rate_limit_per_second
ts = time.monotonic() if now is None else now
window_start = ts - 1.0
with _lock:
q = _hits[ip]
while q and q[0] <= window_start:
q.popleft()
if len(q) >= limit:
return False
q.append(ts)
return True