security: address all findings from vulnerability assessment
Implements the fixes catalogued under Task Master tasks 1-19, derived from the security audit at .claude/plans/adaptive-sleeping-diffie.md. Critical - C-1 SSRF userinfo bypass: PDF proxy now rejects URLs containing userinfo or non-default ports, and asserts netloc == hostname. - C-2 No authentication: optional NADLAN_API_KEY gate guards /api/search/*, /api/appraisals/pdf and /mcp until Traefik mTLS is wired up. Default behavior unchanged when env unset. High - H-1 Error info leak: HTTPException details replaced with generic string; full exceptions still logged via logger.exception. - H-2 Security headers middleware: HSTS, X-Frame-Options, nosniff, Referrer-Policy, Permissions-Policy, CSP, server-banner stripped. - H-3 Swagger gated: /api/docs and /api/openapi.json hidden when ENVIRONMENT=production. - H-4 PDF proxy DoS: streams in 8KB chunks, rejects upstream above 50MB Content-Length, also rejects mid-stream overrun. Magic-byte check moved before StreamingResponse so we can 502 cleanly. - H-5 Per-IP rate limit: token-bucket per (path, IP) on /mcp, /api/search/*, /api/appraisals/pdf. Reads X-Forwarded-For for client IP behind Traefik. Medium - M-1 CORS: ALLOWED_ORIGINS env, allow_headers narrowed, no creds. - M-2 Deal model: extra="ignore" + explicit fields for gushNum, parcelNum, subParcelNum, streetNameHeb, houseNum, deal_source, deal_type, deal_type_description, distance_meters. Frontend now reads snake_case. - M-3 Input validators: regex on block, plot, appraiser names, dates, address, free-text — Hebrew + Latin + safe punctuation. - M-4 COOLIFY_UUID moved out of deploy.yaml into Gitea Actions secret. - M-5 Base images pinned to sha256 digests. - M-6 Runtime deps tightened to ~= compatible-release. - M-7 1MB request-body size limit middleware. Low / privacy - L-1 Heebo self-hosted via @fontsource. Removed Google Fonts links and CSP entries. Added meta referrer + meta robots noindex. - L-2 Cache-Control: immutable for /assets, no-cache for HTML. - L-3 Logger no longer prints user-supplied addresses; emits length only. - L-4 Dockerfile healthcheck uses literal port 8000. - L-5 PDF proxy stream cleanup uses contextlib.suppress. Tests: 331 passed, 17 skipped. Frontend typecheck passes. Findings file: .claude/plans/adaptive-sleeping-diffie.md Tasks: .taskmaster/tasks/tasks.json Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
+339
-23
@@ -11,15 +11,18 @@ Layout:
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from fastapi import FastAPI, HTTPException, Query
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import StreamingResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
from starlette.responses import FileResponse, JSONResponse
|
||||
|
||||
from nadlan_mcp.fastmcp_server import (
|
||||
@@ -33,12 +36,48 @@ from nadlan_mcp.govil.client import ALLOWED_PDF_HOSTS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Environment switch — controls visibility of internal-leak surfaces like
|
||||
# Swagger and openapi.json. Defaults to "production" so an unset env var
|
||||
# is the safe choice.
|
||||
ENVIRONMENT = os.getenv("ENVIRONMENT", "production").lower()
|
||||
IS_PROD = ENVIRONMENT in {"production", "prod"}
|
||||
|
||||
# Hard cap for proxied PDF downloads (bytes). Above this we refuse rather
|
||||
# than buffer/stream. Justice-ministry decisions are rarely >5MB; 50MB is
|
||||
# a generous ceiling that still bounds memory use.
|
||||
MAX_PDF_SIZE = 50 * 1024 * 1024
|
||||
|
||||
# Generic error string returned to clients in lieu of `str(e)` — see H-1.
|
||||
GENERIC_UPSTREAM_ERROR = "Service temporarily unavailable, please retry"
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# Request / response schemas
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
_RE_BLOCK = re.compile(r"^\d{1,8}$")
|
||||
_RE_PLOT = re.compile(r"^\d{1,6}$")
|
||||
_RE_DATE = re.compile(r"^\d{4}-\d{2}-\d{2}$")
|
||||
# Address: Hebrew letters, Latin letters, digits, spaces, common punctuation
|
||||
# (comma, dash, dot, slash, quote, apostrophe). Excludes control chars and
|
||||
# null bytes that could trip downstream APIs.
|
||||
_RE_ADDRESS = re.compile(r"^[-A-Za-z0-9\s,\-./'\"\(\)״׳]+$")
|
||||
# Person/committee names (Hebrew + Latin + spaces + dashes + dots)
|
||||
_RE_NAME = re.compile(r"^[-A-Za-z\s\-.׳״']+$")
|
||||
|
||||
|
||||
def _validate_optional_pattern(value: str | None, pattern: re.Pattern, field: str) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
value = value.strip()
|
||||
if not value:
|
||||
return None
|
||||
if not pattern.match(value):
|
||||
raise ValueError(f"{field}: invalid format")
|
||||
return value
|
||||
|
||||
|
||||
class DealsSearchRequest(BaseModel):
|
||||
address: str = Field(..., min_length=2, max_length=200)
|
||||
years_back: int = Field(2, ge=1, le=10)
|
||||
@@ -46,6 +85,14 @@ class DealsSearchRequest(BaseModel):
|
||||
max_deals: int = Field(50, ge=1, le=200)
|
||||
deal_type: int = Field(2, ge=1, le=2)
|
||||
|
||||
@field_validator("address")
|
||||
@classmethod
|
||||
def _check_address(cls, v: str) -> str:
|
||||
v = v.strip()
|
||||
if not _RE_ADDRESS.match(v):
|
||||
raise ValueError("address contains disallowed characters")
|
||||
return v
|
||||
|
||||
|
||||
class AppraisalsSearchRequest(BaseModel):
|
||||
block: str | None = Field(None, max_length=20)
|
||||
@@ -65,6 +112,43 @@ class AppraisalsSearchRequest(BaseModel):
|
||||
include_nearby: bool = Field(False)
|
||||
nearby_block_radius: int = Field(2, ge=1, le=10)
|
||||
|
||||
@field_validator("block")
|
||||
@classmethod
|
||||
def _check_block(cls, v):
|
||||
return _validate_optional_pattern(v, _RE_BLOCK, "block")
|
||||
|
||||
@field_validator("plot")
|
||||
@classmethod
|
||||
def _check_plot(cls, v):
|
||||
return _validate_optional_pattern(v, _RE_PLOT, "plot")
|
||||
|
||||
@field_validator("decisive_appraiser", "committee")
|
||||
@classmethod
|
||||
def _check_name(cls, v):
|
||||
return _validate_optional_pattern(v, _RE_NAME, "name")
|
||||
|
||||
@field_validator(
|
||||
"decision_date_from",
|
||||
"decision_date_to",
|
||||
"publicity_date_from",
|
||||
"publicity_date_to",
|
||||
)
|
||||
@classmethod
|
||||
def _check_date(cls, v):
|
||||
return _validate_optional_pattern(v, _RE_DATE, "date")
|
||||
|
||||
@field_validator("search_text", "appraisal_header")
|
||||
@classmethod
|
||||
def _check_freetext(cls, v):
|
||||
if v is None:
|
||||
return None
|
||||
v = v.strip()
|
||||
if not v:
|
||||
return None
|
||||
if not _RE_ADDRESS.match(v):
|
||||
raise ValueError("text contains disallowed characters")
|
||||
return v
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# App factory
|
||||
@@ -94,20 +178,178 @@ def create_app() -> FastAPI:
|
||||
"raw FastMCP endpoint at /mcp for LLM clients."
|
||||
),
|
||||
version="2.1.0",
|
||||
docs_url="/api/docs",
|
||||
# Hide API schema in production to avoid handing a roadmap to attackers.
|
||||
docs_url=None if IS_PROD else "/api/docs",
|
||||
redoc_url=None,
|
||||
openapi_url="/api/openapi.json",
|
||||
openapi_url=None if IS_PROD else "/api/openapi.json",
|
||||
lifespan=mcp_app.router.lifespan_context,
|
||||
)
|
||||
|
||||
# CORS — only relevant during local dev when Vite serves on :5173.
|
||||
# CORS — origins overridable via ALLOWED_ORIGINS env (comma-separated).
|
||||
# Default keeps the local-dev Vite origins so contributors aren't broken;
|
||||
# production sets ALLOWED_ORIGINS explicitly (or empty to disable cross-origin).
|
||||
cors_default = "http://localhost:5173,http://127.0.0.1:5173"
|
||||
cors_origins = [
|
||||
o.strip() for o in os.getenv("ALLOWED_ORIGINS", cors_default).split(",") if o.strip()
|
||||
]
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["http://localhost:5173", "http://127.0.0.1:5173"],
|
||||
allow_origins=cors_origins,
|
||||
allow_methods=["GET", "POST", "OPTIONS"],
|
||||
allow_headers=["*"],
|
||||
allow_headers=["Content-Type"],
|
||||
allow_credentials=False,
|
||||
max_age=600,
|
||||
)
|
||||
|
||||
# ── Optional API-key gate (C-2) ─────────────────────────────────
|
||||
# Defense-in-depth complement to Traefik mTLS (which is the canonical
|
||||
# solution but not yet enabled). When NADLAN_API_KEY is set, requests
|
||||
# to non-public paths must present a matching X-API-Key header.
|
||||
# When unset, the gate is bypassed — preserving the current public
|
||||
# behavior so this is safe to ship before mTLS lands.
|
||||
API_KEY = os.getenv("NADLAN_API_KEY", "").strip() or None
|
||||
AUTH_REQUIRED_PATHS = ("/api/search/", "/api/appraisals/pdf", "/mcp")
|
||||
PUBLIC_PATHS = ("/api/health", "/health")
|
||||
|
||||
@app.middleware("http")
|
||||
async def api_key_gate(request, call_next):
|
||||
if API_KEY is None:
|
||||
return await call_next(request)
|
||||
path = request.url.path
|
||||
if path in PUBLIC_PATHS:
|
||||
return await call_next(request)
|
||||
if not any(path.startswith(p) for p in AUTH_REQUIRED_PATHS):
|
||||
return await call_next(request)
|
||||
provided = request.headers.get("x-api-key", "").strip()
|
||||
# Constant-time compare to defeat timing oracles.
|
||||
import hmac
|
||||
|
||||
if not provided or not hmac.compare_digest(provided, API_KEY):
|
||||
return JSONResponse({"detail": "Unauthorized"}, status_code=401)
|
||||
return await call_next(request)
|
||||
|
||||
# ── Per-IP rate limiter (H-5) ───────────────────────────────────
|
||||
# Token bucket per client IP — protects upstream Govmap/gov.il from
|
||||
# bulk scraping via /api/* and /mcp, and prevents one attacker from
|
||||
# exhausting the shared rate-limit budget. Buckets are kept in
|
||||
# process-local memory; a single uvicorn worker is enough at our
|
||||
# scale, and the limit is a defense-in-depth complement to Traefik's
|
||||
# own RateLimit middleware (when added).
|
||||
import asyncio
|
||||
from collections import defaultdict
|
||||
|
||||
RATE_LIMIT_PATHS = (
|
||||
("/mcp", 30, 60.0), # 30 req / 60s on /mcp/*
|
||||
("/api/search/", 60, 60.0), # 60 req / 60s on /api/search/*
|
||||
("/api/appraisals/pdf", 30, 60.0), # 30 PDF fetches / 60s
|
||||
)
|
||||
_rl_buckets: dict[tuple[str, str], list[float]] = defaultdict(list)
|
||||
_rl_lock = asyncio.Lock()
|
||||
|
||||
def _client_ip(request) -> str:
|
||||
# Coolify/Traefik forward the real client IP via X-Forwarded-For.
|
||||
# Take the first entry (left-most), which by convention is the
|
||||
# actual remote client.
|
||||
xff = request.headers.get("x-forwarded-for", "").split(",")[0].strip()
|
||||
if xff:
|
||||
return xff
|
||||
return request.client.host if request.client else "unknown"
|
||||
|
||||
@app.middleware("http")
|
||||
async def rate_limit(request, call_next):
|
||||
path = request.url.path
|
||||
rule = next(
|
||||
(r for r in RATE_LIMIT_PATHS if path.startswith(r[0])),
|
||||
None,
|
||||
)
|
||||
if rule is None:
|
||||
return await call_next(request)
|
||||
prefix, limit, window = rule
|
||||
ip = _client_ip(request)
|
||||
key = (prefix, ip)
|
||||
import time
|
||||
|
||||
now = time.monotonic()
|
||||
async with _rl_lock:
|
||||
bucket = _rl_buckets[key]
|
||||
cutoff = now - window
|
||||
# Drop timestamps that fell out of the window.
|
||||
bucket[:] = [t for t in bucket if t > cutoff]
|
||||
if len(bucket) >= limit:
|
||||
retry_after = max(1, int(window - (now - bucket[0])))
|
||||
resp = JSONResponse({"detail": "Rate limit exceeded"}, status_code=429)
|
||||
resp.headers["Retry-After"] = str(retry_after)
|
||||
return resp
|
||||
bucket.append(now)
|
||||
return await call_next(request)
|
||||
|
||||
# Reject oversized request bodies before parsing — this app's largest
|
||||
# legitimate POST is a few hundred bytes. 1MB is generous.
|
||||
MAX_REQUEST_BODY = 1 * 1024 * 1024
|
||||
|
||||
@app.middleware("http")
|
||||
async def limit_body_size(request, call_next):
|
||||
cl = request.headers.get("content-length")
|
||||
if cl is not None:
|
||||
try:
|
||||
if int(cl) > MAX_REQUEST_BODY:
|
||||
return JSONResponse({"detail": "Request body too large"}, status_code=413)
|
||||
except ValueError:
|
||||
return JSONResponse({"detail": "Invalid Content-Length"}, status_code=400)
|
||||
return await call_next(request)
|
||||
|
||||
# Set HTTP security headers on every response. Done as a single
|
||||
# middleware because we want one place to manage the CSP allowlist
|
||||
# as the UI evolves (e.g. adding Mapbox would require connect-src).
|
||||
@app.middleware("http")
|
||||
async def security_headers(request, call_next):
|
||||
response = await call_next(request)
|
||||
# MIME-sniffing protection — important for the PDF proxy.
|
||||
response.headers.setdefault("X-Content-Type-Options", "nosniff")
|
||||
# Clickjacking — also enforced by CSP frame-ancestors below.
|
||||
response.headers.setdefault("X-Frame-Options", "DENY")
|
||||
# Don't leak query-string PII (searched addresses, gush/חלקה) to
|
||||
# third parties via the Referer header.
|
||||
response.headers.setdefault("Referrer-Policy", "strict-origin-when-cross-origin")
|
||||
# No need for the camera/microphone/geo APIs in this product.
|
||||
response.headers.setdefault(
|
||||
"Permissions-Policy",
|
||||
"geolocation=(), microphone=(), camera=(), payment=(), usb=()",
|
||||
)
|
||||
# CSP — UI is React+Vite (no eval, no inline scripts at build time).
|
||||
# Tailwind injects style attributes at runtime, hence 'unsafe-inline'
|
||||
# for style-src. Heebo is now self-hosted via @fontsource so we no
|
||||
# longer need to allow Google Fonts.
|
||||
response.headers.setdefault(
|
||||
"Content-Security-Policy",
|
||||
"default-src 'self'; "
|
||||
"script-src 'self'; "
|
||||
"style-src 'self' 'unsafe-inline'; "
|
||||
"font-src 'self' data:; "
|
||||
"img-src 'self' data: blob:; "
|
||||
"connect-src 'self'; "
|
||||
"frame-ancestors 'none'; "
|
||||
"form-action 'self'; "
|
||||
"base-uri 'self'; "
|
||||
"object-src 'none'",
|
||||
)
|
||||
# HSTS — only meaningful behind HTTPS, but our Traefik enforces it
|
||||
# in both dev and prod, so always set.
|
||||
response.headers.setdefault(
|
||||
"Strict-Transport-Security", "max-age=31536000; includeSubDomains"
|
||||
)
|
||||
# Drop the server-version banner.
|
||||
response.headers.pop("server", None)
|
||||
# Cache hashed asset bundles aggressively (Vite already adds a hash
|
||||
# to filenames so cache invalidation is a build problem, not a
|
||||
# runtime one). Everything else: don't cache.
|
||||
path = request.url.path
|
||||
if path.startswith("/assets/"):
|
||||
response.headers.setdefault("Cache-Control", "public, max-age=31536000, immutable")
|
||||
elif path == "/" or path.endswith(".html"):
|
||||
response.headers.setdefault("Cache-Control", "no-cache, no-store, must-revalidate")
|
||||
return response
|
||||
|
||||
# ── Health ────────────────────────────────────────────────────────
|
||||
@app.get("/api/health")
|
||||
@app.get("/health")
|
||||
@@ -121,9 +363,9 @@ def create_app() -> FastAPI:
|
||||
available) the gush/חלקה for the top match."""
|
||||
try:
|
||||
response = govmap_client.autocomplete_address(q)
|
||||
except Exception as e:
|
||||
except Exception:
|
||||
logger.exception("autocomplete_address failed")
|
||||
raise HTTPException(status_code=502, detail=str(e))
|
||||
raise HTTPException(status_code=502, detail=GENERIC_UPSTREAM_ERROR)
|
||||
|
||||
results = []
|
||||
for r in response.results:
|
||||
@@ -149,8 +391,10 @@ def create_app() -> FastAPI:
|
||||
coords = response.results[0].coordinates
|
||||
gh = _query_parcel_layer(govmap_client, coords.longitude, coords.latitude)
|
||||
gush_helka = _extract_gush_helka(gh)
|
||||
except Exception as e:
|
||||
logger.warning(f"parcel lookup failed for '{q}': {e}")
|
||||
except Exception:
|
||||
# Don't log the user-supplied query — it's PII (address)
|
||||
# and goes to whoever has access to the container logs.
|
||||
logger.warning("parcel lookup failed (query length=%d)", len(q))
|
||||
|
||||
return {"query": q, "results": results, "gush_helka": gush_helka}
|
||||
|
||||
@@ -166,9 +410,9 @@ def create_app() -> FastAPI:
|
||||
req.max_deals,
|
||||
req.deal_type,
|
||||
)
|
||||
except Exception as e:
|
||||
except Exception:
|
||||
logger.exception("find_recent_deals_for_address failed")
|
||||
raise HTTPException(status_code=502, detail=str(e))
|
||||
raise HTTPException(status_code=502, detail=GENERIC_UPSTREAM_ERROR)
|
||||
|
||||
# Search-center coords + resolved address text for the UI (best-effort).
|
||||
search_coords = None
|
||||
@@ -243,9 +487,9 @@ def create_app() -> FastAPI:
|
||||
|
||||
try:
|
||||
exact_response = _run(req.block, req.plot)
|
||||
except Exception as e:
|
||||
except Exception:
|
||||
logger.exception("search_decisions_paged (exact) failed")
|
||||
raise HTTPException(status_code=502, detail=str(e))
|
||||
raise HTTPException(status_code=502, detail=GENERIC_UPSTREAM_ERROR)
|
||||
|
||||
# Dedupe across buckets by a composite key (gov.il doesn't expose a
|
||||
# single stable id we can rely on).
|
||||
@@ -333,7 +577,7 @@ def create_app() -> FastAPI:
|
||||
|
||||
# ── /api/appraisals/pdf ───────────────────────────────────────────
|
||||
@app.get("/api/appraisals/pdf")
|
||||
async def proxy_appraisal_pdf(url: str = Query(..., min_length=20)):
|
||||
async def proxy_appraisal_pdf(url: str = Query(..., min_length=20, max_length=2048)):
|
||||
"""
|
||||
Stream a decisive-appraiser PDF through this server.
|
||||
|
||||
@@ -342,12 +586,36 @@ def create_app() -> FastAPI:
|
||||
"""
|
||||
from urllib.parse import urlparse
|
||||
|
||||
parsed = urlparse(url)
|
||||
if parsed.scheme != "https" or parsed.hostname not in ALLOWED_PDF_HOSTS:
|
||||
# ── URL validation (C-1: defeat userinfo-bypass SSRF) ───────────
|
||||
# urlparse(...).hostname returns the value after the optional
|
||||
# `user:pass@` userinfo, so a URL like
|
||||
# https://attacker.com@pub-justice.openapi.gov.il/x.pdf
|
||||
# has hostname == pub-justice.openapi.gov.il *and* would be fetched
|
||||
# from attacker.com by curl. Reject any URL with userinfo, and any
|
||||
# URL whose port differs from the default for its scheme.
|
||||
try:
|
||||
parsed = urlparse(url)
|
||||
except ValueError:
|
||||
raise HTTPException(status_code=400, detail="Malformed URL")
|
||||
|
||||
if parsed.scheme != "https":
|
||||
raise HTTPException(status_code=400, detail="Only https URLs are allowed")
|
||||
if parsed.hostname not in ALLOWED_PDF_HOSTS:
|
||||
raise HTTPException(status_code=400, detail="URL host not allowed")
|
||||
if parsed.username is not None or parsed.password is not None:
|
||||
raise HTTPException(status_code=400, detail="URL must not contain userinfo")
|
||||
if parsed.port not in (None, 443):
|
||||
raise HTTPException(status_code=400, detail="URL port not allowed")
|
||||
# Final sanity: the netloc the URL library will dial must be exactly
|
||||
# the hostname (no leftover @ or port). This catches edge cases the
|
||||
# individual checks above miss.
|
||||
if parsed.netloc.lower() != parsed.hostname.lower():
|
||||
raise HTTPException(status_code=400, detail="URL netloc mismatch")
|
||||
|
||||
try:
|
||||
# Reuse the curl_cffi session that already carries x-client-id.
|
||||
# stream=True so we read in chunks rather than buffering the
|
||||
# whole body in memory (H-4: defeat OOM via large upstream).
|
||||
response = decisive_appraiser_client._session.get(
|
||||
url,
|
||||
timeout=(
|
||||
@@ -355,26 +623,74 @@ def create_app() -> FastAPI:
|
||||
decisive_appraiser_client.config.read_timeout * 4,
|
||||
),
|
||||
headers={"Accept": "application/pdf,*/*"},
|
||||
stream=True,
|
||||
allow_redirects=False,
|
||||
)
|
||||
except Exception as e:
|
||||
except Exception:
|
||||
logger.exception("PDF proxy upstream call failed")
|
||||
raise HTTPException(status_code=502, detail=str(e))
|
||||
raise HTTPException(status_code=502, detail=GENERIC_UPSTREAM_ERROR)
|
||||
|
||||
if response.status_code != 200:
|
||||
with contextlib.suppress(Exception):
|
||||
response.close()
|
||||
raise HTTPException(
|
||||
status_code=response.status_code,
|
||||
status_code=502,
|
||||
detail=f"Upstream returned {response.status_code}",
|
||||
)
|
||||
if not response.content.startswith(b"%PDF"):
|
||||
raise HTTPException(status_code=502, detail="Upstream did not return a PDF")
|
||||
|
||||
# Refuse before reading anything if upstream advertises an oversized body.
|
||||
try:
|
||||
advertised = int(response.headers.get("Content-Length", "0") or 0)
|
||||
except ValueError:
|
||||
advertised = 0
|
||||
if advertised > MAX_PDF_SIZE:
|
||||
with contextlib.suppress(Exception):
|
||||
response.close()
|
||||
raise HTTPException(status_code=413, detail="Upstream PDF too large")
|
||||
|
||||
# Suggest a reasonable filename for the browser.
|
||||
filename = parsed.path.rstrip("/").split("/")[-1] or "appraisal"
|
||||
if not filename.lower().endswith(".pdf"):
|
||||
filename = f"{filename}.pdf"
|
||||
|
||||
# Peek the first chunk synchronously so we can verify the PDF
|
||||
# magic header before the StreamingResponse takes over (we can't
|
||||
# raise HTTPException from inside the generator).
|
||||
chunks = response.iter_content(chunk_size=8192)
|
||||
try:
|
||||
first_chunk = next(chunks)
|
||||
except StopIteration:
|
||||
with contextlib.suppress(Exception):
|
||||
response.close()
|
||||
raise HTTPException(status_code=502, detail="Empty upstream response")
|
||||
except Exception:
|
||||
logger.exception("PDF proxy: upstream stream failed")
|
||||
with contextlib.suppress(Exception):
|
||||
response.close()
|
||||
raise HTTPException(status_code=502, detail=GENERIC_UPSTREAM_ERROR)
|
||||
|
||||
if not first_chunk.startswith(b"%PDF"):
|
||||
with contextlib.suppress(Exception):
|
||||
response.close()
|
||||
raise HTTPException(status_code=502, detail="Upstream did not return a PDF")
|
||||
|
||||
def _iter() -> Any:
|
||||
yield response.content
|
||||
seen = len(first_chunk)
|
||||
yield first_chunk
|
||||
try:
|
||||
for chunk in chunks:
|
||||
if not chunk:
|
||||
continue
|
||||
seen += len(chunk)
|
||||
if seen > MAX_PDF_SIZE:
|
||||
logger.warning(
|
||||
f"PDF proxy: upstream exceeded {MAX_PDF_SIZE} bytes; truncating"
|
||||
)
|
||||
return
|
||||
yield chunk
|
||||
finally:
|
||||
with contextlib.suppress(Exception):
|
||||
response.close()
|
||||
|
||||
return StreamingResponse(
|
||||
_iter(),
|
||||
|
||||
Reference in New Issue
Block a user