refactor(llm-access): domain-split provider.rs into concern submodules (PR-B)#20
Conversation
Second of two PRs splitting the 14.3k-line provider.rs (PR-A #19 extracted the 106 tests; this PR-B domain-splits the remaining ~7300 code lines). Free-fn DAG shape (like the duckdb split), already mod.rs-style. provider.rs: 7301 -> 603 lines. 182 free fns + 16 impls (63 methods) move into 19 concern submodules: usage_meta, state, limiter, route_selection, codex_dispatch, kiro_media, kiro_dispatch, stream_guards, kiro_protocol, client, errors, kiro_model, kiro_usage, entry, codex_auth, codex_sse, kiro_summary, util, codex_models. Behavior-preserving structural move: identical top-level name multiset (336), same 106 tests. Fat facade keeps all 47 types + 2 traits + 24 consts/statics (submodules read private fields as ancestor-privates, no field-vis change). Moved free fns -> pub only where cross-region referenced (102/182); inherent methods -> pub(super) only where a foreign region calls them (25/46); trait-impl methods never bumped. Facade re-imports in 3 categories: external API preserved via pub/pub(crate) use (provider_entry*, codex_upstream_base_url, compute_codex_upstream_url, resolve_codex_client_version, codex_public_model_catalog_response, default_codex_public_model_catalog_response, call_kiro_generate_for_route, decode_kiro_events_from_bytes); unconditional `use` for fns the facade's LazyLock statics call (build_provider_client + client-pool tunables); and `#[cfg(test)] use` for fns/types tests reach via super:: (incl. axum Method). Named the kiro-header module `kiro_protocol` to avoid colliding with the `use crate::kiro_headers` binding (E0255). Pre-existing kiro_error.rs calls super::{anthropic_json_error, summarize_error_bytes} -> facade re-imports them so those super:: paths still resolve. Verification: - cargo clippy -p llm-access --all-targets -- -D warnings -> clean - cargo test -p llm-access provider:: -> 111 passed - rustfmt on the 20 changed files only; deps/lance, deps/lancedb untouched Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
/gemini review |
|
@codex review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive proxy and gateway system for Codex and Kiro providers, implementing request dispatching, authentication, route selection, rate limiting, streaming, and usage/preflight logging. During the review, two critical issues were identified: first, a bug in strip_codex_encrypted_reasoning_items that silently wipes out the entire input array for standard requests when no encrypted reasoning is modified; second, a high-severity DNS rebinding SSRF vulnerability in the remote media fetcher due to a Time-of-Check to Time-of-Use (TOCTOU) gap between IP validation and the actual request execution.
| if changed { | ||
| if filtered.is_empty() { | ||
| remove_input = true; | ||
| } else { | ||
| *items = filtered; | ||
| } | ||
| } |
There was a problem hiding this comment.
This is a critical bug. When changed is false (i.e., no encrypted reasoning items were found or sanitized), the if changed block is skipped, leaving *items as an empty vector (due to the std::mem::take(items) call on line 291). This silently wipes out the entire input array for all standard requests that do not contain encrypted reasoning, causing them to fail or be processed as empty requests.
if changed && filtered.is_empty() {
remove_input = true;
} else {
*items = filtered;
}There was a problem hiding this comment.
Verified against the full call chain — this is a false positive; the current code is correct.
The observation is accurate in isolation (when changed == false, std::mem::take(items) on line 291 leaves *items empty and the if changed block doesn't restore it), but the conclusion doesn't hold once you follow how the map is used:
strip_codex_encrypted_reasoning_itemshas exactly one caller —retry_codex_without_encrypted_reasoning(this file). That caller operates on a fresh local clonevalueparsed fromprepared.request_body(preparedis an immutable&PreparedGatewayRequest), so the emptiedinputis never the live request.- On the
changed == falsepath,stripreturnsfalse, and the caller doesif !strip(...) { return None; }— it discardsvalueentirely and the single caller of that (codex_dispatch.rs, theif let Some(retry_prepared) = ...site) reuses the originalpreparedunchanged. valueis serialized (serde_json::to_vec(&value)) only whenstripreturnstrue(i.e.changed == true) — and on that path theif changedblock has already set*items = filtered(or removedinputwhenfiltered.is_empty()). So whenever the map is actually used,inputis correctly populated.
The empty-items state therefore only exists on the no-op path whose result is thrown away. No standard request is affected. (The suggested changed && filtered.is_empty() rewrite is behavior-equivalent and harmless, but it's a fragile-invariant tidy-up, not a bugfix — and this is a byte-identical structural-move PR, so I'm leaving the body untouched.)
| let url = validate_kiro_remote_media_url(request.url)?; | ||
| validate_kiro_remote_media_resolved_addresses(&url).await?; |
There was a problem hiding this comment.
There is a Time-of-Check to Time-of-Use (TOCTOU) vulnerability here that allows DNS Rebinding SSRF.
Although validate_kiro_remote_media_resolved_addresses resolves the host and verifies that the IP addresses are not private, reqwest will perform its own DNS resolution when self.client.get(url.clone()).send().await is called. An attacker can configure a DNS server to return a safe public IP during the first resolution (the check) and a private/local IP (e.g., 127.0.0.1 or 10.0.0.1) during the second resolution (the use), bypassing the private IP check.
To prevent DNS Rebinding SSRF, you should configure a custom DNS resolver on the reqwest::Client (when building KIRO_REMOTE_MEDIA_CLIENT) that filters out private IP addresses at the socket connection level, or resolve the IP address once and rewrite the request URL to use the resolved IP directly (while setting the Host header to the original host name, though this requires custom TLS verification for HTTPS).
There was a problem hiding this comment.
Confirmed real — and good catch. The two-resolution gap is exactly as described: validate_kiro_remote_media_resolved_addresses does its own tokio::net::lookup_host(...) and rejects private IPs (resolution #1), then self.client.get(url.clone()).send() makes reqwest perform an independent DNS resolution #2. A hostile authoritative DNS server can answer #1 with a public IP and #2 with 127.0.0.1 / 169.254.169.254 / RFC1918, bypassing the check → TOCTOU DNS-rebinding SSRF (reachable when kiro_remote_media_resolution_enabled is set, since the source URLs come from the request payload).
Deferring to a dedicated security PR rather than this one, because: (a) it's pre-existing code moved verbatim here — this PR is a behavior-preserving structural split verified by identical top-level name-multiset + byte-identical bodies, and folding a security fix in would void that guarantee; (b) the fix is genuinely behavior-changing and non-trivial — either a custom connection-level resolver on KIRO_REMOTE_MEDIA_CLIENT that re-validates the socket peer IP at connect time, or resolve-once-then-connect-to-the-pinned-IP (with Host/SNI preserved for TLS). Will open a follow-up with a rebinding regression test, same as the duckdb write-race fix (#18) found during the #16 review.
|
Codex Review: Didn't find any major issues. Keep it up! ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
There was a problem hiding this comment.
Code Review
This pull request introduces comprehensive proxy dispatching, routing, session management, and usage tracking for both Codex and Kiro providers. The review feedback highlights several critical security vulnerabilities in kiro_media.rs, including a potential Denial of Service (DoS) from reading unbounded response bodies, and multiple Server-Side Request Forgery (SSRF) bypasses (IPv4-mapped IPv6 addresses, DNS rebinding, and multicast IP addresses). Additionally, a performance improvement was suggested in entry.rs to avoid an unnecessary string allocation when handling bearer tokens.
| let bytes = response.bytes().await.map_err(|err| { | ||
| KiroRemoteMediaResolutionError::new(format!("failed to read URL source body: {err}")) | ||
| })?; |
There was a problem hiding this comment.
The response.bytes().await call reads the entire response body into memory without enforcing a maximum size limit when the Content-Length header is missing or spoofed. Since this fetcher retrieves arbitrary user-provided URLs, a malicious server could return an extremely large or infinite stream of bytes, leading to process Out-Of-Memory (OOM) crashes. To prevent this, read the response body as a stream and enforce a strict limit on the number of bytes read.
| let bytes = response.bytes().await.map_err(|err| { | |
| KiroRemoteMediaResolutionError::new(format!("failed to read URL source body: {err}")) | |
| })?; | |
| let mut bytes_vec = Vec::new(); | |
| let mut stream = response.bytes_stream(); | |
| use futures_util::StreamExt; | |
| while let Some(chunk) = stream.next().await { | |
| let chunk = chunk.map_err(|err| { | |
| KiroRemoteMediaResolutionError::new(format!("failed to read URL source body: {err}")) | |
| })?; | |
| if bytes_vec.len() + chunk.len() > max_bytes { | |
| return Err(KiroRemoteMediaResolutionError::new(format!( | |
| "URL source exceeds {} byte limit", | |
| max_bytes | |
| ))); | |
| } | |
| bytes_vec.extend_from_slice(&chunk); | |
| } | |
| let bytes = bytes::Bytes::from(bytes_vec); |
| match ip { | ||
| IpAddr::V4(ip) => is_private_kiro_remote_media_ipv4(ip), | ||
| IpAddr::V6(ip) => is_private_kiro_remote_media_ipv6(ip), | ||
| } | ||
| } | ||
| fn is_private_kiro_remote_media_ipv4(ip: Ipv4Addr) -> bool { |
There was a problem hiding this comment.
The private IP check for IPv6 addresses does not account for IPv4-mapped IPv6 addresses (e.g., ::ffff:127.0.0.1). Since Ipv6Addr::is_loopback() only returns true for ::1, an attacker could bypass the SSRF protection by using an IPv4-mapped IPv6 address pointing to a private or loopback IPv4 address. To fix this, check if the IPv6 address can be converted to an IPv4 address, and if so, validate it against the IPv4 private/loopback rules.
| match ip { | |
| IpAddr::V4(ip) => is_private_kiro_remote_media_ipv4(ip), | |
| IpAddr::V6(ip) => is_private_kiro_remote_media_ipv6(ip), | |
| } | |
| } | |
| fn is_private_kiro_remote_media_ipv4(ip: Ipv4Addr) -> bool { | |
| fn is_private_kiro_remote_media_ip(ip: IpAddr) -> bool { | |
| match ip { | |
| IpAddr::V4(ip) => is_private_kiro_remote_media_ipv4(ip), | |
| IpAddr::V6(ip) => { | |
| if is_private_kiro_remote_media_ipv6(ip) { | |
| return true; | |
| } | |
| if let Some(ipv4) = ip.to_ipv4() { | |
| if is_private_kiro_remote_media_ipv4(ipv4) { | |
| return true; | |
| } | |
| } | |
| false | |
| } | |
| } | |
| } |
| async fn validate_kiro_remote_media_resolved_addresses( | ||
| url: &url::Url, | ||
| ) -> Result<(), KiroRemoteMediaResolutionError> { | ||
| let host = url | ||
| .host_str() | ||
| .ok_or_else(|| KiroRemoteMediaResolutionError::new("URL source is missing host"))?; | ||
| if host.parse::<IpAddr>().is_ok() { | ||
| return Ok(()); | ||
| } | ||
| let port = url | ||
| .port_or_known_default() | ||
| .ok_or_else(|| KiroRemoteMediaResolutionError::new("URL source is missing port"))?; | ||
| let addresses = tokio::net::lookup_host((host, port)).await.map_err(|err| { | ||
| KiroRemoteMediaResolutionError::new(format!("failed to resolve URL source host: {err}")) | ||
| })?; | ||
| let mut resolved_any = false; | ||
| for address in addresses { | ||
| resolved_any = true; | ||
| reject_private_kiro_remote_media_ip(address.ip())?; | ||
| } | ||
| if !resolved_any { | ||
| return Err(KiroRemoteMediaResolutionError::new( | ||
| "URL source host resolved to no addresses", | ||
| )); | ||
| } | ||
| Ok(()) | ||
| } |
There was a problem hiding this comment.
The DNS resolution check in validate_kiro_remote_media_resolved_addresses is vulnerable to DNS Rebinding attacks. A malicious user can configure a domain name with a very low TTL that initially resolves to a public IP (passing the validation check), but subsequently resolves to a private IP (e.g., 127.0.0.1 or 169.254.169.254) when reqwest resolves it again to send the actual HTTP request. To prevent DNS Rebinding, resolve the domain name once, validate the IP, and then pin the domain to that specific IP address using reqwest::ClientBuilder::resolve.
| fn is_private_kiro_remote_media_ipv4(ip: Ipv4Addr) -> bool { | ||
| ip.is_private() | ||
| || ip.is_loopback() | ||
| || ip.is_link_local() | ||
| || ip.is_broadcast() | ||
| || ip.is_documentation() | ||
| || ip == Ipv4Addr::UNSPECIFIED | ||
| } | ||
| fn is_private_kiro_remote_media_ipv6(ip: Ipv6Addr) -> bool { | ||
| ip.is_loopback() | ||
| || ip.is_unique_local() | ||
| || ip.is_unicast_link_local() | ||
| || ip.is_unspecified() | ||
| || matches!(ip.segments(), [0x2001, 0x0db8, _, _, _, _, _, _]) | ||
| } |
There was a problem hiding this comment.
Multicast IP addresses (e.g., 224.0.0.0/4 for IPv4 and ff00::/8 for IPv6) are not blocked by the private IP checks. Multicast addresses are used for local network routing and should be blocked to prevent local network probing or other unexpected behavior. Add ip.is_multicast() checks to both the IPv4 and IPv6 private IP validation functions.
| fn is_private_kiro_remote_media_ipv4(ip: Ipv4Addr) -> bool { | |
| ip.is_private() | |
| || ip.is_loopback() | |
| || ip.is_link_local() | |
| || ip.is_broadcast() | |
| || ip.is_documentation() | |
| || ip == Ipv4Addr::UNSPECIFIED | |
| } | |
| fn is_private_kiro_remote_media_ipv6(ip: Ipv6Addr) -> bool { | |
| ip.is_loopback() | |
| || ip.is_unique_local() | |
| || ip.is_unicast_link_local() | |
| || ip.is_unspecified() | |
| || matches!(ip.segments(), [0x2001, 0x0db8, _, _, _, _, _, _]) | |
| } | |
| fn is_private_kiro_remote_media_ipv4(ip: Ipv4Addr) -> bool { | |
| ip.is_private() | |
| || ip.is_loopback() | |
| || ip.is_link_local() | |
| || ip.is_broadcast() | |
| || ip.is_documentation() | |
| || ip.is_multicast() | |
| || ip == Ipv4Addr::UNSPECIFIED | |
| } | |
| fn is_private_kiro_remote_media_ipv6(ip: Ipv6Addr) -> bool { | |
| ip.is_loopback() | |
| || ip.is_unique_local() | |
| || ip.is_unicast_link_local() | |
| || ip.is_unspecified() | |
| || ip.is_multicast() | |
| || matches!(ip.segments(), [0x2001, 0x0db8, _, _, _, _, _, _]) | |
| } |
| let Some(secret) = presented_secret(request.headers(), request.uri().path()).map(str::to_owned) | ||
| else { | ||
| return (StatusCode::UNAUTHORIZED, "missing bearer token").into_response(); | ||
| }; | ||
| let key = match state | ||
| .control_store | ||
| .authenticate_bearer_secret(&secret) |
There was a problem hiding this comment.
The bearer token secret is cloned/allocated into an owned String using .map(str::to_owned) unnecessarily. Since presented_secret borrows from the request headers which live for the duration of the function, we can keep it as a borrowed &str and avoid the redundant allocation.
| let Some(secret) = presented_secret(request.headers(), request.uri().path()).map(str::to_owned) | |
| else { | |
| return (StatusCode::UNAUTHORIZED, "missing bearer token").into_response(); | |
| }; | |
| let key = match state | |
| .control_store | |
| .authenticate_bearer_secret(&secret) | |
| let Some(secret) = presented_secret(request.headers(), request.uri().path()) else { | |
| return (StatusCode::UNAUTHORIZED, "missing bearer token").into_response(); | |
| }; | |
| let key = match state | |
| .control_store | |
| .authenticate_bearer_secret(secret) | |
| .await |
…21) * fix(llm-access): block DNS-rebinding SSRF in kiro remote-media fetch The kiro remote-media fetcher validated resolved IPs with its own `tokio::net::lookup_host` (rejecting private/local addresses), then let reqwest perform an INDEPENDENT DNS resolution at connect time. A hostile authoritative DNS server could answer the pre-flight check with a public IP and the connect-time resolution with 127.0.0.1 / 169.254.169.254 / RFC1918, bypassing the guard — a TOCTOU DNS-rebinding SSRF (reachable when a key has kiro_remote_media_resolution_enabled, since the source URLs come from the request payload). Found in PR #20 review (gemini). Fix: a custom `reqwest::dns::Resolve` (PrivateFilteringDnsResolver) installed on KIRO_REMOTE_MEDIA_CLIENT via `dns_resolver`. It resolves the host and filters out private/local addresses, returning only public ones (erroring if none remain). Because reqwest dials exactly the addresses the resolver returns, the IP that is vetted IS the IP connected to — closing the check/use gap. IP-literal URLs (never sent to a resolver) remain covered by the existing literal-IP gate in validate_kiro_remote_media_url; the strict pre-flight check is kept for early, clear errors (defense in depth). Tests: pure `filter_public_kiro_remote_media_addrs` cases (keeps public / drops private / rejects all-private / rejects empty) + a resolver test that `localhost` (RFC 6761 loopback) is rejected end-to-end. Verification: - cargo clippy -p llm-access --all-targets -- -D warnings -> clean - cargo test -p llm-access provider:: -> 115 passed (111 prior + 4 new) - rustfmt on the 2 changed files only; deps/lance, deps/lancedb untouched Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * fix(llm-access): close IPv4-mapped and literal-IPv6 SSRF gate holes CR on PR #21 surfaced two residual bypasses of the kiro remote-media private-address gate: - IPv4-mapped (`::ffff:a.b.c.d`) / IPv4-compatible (`::a.b.c.d`) IPv6 addresses are dialed as their embedded IPv4 by dual-stack sockets, so e.g. `::ffff:127.0.0.1` slipped past the native-IPv6 checks. Re-check the embedded IPv4 (via `Ipv6Addr::to_ipv4`) against the v4 rules in both the literal gate and the resolving filter. - `url.host_str()` brackets IPv6 literals (`[::1]`), so the `parse::<IpAddr>()` literal gate never matched them and a literal-IP URL skipped the private check entirely (reqwest dials the literal without ever consulting the filtering resolver). Match the parsed `url.host()` enum (Domain/Ipv4/Ipv6) instead, in both the pre-flight validate and the resolve-addresses pre-check. Tests: IPv4-mapped private rejection + mapped-public retention; literal loopback / mapped-private / url-normalized decimal-IPv4 URL rejection; public literal IPv6 allowed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
1 participant
What
Second and final PR splitting the 14.3k-line
llm-access/src/provider.rs(PR-A #19 extracted the 106 tests intoprovider/tests.rs; this PR-B domain-splits the remaining ~7300 code lines). This is the finale of thellm-access*big-file refactor sequence.provider.rs: 7301 → 603 lines. 182 free fns + 16 impls (63 methods) move into 19 concern submodules:usage_meta, state, limiter, route_selection, codex_dispatch, kiro_media, kiro_dispatch, stream_guards, kiro_protocol, client, errors, kiro_model, kiro_usage, entry, codex_auth, codex_sse, kiro_summary, util, codex_models.Same free-fn-DAG shape as the duckdb split (#16);
provider.rswas alreadymod.rs-style (it declaredmod kiro_error; mod kiro_session_affinity;).Behavior-preserving
Structural move only — fn/method bodies byte-identical.
#[test]/#[tokio::test].Fat facade keeps all 47 types + 2 traits + 24 consts/statics (submodules read private fields as ancestor-privates — no field-visibility change). Moved free fns →
pubonly where cross-region referenced (102/182); inherent methods →pub(super)only where a foreign region calls them (25/46); trait-impl methods never bumped.Facade re-imports (3 categories)
pub/pub(crate) use—provider_entry,provider_entry_handler,codex_upstream_base_url,compute_codex_upstream_url,resolve_codex_client_version,codex_public_model_catalog_response,default_codex_public_model_catalog_response,call_kiro_generate_for_route,decode_kiro_events_from_bytes. (codex_openai_models_responsewas not re-exported — it's only called by a sibling, never viacrate::provider::, so its facade re-export was redundant.)usefor fns the facade'sLazyLockstatics call (build_provider_client+ client-pool tunables).#[cfg(test)] usefor fns/types tests reach viasuper::(incl. externalaxum::http::Method).Provider-specific gotchas handled
kiro_protocol(notkiro_headers) to avoid colliding with theuse crate::kiro_headersbinding → E0255.kiro_error.rs(a facade child) callssuper::{anthropic_json_error, summarize_error_bytes}→ the facade re-imports them so thosesuper::paths still resolve (ref-scan corpus had to include the pre-existing submodules).self::→super::rewrite on pasted head-imports.Verification
cargo clippy -p llm-access --all-targets -- -D warnings→ cleancargo test -p llm-access provider::→ 111 passedrustfmton the 20 changed files only;deps/lance/deps/lancedbuntouched🤖 Generated with Claude Code