Skip to content

fix(network): preserve encoded request bodies during secret substitution #768

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
toksdotdev wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

fix(network): preserve encoded request bodies during secret substitution #768

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
169 changes: 152 additions & 17 deletions crates/network/lib/secrets/handler.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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,
Expand All @@ -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.
Expand Down Expand Up @@ -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<Vec<u8>> {
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<u8> = percent_decode(haystack).collect();
Expand Down Expand Up @@ -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![]);
Expand Down
Loading