"""SSRF guard for the ML service — image URL validation + safe download. 【這個模組在做什麼？】 判斷「這個 URL 可以 fetch 嗎」並做實際下載，是 async FastAPI handler 的 SSRF 防線。哪些 host / scheme 被擋掉的「策略」住在 `ssrf_policy.py`（由 packages/ssrf-policy/policy.yaml 驅動，Node 端共用）； 本檔案只負責 transport 層的硬性限制 — follow_redirects / body cap / timeout / PIL offload — 這四個 caller 不可以 override。 """ from __future__ import annotations import asyncio import ipaddress import socket from io import BytesIO from urllib.parse import urlparse import httpx from PIL import Image from src.ssrf_policy import is_allowed_scheme, is_blocked_host # 10 MB per image download — 比典型電商圖（~1 MB）寬，但擋掉惡意超大檔案 # 讓 handler 在記憶體載入前就失敗。 _MAX_IMAGE_BYTES = 10 * 1024 * 1024 # PIL 解壓縮炸彈防護 — 4 KB 壓縮 PNG 可宣告 50000×50000 RGB → 解壓 7.5 GB # OOM-kill 整個 service。10 MB body cap 不擋這條，因為惡意檔案壓縮後很小。 # 在本模組頂端設一次（非靠 main.py 的 import 順序）— ssrf.py 是 download_image # 唯一的 entry point，任何 caller 拿到 PIL Image 之前一定先過這條 import。 # 對齊 paste_parser.py + main.py 的 40_000_000 上限（feedback_pil_decompression_bomb_cap.md）。 Image.MAX_IMAGE_PIXELS = 40_000_000 # Slowloris defence layer 1 — granular per-recv httpx.Timeout # Pre-fix code passed a scalar `timeout=timeout_s` which httpx applies to each # socket op (connect / read / write / pool) **separately**. Attacker streaming # 1 byte every (read - 0.1)s within a chunk loop holds the connection across # many recv calls — `len(chunks) × read_timeout` worst case wall clock # (10 MB cap / 65 KB chunks = 160 chunks, 160 × 10s = 1600s on a single # download). Under `asyncio.to_thread` worker concurrency (`min(32, cpu+4)`) # 32 simultaneous slowloris connections starve all other endpoints. # # Granular tiers split connect/pool (TCP handshake — should be fast) from # read/write (per-recv data ops — can be slower for large bodies). Layer 2 # (asyncio.wait_for total ceiling) catches drip feed within per-recv budget. _PER_RECV_TIMEOUT = httpx.Timeout(connect=5.0, read=10.0, write=10.0, pool=5.0) # 預設總 wall-clock ceiling — F4 v2 valuation_observations 同步 await 路徑可 # 縮到 3s 做 fast-fail，新增 caller 也可 override 用更短上限。30s 給標準 # 10 MB image 充足餘裕（典型 1 MB CDN image 通常 < 1s 完成）。 _DEFAULT_TIMEOUT_S = 30.0 def validate_image_url(url: str) -> None: """Validate a URL before fetching. Raises ValueError on SSRF indicators. 檢查順序： 1. Scheme 必須在 allowed 清單（目前只有 https）。 2. hostname 非空。 3. hostname 不在 block list（精確或 CIDR）。 4. DNS resolve hostname → 每個解析到的 IP 也得通過 block list （codebase review v19 ML-03，DNS rebinding mitigation）。 【DNS rebinding 攻擊面】 Attacker 註冊 `img.attacker.com` → DNS 先回 public IP → static hostname check 通過 → httpx 連線時 attacker 切換 DNS 回 `169.254.169.254`（AWS metadata）/ `10.0.0.1`（內部服務）。 `follow_redirects=False` 救不到，因為是 DNS-level 切換。 這層 mitigation 做的是 **解析後立刻 re-check IP**：static check 通過、 DNS resolve、每個 IP 也跑 `is_blocked_host`。**TOCTOU 仍存在** — 真正零窗口要 socket-level peer IP 驗證（httpx transport hook 或 自訂 resolver）— 但 static-resolve 已關掉絕大部分 rebind window， 留待後續 sprint 補完。 DNS resolve 失敗 → raise（fail-closed，跟既有 SSRF 嚴格一致）。 """ parsed = urlparse(url) if not is_allowed_scheme(parsed.scheme): raise ValueError(f"Only https URLs allowed, got: {parsed.scheme or '(none)'}") hostname = (parsed.hostname or "").lower() if not hostname: raise ValueError("URL missing hostname") if is_blocked_host(hostname): raise ValueError(f"Internal host not allowed: {hostname}") # DNS resolve + per-IP block check (ML-03). Skip when hostname is already # an IP literal — `is_blocked_host` already exercised the literal/CIDR # paths, and `getaddrinfo` on a literal just round-trips it (extra DNS # latency for no security gain). The literal/CIDR coverage above is the # canonical path for IP-form URLs. try: ipaddress.ip_address(hostname.rstrip(".").strip("[]")) is_ip_literal = True except ValueError: is_ip_literal = False if is_ip_literal: return try: addrs = socket.getaddrinfo(hostname, None) except socket.gaierror as exc: # Fail-closed: a hostname that won't resolve is not a hostname we # are willing to fetch. Matches the strict SSRF posture elsewhere # in this module. raise ValueError(f"DNS resolution failed for: {hostname}") from exc for entry in addrs: sockaddr = entry[4] # (host, port[, flowinfo, scopeid]) ip_str = sockaddr[0] if is_blocked_host(ip_str): raise ValueError(f"Resolved IP {ip_str} for {hostname} is blocked") def _decode_image(buf: BytesIO) -> Image.Image: """Synchronous PIL decode — call via asyncio.to_thread to keep event loop free. 【為什麼接 BytesIO 不接 bytes？】 `_do_download` 把 chunk 寫進單一 BytesIO accumulator（取代舊的 `chunks: list[bytes] + b"".join`），少一份完整 bytes 副本 — 10 MB cap × 32 並發 worker peak 從 ~640 MB 降回 ~320 MB。PIL 的 `Image.open` 直 接吃 file-like，零成本接管。Buf 必須活到 `.convert("RGB")` 完成 — Image.open 是 lazy 只 parse header；img.load() 顯式 force decode 把 data 拉進新 PIL object，之後 buf 可以被 GC。本函式 frame 持有 buf 到 return 為止，符合此約束；未來 refactor 不可把 buf 提早 close。 【為什麼要 offload？】 `Image.open` 是 lazy（只 parse header），但 `.convert("RGB")` 會 force full decode + color transform：10 MB JPEG 約 100-300ms 純 CPU。跑在 event loop 上會把所有其他 coroutine 都凍住——4 個 async caller (`/verify/local` 雙圖 + outfit_match + valuation observation) 全部 serialise 在 PIL CPU 後面。`img.load()` 顯式觸發 full decode，避免 `.convert` 內部 lazy IO 還埋在這條 thread 之外。 """ # buf 寫完 position 在尾端；Image.open 內部會 seek 但顯式重置防 PIL 行為變動 buf.seek(0) img = Image.open(buf) img.load() # force full decode (else .convert lazy-reads from BytesIO) return img.convert("RGB") async def download_image( url: str, timeout_s: float = _DEFAULT_TIMEOUT_S ) -> Image.Image: """Download and decode an image with SSRF validation + size cap + slowloris defence. Use this instead of `requests.get` or `httpx.get` directly from any async handler — it enforces the allow list, disables redirects (so an allow-listed host can't bounce to an internal IP), caps the body at 10 MB during streaming, applies granular per-recv timeouts, wraps the whole download in an asyncio.wait_for total ceiling, and offloads PIL decode to a thread. `timeout_s` 是**總 wall-clock ceiling**（asyncio.wait_for 上限）。預設 30s 給標準 10 MB image 充足餘裕；valuation_observations 在 predict 同步 await 路徑上用 3s 做 fast-fail。Per-recv 粒度（connect/read/write/pool） 住在 `_PER_RECV_TIMEOUT` 常數（5/10/10/5）—— caller 不可 override， 維持 sibling parity 與 sync `lgbm_valuation_image._fetch_and_embed` 一致。 【為什麼 stream 而非 get？】 sync sibling `_fetch_and_embed` post-review (security conf 80) 教訓： `await client.get(url)` 會把整個 body 讀進記憶體才回。攻擊者宣告 Content-Length: 1024 但 server 實際 stream 100MB，post-hoc `len(resp.content)` 已經 OOM；asyncio.to_thread 並發場景下多 request 各 buffer 10MB+ 會耗 thread pool memory。stream + chunk-loop 累計，超過 cap 立即 raise — 硬上限就是 cap，不是事後檢查 （feedback_streaming_body_cap_before_full_read.md）。 【為什麼分 per-recv granularity 與 total ceiling 兩層？】 Pre-fix 用 scalar `httpx.AsyncClient(timeout=timeout_s)`，httpx 把這個 值套到每個 socket op (connect / read / write / pool) **各自**做計時。 攻擊者每 (read - 0.1)s 送 1 byte 進 chunk loop，每個 recv 都剛好不超 per-recv 上限，但連線可被拉到 `len(chunks) × read_timeout` 才斷。 10MB cap / 65KB chunk = 160 chunk，1600s 拖一條連線 — 32 條並發直接 塞滿 worker pool。修法兩層：(1) `_PER_RECV_TIMEOUT` 拆 connect/pool (TCP 握手) vs read/write (data 傳輸) 各自配額；(2) `asyncio.wait_for` 包整個 download 過程 — 即使每 recv 都在 per-recv 預算內，總 wall clock 超 ceiling 也立刻中止。兩層必須 coexist。 """ # offload to thread: validate_image_url runs blocking socket.getaddrinfo, # which freezes the event loop and starves every Gemma capability sharing # the loop when DNS is slow (uncached lookup / attacker-controlled slow # auth NS). Per .claude/rules/ml-service.md §Async discipline. await asyncio.to_thread(validate_image_url, url) async def _do_download() -> BytesIO: # BytesIO accumulator 取代 `chunks: list[bytes] + b"".join`：少一份 # 完整 bytes 副本，10 MB cap × 32 並發 worker peak 從 ~640 MB 降到 # ~320 MB。BytesIO.write 累進到單一內部緩衝（CPython 實作細節 — # 用 C-level over-allocation，比每次重新 alloc 一份新 bytes 來得便宜， # 但具體 growth 策略不該被本層 code 假設）；_decode_image 直接收 # BytesIO。Sibling parity 與 sync `lgbm_valuation_image._fetch_and_embed` # 一致。 buf = BytesIO() total = 0 async with httpx.AsyncClient( timeout=_PER_RECV_TIMEOUT, follow_redirects=False ) as client: async with client.stream("GET", url) as resp: resp.raise_for_status() # 重要：aiter_bytes() 回傳「已解壓」bytes（httpx 透明解碼 # Content-Encoding: gzip / brotli）。禁止改為 aiter_raw() — # raw 是「壓縮前」bytes，會讓 gzip bomb 用 10KB 壓縮資料宣告 # 1GB 解壓後內容繞過 _MAX_IMAGE_BYTES cap（post-review security # P2，conf 80）。10MB cap 套在解壓後大小才正確。 async for chunk in resp.aiter_bytes(chunk_size=65536): total += len(chunk) if total > _MAX_IMAGE_BYTES: raise ValueError( f"Image too large: > {_MAX_IMAGE_BYTES} bytes " "(streaming cap)" ) buf.write(chunk) return buf try: buf = await asyncio.wait_for(_do_download(), timeout=timeout_s) except asyncio.TimeoutError as exc: # Convert to ValueError so caller's uniform `except Exception` log # path stays consistent with body-cap + SSRF reject (1-line # `safe_url, exc` warning, no FastAPI default 500 stack trace). raise ValueError( f"Image download exceeded total ceiling: > {timeout_s}s" ) from exc # PIL decode is CPU-bound (100-300ms for 10MB JPEG). Offload to thread # to keep the event loop free for other coroutines — under 4-caller # async load (`/verify/local` dual + outfit_match + valuation observation) # PIL on the loop serialises every other inference. See _decode_image # docstring for the load()-then-convert ordering rationale + BytesIO # ownership lifetime constraint (buf must outlive .convert). return await asyncio.to_thread(_decode_image, buf)