//! Defensive helpers for the image-download paths.
//!
//! Two threats this module addresses:
//!
//! - **SSRF**: a scraped chapter or manga page can embed an absolute
//!   `<img src="http://10.0.0.1/...">`. The crawler runs inside the
//!   backend container with intra-compose access to `postgres:5432`
//!   and possibly other internal services; without a host check the
//!   crawler would happily probe them. [`is_safe_url`] rejects
//!   anything whose host isn't on the operator-configured allowlist,
//!   plus any IP literal in RFC1918 / loopback / link-local / unique-
//!   local space (including IPv4-mapped IPv6 like `::ffff:127.0.0.1`)
//!   as a second defence for the case where an allowlisted hostname's
//!   DNS happens to resolve to a literal private address.
//!
//!   **DNS rebinding is not covered.** A hostname like `cdn.allowed.com`
//!   that *resolves* to `127.0.0.1` via hostile DNS bypasses the IP
//!   check entirely — `is_safe_url` only inspects URL strings, not
//!   resolved IPs. Mitigating that requires a custom reqwest resolver
//!   that filters IPs after DNS, which would mean rebuilding reqwest's
//!   connector. The allowlist + good operator DNS hygiene is the
//!   realistic mitigation today.
//!
//! - **Unbounded download**: `Response::bytes().await` reads the full
//!   body before returning. A malicious source serving a 10 GiB image
//!   would fill memory and then disk. [`accumulate_capped`] streams
//!   the body chunk-by-chunk into a [`bytes::BytesMut`] and bails as
//!   soon as the running total exceeds the cap.
//!
//! Both helpers are pure-data: the SSRF check is keyed off a parsed
//! URL string, and the byte accumulator is keyed off a generic stream.
//! Easy to unit-test without a live network or browser.

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::Arc;

use anyhow::{bail, Context};
use bytes::BytesMut;
use futures_util::StreamExt;
use reqwest::dns::{Addrs, Name, Resolve, Resolving};
use reqwest::Url;

/// Default per-image download cap. A page image is generally <2 MiB;
/// 32 MiB leaves headroom for high-resolution covers while still
/// stopping a misbehaving CDN dead. Override via `CRAWLER_MAX_IMAGE_BYTES`.
pub const DEFAULT_MAX_IMAGE_BYTES: usize = 32 * 1024 * 1024;

/// Hosts that are always allowed in addition to the operator's
/// configured allowlist. None by default — keeping the surface area
/// minimal so the only way a URL gets through is if it matches an
/// explicit catalog/CDN entry.
///
/// `allow_any` flips the host check off entirely (private-IP and
/// scheme checks still apply). It exists for operators whose sources
/// shard images across numbered CDN subdomains (`cdn1`, `cdn2`, …)
/// where enumerating each host upfront is impractical. Off by default.
#[derive(Clone, Debug, Default)]
pub struct DownloadAllowlist {
    hosts: Vec<String>,
    allow_any: bool,
}

impl DownloadAllowlist {
    pub fn new() -> Self {
        Self {
            hosts: Vec::new(),
            allow_any: false,
        }
    }

    /// Bypass the host allowlist. Scheme, localhost, and private-IP
    /// checks in [`is_safe_url`] continue to apply — this only opens
    /// up public hosts that weren't pre-enumerated.
    pub fn allow_any() -> Self {
        Self {
            hosts: Vec::new(),
            allow_any: true,
        }
    }

    /// Add a host (case-insensitive match). Sub-domains are *not*
    /// implied: pass `cdn.example.com` and `example.com` separately
    /// if both should be reachable.
    pub fn allow(mut self, host: impl Into<String>) -> Self {
        let h = host.into().to_ascii_lowercase();
        if !h.is_empty() && !self.hosts.iter().any(|existing| existing == &h) {
            self.hosts.push(h);
        }
        self
    }

    pub fn is_empty(&self) -> bool {
        self.hosts.is_empty()
    }

    /// The explicitly-allowed hosts (lowercased). Empty when `allow_any`.
    pub fn hosts(&self) -> &[String] {
        &self.hosts
    }

    /// Whether the host check is bypassed entirely (`CRAWLER_ALLOW_ANY_HOST`).
    pub fn is_allow_any(&self) -> bool {
        self.allow_any
    }

    pub fn contains(&self, host: &str) -> bool {
        if self.allow_any {
            return true;
        }
        let lower = host.to_ascii_lowercase();
        self.hosts.iter().any(|h| h == &lower)
    }
}

