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 import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile
from fastapi.responses import FileResponse from fastapi.responses import FileResponse
from sqlalchemy.orm import Session from sqlalchemy.orm import Session
from starlette.background import BackgroundTask from starlette.background import BackgroundTask
from app import db as db_module
from app.db import get_db from app.db import get_db
from app.models import Item
from app.schemas import ItemCreate, ItemOut, TTLOption, WallResponse 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 items as items_service
from app.services import rate_limit as rate_limit_service
from app.services import storage as storage_service from app.services import storage as storage_service
router = APIRouter(prefix="/api/items", tags=["items"]) router = APIRouter(prefix="/api/items", tags=["items"])
@@ -16,21 +18,47 @@ _INLINE_IMAGE_PREFIXES = ("image/",)
def _client_ip(request: Request) -> str: 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") forwarded = request.headers.get("x-forwarded-for")
if forwarded: if forwarded:
return forwarded.split(",")[0].strip() return forwarded.split(",")[0].strip()
if request.client and request.client.host:
return request.client.host
return "" 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) @router.post("", response_model=ItemOut)
def create_item( def create_item(
payload: ItemCreate, payload: ItemCreate,
request: Request, request: Request,
db: Session = Depends(get_db), 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) return items_service.item_to_dict(item)
@@ -44,6 +72,7 @@ async def upload_item(
ttl: TTLOption = Form("24h"), ttl: TTLOption = Form("24h"),
title: str | None = Form(None), title: str | None = Form(None),
): ):
ip = _enforce_write_guards(request, db)
rel_path, size, mime, safe_name = await storage_service.save_upload(file) rel_path, size, mime, safe_name = await storage_service.save_upload(file)
item = items_service.create_file_item( item = items_service.create_file_item(
db, db,
@@ -54,7 +83,7 @@ async def upload_item(
is_public=is_public, is_public=is_public,
burn_after_read=burn_after_read, burn_after_read=burn_after_read,
ttl=ttl, ttl=ttl,
created_ip=_client_ip(request), created_ip=ip,
title=title, title=title,
) )
return items_service.item_to_dict(item) return items_service.item_to_dict(item)
@@ -75,14 +104,11 @@ def wall(
) )
def _burn_file_after_response(item_id: int) -> None: def _delete_file_path(path: str) -> None:
db = db_module.SessionLocal()
try: try:
item = db.get(Item, item_id) Path(path).unlink(missing_ok=True)
if item is not None and item.burn_after_read and not item.burned: except OSError:
items_service.burn_file_item(db, item) pass
finally:
db.close()
@router.get("/{slug}/file") @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" media_type = item.mime or "application/octet-stream"
filename = item.file_name or path.name filename = item.file_name or path.name
inline = media_type.startswith(_INLINE_IMAGE_PREFIXES) disposition = _content_disposition_type(media_type, filename)
background = None background = None
if item.burn_after_read: 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( return FileResponse(
path, path,
media_type=media_type, media_type=media_type,
filename=filename, filename=filename,
content_disposition_type="inline" if inline else "attachment", content_disposition_type=disposition,
background=background, 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 return item
def burn_file_item(db: Session, item: Item) -> None: def claim_burn(db: Session, item: Item) -> None:
"""Mark a file item burned and remove its bytes from disk.""" """Mark burn consumed in DB before streaming bytes."""
item.burned = True item.burned = True
path = resolve_file_path(item)
db.commit() 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: if path is not None:
try: try:
path.unlink(missing_ok=True) path.unlink(missing_ok=True)
@@ -146,6 +150,12 @@ def burn_file_item(db: Session, item: Item) -> None:
pass 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: def get_item(db: Session, slug: str) -> Item | None:
item = db.scalar(select(Item).where(Item.slug == slug)) item = db.scalar(select(Item).where(Item.slug == slug))
if item is None or _is_unavailable(item): 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
+4
View File
@@ -14,9 +14,13 @@ def client(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
import app.config as config_module import app.config as config_module
import app.db as db_module import app.db as db_module
import app.main as main_module import app.main as main_module
from app.services import rate_limit as rate_limit_module
settings = Settings(_env_file=None, DATA_DIR=tmp_path) settings = Settings(_env_file=None, DATA_DIR=tmp_path)
config_module.settings = settings config_module.settings = settings
# Multi-create tests are not about rate limits; the dedicated test sets 2/s.
monkeypatch.setattr(config_module.settings, "rate_limit_per_second", 10_000)
rate_limit_module.reset()
engine = create_engine( engine = create_engine(
settings.db_url, settings.db_url,
+80
View File
@@ -119,3 +119,83 @@ def test_file_burn_after_download_not_metadata(client):
assert client.get(f"/api/items/{slug}/file").status_code == 404 assert client.get(f"/api/items/{slug}/file").status_code == 404
assert client.get(f"/api/items/{slug}").status_code == 404 assert client.get(f"/api/items/{slug}").status_code == 404
def test_svg_forced_attachment_disposition(client):
files = {"file": ("xss.svg", b"<svg xmlns='http://www.w3.org/2000/svg'></svg>", "image/svg+xml")}
r = client.post(
"/api/items/upload",
files=files,
data={"is_public": "true", "ttl": "24h"},
)
assert r.status_code == 200
slug = r.json()["slug"]
f = client.get(f"/api/items/{slug}/file")
assert f.status_code == 200
cd = f.headers.get("content-disposition", "")
assert "attachment" in cd
assert "inline" not in cd.split(";")[0]
def test_burn_claims_before_stream(client, monkeypatch):
"""burned=True is set before streaming, not only in the post-response BackgroundTask."""
import app.db as db_module
import app.api.public as public_api
from app.models import Item
from sqlalchemy import select
# Background only deletes bytes; claiming is done in the request handler.
monkeypatch.setattr(public_api, "_delete_file_path", lambda *_a, **_k: None)
files = {"file": ("once.txt", b"payload", "text/plain")}
r = client.post(
"/api/items/upload",
files=files,
data={"is_public": "true", "ttl": "never", "burn_after_read": "true"},
)
assert r.status_code == 200
slug = r.json()["slug"]
first = client.get(f"/api/items/{slug}/file")
assert first.status_code == 200
assert first.content == b"payload"
db = db_module.SessionLocal()
try:
item = db.scalar(select(Item).where(Item.slug == slug))
assert item is not None
assert item.burned is True
finally:
db.close()
assert client.get(f"/api/items/{slug}/file").status_code == 404
def test_rate_limit_third_request_in_same_second(client, monkeypatch):
import app.config as config_module
from app.services import rate_limit as rate_limit_module
monkeypatch.setattr(config_module.settings, "rate_limit_per_second", 2)
rate_limit_module.reset()
r1 = client.post("/api/items", json={"body": "one"})
r2 = client.post("/api/items", json={"body": "two"})
r3 = client.post("/api/items", json={"body": "three"})
assert r1.status_code == 200
assert r2.status_code == 200
assert r3.status_code == 429
def test_banned_ip_cannot_create(client):
import app.db as db_module
from app.models import BannedIP
db = db_module.SessionLocal()
try:
db.add(BannedIP(ip="testclient", reason="test"))
db.commit()
finally:
db.close()
r = client.post("/api/items", json={"body": "x"})
assert r.status_code == 403