From 022df8771562c93c6192b03b35c97e290ed16e63 Mon Sep 17 00:00:00 2001 From: Tochukwu Nkemdilim <11903253+toksdotdev@users.noreply.github.com> Date: Wed, 20 May 2026 09:16:33 -0400 Subject: [PATCH] fix(network): preserve encoded request bodies during secret substitution --- crates/network/lib/secrets/handler.rs | 169 +++++++++++++++++++++++--- 1 file changed, 152 insertions(+), 17 deletions(-) diff --git a/crates/network/lib/secrets/handler.rs b/crates/network/lib/secrets/handler.rs index aebf74b4..da4da399 100644 --- a/crates/network/lib/secrets/handler.rs +++ b/crates/network/lib/secrets/handler.rs @@ -59,6 +59,28 @@ impl EligibleSecret { self.inject_headers || self.inject_basic_auth || self.inject_query_params } + /// Returns true when the current header bytes contain this secret's + /// placeholder in a header-substitution scope. + fn may_substitute_in_headers(&self, headers: &[u8]) -> bool { + if !self.wants_header_injection() { + return false; + } + + let needle = self.placeholder.as_bytes(); + if (self.inject_headers || self.inject_query_params) && contains_bytes(headers, needle) { + return true; + } + + if self.inject_basic_auth { + return basic_auth_decoded_contains( + String::from_utf8_lossy(headers).as_ref(), + &self.placeholder, + ); + } + + false + } + /// Substitute this secret's placeholder in the headers portion, scoped by /// the secret's `headers` / `basic_auth` / `query_params` flags. fn substitute_in_headers(&self, headers: &str) -> String { @@ -175,15 +197,11 @@ impl SecretsHandler { Some(pos) => (&data[..pos], &data[pos..]), None => (data, &[] as &[u8]), }; - let mut header_str = String::from_utf8_lossy(header_bytes).into_owned(); - let mut body_str = if boundary.is_some() { - String::from_utf8_lossy(body_bytes).into_owned() - } else { - String::new() - }; // Fast path: skip violation check when no ineligible secrets exist. - if self.has_ineligible && self.has_violation(data, &header_str) { + if self.has_ineligible + && self.has_violation(data, String::from_utf8_lossy(header_bytes).as_ref()) + { self.update_tail(data); match self.on_violation { ViolationAction::Block => return None, @@ -206,27 +224,67 @@ impl SecretsHandler { return Some(Cow::Borrowed(data)); } + let mut header_str = None; + let mut body = None; + for secret in &self.eligible { // Skip secrets that require TLS identity on non-intercepted connections. if secret.require_tls_identity && !self.tls_intercepted { continue; } - if secret.wants_header_injection() { - header_str = secret.substitute_in_headers(&header_str); + + if secret.may_substitute_in_headers(header_bytes) { + let current = header_str + .get_or_insert_with(|| String::from_utf8_lossy(header_bytes).into_owned()); + *current = secret.substitute_in_headers(current); } - if boundary.is_some() && secret.inject_body && body_str.contains(&secret.placeholder) { - body_str = body_str.replace(&secret.placeholder, &secret.value); + + if boundary.is_some() && secret.inject_body { + let source = body.as_deref().unwrap_or(body_bytes); + if let Some(replaced) = replace_bytes( + source, + secret.placeholder.as_bytes(), + secret.value.as_bytes(), + ) { + body = Some(replaced); + } } } - // If body substitution changed the length, update Content-Length. - if boundary.is_some() && body_str.len() != body_bytes.len() { - header_str = update_content_length(&header_str, body_str.len()); + let header_changed = header_str + .as_ref() + .is_some_and(|headers| headers.as_bytes() != header_bytes); + let body_changed = body.is_some(); + + if !header_changed && !body_changed { + return Some(Cow::Borrowed(data)); + } + + let mut output = Vec::with_capacity( + header_str + .as_ref() + .map_or(header_bytes.len(), |headers| headers.len()) + + body.as_ref().map_or(body_bytes.len(), Vec::len), + ); + + let body_bytes_out = body.as_deref().unwrap_or(body_bytes); + if body_changed && body_bytes_out.len() != body_bytes.len() { + let headers = match header_str { + Some(headers) => update_content_length(&headers, body_bytes_out.len()), + None => update_content_length( + String::from_utf8_lossy(header_bytes).as_ref(), + body_bytes_out.len(), + ), + }; + output.extend_from_slice(headers.as_bytes()); + } else if let Some(headers) = header_str { + output.extend_from_slice(headers.as_bytes()); + } else { + output.extend_from_slice(header_bytes); } - let mut output = header_str; - output.push_str(&body_str); - Some(Cow::Owned(output.into_bytes())) + output.extend_from_slice(body_bytes_out); + Some(Cow::Owned(output)) } /// Returns true if no secrets are configured. @@ -351,6 +409,29 @@ fn contains_bytes(haystack: &[u8], needle: &[u8]) -> bool { haystack.windows(needle.len()).any(|w| w == needle) } +/// Replace all occurrences of `needle` in `haystack`. +/// +/// Returns `None` when no replacement is needed so callers can preserve the +/// original byte slice without rebuilding arbitrary binary payloads. +fn replace_bytes(haystack: &[u8], needle: &[u8], replacement: &[u8]) -> Option> { + if !contains_bytes(haystack, needle) { + return None; + } + + let mut result = Vec::with_capacity(haystack.len()); + let mut cursor = 0; + while cursor < haystack.len() { + if haystack[cursor..].starts_with(needle) { + result.extend_from_slice(replacement); + cursor += needle.len(); + } else { + result.push(haystack[cursor]); + cursor += 1; + } + } + Some(result) +} + /// Returns true if `haystack`, after URL percent-decoding, contains `needle`. fn url_decoded_contains(haystack: &[u8], needle: &[u8]) -> bool { let decoded: Vec = percent_decode(haystack).collect(); @@ -559,6 +640,60 @@ mod tests { assert!(result.ends_with("hello")); } + #[test] + fn eligible_secret_preserves_binary_body_without_placeholder() { + let config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]); + let mut handler = SecretsHandler::new(&config, "api.openai.com", true); + + let body = vec![0x1f, 0x8b, 0x08, 0x00, 0xff, 0x00, 0x80, 0xfe]; + let mut input = format!( + "POST /git-upload-pack HTTP/1.1\r\nContent-Encoding: gzip\r\nContent-Length: {}\r\n\r\n", + body.len() + ) + .into_bytes(); + input.extend_from_slice(&body); + + let output = handler.substitute(&input).unwrap(); + assert_eq!(&*output, input.as_slice()); + } + + #[test] + fn eligible_secret_preserves_binary_chunk_without_placeholder() { + let config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]); + let mut handler = SecretsHandler::new(&config, "api.openai.com", true); + + let input = [0x1f, 0x8b, 0x08, 0x00, 0xff, 0x00, 0x80, 0xfe]; + let output = handler.substitute(&input).unwrap(); + assert_eq!(&*output, input.as_slice()); + } + + #[test] + fn body_injection_preserves_non_utf8_bytes() { + let mut secret = make_secret("$KEY", "real-secret", "api.openai.com"); + secret.injection.body = true; + let config = make_config(vec![secret]); + let mut handler = SecretsHandler::new(&config, "api.openai.com", true); + + let body = [0xff, b'$', b'K', b'E', b'Y', 0xfe]; + let mut input = + format!("POST / HTTP/1.1\r\nContent-Length: {}\r\n\r\n", body.len()).into_bytes(); + input.extend_from_slice(&body); + + let output = handler.substitute(&input).unwrap().into_owned(); + let expected_body = [b"\xffreal-secret".as_slice(), &[0xfe]].concat(); + let expected = [ + format!( + "POST / HTTP/1.1\r\nContent-Length: {}\r\n\r\n", + expected_body.len() + ) + .as_bytes(), + expected_body.as_slice(), + ] + .concat(); + + assert_eq!(output, expected); + } + #[test] fn no_secrets_passthrough() { let config = make_config(vec![]);