/// Verify a URL is safe for the crawler to fetch.
///
/// Rejects:
/// - non-http(s) schemes (file://, gopher://, …),
/// - any IP literal in private / loopback / link-local / unique-local
///   space (defense in depth — a DNS allowlist alone wouldn't cover an
///   attacker that places an entry like `cdn.evil` pointing at
///   `192.168.1.1`),
/// - the literal hostname `localhost`,
/// - hosts that aren't on the supplied allowlist.
///
/// An empty allowlist rejects everything (the conservative default —
/// callers must explicitly allow the catalog and CDN hosts).
pub fn is_safe_url(raw_url: &str, allow: &DownloadAllowlist) -> Result<(), UrlSafetyError> {
    let url = ensure_public_target_inner(raw_url)?;
    let lower_host = url
        .host_str()
        .expect("host validated by ensure_public_target_inner")
        .to_ascii_lowercase();
    if !allow.contains(&lower_host) {
        return Err(UrlSafetyError::HostNotAllowed(lower_host));
    }
    Ok(())
}

/// Validate that an admin-supplied URL points at a publicly-routable target:
/// scheme is http or https, host is present, host isn't `localhost`, and (if
/// the host is an IP literal) it isn't loopback/private/link-local/CGNAT/etc.
///
/// Used to validate `endpoint`-style admin settings (crawler `start_url`,
/// analysis `endpoint`) where there is no per-deployment allowlist to consult,
/// but where a hostile or careless admin value (e.g. `http://169.254.169.254/`,
/// `http://127.0.0.1:5432/`) would let the worker pivot inside the deployment.
///
/// Note: DNS hostnames (`mangalord-vision`, `vision.internal.corp`) pass — the
/// check is on *literal* IP private-range strings only, so the documented
/// docker-internal vision endpoint keeps working.
pub fn ensure_public_target(raw_url: &str) -> Result<(), UrlSafetyError> {
    ensure_public_target_inner(raw_url).map(|_| ())
}

fn ensure_public_target_inner(raw_url: &str) -> Result<Url, UrlSafetyError> {
    let url = Url::parse(raw_url).map_err(|_| UrlSafetyError::Unparseable)?;
    let scheme = url.scheme();
    if scheme != "http" && scheme != "https" {
        return Err(UrlSafetyError::BadScheme(scheme.to_string()));
    }
    let host = url.host_str().ok_or(UrlSafetyError::NoHost)?;
    let lower_host = host.to_ascii_lowercase();
    if lower_host == "localhost" {
        return Err(UrlSafetyError::Loopback);
    }
    // Reject IP literals in private/loopback ranges regardless of the
    // allowlist — if someone puts an IP literal on the allowlist they
    // almost certainly didn't mean a private range.
    // reqwest::Url normalises IPv6 literals as `[::1]` (brackets
    // included) in `host_str()`. Strip the brackets before parsing.
    let ip_candidate = lower_host
        .strip_prefix('[')
        .and_then(|s| s.strip_suffix(']'))
        .unwrap_or(&lower_host);
    if let Ok(ip) = ip_candidate.parse::<IpAddr>() {
        if is_private_ip(&ip) {
            return Err(UrlSafetyError::PrivateIp(ip));
        }
    }
    Ok(url)
}

pub(crate) fn is_private_ip(ip: &IpAddr) -> bool {
    match ip {
        IpAddr::V4(v4) => {
            v4.is_loopback()
                || v4.is_private()
                || v4.is_link_local()
                || v4.is_unspecified()
                || v4.is_broadcast()
                // CGNAT 100.64.0.0/10
                || (v4.octets()[0] == 100 && (v4.octets()[1] & 0xC0) == 64)
                // 169.254/16 link-local already covered, but 0.0.0.0/8 is special-use
                || v4.octets()[0] == 0
        }
        IpAddr::V6(v6) => {
            // Any IPv6 form that *embeds* an IPv4 address (mapped
            // `::ffff:0:0/96`, compatible `::/96`, NAT64 `64:ff9b::/96`,
            // 6to4 `2002::/16`) is unwrapped and re-checked as its IPv4 —
            // otherwise `::127.0.0.1` / `2002:7f00:1::` / `64:ff9b::7f00:1`
            // would smuggle an internal IPv4 past the check (the audit's
            // IPv6-embedding gap). `Ipv6Addr::is_loopback()` only matches
            // `::1` exactly, so these embeddings need explicit handling.
            if let Some(v4) = embedded_ipv4(v6) {
                return is_private_ip(&IpAddr::V4(v4));
            }
            v6.is_loopback()
                || v6.is_unspecified()
                // fc00::/7 unique-local
                || (v6.segments()[0] & 0xfe00) == 0xfc00
                // fe80::/10 link-local
                || (v6.segments()[0] & 0xffc0) == 0xfe80
        }
    }
}

