src/ssrf.py

"""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)