fix(provider-utils,providers): validate and pin provider-supplied download URLs

Providers fetch URLs taken from response bodies and headers — generated
assets, polling and result URLs, the Google Files upload URL — through
the shared auto-redirecting client with no target validation, so a
compromised or spoofed provider response could steer authenticated
requests at internal services (cloud metadata, loopback, RFC1918).

The new send_validated_download applies AI SDK-aligned SSRF protection:

- URL validation against AI SDK's validateUrl blocklists: http/https
  only, localhost/.local/.localhost rejected, literal-IP checks
  including the IPv4-embedded IPv6 forms (mapped, compatible, SIIT,
  NAT64), so [::ffff:169.254.169.254] is caught.
- Every DNS answer is validated and the connection pinned to exactly
  those addresses (resolve overrides), defeating TTL-0 rebinding. The
  proxy configuration is applied to the pinned client so reqwest makes
  the per-URL routing decision: proxied requests are resolved by that
  trusted transport, direct ones (NO_PROXY match, no proxy for the
  scheme) stay pinned.
- Redirects are followed manually (RedirectMode::Manual), at most
  MAX_REDIRECTS hops, each hop re-validated and re-pinned; redirects to
  non-HTTP(S) schemes are rejected, matching fetch.
- Caller headers are sanitized (hop-by-hop, forwarding, and
  metadata-service headers dropped) and sent only when the target is
  strictly same-origin with the configured trusted origin — AI SDK's
  credentialedOrigin semantics — checked before the first request and
  on every redirect hop.
- URLs same-origin with the configured base_url skip address validation
  so self-hosted deployments serving assets from their own (possibly
  private) origin keep working.

Wired into every response-supplied-URL fetch: Black Forest Labs
(polling + download), Gladia (result polling), Luma, Recraft,
Replicate, and fal downloads, and the Google Files upload.

Author: cunninghamcard-bit

aimux-provider-utils/src/download_guard.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};

aimux-provider-utils/src/http.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::<serde_json::Value>(&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<String> = 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()));
+}

aimux-provider-utils/src/lib.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));
+    }
+}