/// Extract the IPv4 address embedded in an IPv6 literal, for every
/// transitional encoding that can carry one: IPv4-mapped (`::ffff:0:0/96`),
/// IPv4-compatible (`::/96`, deprecated but still routable via some stacks),
/// NAT64 (`64:ff9b::/96`), and 6to4 (`2002::/16`). Returns `None` for a
/// native IPv6 address. Callers recurse into [`is_private_ip`] on the result
/// so a private IPv4 can't hide inside an IPv6 literal.
fn embedded_ipv4(v6: &Ipv6Addr) -> Option<Ipv4Addr> {
    let seg = v6.segments();
    let low32 = |a: u16, b: u16| Ipv4Addr::new((a >> 8) as u8, (a & 0xff) as u8, (b >> 8) as u8, (b & 0xff) as u8);
    // ::ffff:0:0/96 (mapped) and ::/96 (compatible) — top 96 bits zero,
    // except mapped which has 0xffff at seg[5]. to_ipv4() covers both.
    if seg[0..5] == [0, 0, 0, 0, 0] && (seg[5] == 0 || seg[5] == 0xffff) {
        return Some(low32(seg[6], seg[7]));
    }
    // NAT64 64:ff9b::/96
    if seg[0] == 0x0064 && seg[1] == 0xff9b && seg[2..6] == [0, 0, 0, 0] {
        return Some(low32(seg[6], seg[7]));
    }
    // 6to4 2002::/16 — embedded IPv4 is bits 16..48 (seg[1], seg[2]).
    if seg[0] == 0x2002 {
        return Some(low32(seg[1], seg[2]));
    }
    None
}

/// A `reqwest::dns::Resolve` that performs normal system resolution, then
/// drops any resolved address in a private / loopback / link-local / metadata
/// range. Installed on every crawler + analysis reqwest client so a hostname
/// that resolves to an internal IP (DNS rebinding: an attacker-owned domain
/// with an `A` record for `169.254.169.254` or `10.x`) can never be connected
/// to — closing the TOCTOU gap that the string-only `ensure_public_target`
/// check leaves open. Fires per connection, so it also guards redirect hops.
#[derive(Debug, Default)]
pub struct SafeResolver;

/// Partition resolved addresses into the public ones, rejecting when nothing
/// survives. Split out from the async `resolve` so the security-critical
/// filter is unit-testable without real DNS.
fn retain_public_addrs(
    host: &str,
    addrs: impl Iterator<Item = SocketAddr>,
) -> Result<Vec<SocketAddr>, BlockedResolution> {
    let public: Vec<SocketAddr> = addrs.filter(|a| !is_private_ip(&a.ip())).collect();
    if public.is_empty() {
        return Err(BlockedResolution(host.to_string()));
    }
    Ok(public)
}

#[derive(Debug, thiserror::Error)]
#[error("host {0} resolved only to private/blocked addresses")]
struct BlockedResolution(String);

impl Resolve for SafeResolver {
    fn resolve(&self, name: Name) -> Resolving {
        Box::pin(async move {
            let host = name.as_str().to_string();
            // Port 0: reqwest overrides it with the URL's port after
            // resolution (same convention as reqwest's default GaiResolver).
            let resolved = tokio::net::lookup_host((host.as_str(), 0))
                .await
                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
            let public = retain_public_addrs(&host, resolved)
                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
            Ok(Box::new(public.into_iter()) as Addrs)
        })
    }
}

/// Shared [`SafeResolver`] for wiring into `ClientBuilder::dns_resolver`.
pub fn safe_dns_resolver() -> Arc<SafeResolver> {
    Arc::new(SafeResolver)
}

#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum UrlSafetyError {
    #[error("URL is not parseable")]
    Unparseable,
    #[error("scheme {0:?} is not http or https")]
    BadScheme(String),
    #[error("URL is missing a host")]
    NoHost,
    #[error("host points at the loopback interface")]
    Loopback,
    #[error("host is a private/internal IP: {0}")]
    PrivateIp(IpAddr),
    #[error("host {0:?} is not on the crawler download allowlist")]
    HostNotAllowed(String),
}

/// Maximum number of redirects the crawler will follow before giving up.
/// Matches reqwest's historical default; every hop is re-validated by
/// [`check_redirect_hop`], so the cap is a belt-and-braces stop against a
/// redirect loop rather than the primary SSRF defence.
pub const MAX_REDIRECTS: usize = 10;

