diff --git a/aimux-core/src/trace/hash.rs b/aimux-core/src/trace/hash.rs index e0b25273..5d4de58a 100644 --- a/aimux-core/src/trace/hash.rs +++ b/aimux-core/src/trace/hash.rs @@ -23,7 +23,7 @@ pub fn mix(mut x: u64) -> u64 { /// prefix included). // Rust 1.98 clippy suggests `as_chunks::<8>()`, which stabilized in 1.88; // the workspace MSRV is 1.85. Drop this allow when the MSRV moves past 1.88. -#[allow(clippy::chunks_exact_to_as_chunks)] +#[allow(unknown_lints, clippy::chunks_exact_to_as_chunks)] #[must_use] pub fn hash64(key: u64, data: &[u8]) -> u64 { let mut h = key ^ 0x9e37_79b9_7f4a_7c15; diff --git a/aimux-provider-utils/src/download_guard.rs b/aimux-provider-utils/src/download_guard.rs new file mode 100644 index 00000000..c009a608 --- /dev/null +++ b/aimux-provider-utils/src/download_guard.rs @@ -0,0 +1,386 @@ +//! SSRF guard for downloads of provider-supplied URLs. +//! +//! Providers return asset URLs, polling URLs, and result URLs in response +//! bodies; fetching them blindly lets a compromised or spoofed response steer +//! authenticated requests at internal services (cloud metadata, loopback, +//! RFC1918 space). This module validates such URLs, resolves and validates +//! every DNS answer, and hands the validated addresses back so the transport +//! can pin the actual connection to them (defeating TTL-0 DNS rebinding). +//! The address policy mirrors AI SDK's `validateUrl` blocklists. + +use std::net::IpAddr; + +use aimux_core::error::AiMuxError; + +/// Compare scheme, host, and effective port. Unparseable inputs are never +/// same-origin. Public so providers can build credential allowlists on top +/// of it (AI SDK's `isSameOrigin`). +#[must_use] +pub fn same_origin(url: &str, origin: &str) -> bool { + let (Ok(url), Ok(origin)) = (url::Url::parse(url), url::Url::parse(origin)) else { + return false; + }; + url.scheme() == origin.scheme() + && url.host_str() == origin.host_str() + && url.port_or_known_default() == origin.port_or_known_default() +} + +/// Trust flows only within the trusted origin: the exemption applies to a +/// redirect target only when the redirecting URL is itself on the trusted +/// origin, so a foreign hop cannot launder a request into it. +pub(crate) fn hop_trusted_origin<'a>( + trusted_origin: Option<&'a str>, + current_url: &str, +) -> Option<&'a str> { + trusted_origin.filter(|origin| same_origin(current_url, origin)) +} + +fn without_query(parsed: &url::Url) -> String { + let mut redacted = parsed.clone(); + redacted.set_query(None); + redacted.to_string() +} + +/// Parse a host as `url::Url::host_str` yields it into an IP literal. The url +/// crate keeps IPv6 hosts bracketed (and may render mapped addresses in hex +/// form, e.g. "[::ffff:7f00:1]"), so brackets are stripped before parsing. +fn host_ip_literal(host: &str) -> Option { + host.trim_start_matches('[') + .trim_end_matches(']') + .parse() + .ok() +} + +/// Syntactic checks: scheme, disallowed hostnames, and literal-IP publicness. +/// +/// Callers that resolve DNS afterwards rely on this having rejected every +/// non-public literal, so keep this the single literal-check site. +pub(crate) fn validate_download_url(url: &str) -> Result { + let parsed = url::Url::parse(url) + .map_err(|error| AiMuxError::InvalidArgument(format!("invalid download URL: {error}")))?; + if !matches!(parsed.scheme(), "http" | "https") { + return Err(AiMuxError::InvalidArgument(format!( + "download URL must use HTTP or HTTPS: {}", + without_query(&parsed) + ))); + } + let host = parsed + .host_str() + .ok_or_else(|| AiMuxError::InvalidArgument("download URL has no host".into()))?; + let normalized_host = host.to_ascii_lowercase(); + let normalized_host = normalized_host.trim_end_matches('.'); + if normalized_host == "localhost" + || normalized_host.ends_with(".local") + || normalized_host.ends_with(".localhost") + { + return Err(AiMuxError::InvalidArgument(format!( + "download URL targets a disallowed hostname: {normalized_host}" + ))); + } + if let Some(address) = host_ip_literal(normalized_host) + && !is_public_download_address(address) + { + return Err(AiMuxError::InvalidArgument(format!( + "download URL targets a non-public address: {address}" + ))); + } + Ok(parsed) +} + +/// Validate a download URL and resolve its host, returning the DNS answers +/// that passed the guard so the connection can be pinned to them. +/// +/// A URL same-origin with `trusted_origin` (normally the configured +/// `base_url`) is exempt and returns no addresses; self-hosted deployments +/// legitimately serve assets from private space on their own origin. +pub(crate) async fn validate_download_target( + url: &str, + trusted_origin: Option<&str>, +) -> Result, AiMuxError> { + if trusted_origin.is_some_and(|origin| same_origin(url, origin)) { + return Ok(Vec::new()); + } + let parsed = validate_download_url(url)?; + let host = parsed + .host_str() + .ok_or_else(|| AiMuxError::InvalidArgument("download URL has no host".into()))?; + if let Some(address) = host_ip_literal(host) { + // validate_download_url above already rejected non-public literals. + return Ok(vec![address]); + } + let port = parsed + .port_or_known_default() + .ok_or_else(|| AiMuxError::InvalidArgument("download URL has no known port".into()))?; + let addresses: Vec<_> = tokio::net::lookup_host((host, port)) + .await + .map_err(|error| { + AiMuxError::InvalidArgument(format!("download URL host could not be resolved: {error}")) + })? + .map(|address| address.ip()) + .collect(); + validate_resolved_download_addresses(host, addresses) +} + +fn validate_resolved_download_addresses( + host: &str, + addresses: Vec, +) -> Result, AiMuxError> { + if addresses.is_empty() { + return Err(AiMuxError::InvalidArgument(format!( + "download URL host did not resolve to an address: {host}" + ))); + } + let mut validated = Vec::with_capacity(addresses.len()); + for address in addresses { + if !is_public_download_address(address) { + return Err(AiMuxError::InvalidArgument(format!( + "download URL resolves to a non-public address: {address}" + ))); + } + if !validated.contains(&address) { + validated.push(address); + } + } + Ok(validated) +} + +fn is_public_download_address(address: IpAddr) -> bool { + match address { + IpAddr::V4(address) => { + let [a, b, c, _] = address.octets(); + !(a == 0 + || a == 10 + || (a == 100 && (64..=127).contains(&b)) + || a == 127 + || (a == 169 && b == 254) + || (a == 172 && (16..=31).contains(&b)) + || (a == 192 && b == 0 && matches!(c, 0 | 2)) + || (a == 192 && b == 168) + || (a == 198 && matches!(b, 18 | 19)) + || (a == 198 && b == 51 && c == 100) + || (a == 203 && b == 0 && c == 113) + || a >= 224) + } + IpAddr::V6(address) => { + let groups = address.segments(); + let top_zero = |count: usize| groups[..count].iter().all(|group| *group == 0); + if (top_zero(7) && matches!(groups[7], 0 | 1)) + || (groups[0] & 0xfe00) == 0xfc00 + || (groups[0] & 0xffc0) == 0xfe80 + || (groups[0] & 0xffc0) == 0xfec0 + || (groups[0] & 0xff00) == 0xff00 + || (groups[0] == 0x2001 && groups[1] == 0x0db8) + || (groups[0] == 0x3fff && (groups[1] & 0xf000) == 0) + { + return false; + } + + // Transition prefixes that embed an IPv4 address are judged by + // the embedded IPv4 bits: IPv4-compatible (::a.b.c.d), mapped + // (::ffff:a.b.c.d), SIIT (::ffff:0:a.b.c.d), and NAT64 + // (64:ff9b::/96, 64:ff9b:1::/48). 6to4/Teredo are deliberately + // omitted for parity with AI SDK's isPrivateIPv6. + let embeds_ipv4 = top_zero(6) + || (top_zero(5) && groups[5] == 0xffff) + || (top_zero(4) && groups[4] == 0xffff && groups[5] == 0) + || (groups[0] == 0x0064 + && groups[1] == 0xff9b + && groups[2..6].iter().all(|group| *group == 0)) + || (groups[0] == 0x0064 && groups[1] == 0xff9b && groups[2] == 1); + if !embeds_ipv4 { + return true; + } + + let embedded = std::net::Ipv4Addr::new( + (groups[6] >> 8) as u8, + groups[6] as u8, + (groups[7] >> 8) as u8, + groups[7] as u8, + ); + is_public_download_address(IpAddr::V4(embedded)) + } + } +} + +/// Drop hop-by-hop, forwarding, and metadata-service headers from a download +/// request; they leak deployment topology or unlock metadata endpoints when +/// forwarded to a provider-supplied host. Mirrors AI SDK's download header +/// policy (auth headers survive; the redirect loop clears them cross-origin). +pub(crate) fn sanitize_download_headers(headers: &mut Vec<(String, String)>) { + const BLOCKED: &[&str] = &[ + "connection", + "keep-alive", + "te", + "trailer", + "transfer-encoding", + "upgrade", + "host", + "forwarded", + "proxy-authorization", + "via", + "x-forwarded-for", + "x-forwarded-host", + "x-forwarded-proto", + "x-real-ip", + "metadata", + "metadata-flavor", + "x-aws-ec2-metadata-token", + "x-metadata-token", + "cookie", + "set-cookie", + ]; + headers.retain(|(name, _)| { + !BLOCKED + .iter() + .any(|blocked| name.eq_ignore_ascii_case(blocked)) + }); +} + +/// Drop caller headers except `User-Agent` when a redirect crosses origin. +pub(crate) fn retain_user_agent(headers: &mut Vec<(String, String)>) { + let user_agent = headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("user-agent")) + .cloned(); + headers.clear(); + if let Some(header) = user_agent { + headers.push(header); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn public(address: &str) -> bool { + is_public_download_address(address.parse().unwrap()) + } + + #[test] + fn address_policy_matches_ai_sdk_ranges() { + for address in [ + "0.1.2.3", + "10.0.0.1", + "100.64.0.1", + "100.127.255.255", + "127.255.0.1", + "169.254.169.254", + "172.16.0.1", + "172.31.255.255", + "192.0.0.1", + "192.0.2.1", + "192.168.1.1", + "198.18.0.1", + "198.19.255.255", + "198.51.100.1", + "203.0.113.1", + "224.0.0.1", + "255.255.255.255", + "::", + "::1", + "fc00::1", + "fd12::1", + "fe80::1", + "fec0::1", + "ff02::1", + "2001:db8::1", + "3fff::1", + "3fff:fff::1", + "::7f00:1", + "::ffff:7f00:1", + "::ffff:0:7f00:1", + "64:ff9b::7f00:1", + "64:ff9b::a9fe:a9fe", + "64:ff9b:1::a9fe:a9fe", + ] { + assert!(!public(address), "{address} must be blocked"); + } + + for address in [ + "8.8.8.8", + "100.63.0.1", + "100.128.0.1", + "172.15.0.1", + "172.32.0.1", + "192.0.3.1", + "198.51.101.1", + "203.0.114.1", + "2606:4700::1", + "3fff:1000::1", + "::ffff:808:808", + "64:ff9b::808:808", + ] { + assert!(public(address), "{address} must be allowed"); + } + } + + #[test] + fn url_policy_matches_ai_sdk_hostname_and_scheme_rules() { + for url in [ + "http://localhost/file", + "http://localhost./file", + "http://myhost.local/file", + "http://myhost.local./file", + "http://app.localhost/file", + "http://app.localhost./file", + "http://2130706433/file", + "http://0x7f000001/file", + "http://0177.0.0.1/file", + "http://[::127.0.0.1]/file", + "http://[::ffff:0:127.0.0.1]/file", + "http://[64:ff9b::169.254.169.254]/file", + "http://[64:ff9b:1::169.254.169.254]/file", + "file:///etc/passwd", + "ftp://example.com/file", + "data:text/plain;base64,aGVsbG8=", + "javascript:alert(1)", + ] { + assert!(validate_download_url(url).is_err(), "{url} must be blocked"); + } + + for url in [ + "https://example.com/file", + "https://example.com./file", + "http://8.8.8.8/file", + ] { + validate_download_url(url).unwrap_or_else(|error| { + panic!("{url} must be allowed: {error:?}"); + }); + } + } + + #[test] + fn resolved_address_validation_fails_closed_and_preserves_all_answers() { + let empty = validate_resolved_download_addresses("empty.example", Vec::new()) + .expect_err("an empty DNS answer must fail closed"); + assert!( + matches!(empty, AiMuxError::InvalidArgument(ref message) if message.contains("did not resolve")) + ); + + let mixed = validate_resolved_download_addresses( + "mixed.example", + vec!["8.8.8.8".parse().unwrap(), "127.0.0.1".parse().unwrap()], + ) + .expect_err("one private answer must reject the entire DNS result"); + assert!( + matches!(mixed, AiMuxError::InvalidArgument(ref message) if message.contains("non-public")) + ); + + let all = validate_resolved_download_addresses( + "public.example", + vec![ + "8.8.8.8".parse().unwrap(), + "1.1.1.1".parse().unwrap(), + "8.8.8.8".parse().unwrap(), + ], + ) + .unwrap(); + assert_eq!( + all, + vec![ + "8.8.8.8".parse::().unwrap(), + "1.1.1.1".parse::().unwrap() + ] + ); + } +} diff --git a/aimux-provider-utils/src/http.rs b/aimux-provider-utils/src/http.rs index 9d717f63..37de35a3 100644 --- a/aimux-provider-utils/src/http.rs +++ b/aimux-provider-utils/src/http.rs @@ -166,6 +166,27 @@ pub fn shared_client() -> Result<&'static Client, AiMuxError> { PoolConfig::default(), TimeoutConfig::default(), global_proxy(), + RedirectMode::Automatic, + ) + }) + .as_ref() + .map_err(|e| client_init_error(e)) +} + +/// 校验下载共享 Client（带 30s 整体超时，禁自动重定向——逐跳手动校验）。 +static DOWNLOAD: OnceLock> = OnceLock::new(); + +/// Shared client for validated downloads whose hop is not pinned (trusted +/// origin, or a proxy owns resolution). Redirects are followed manually by +/// [`send_validated_redirects`] so each hop is re-validated. +fn download_client() -> Result<&'static Client, AiMuxError> { + DOWNLOAD + .get_or_init(|| { + build_client( + PoolConfig::default(), + TimeoutConfig::default(), + global_proxy(), + RedirectMode::Manual, ) }) .as_ref() @@ -185,6 +206,7 @@ pub fn shared_streaming_client() -> Result<&'static Client, AiMuxError> { PoolConfig::default(), TimeoutConfig::streaming(), global_proxy(), + RedirectMode::Automatic, ) }) .as_ref() @@ -194,12 +216,46 @@ pub fn shared_streaming_client() -> Result<&'static Client, AiMuxError> { /// 用给定配置构建一个 reqwest Client。构建失败返回错误字符串（reqwest 仅 /// 提供 Display），由调用方决定映射——不再 `expect`（issue #115：受限环境 /// 下 TLS/资源初始化失败不应在 panic=abort 产物中直接终止宿主进程）。 +/// One redirect ceiling for both modes: reqwest's automatic following and +/// the manual validated-download loop. +const MAX_REDIRECTS: usize = 10; + +/// How a client treats redirects. `Manual` is for the validated-download +/// path, which follows redirects itself so every hop can be re-validated +/// and re-pinned. +#[derive(Clone, Copy)] +enum RedirectMode { + Automatic, + Manual, +} + +impl RedirectMode { + fn policy(self) -> reqwest::redirect::Policy { + match self { + Self::Automatic => reqwest::redirect::Policy::limited(MAX_REDIRECTS), + Self::Manual => reqwest::redirect::Policy::none(), + } + } +} + fn build_client( pool: PoolConfig, timeout: TimeoutConfig, proxy: ProxyConfig, + redirects: RedirectMode, ) -> Result { + apply_proxy(client_builder(&pool, &timeout, redirects), &proxy) + .build() + .map_err(|e| e.to_string()) +} + +fn client_builder( + pool: &PoolConfig, + timeout: &TimeoutConfig, + redirects: RedirectMode, +) -> reqwest::ClientBuilder { let mut b = Client::builder() + .redirect(redirects.policy()) .connect_timeout(Duration::from_millis(timeout.connect_timeout_ms)) .pool_max_idle_per_host(pool.max_idle_per_host) .pool_idle_timeout(Some(Duration::from_secs(pool.idle_timeout_secs))); @@ -209,8 +265,51 @@ fn build_client( if timeout.response_timeout_ms > 0 { b = b.timeout(Duration::from_millis(timeout.response_timeout_ms)); } - b = apply_proxy(b, &proxy); - b.build().map_err(|e| e.to_string()) + b +} + +/// Build a one-off client whose resolver only ever answers with the +/// pre-validated addresses. Used for downloads whose DNS results passed the +/// SSRF guard; prevents rebinding between validation and connection. +/// +/// `reqwest` matches `resolve` entries by exact host, so the URL's host (and +/// port, which `reqwest` reuses for the socket) must be supplied explicitly; +/// port 0 would make it dial port 0. +fn pinned_client(url: &str, addresses: &[std::net::IpAddr]) -> Result { + if addresses.is_empty() { + return Err(AiMuxError::Other( + "cannot build a pinned download client without an address".into(), + )); + } + let parsed = reqwest::Url::parse(url) + .map_err(|e| AiMuxError::Other(format!("invalid download url: {e}")))?; + let host = parsed + .host_str() + .ok_or_else(|| AiMuxError::Other("download url has no host".to_string()))? + .to_string(); + let port = parsed.port_or_known_default().unwrap_or(80); + // The proxy configuration is applied so reqwest makes the per-URL + // routing decision itself: when a proxy carries the request the proxy + // resolves the target (a trusted transport, the override below is + // unused), and any request the proxy rules send DIRECT — a NO_PROXY + // match, or no proxy configured for the URL's scheme — still connects + // only through the validated, pinned addresses. + let mut b = apply_proxy( + client_builder( + &PoolConfig::default(), + &TimeoutConfig::default(), + RedirectMode::Manual, + ), + &global_proxy(), + ); + let socket_addresses: Vec<_> = addresses + .iter() + .map(|address| std::net::SocketAddr::new(*address, port)) + .collect(); + b = b.resolve_to_addrs(&host, &socket_addresses); + b.build().map_err(|e| { + AiMuxError::Other(format!("pinned download client initialization failed: {e}")) + }) } /// Apply proxy configuration to a reqwest client builder (by-value chain). @@ -376,9 +475,69 @@ pub async fn send( ) -> Result { auto_init_from_env(); let client = shared_client()?; + send_via( + RequestTransport::Direct(client), + request, + retry_config, + error_structure, + ) + .await +} + +/// Fetch a **provider-supplied URL** (a generated asset, polling URL, result +/// URL, or upload URL taken from a response body or header) with SSRF +/// protection: the URL is validated against AI SDK's address blocklists, +/// every DNS answer is checked and the connection pinned to the validated +/// addresses (defeating TTL-0 rebinding), every redirect hop is re-validated, +/// and headers are sanitized (hop-by-hop, forwarding, and metadata-service +/// headers dropped). +/// +/// Both origins mirror AI SDK's options and must come from developer +/// configuration or a provider's own allowlist — never from a response: +/// +/// - `trusted_origin` (normally the configured `base_url`) exempts +/// same-origin URLs from the address blocklist so self-hosted deployments +/// keep working. It is only about reachability, never about headers. +/// - `credentialed_origin` — AI SDK's `credentialedOrigin` — confines caller +/// headers (which may carry the provider API key) to that origin, from the +/// first request and on every redirect hop; once stripped they are never +/// restored. `None` means the caller gates its own headers (e.g. BFL's +/// host allowlist); headers then still strip on any redirect leaving the +/// request's origin. +/// +/// # Errors +/// +/// Returns [`AiMuxError::InvalidArgument`] when the URL or a DNS answer +/// fails validation, plus the same transport/response errors as [`send`]. +pub async fn send_validated( + request: HttpRequest, + trusted_origin: Option<&str>, + credentialed_origin: Option<&str>, + retry_config: RetryConfig, + error_structure: &ErrorStructure, +) -> Result { + auto_init_from_env(); + send_via( + RequestTransport::Validated { + trusted_origin, + credentialed_origin, + }, + request, + retry_config, + error_structure, + ) + .await +} + +async fn send_via( + transport: RequestTransport<'_>, + request: HttpRequest, + retry_config: RetryConfig, + error_structure: &ErrorStructure, +) -> Result { let started = Instant::now(); let (resp, attempt) = - send_with_retry_raw(client, &request, retry_config, error_structure).await?; + send_with_retry_raw(&transport, &request, retry_config, error_structure).await?; let status = resp.status().as_u16(); let headers = collect_headers(resp.headers()); @@ -494,8 +653,13 @@ pub async fn send_stream( auto_init_from_env(); let client = shared_streaming_client()?; let started = Instant::now(); - let (resp, attempt) = - send_with_retry_raw(client, &request, retry_config, error_structure).await?; + let (resp, attempt) = send_with_retry_raw( + &RequestTransport::Direct(client), + &request, + retry_config, + error_structure, + ) + .await?; let status = resp.status().as_u16(); let headers = collect_headers(resp.headers()); @@ -1326,8 +1490,143 @@ fn record_failed_exchange( /// /// 这是 http 层内部函数——`reqwest::Response` 不外泄。每次重试从 `&request` /// 重建 `RequestBuilder`（HttpRequest 是纯数据，可重复读）。 +/// How one attempt of the retry loop reaches the network: directly on a +/// client, or through the validated-download redirect loop. +enum RequestTransport<'a> { + Direct(&'a Client), + Validated { + trusted_origin: Option<&'a str>, + credentialed_origin: Option<&'a str>, + }, +} + +impl RequestTransport<'_> { + async fn send(&self, request: &HttpRequest) -> Result { + match self { + Self::Direct(client) => send_request(client, request).await, + Self::Validated { + trusted_origin, + credentialed_origin, + } => send_validated_redirects(request, *trusted_origin, *credentialed_origin).await, + } + } +} + +/// The redirect statuses fetch follows; 300/304 are responses, not hops. +fn is_redirect_status(status: reqwest::StatusCode) -> bool { + matches!( + status, + reqwest::StatusCode::MOVED_PERMANENTLY + | reqwest::StatusCode::FOUND + | reqwest::StatusCode::SEE_OTHER + | reqwest::StatusCode::TEMPORARY_REDIRECT + | reqwest::StatusCode::PERMANENT_REDIRECT + ) +} + +fn redirect_error(message: impl Into, status: reqwest::StatusCode) -> AiMuxError { + AiMuxError::ApiCall(ApiCallError { + message: message.into(), + status_code: Some(status.as_u16()), + is_retryable: false, + ..Default::default() + }) +} + +/// Send one attempt of a validated download, following redirects manually so +/// every hop is validated and its connection pinned to the DNS answers that +/// passed the guard. The whole chain counts as one attempt to the retry +/// layer, matching how the auto-following shared client behaves. +async fn send_validated_redirects( + request: &HttpRequest, + trusted_origin: Option<&str>, + credentialed_origin: Option<&str>, +) -> Result { + let mut current = request.clone(); + crate::download_guard::sanitize_download_headers(&mut current.headers); + // AI SDK's credentialedOrigin: when set, caller headers (which may carry + // the provider API key) are confined to that exact origin from the very + // first request. When unset the caller has already gated its own headers + // (e.g. BFL's host allowlist); redirects still strip below either way. + if let Some(origin) = credentialed_origin + && !crate::download_guard::same_origin(¤t.url, origin) + { + crate::download_guard::retain_user_agent(&mut current.headers); + } + // Headers never travel past this origin on a redirect; stripping is + // one-way, so a hop back onto it cannot restore them. + let credential_anchor = credentialed_origin.unwrap_or(&request.url); + let mut pinned = + crate::download_guard::validate_download_target(¤t.url, trusted_origin).await?; + for redirect_count in 0..=MAX_REDIRECTS { + let hop_client; + let client: &Client = if pinned.is_empty() { + // Empty pins mean the hop is on the trusted origin. + download_client()? + } else { + // DNS answers were validated; pin them so a direct connection + // can only reach the addresses that passed the guard. When a + // configured proxy carries the request instead, the proxy + // resolves the target and the pin is deliberately unused. + hop_client = pinned_client(¤t.url, &pinned)?; + &hop_client + }; + let response = send_request(client, ¤t).await?; + let status = response.status(); + if !is_redirect_status(status) { + return Ok(response); + } + let Some(location) = response.headers().get(reqwest::header::LOCATION) else { + return Ok(response); + }; + if redirect_count == MAX_REDIRECTS { + return Err(redirect_error("too many redirects", status)); + } + let location = location + .to_str() + .map(str::to_owned) + .map_err(|_| redirect_error("redirect location is not valid UTF-8", status))?; + // Dropping the unconsumed 3xx body releases its connection before a + // potentially slow DNS check for the next hop. + drop(response); + let base = url::Url::parse(¤t.url) + .map_err(|e| AiMuxError::InvalidArgument(format!("invalid request URL: {e}")))?; + let next = base + .join(&location) + .map_err(|e| redirect_error(format!("invalid redirect URL: {e}"), status))?; + // Fetch treats a redirect to a non-HTTP(S) scheme (data:, file:, ...) + // as a network error; following one would let the redirecting server + // fabricate a response outside the transport. + if !matches!(next.scheme(), "http" | "https") { + return Err(redirect_error( + format!("redirect to non-HTTP scheme: {}", next.scheme()), + status, + )); + } + let next = next.to_string(); + let hop_trusted = crate::download_guard::hop_trusted_origin(trusted_origin, ¤t.url); + pinned = crate::download_guard::validate_download_target(&next, hop_trusted).await?; + // Credentials are scoped to the credential anchor, not the previous + // hop: once a redirect leaves it, headers cannot come back. + if !crate::download_guard::same_origin(&next, credential_anchor) { + crate::download_guard::retain_user_agent(&mut current.headers); + } + if status == reqwest::StatusCode::SEE_OTHER + || (matches!( + status, + reqwest::StatusCode::MOVED_PERMANENTLY | reqwest::StatusCode::FOUND + ) && matches!(current.method, HttpMethod::Post)) + { + current.method = HttpMethod::Get; + current.body = HttpBody::Empty; + } + current.url = next; + } + unreachable!("redirect loop returns on response or error") +} + async fn send_with_retry_raw( - client: &Client, + transport: &RequestTransport<'_>, request: &HttpRequest, retry_config: RetryConfig, error_structure: &ErrorStructure, @@ -1347,7 +1646,7 @@ async fn send_with_retry_raw( for attempt in 0..=retry_config.max_retries { let attempt_start = Instant::now(); - let resp = send_request(client, request).await; + let resp = transport.send(request).await; let latency_ms = attempt_start.elapsed().as_millis() as u64; match resp { @@ -2087,4 +2386,35 @@ mod tests { // both are valid. We just assert it doesn't panic. let _ = result; } + + #[tokio::test] + async fn pinned_client_resolves_only_through_validated_addresses() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = [0; 1024]; + let _ = stream.read(&mut request).await.unwrap(); + stream + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok") + .await + .unwrap(); + }); + // .example never resolves in real DNS, so success proves the resolve + // override was used rather than the system resolver. + let url = format!("http://download.example:{port}/file"); + let client = + pinned_client(&url, &["127.0.0.1".parse().unwrap()]).expect("pinned client builds"); + let response = tokio::time::timeout(Duration::from_secs(2), client.get(url).send()) + .await + .expect("request must not hang") + .expect("the validated address must be used"); + assert_eq!(response.status(), reqwest::StatusCode::OK); + tokio::time::timeout(Duration::from_secs(2), server) + .await + .expect("test server must finish") + .unwrap(); + } } diff --git a/aimux-provider-utils/src/lib.rs b/aimux-provider-utils/src/lib.rs index d7d4f412..8d89e6cf 100644 --- a/aimux-provider-utils/src/lib.rs +++ b/aimux-provider-utils/src/lib.rs @@ -6,6 +6,7 @@ //! and retry logic — the Rust equivalents of `@ai-sdk/provider-utils`. pub mod api_key; +mod download_guard; pub mod headers; pub mod http; pub mod logging; @@ -19,11 +20,12 @@ pub mod url; pub mod ws; pub use api_key::load_api_key; +pub use download_guard::same_origin; pub use headers::with_user_agent_suffix; pub use http::{ HttpBody, HttpMethod, HttpRequest, HttpResponse, HttpStreamResponse, PoolConfig, ProxyConfig, RequestTimeout, TimeoutConfig, init_proxy, send, send_stream, send_stream_timed, send_timed, - shared_client, shared_streaming_client, sleep_or_abort, + send_validated, shared_client, shared_streaming_client, sleep_or_abort, }; pub use logging::{body_logging_enabled, init_logging, redact_body}; pub use multipart::{MultipartForm, media_type_to_extension}; diff --git a/aimux-provider-utils/tests/download_guard_test.rs b/aimux-provider-utils/tests/download_guard_test.rs new file mode 100644 index 00000000..a6478bf7 --- /dev/null +++ b/aimux-provider-utils/tests/download_guard_test.rs @@ -0,0 +1,232 @@ +//! SSRF guard wiring tests for `send_validated`. +//! +//! The guard must reject provider-supplied URLs that point at +//! private/loopback/link-local space (before any connection is attempted), +//! re-validate every redirect hop, and keep same-origin traffic against the +//! configured `trusted_origin` working — mirroring AI SDK's `validateUrl` + +//! `trustedOrigin` semantics. The mock server runs on loopback, so it is only +//! reachable through the trusted-origin exemption; anything the guard treats +//! as foreign is rejected as a non-public literal. + +use serde_json::json; +use wiremock::matchers::{header, method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +use aimux_core::error::AiMuxError; +use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; +use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send_validated}; + +fn request(url: String) -> HttpRequest { + HttpRequest { + method: HttpMethod::Get, + url, + headers: vec![("authorization".into(), "Bearer test".into())], + body: HttpBody::Empty, + abort_signal: None, + call_id: None, + recording_context: None, + } +} + +#[tokio::test] +async fn rejects_non_public_literal_urls_before_connecting() { + for url in [ + "http://169.254.169.254/latest/meta-data", + "http://127.0.0.1:9/file", + "http://[::ffff:169.254.169.254]/meta", + "http://10.1.2.3/file", + ] { + let error = send_validated( + request(url.into()), + None, + None, + RetryConfig::default(), + &DEFAULT_ERROR_STRUCTURE, + ) + .await + .expect_err("non-public literal must be rejected"); + assert!( + matches!(error, AiMuxError::InvalidArgument(ref m) if m.contains("non-public")), + "unexpected error for {url}: {error:?}" + ); + } +} + +#[tokio::test] +async fn trusted_origin_download_succeeds_and_keeps_auth_headers() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/generated/file")) + .and(header("authorization", "Bearer test")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "ok"}))) + .expect(1) + .mount(&server) + .await; + + let response = send_validated( + request(format!("{}/generated/file", server.uri())), + Some(&server.uri()), + Some(&server.uri()), + RetryConfig::default(), + &DEFAULT_ERROR_STRUCTURE, + ) + .await + .unwrap(); + assert_eq!(response.status, 200); + assert_eq!( + serde_json::from_slice::(&response.body).unwrap(), + json!({"value": "ok"}) + ); +} + +#[tokio::test] +async fn follows_a_relative_redirect_within_the_trusted_origin() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/start")) + .respond_with(ResponseTemplate::new(302).insert_header("location", "/final")) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/final")) + .respond_with(ResponseTemplate::new(200).set_body_bytes(b"done")) + .expect(1) + .mount(&server) + .await; + + let response = send_validated( + request(format!("{}/start", server.uri())), + Some(&server.uri()), + Some(&server.uri()), + RetryConfig::default(), + &DEFAULT_ERROR_STRUCTURE, + ) + .await + .unwrap(); + assert_eq!(response.status, 200); + assert_eq!(response.body.as_ref(), b"done"); +} + +#[tokio::test] +async fn rejects_a_redirect_to_a_non_public_target() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/start")) + .respond_with( + ResponseTemplate::new(302).insert_header("location", "http://169.254.169.254/meta"), + ) + .expect(1) + .mount(&server) + .await; + + let error = send_validated( + request(format!("{}/start", server.uri())), + Some(&server.uri()), + Some(&server.uri()), + RetryConfig::default(), + &DEFAULT_ERROR_STRUCTURE, + ) + .await + .expect_err("redirect to metadata IP must be rejected"); + assert!( + matches!(error, AiMuxError::InvalidArgument(ref m) if m.contains("non-public")), + "unexpected error: {error:?}" + ); +} + +#[tokio::test] +async fn rejects_a_redirect_onto_a_foreign_loopback_origin() { + // The trusted origin is one loopback server; a redirect to a DIFFERENT + // loopback origin must not inherit the exemption. + let server = MockServer::start().await; + let other = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/start")) + .respond_with( + ResponseTemplate::new(302) + .insert_header("location", format!("{}/private", other.uri()).as_str()), + ) + .expect(1) + .mount(&server) + .await; + + let error = send_validated( + request(format!("{}/start", server.uri())), + Some(&server.uri()), + Some(&server.uri()), + RetryConfig::default(), + &DEFAULT_ERROR_STRUCTURE, + ) + .await + .expect_err("foreign loopback origin must be rejected"); + assert!( + matches!(error, AiMuxError::InvalidArgument(ref m) if m.contains("non-public")), + "unexpected error: {error:?}" + ); +} + +#[tokio::test] +async fn rejects_a_redirect_to_a_data_url() { + // Fetch treats a redirect to a non-HTTP(S) scheme as a network error; a + // server must not be able to fabricate a response via Location: data:. + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/start")) + .respond_with( + ResponseTemplate::new(302).insert_header("location", "data:text/plain,forged"), + ) + .expect(1) + .mount(&server) + .await; + + let error = send_validated( + request(format!("{}/start", server.uri())), + Some(&server.uri()), + Some(&server.uri()), + RetryConfig::default(), + &DEFAULT_ERROR_STRUCTURE, + ) + .await + .expect_err("redirect to a data: URL must be rejected"); + assert!( + matches!(error, AiMuxError::ApiCall(ref e) if e.message.contains("non-HTTP scheme")), + "unexpected error: {error:?}" + ); +} + +#[tokio::test] +async fn sanitizes_metadata_and_cookie_headers_from_download_requests() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/file")) + .respond_with(ResponseTemplate::new(200).set_body_bytes(b"ok")) + .expect(1) + .mount(&server) + .await; + + let mut req = request(format!("{}/file", server.uri())); + req.headers.push(("Cookie".into(), "session=secret".into())); + req.headers + .push(("Metadata-Flavor".into(), "Google".into())); + send_validated( + req, + Some(&server.uri()), + Some(&server.uri()), + RetryConfig::default(), + &DEFAULT_ERROR_STRUCTURE, + ) + .await + .unwrap(); + + let received = server.received_requests().await.unwrap(); + assert_eq!(received.len(), 1); + let names: Vec = received[0] + .headers + .keys() + .map(|name| name.as_str().to_ascii_lowercase()) + .collect(); + assert!(!names.contains(&"cookie".to_string())); + assert!(!names.contains(&"metadata-flavor".to_string())); + assert!(names.contains(&"authorization".to_string())); +} diff --git a/aimux-providers/src/black_forest_labs.rs b/aimux-providers/src/black_forest_labs.rs index e0557c55..04019dec 100644 --- a/aimux-providers/src/black_forest_labs.rs +++ b/aimux-providers/src/black_forest_labs.rs @@ -20,8 +20,8 @@ use aimux_core::image_model::{ use aimux_core::shared::Warning; use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, sleep_or_abort, - without_trailing_slash, + HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, send_validated, + sleep_or_abort, without_trailing_slash, }; const DEFAULT_POLL_INTERVAL_MS: u64 = 500; @@ -140,6 +140,23 @@ fn gcd(a: u32, b: u32) -> u32 { if b == 0 { a } else { gcd(b, a % b) } } +/// AI SDK's `isTrustedUrl` (black-forest-labs-api.ts): credentials may go to +/// the configured origin, or over HTTPS to `bfl.ai` and its subdomains — BFL +/// serves polling and asset URLs from regional clusters. Allowlisted hosts +/// still get full URL/DNS validation; this gates only the headers. +fn bfl_trusted_url(url: &str, base_url: &str) -> bool { + if aimux_provider_utils::same_origin(url, base_url) { + return true; + } + let Ok(parsed) = url::Url::parse(url) else { + return false; + }; + parsed.scheme() == "https" + && parsed + .host_str() + .is_some_and(|host| host == "bfl.ai" || host.ends_with(".bfl.ai")) +} + #[async_trait] impl ImageModel for BlackForestLabsImageModel { fn provider(&self) -> &str { @@ -285,6 +302,15 @@ impl ImageModel for BlackForestLabsImageModel { let headers = self.build_headers(options.headers.as_ref()); let header_list: Vec<(String, String)> = headers.into_iter().collect(); + // AI SDK gates headers per URL via isTrustedUrl; response-supplied + // targets outside the BFL allowlist get none. + let gated_headers = |url: &str| -> Vec<(String, String)> { + if bfl_trusted_url(url, &self.config.base_url) { + header_list.clone() + } else { + vec![] + } + }; // Submit let resp = send( @@ -343,17 +369,23 @@ impl ImageModel for BlackForestLabsImageModel { let mut result_duration = None; for _ in 0..max_attempts { - let pr = send( + // AI SDK polls polling_url with validateUrl: true and gates the + // headers itself via isTrustedUrl (base_url origin or HTTPS + // *.bfl.ai), so credentialed_origin is None here. + let poll_url = poll_url_with_id.to_string(); + let pr = send_validated( HttpRequest { method: HttpMethod::Get, - url: poll_url_with_id.to_string(), - headers: header_list.clone(), + headers: gated_headers(&poll_url), + url: poll_url, body: HttpBody::Empty, abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, + Some(&self.config.base_url), + None, RetryConfig::default(), &DEFAULT_ERROR_STRUCTURE, ) @@ -405,18 +437,22 @@ impl ImageModel for BlackForestLabsImageModel { )) })?; - // Download image - let ir = send( + // Download image; result.sample is a URL from the poll response body. + // AI SDK sends its headers to trusted BFL hosts on the download too + // (isTrustedUrl-gated), so mirror the poll's header policy. + let ir = send_validated( HttpRequest { method: HttpMethod::Get, + headers: gated_headers(&image_url), url: image_url, - headers: vec![], body: HttpBody::Empty, abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, + Some(&self.config.base_url), + None, RetryConfig::default(), &DEFAULT_ERROR_STRUCTURE, ) @@ -465,3 +501,22 @@ impl ImageModel for BlackForestLabsImageModel { }) } } + +#[cfg(test)] +mod tests { + use super::bfl_trusted_url; + + #[test] + fn credentials_go_to_the_base_url_and_bfl_hosts_only() { + let base = "https://api.bfl.ai"; + assert!(bfl_trusted_url("https://api.bfl.ai/v1/get_result", base)); + assert!(bfl_trusted_url( + "https://api.us1.bfl.ai/v1/get_result", + base + )); + assert!(bfl_trusted_url("https://bfl.ai/x", base)); + assert!(!bfl_trusted_url("http://api.us1.bfl.ai/x", base)); + assert!(!bfl_trusted_url("https://evil-bfl.ai/x", base)); + assert!(!bfl_trusted_url("https://attacker.example/x", base)); + } +} diff --git a/aimux-providers/src/fal.rs b/aimux-providers/src/fal.rs index 6fc0e2d2..3bf570b4 100644 --- a/aimux-providers/src/fal.rs +++ b/aimux-providers/src/fal.rs @@ -20,8 +20,8 @@ use aimux_core::transcription_model::{ }; use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, sleep_or_abort, - without_trailing_slash, + HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, send_validated, + sleep_or_abort, without_trailing_slash, }; // ── Config ────────────────────────────────────────────────────────────────── @@ -551,7 +551,8 @@ impl ImageModel for FalImageModel { let mut downloaded: Vec> = Vec::new(); for img in &target_images { if let Some(url) = img.get("url").and_then(|v| v.as_str()) { - let ir = send( + // images[].url comes from the queue result response body. + let ir = send_validated( HttpRequest { method: HttpMethod::Get, url: url.to_string(), @@ -562,6 +563,8 @@ impl ImageModel for FalImageModel { call_id: None, recording_context: None, }, + Some(&self.config.base_url), + Some(&self.config.base_url), RetryConfig::default(), &DEFAULT_ERROR_STRUCTURE, ) diff --git a/aimux-providers/src/gladia.rs b/aimux-providers/src/gladia.rs index 646b2375..7f339c96 100644 --- a/aimux-providers/src/gladia.rs +++ b/aimux-providers/src/gladia.rs @@ -23,7 +23,7 @@ use aimux_core::transcription_model::{ use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; use aimux_provider_utils::{ HttpBody, HttpMethod, HttpRequest, MultipartForm, RetryConfig, load_api_key, - media_type_to_extension, send, sleep_or_abort, without_trailing_slash, + media_type_to_extension, send, send_validated, sleep_or_abort, without_trailing_slash, }; // ── Config ────────────────────────────────────────────────────────────────── @@ -274,7 +274,10 @@ impl TranscriptionModel for GladiaTranscriptionModel { ) .await?; - let resp = send( + // AI SDK polls result_url with validateUrl: true and + // credentialedOrigin = the API origin: the target is validated + // and headers survive only while it stays on base_url's origin. + let resp = send_validated( HttpRequest { method: HttpMethod::Get, url: init.result_url.clone(), @@ -288,6 +291,8 @@ impl TranscriptionModel for GladiaTranscriptionModel { call_id: None, recording_context: None, }, + Some(&self.config.base_url), + Some(&self.config.base_url), RetryConfig::default(), &DEFAULT_ERROR_STRUCTURE, ) diff --git a/aimux-providers/src/google/files.rs b/aimux-providers/src/google/files.rs index e4d68a15..d0e420c5 100644 --- a/aimux-providers/src/google/files.rs +++ b/aimux-providers/src/google/files.rs @@ -26,7 +26,9 @@ use aimux_core::shared::FileBytes; use aimux_core::types::Warning; use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send, sleep_or_abort}; +use aimux_provider_utils::{ + HttpBody, HttpMethod, HttpRequest, RetryConfig, send, send_validated, sleep_or_abort, +}; use super::GoogleConfig; @@ -237,7 +239,10 @@ impl Files for GoogleFiles { ), ]; - let upload_resp = send( + // The upload URL comes from the init response's x-goog-upload-url + // header and receives the user's file bytes; validate it. (AI SDK + // fetches this URL unvalidated — kept stricter here deliberately.) + let upload_resp = send_validated( HttpRequest { method: HttpMethod::Post, url: upload_url, @@ -248,6 +253,8 @@ impl Files for GoogleFiles { call_id: None, recording_context: None, }, + Some(&self.config.base_url), + Some(&self.config.base_url), RetryConfig::default(), &GOOGLE_ERROR_STRUCTURE, ) diff --git a/aimux-providers/src/luma.rs b/aimux-providers/src/luma.rs index a40364f4..2ba2403a 100644 --- a/aimux-providers/src/luma.rs +++ b/aimux-providers/src/luma.rs @@ -19,8 +19,8 @@ use aimux_core::image_model::{ use aimux_core::shared::Warning; use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, sleep_or_abort, - without_trailing_slash, + HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, send_validated, + sleep_or_abort, without_trailing_slash, }; const DEFAULT_POLL_INTERVAL_MS: u64 = 500; @@ -386,8 +386,8 @@ impl ImageModel for LumaImageModel { )) })?; - // Download image - let ir = send( + // Download image; assets.image is a URL from the poll response body. + let ir = send_validated( HttpRequest { method: HttpMethod::Get, url: image_url, @@ -398,6 +398,8 @@ impl ImageModel for LumaImageModel { call_id: None, recording_context: None, }, + Some(&self.config.base_url), + Some(&self.config.base_url), RetryConfig::default(), &DEFAULT_ERROR_STRUCTURE, ) diff --git a/aimux-providers/src/recraft.rs b/aimux-providers/src/recraft.rs index f3b820ed..e9f52e8e 100644 --- a/aimux-providers/src/recraft.rs +++ b/aimux-providers/src/recraft.rs @@ -22,7 +22,8 @@ use aimux_core::provider::Provider; use aimux_core::shared::Warning; use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, without_trailing_slash, + HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, send_validated, + without_trailing_slash, }; const DEFAULT_BASE_URL: &str = "https://external.api.recraft.ai/v1"; @@ -211,6 +212,7 @@ fn build_generation_body( async fn extract_images( response: &Value, abort_signal: Option, + base_url: &str, ) -> Result { let items = response.get("data").and_then(|d| d.as_array()); @@ -240,7 +242,8 @@ async fn extract_images( let mut binaries = Vec::with_capacity(urls.len()); for url in &urls { - let resp = send( + // data[].url is a generated-image URL from the response body. + let resp = send_validated( HttpRequest { method: HttpMethod::Get, url: url.clone(), @@ -251,6 +254,8 @@ async fn extract_images( call_id: None, recording_context: None, }, + Some(base_url), + Some(base_url), RetryConfig::default(), &DEFAULT_ERROR_STRUCTURE, ) @@ -312,7 +317,8 @@ impl ImageModel for RecraftImageModel { let response_headers = resp.headers; let value: Value = serde_json::from_slice(&resp.body)?; - let images = extract_images(&value, options.abort_signal.clone()).await?; + let images = + extract_images(&value, options.abort_signal.clone(), &self.config.base_url).await?; Ok(ImageResult { images, diff --git a/aimux-providers/src/replicate.rs b/aimux-providers/src/replicate.rs index b8ca85a2..f6887869 100644 --- a/aimux-providers/src/replicate.rs +++ b/aimux-providers/src/replicate.rs @@ -16,8 +16,8 @@ use aimux_core::image_model::{ use aimux_core::shared::Warning; use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, sleep_or_abort, - without_trailing_slash, + HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, send_validated, + sleep_or_abort, without_trailing_slash, }; /// Configuration for the Replicate provider. @@ -303,7 +303,8 @@ impl ImageModel for ReplicateImageModel { // Download images let mut downloaded: Vec> = Vec::new(); for url in &urls { - let ir = send( + // output URLs come from the prediction response body. + let ir = send_validated( HttpRequest { method: HttpMethod::Get, url: url.clone(), @@ -314,6 +315,8 @@ impl ImageModel for ReplicateImageModel { call_id: None, recording_context: None, }, + Some(&self.config.base_url), + Some(&self.config.base_url), RetryConfig::default(), &DEFAULT_ERROR_STRUCTURE, ) diff --git a/aimux-providers/tests/download_ssrf_guard_test.rs b/aimux-providers/tests/download_ssrf_guard_test.rs new file mode 100644 index 00000000..133a9501 --- /dev/null +++ b/aimux-providers/tests/download_ssrf_guard_test.rs @@ -0,0 +1,81 @@ +//! SSRF guard wiring tests for the provider download path. +//! +//! Providers must fetch response-body URLs through `send_validated_download`, +//! which rejects private/loopback/link-local targets while same-origin URLs +//! against the configured `base_url` (the trusted origin) stay allowed — +//! mirroring AI SDK's `validateUrl` + `trustedOrigin` semantics. +//! +//! These tests use a wiremock server as the "provider endpoint" and point the +//! response-body URL at literal loopback/private IPs that the guard blocks +//! before any connection is attempted. + +use serde_json::json; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +use aimux_core::image_model::{ImageCallOptions, ImageModel}; +use aimux_providers::{RecraftConfig, RecraftProvider}; + +const PROMPT: &str = "A cute baby sea otter"; + +fn options(prompt: &str) -> ImageCallOptions { + ImageCallOptions::new(prompt.to_string()) +} + +/// Mount a generations response whose `data[0].url` points at `url`. +async fn mock_generations_with_url(server: &MockServer, url: &str) { + Mock::given(method("POST")) + .and(path("/images/generations")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "data": [ { "url": url } ] + }))) + .mount(server) + .await; +} + +fn make_model(server: &MockServer) -> impl ImageModel { + let config = RecraftConfig::new("test-recraft-key").with_base_url(server.uri()); + RecraftProvider::new(config).image("recraftv3") +} + +#[tokio::test] +async fn rejects_response_body_url_pointing_at_loopback_metadata() { + let server = MockServer::start().await; + // 169.254.169.254 is the classic cloud-metadata SSRF target. + mock_generations_with_url(&server, "http://169.254.169.254/latest/meta-data").await; + + let model = make_model(&server); + let err = model + .do_generate(&options(PROMPT)) + .await + .expect_err("guard must block metadata IP"); + match err { + aimux_core::AiMuxError::InvalidArgument(msg) => { + assert!(msg.contains("non-public"), "unexpected message: {msg}"); + } + other => panic!("expected InvalidArgument, got {other:?}"), + } +} + +#[tokio::test] +async fn allows_same_origin_download_against_configured_base_url() { + // A self-hosted deployment legitimately returns URLs on its own host. + // trustedOrigin (the configured base_url) must keep this working. + let server = MockServer::start().await; + mock_generations_with_url(&server, &format!("{}/generated/x.png", server.uri())).await; + Mock::given(method("GET")) + .and(path("/generated/x.png")) + .respond_with(ResponseTemplate::new(200).set_body_bytes(b"fake-image-data")) + .mount(&server) + .await; + + let model = make_model(&server); + let result = model.do_generate(&options(PROMPT)).await.unwrap(); + match result.images { + aimux_core::image_model::ImageOutputs::Binary(imgs) => { + assert_eq!(imgs.len(), 1); + assert_eq!(imgs[0], b"fake-image-data"); + } + other => panic!("expected Binary outputs, got {other:?}"), + } +}