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:
+46
-16
@@ -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,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -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):
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -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,
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user