/// Why a redirect hop was refused.
#[derive(Debug, thiserror::Error)]
pub enum RedirectError {
    #[error("redirect chain exceeded {0} hops")]
    TooManyHops(usize),
    #[error("redirect target rejected: {0}")]
    Unsafe(#[from] UrlSafetyError),
}

/// Decide whether a single redirect hop is safe to follow.
///
/// `is_safe_url` only inspects the *initial* URL a caller hands to reqwest;
/// without re-validation an allowlisted CDN that answers `302 ->
/// http://169.254.169.254/...` or `-> http://127.0.0.1:5432/` would be
/// followed transparently (reqwest's default policy follows up to 10
/// redirects). This re-runs the full allowlist + private-IP + scheme check on
/// each hop and enforces [`MAX_REDIRECTS`]. `completed_hops` is the number of
/// URLs already visited in the chain (reqwest's `attempt.previous().len()`).
pub fn check_redirect_hop(
    next_url: &str,
    completed_hops: usize,
    allow: &DownloadAllowlist,
) -> Result<(), RedirectError> {
    if completed_hops >= MAX_REDIRECTS {
        return Err(RedirectError::TooManyHops(completed_hops));
    }
    is_safe_url(next_url, allow)?;
    Ok(())
}

/// Build a reqwest redirect policy that re-validates every hop against the
/// download allowlist (see [`check_redirect_hop`]). Use for the crawler image
/// clients, which fetch attacker-influenced URLs.
pub fn safe_redirect_policy(allow: DownloadAllowlist) -> reqwest::redirect::Policy {
    reqwest::redirect::Policy::custom(move |attempt| {
        let hops = attempt.previous().len();
        match check_redirect_hop(attempt.url().as_str(), hops, &allow) {
            Ok(()) => attempt.follow(),
            Err(e) => attempt.error(e),
        }
    })
}

/// Build a reqwest redirect policy that re-validates every hop with
/// [`ensure_public_target`] (scheme + private-IP, no allowlist). Use for
/// single-endpoint clients (the analysis vision endpoint / its probe) where
/// there is no per-deployment allowlist but a redirect into the deployment's
/// internal network must still be refused.
pub fn public_redirect_policy() -> reqwest::redirect::Policy {
    reqwest::redirect::Policy::custom(move |attempt| {
        let hops = attempt.previous().len();
        if hops >= MAX_REDIRECTS {
            return attempt.error(RedirectError::TooManyHops(hops));
        }
        match ensure_public_target(attempt.url().as_str()) {
            Ok(()) => attempt.follow(),
            Err(e) => attempt.error(RedirectError::Unsafe(e)),
        }
    })
}

/// Drain a byte stream into a single buffer, bailing out as soon as
/// the running total exceeds `max_bytes`. Generic over the stream so
/// it's testable without a live HTTP response.
pub async fn accumulate_capped<S, E>(stream: S, max_bytes: usize) -> anyhow::Result<bytes::Bytes>
where
    S: futures_core::Stream<Item = Result<bytes::Bytes, E>>,
    E: std::error::Error + Send + Sync + 'static,
{
    let mut buf = BytesMut::new();
    let mut stream = std::pin::pin!(stream);
    while let Some(chunk) = stream.next().await {
        let chunk = chunk.map_err(|e| anyhow::anyhow!("stream chunk: {e}"))?;
        if buf.len().saturating_add(chunk.len()) > max_bytes {
            bail!(
                "response exceeds {max_bytes}-byte cap (received >{}+{})",
                buf.len(),
                chunk.len()
            );
        }
        buf.extend_from_slice(&chunk);
    }
    Ok(buf.freeze())
}

/// Send `req` and stream the response into a length-limited buffer.
/// Combines [`is_safe_url`] check + [`accumulate_capped`] so each
/// call-site is one line.
pub async fn fetch_bytes_capped(
    http: &reqwest::Client,
    url: &str,
    referer: Option<&str>,
    allow: &DownloadAllowlist,
    max_bytes: usize,
) -> anyhow::Result<bytes::Bytes> {
    is_safe_url(url, allow).with_context(|| format!("reject unsafe URL {url}"))?;
    let mut req = http.get(url);
    if let Some(r) = referer {
        req = req.header(reqwest::header::REFERER, r);
    }
    let resp = req
        .send()
        .await
        .with_context(|| format!("GET {url}"))?
        .error_for_status()
        .with_context(|| format!("non-2xx for {url}"))?;
    accumulate_capped(resp.bytes_stream(), max_bytes)
        .await
        .with_context(|| format!("download body for {url}"))
}

/// Send `req` and return the response body as a stream after the
/// safety check + 2xx status check. Caller owns chunking, capping, and
/// piping to storage. Used by `download_and_store_page` so peak memory
/// stays at one chunk per concurrent dispatch instead of one full
/// image.
pub async fn fetch_stream(
    http: &reqwest::Client,
    url: &str,
    referer: Option<&str>,
    allow: &DownloadAllowlist,
) -> anyhow::Result<reqwest::Response> {
    is_safe_url(url, allow).with_context(|| format!("reject unsafe URL {url}"))?;
    let mut req = http.get(url);
    if let Some(r) = referer {
        req = req.header(reqwest::header::REFERER, r);
    }
    let resp = req
        .send()
        .await
        .with_context(|| format!("GET {url}"))?
        .error_for_status()
        .with_context(|| format!("non-2xx for {url}"))?;
    Ok(resp)
}

/// True when `bytes` sniffs as one of the *renderable* image formats
/// the `/files/*key` endpoint can serve with a correct Content-Type:
/// JPEG, PNG, WebP, GIF, AVIF. Matches the upload pipeline's
/// whitelist in `upload::parse_image`.
///
/// `infer::MatcherType::Image` is intentionally NOT used — it also
/// matches BMP, TIFF, HEIF, ICO, PSD, and JP2. Those would sniff as
/// "image" here but [`api::files::content_type_for`] would fall back
/// to `application/octet-stream`, prompting browsers to download
/// instead of render. Keep the two layers aligned.
pub fn looks_like_image(bytes: &[u8]) -> bool {
    matches!(
        infer::get(bytes).map(|k| k.mime_type()),
        Some("image/jpeg" | "image/png" | "image/webp" | "image/gif" | "image/avif")
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures_util::stream;

    fn allow_just(host: &str) -> DownloadAllowlist {
        DownloadAllowlist::new().allow(host)
    }

    #[test]
    fn allow_any_admits_arbitrary_public_host() {
        // Operators who can't pre-enumerate a numbered-CDN fleet
        // (cdn1, cdn2, …) opt into allow_any. Any public host passes.
        let allow = DownloadAllowlist::allow_any();
        assert!(is_safe_url("https://cdn7.random.tld/x.jpg", &allow).is_ok());
        assert!(is_safe_url("https://anything-goes.example/", &allow).is_ok());
    }

    #[test]
    fn allow_any_still_blocks_private_ips() {
        // The point of the bypass is the host-allowlist check, not the
        // SSRF defense. Private/loopback IPs stay refused.
        let allow = DownloadAllowlist::allow_any();
        for url in [
            "http://10.0.0.1/",
            "http://192.168.1.1/",
            "http://169.254.169.254/",
            "http://127.0.0.1/",
            "http://[::1]/",
            "http://[::ffff:127.0.0.1]/",
        ] {
            assert!(
                matches!(
                    is_safe_url(url, &allow).unwrap_err(),
                    UrlSafetyError::PrivateIp(_)
                ),
                "allow_any must still reject {url}"
            );
        }
    }

    #[test]
    fn allow_any_still_blocks_localhost() {
        let allow = DownloadAllowlist::allow_any();
        assert!(matches!(
            is_safe_url("http://localhost:8080/", &allow).unwrap_err(),
            UrlSafetyError::Loopback
        ));
    }

    #[test]
    fn allow_any_still_blocks_non_http_schemes() {
        let allow = DownloadAllowlist::allow_any();
        assert!(matches!(
            is_safe_url("file:///etc/passwd", &allow).unwrap_err(),
            UrlSafetyError::BadScheme(_)
        ));
    }

    #[test]
    fn safe_url_allows_listed_host() {
        let allow = allow_just("cdn.example.com");
        assert!(is_safe_url("https://cdn.example.com/img.jpg", &allow).is_ok());
    }

    #[test]
    fn safe_url_blocks_unlisted_host() {
        let allow = allow_just("cdn.example.com");
        let err = is_safe_url("https://evil.example.org/img.jpg", &allow).unwrap_err();
        assert!(matches!(err, UrlSafetyError::HostNotAllowed(h) if h == "evil.example.org"));
    }

    #[test]
    fn safe_url_blocks_localhost_even_if_allowlisted() {
        let allow = allow_just("localhost");
        assert!(matches!(
            is_safe_url("http://localhost:8080/", &allow).unwrap_err(),
            UrlSafetyError::Loopback
        ));
    }

    #[test]
    fn safe_url_blocks_loopback_ipv4() {
        let allow = allow_just("127.0.0.1");
        assert!(matches!(
            is_safe_url("http://127.0.0.1/", &allow).unwrap_err(),
            UrlSafetyError::PrivateIp(_)
        ));
    }

    // --- ensure_public_target (allowlist-free admin URL validation) ---

    #[test]
    fn public_target_allows_dns_hostnames() {
        // The whole point of this helper is to validate admin-supplied
        // endpoint URLs (analysis vision endpoint, crawler start_url)
        // WITHOUT consulting an allowlist. Docker-internal DNS names
        // (which resolve at runtime to private IPs) MUST pass.
        assert!(ensure_public_target("http://mangalord-vision:8000/v1/chat/completions").is_ok());
        assert!(ensure_public_target("https://api.openai.com/v1/chat/completions").is_ok());
        assert!(ensure_public_target("https://example.com/").is_ok());
    }

    #[test]
    fn public_target_blocks_ip_literal_attacks() {
        // Literal IPs in private/loopback/link-local ranges — the routes
        // an attacker would actually use (AWS IMDS, postgres on
        // 127.0.0.1, RFC1918 inside a corp network).
        for url in [
            "http://169.254.169.254/latest/meta-data/",
            "http://127.0.0.1:5432/",
            "http://10.0.0.1/",
            "http://192.168.1.1/",
            "http://[::1]/",
            "http://[::ffff:127.0.0.1]/",
            "http://0.0.0.0/",
        ] {
            assert!(
                matches!(
                    ensure_public_target(url).unwrap_err(),
                    UrlSafetyError::PrivateIp(_)
                ),
                "must reject {url}"
            );
        }
    }

    #[test]
    fn public_target_blocks_localhost_hostname() {
        assert!(matches!(
            ensure_public_target("http://localhost:5432/").unwrap_err(),
            UrlSafetyError::Loopback
        ));
    }

    #[test]
    fn public_target_blocks_non_http_schemes() {
        for url in ["file:///etc/passwd", "gopher://x.example/", "ftp://x/"] {
            assert!(matches!(
                ensure_public_target(url).unwrap_err(),
                UrlSafetyError::BadScheme(_)
            ));
        }
    }

    #[test]
    fn public_target_rejects_unparseable() {
        assert!(matches!(
            ensure_public_target("not a url").unwrap_err(),
            UrlSafetyError::Unparseable
        ));
    }

    // --- back to is_safe_url ---

    #[test]
    fn safe_url_blocks_rfc1918() {
        let allow = allow_just("10.0.0.1");
        for url in [
            "http://10.0.0.1/",
            "http://192.168.1.1/",
            "http://172.16.0.5/",
            "http://172.31.255.255/",
        ] {
            assert!(
                matches!(
                    is_safe_url(url, &allow).unwrap_err(),
                    UrlSafetyError::PrivateIp(_)
                ),
                "should reject {url}"
            );
        }
    }

    #[test]
    fn safe_url_blocks_link_local() {
        let allow = allow_just("169.254.169.254");
        // 169.254.169.254 is the AWS/GCP metadata service — the most
        // dangerous SSRF target on a default cloud VM.
        assert!(matches!(
            is_safe_url("http://169.254.169.254/", &allow).unwrap_err(),
            UrlSafetyError::PrivateIp(_)
        ));
    }

    #[test]
    fn safe_url_blocks_ipv6_loopback_and_ula() {
        // Debug what host_str returns first — reqwest::Url normalises
        // IPv6 literals as `[::1]` with brackets, which doesn't parse
        // as `IpAddr` directly. The implementation strips them.
        let allow = allow_just("[::1]");
        let err = is_safe_url("http://[::1]/", &allow).unwrap_err();
        assert!(
            matches!(err, UrlSafetyError::PrivateIp(_)),
            "expected PrivateIp, got {err:?}"
        );
        let allow = allow_just("[fd00::1]");
        let err = is_safe_url("http://[fd00::1]/", &allow).unwrap_err();
        assert!(
            matches!(err, UrlSafetyError::PrivateIp(_)),
            "expected PrivateIp, got {err:?}"
        );
    }

    #[test]
    fn safe_url_blocks_ipv4_mapped_ipv6_loopback() {
        // `Ipv6Addr::is_loopback()` only matches `::1` exactly, so
        // `::ffff:127.0.0.1` would slip through without the
        // to_ipv4_mapped() unwrap in is_private_ip.
        let allow = allow_just("[::ffff:127.0.0.1]");
        let err = is_safe_url("http://[::ffff:127.0.0.1]/", &allow).unwrap_err();
        assert!(
            matches!(err, UrlSafetyError::PrivateIp(_)),
            "expected PrivateIp, got {err:?}"
        );
    }

    #[test]
    fn safe_url_blocks_ipv4_mapped_ipv6_rfc1918() {
        let allow = allow_just("[::ffff:10.0.0.1]");
        let err = is_safe_url("http://[::ffff:10.0.0.1]/", &allow).unwrap_err();
        assert!(matches!(err, UrlSafetyError::PrivateIp(_)));
    }

    #[test]
    fn is_private_ip_unwraps_embedded_ipv4_encodings() {
        // Every IPv6 encoding that can smuggle an internal IPv4 must be
        // caught. The audit flagged compatible ::/96, NAT64, and 6to4 as
        // gaps past the original mapped-only handling.
        for s in [
            "::ffff:127.0.0.1", // IPv4-mapped loopback
            "::127.0.0.1",      // IPv4-compatible loopback (was a gap)
            "::ffff:10.1.2.3",  // mapped RFC1918
            "::10.1.2.3",       // compatible RFC1918 (was a gap)
            "64:ff9b::7f00:1",  // NAT64 of 127.0.0.1 (was a gap)
            "64:ff9b::a01:203", // NAT64 of 10.1.2.3
            "2002:7f00:1::",    // 6to4 of 127.0.0.1 (was a gap)
            "2002:a01:203::",   // 6to4 of 10.1.2.3
            "2002:a9fe:a9fe::",  // 6to4 of 169.254.169.254 (metadata)
        ] {
            let ip: IpAddr = s.parse().unwrap();
            assert!(is_private_ip(&ip), "{s} must be flagged private");
        }
    }

    #[test]
    fn is_private_ip_allows_public_embedded_and_native_ipv6() {
        // A public IPv4 embedded in IPv6, and a native public IPv6, must
        // NOT be flagged — the unwrap only blocks when the embedded v4 is
        // itself private.
        for s in [
            "::ffff:8.8.8.8",  // mapped public
            "2002:808:808::",  // 6to4 of 8.8.8.8 (public)
            "2606:4700:4700::1111", // native public (Cloudflare)
        ] {
            let ip: IpAddr = s.parse().unwrap();
            assert!(!is_private_ip(&ip), "{s} must be allowed");
        }
    }

    #[test]
    fn retain_public_addrs_drops_private_and_errors_when_all_private() {
        use std::net::{Ipv4Addr, SocketAddr};
        let pub_addr = SocketAddr::from((Ipv4Addr::new(93, 184, 216, 34), 0));
        let loopback = SocketAddr::from((Ipv4Addr::new(127, 0, 0, 1), 0));
        let metadata = SocketAddr::from((Ipv4Addr::new(169, 254, 169, 254), 0));

        // Mixed result keeps only the public address.
        let kept =
            retain_public_addrs("mixed.example", [pub_addr, loopback, metadata].into_iter())
                .expect("public address survives");
        assert_eq!(kept, vec![pub_addr]);

        // All-private (DNS rebinding to internal) is rejected outright.
        let err =
            retain_public_addrs("rebind.attacker", [loopback, metadata].into_iter()).unwrap_err();
        assert!(err.to_string().contains("rebind.attacker"));
    }

    #[test]
    fn safe_url_blocks_non_http_schemes() {
        let allow = allow_just("anywhere");
        assert!(matches!(
            is_safe_url("file:///etc/passwd", &allow).unwrap_err(),
            UrlSafetyError::BadScheme(_)
        ));
        assert!(matches!(
            is_safe_url("gopher://anywhere:70/", &allow).unwrap_err(),
            UrlSafetyError::BadScheme(_)
        ));
    }

    #[test]
    fn safe_url_rejects_unparseable() {
        let allow = allow_just("anywhere");
        assert!(matches!(
            is_safe_url("not a url", &allow).unwrap_err(),
            UrlSafetyError::Unparseable
        ));
    }

    #[test]
    fn safe_url_empty_allowlist_rejects_everything() {
        let allow = DownloadAllowlist::new();
        let err = is_safe_url("https://cdn.example.com/img.jpg", &allow).unwrap_err();
        assert!(matches!(err, UrlSafetyError::HostNotAllowed(_)));
    }

    #[test]
    fn allowlist_matches_case_insensitively() {
        let allow = DownloadAllowlist::new().allow("CDN.Example.COM");
        assert!(is_safe_url("https://cdn.example.com/x.jpg", &allow).is_ok());
        assert!(is_safe_url("https://CDN.EXAMPLE.com/x.jpg", &allow).is_ok());
    }

    // --- redirect-hop re-validation (SSRF via 3xx) ---

    #[test]
    fn redirect_hop_allows_listed_public_target() {
        let allow = allow_just("cdn.example.com");
        assert!(check_redirect_hop("https://cdn.example.com/next.jpg", 1, &allow).is_ok());
    }

    #[test]
    fn redirect_hop_blocks_private_ip_target() {
        // The core SSRF case: an allowlisted CDN 302s to the cloud metadata
        // service / an intra-compose port. Must be refused mid-chain.
        let allow = allow_just("cdn.example.com");
        for url in ["http://169.254.169.254/", "http://127.0.0.1:5432/", "http://10.0.0.1/"] {
            let err = check_redirect_hop(url, 1, &allow).unwrap_err();
            assert!(
                matches!(err, RedirectError::Unsafe(UrlSafetyError::PrivateIp(_))),
                "expected PrivateIp for {url}, got {err:?}"
            );
        }
    }

    #[test]
    fn redirect_hop_blocks_off_allowlist_public_host() {
        // Per the strict policy: a redirect to an unlisted *public* host is
        // also refused (allow_any covers the numbered-CDN case instead).
        let allow = allow_just("cdn.example.com");
        let err = check_redirect_hop("https://evil.example.org/x", 1, &allow).unwrap_err();
        assert!(matches!(err, RedirectError::Unsafe(UrlSafetyError::HostNotAllowed(_))));
    }

    #[test]
    fn redirect_hop_blocks_bad_scheme_target() {
        let allow = DownloadAllowlist::allow_any();
        let err = check_redirect_hop("file:///etc/passwd", 1, &allow).unwrap_err();
        assert!(matches!(err, RedirectError::Unsafe(UrlSafetyError::BadScheme(_))));
    }

    #[test]
    fn redirect_hop_caps_chain_length() {
        let allow = allow_just("cdn.example.com");
        // A safe target is still refused once the hop cap is reached, so a
        // redirect loop can't spin forever.
        let err = check_redirect_hop("https://cdn.example.com/x", MAX_REDIRECTS, &allow)
            .unwrap_err();
        assert!(matches!(err, RedirectError::TooManyHops(n) if n == MAX_REDIRECTS));
    }

    #[tokio::test]
    async fn accumulate_capped_returns_full_body_under_cap() {
        let chunks: Vec<Result<bytes::Bytes, std::io::Error>> = vec![
            Ok(bytes::Bytes::from_static(b"hello ")),
            Ok(bytes::Bytes::from_static(b"world")),
        ];
        let s = stream::iter(chunks);
        let out = accumulate_capped(s, 100).await.unwrap();
        assert_eq!(out.as_ref(), b"hello world");
    }

    #[tokio::test]
    async fn accumulate_capped_bails_past_cap() {
        let chunks: Vec<Result<bytes::Bytes, std::io::Error>> = vec![
            Ok(bytes::Bytes::from(vec![0u8; 50])),
            Ok(bytes::Bytes::from(vec![0u8; 60])),
        ];
        let s = stream::iter(chunks);
        let err = accumulate_capped(s, 100).await.unwrap_err();
        assert!(err.to_string().contains("100-byte cap"));
    }

    #[tokio::test]
    async fn accumulate_capped_surfaces_stream_errors() {
        let chunks: Vec<Result<bytes::Bytes, std::io::Error>> = vec![
            Ok(bytes::Bytes::from_static(b"ok")),
            Err(std::io::Error::other("network blip")),
        ];
        let s = stream::iter(chunks);
        let err = accumulate_capped(s, 100).await.unwrap_err();
        assert!(err.to_string().contains("network blip"));
    }

    #[test]
    fn looks_like_image_accepts_jpeg() {
        // JPEG SOI + APP0 segment.
        let jpeg = [0xff, 0xd8, 0xff, 0xe0, 0, 0x10, b'J', b'F', b'I', b'F'];
        assert!(looks_like_image(&jpeg));
    }

    #[test]
    fn looks_like_image_accepts_png() {
        let png = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0, 0, 0, 0];
        assert!(looks_like_image(&png));
    }

    #[test]
    fn looks_like_image_rejects_html_disguised_as_image() {
        let html = b"<html><body>not an image</body></html>";
        assert!(!looks_like_image(html));
    }

    #[test]
    fn looks_like_image_rejects_empty() {
        assert!(!looks_like_image(&[]));
    }

    #[test]
    fn looks_like_image_rejects_renderable_but_unsupported_formats() {
        // BMP, TIFF, ICO, PSD are `infer::MatcherType::Image` but the
        // /files/*key handler doesn't have Content-Type mappings for
        // them, so they'd be served as application/octet-stream and
        // download instead of render. Reject at the crawler so we
        // never land them in storage.
        // BMP magic: "BM" + 4-byte size.
        let bmp = [b'B', b'M', 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert!(!looks_like_image(&bmp), "BMP must be rejected (not renderable by /files)");

        // TIFF little-endian magic: "II" + 42.
        let tiff = [0x49, 0x49, 0x2a, 0x00, 0, 0, 0, 0];
        assert!(!looks_like_image(&tiff), "TIFF must be rejected");

        // ICO magic: 0x00,0x00,0x01,0x00.
        let ico = [0x00, 0x00, 0x01, 0x00, 1, 0, 16, 16, 0, 0, 1, 0, 0x18, 0, 0x40, 0, 0, 0, 0x16, 0, 0, 0];
        assert!(!looks_like_image(&ico), "ICO must be rejected");
    }

    #[test]
    fn looks_like_image_accepts_webp_gif_avif() {
        // Cover the three remaining whitelisted formats so a future
        // tightening that drops one would fail noisily.
        let webp = [
            b'R', b'I', b'F', b'F',
            0, 0, 0, 0,
            b'W', b'E', b'B', b'P',
            b'V', b'P', b'8', b' ',
        ];
        assert!(looks_like_image(&webp));

        let gif = [b'G', b'I', b'F', b'8', b'7', b'a', 0, 0, 0, 0];
        assert!(looks_like_image(&gif));

        let avif = [
            0x00, 0x00, 0x00, 0x18,
            b'f', b't', b'y', b'p',
            b'a', b'v', b'i', b'f',
            0x00, 0x00, 0x00, 0x00,
            b'm', b'i', b'f', b'1',
            b'a', b'v', b'i', b'f',
        ];
        assert!(looks_like_image(&avif));
    }
}