1use std::borrow::Cow;
7
8use base64::{Engine, engine::general_purpose::STANDARD as BASE64};
9use percent_encoding::percent_decode;
10
11use super::config::{SecretsConfig, ViolationAction};
12
13pub struct SecretsHandler {
22 eligible_for_substitution: Vec<EligibleSecret>,
24 ineligible_for_substitution: Vec<IneligibleSecret>,
26 tls_intercepted: bool,
28 max_placeholder_len: usize,
30 prev_tail: Vec<u8>,
34}
35
36struct EligibleSecret {
38 placeholder: String,
39 value: String,
40 inject_headers: bool,
41 inject_basic_auth: bool,
42 inject_query_params: bool,
43 inject_body: bool,
44 require_tls_identity: bool,
45}
46
47struct IneligibleSecret {
49 placeholder: String,
50 action: BlockingAction,
51}
52
53#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
55enum BlockingAction {
56 Block,
57 #[default]
58 BlockAndLog,
59 BlockAndTerminate,
60}
61
62impl EligibleSecret {
67 fn wants_header_injection(&self) -> bool {
70 self.inject_headers || self.inject_basic_auth || self.inject_query_params
71 }
72
73 fn may_substitute_in_headers(&self, headers: &[u8]) -> bool {
76 if !self.wants_header_injection() {
77 return false;
78 }
79
80 let needle = self.placeholder.as_bytes();
81 if (self.inject_headers || self.inject_query_params) && contains_bytes(headers, needle) {
82 return true;
83 }
84
85 if self.inject_basic_auth {
87 return basic_auth_decoded_contains(
88 String::from_utf8_lossy(headers).as_ref(),
89 &self.placeholder,
90 );
91 }
92
93 false
94 }
95
96 fn substitute_in_headers(&self, headers: &str) -> String {
99 let mut result = String::with_capacity(headers.len());
100 for (i, line) in headers.split("\r\n").enumerate() {
101 if i > 0 {
102 result.push_str("\r\n");
103 }
104 match self.substitute_in_header_line(line, i == 0) {
105 Some(s) => result.push_str(&s),
106 None => result.push_str(line),
107 }
108 }
109 result
110 }
111
112 fn substitute_in_header_line(&self, line: &str, is_request_line: bool) -> Option<String> {
116 if self.inject_basic_auth
117 && is_authorization_header(line)
118 && let Some(replaced) = self.substitute_basic_auth_header(line)
119 {
120 return Some(replaced);
121 }
122 if self.inject_headers {
123 return Some(line.replace(&self.placeholder, &self.value));
124 }
125 if is_request_line && self.inject_query_params {
126 return Some(line.replace(&self.placeholder, &self.value));
127 }
128 None
129 }
130
131 fn substitute_basic_auth_header(&self, line: &str) -> Option<String> {
137 let decoded = decode_basic_credentials(line)?;
138 if !decoded.contains(&self.placeholder) {
139 return None;
140 }
141 let (name, _) = line.split_once(':')?;
142 let replaced = decoded.replace(&self.placeholder, &self.value);
143 Some(format!(
144 "{name}: Basic {}",
145 BASE64.encode(replaced.as_bytes())
146 ))
147 }
148}
149
150impl BlockingAction {
151 fn from_violation_action(action: &ViolationAction) -> Option<Self> {
152 match action {
153 ViolationAction::Block => Some(Self::Block),
154 ViolationAction::BlockAndLog => Some(Self::BlockAndLog),
155 ViolationAction::BlockAndTerminate => Some(Self::BlockAndTerminate),
156 ViolationAction::Passthrough(_) => None,
157 }
158 }
159
160 fn into_violation_action(self) -> ViolationAction {
161 match self {
162 Self::Block => ViolationAction::Block,
163 Self::BlockAndLog => ViolationAction::BlockAndLog,
164 Self::BlockAndTerminate => ViolationAction::BlockAndTerminate,
165 }
166 }
167}
168
169impl SecretsHandler {
170 pub fn new(config: &SecretsConfig, sni: &str, tls_intercepted: bool) -> Self {
177 let mut eligible_for_substitution = Vec::new();
178 let mut ineligible_for_substitution = Vec::new();
179 let mut max_placeholder_len = 0;
180
181 for secret in &config.secrets {
182 max_placeholder_len = max_placeholder_len.max(secret.placeholder.len());
183
184 let host_allowed = secret.allowed_hosts.is_empty()
185 || secret.allowed_hosts.iter().any(|p| p.matches(sni));
186
187 if host_allowed {
190 eligible_for_substitution.push(EligibleSecret {
191 placeholder: secret.placeholder.clone(),
192 value: secret.value.clone(),
193 inject_headers: secret.injection.headers,
194 inject_basic_auth: secret.injection.basic_auth,
195 inject_query_params: secret.injection.query_params,
196 inject_body: secret.injection.body,
197 require_tls_identity: secret.require_tls_identity,
198 });
199
200 continue;
201 }
202
203 let action = secret.on_violation.as_ref().unwrap_or(&config.on_violation);
204
205 if let ViolationAction::Passthrough(hosts) = action
207 && hosts.iter().any(|p| p.matches(sni))
208 {
209 continue;
210 }
211
212 ineligible_for_substitution.push(IneligibleSecret {
214 placeholder: secret.placeholder.clone(),
215 action: BlockingAction::from_violation_action(action).unwrap_or_default(),
216 });
217 }
218
219 Self {
220 eligible_for_substitution,
221 ineligible_for_substitution,
222 tls_intercepted,
223 max_placeholder_len,
224 prev_tail: Vec::new(),
225 }
226 }
227
228 pub fn substitute<'a>(&mut self, data: &'a [u8]) -> Result<Cow<'a, [u8]>, ViolationAction> {
239 let boundary = find_header_boundary(data);
242 let (header_bytes, body_bytes) = match boundary {
243 Some(pos) => (&data[..pos], &data[pos..]),
244 None => (data, &[] as &[u8]),
245 };
246
247 if let Some(action) =
249 self.detect_blocking_action(data, String::from_utf8_lossy(header_bytes).as_ref())
250 {
251 match action {
252 BlockingAction::Block => return Err(action.into_violation_action()),
253 BlockingAction::BlockAndLog => {
254 tracing::warn!("secret violation: placeholder detected for disallowed host");
255 return Err(action.into_violation_action());
256 }
257 BlockingAction::BlockAndTerminate => {
258 tracing::error!(
259 "secret violation: placeholder detected for disallowed host — terminating"
260 );
261 return Err(action.into_violation_action());
262 }
263 }
264 }
265 self.update_tail(data);
266
267 if self.eligible_for_substitution.is_empty() {
268 return Ok(Cow::Borrowed(data));
270 }
271
272 let mut header_str = None;
274 let mut body = None;
275
276 for secret in &self.eligible_for_substitution {
277 if secret.require_tls_identity && !self.tls_intercepted {
279 continue;
280 }
281
282 if secret.may_substitute_in_headers(header_bytes) {
284 let current = header_str
285 .get_or_insert_with(|| String::from_utf8_lossy(header_bytes).into_owned());
286 *current = secret.substitute_in_headers(current);
287 }
288
289 if boundary.is_some() && secret.inject_body {
291 let source = body.as_deref().unwrap_or(body_bytes);
292 if let Some(replaced) = replace_bytes(
293 source,
294 secret.placeholder.as_bytes(),
295 secret.value.as_bytes(),
296 ) {
297 body = Some(replaced);
298 }
299 }
300 }
301
302 let header_changed = header_str
303 .as_ref()
304 .is_some_and(|headers| headers.as_bytes() != header_bytes);
305 let body_changed = body.is_some();
306
307 if !header_changed && !body_changed {
309 return Ok(Cow::Borrowed(data));
310 }
311
312 let header_len = header_str
313 .as_ref()
314 .map_or(header_bytes.len(), |headers| headers.len());
315 let body_len = body.as_ref().map_or(body_bytes.len(), Vec::len);
316 let mut output = Vec::with_capacity(header_len + body_len);
317
318 let body_bytes_out = body.as_deref().unwrap_or(body_bytes);
319 if body_changed && body_bytes_out.len() != body_bytes.len() {
321 let headers = match header_str {
322 Some(headers) => update_content_length(&headers, body_bytes_out.len()),
323 None => update_content_length(
324 String::from_utf8_lossy(header_bytes).as_ref(),
325 body_bytes_out.len(),
326 ),
327 };
328 output.extend_from_slice(headers.as_bytes());
329 } else if let Some(headers) = header_str {
330 output.extend_from_slice(headers.as_bytes());
331 } else {
332 output.extend_from_slice(header_bytes);
333 }
334
335 output.extend_from_slice(body_bytes_out);
336 Ok(Cow::Owned(output))
337 }
338
339 pub fn is_empty(&self) -> bool {
341 self.eligible_for_substitution.is_empty() && self.ineligible_for_substitution.is_empty()
342 }
343
344 fn detect_blocking_action(&self, data: &[u8], headers: &str) -> Option<BlockingAction> {
352 if self.ineligible_for_substitution.is_empty() {
353 return None;
354 }
355
356 let scan_buf: Cow<[u8]> = if self.prev_tail.is_empty() {
357 Cow::Borrowed(data)
358 } else {
359 let mut stitched = Vec::with_capacity(self.prev_tail.len() + data.len());
360 stitched.extend_from_slice(&self.prev_tail);
361 stitched.extend_from_slice(data);
362 Cow::Owned(stitched)
363 };
364 let scan = scan_buf.as_ref();
365
366 let mut detected = None;
367 for secret in &self.ineligible_for_substitution {
368 let needle = secret.placeholder.as_bytes();
369 if contains_bytes(scan, needle)
370 || url_decoded_contains(scan, needle)
371 || json_escaped_contains(scan, needle)
372 || basic_auth_decoded_contains(headers, &secret.placeholder)
373 {
374 detected = Some(strictest_violation_action(detected, secret.action));
375 }
376 }
377
378 detected
379 }
380
381 fn update_tail(&mut self, data: &[u8]) {
385 let tail_size = self.max_placeholder_len.saturating_sub(1);
386 if tail_size == 0 {
387 return;
388 }
389 if data.len() >= tail_size {
390 self.prev_tail.clear();
391 self.prev_tail
392 .extend_from_slice(&data[data.len() - tail_size..]);
393 return;
394 }
395 self.prev_tail.extend_from_slice(data);
396 let overflow = self.prev_tail.len().saturating_sub(tail_size);
397 if overflow > 0 {
398 self.prev_tail.drain(..overflow);
399 }
400 }
401}
402
403fn is_authorization_header(line: &str) -> bool {
410 line.as_bytes()
411 .get(..14)
412 .is_some_and(|b| b.eq_ignore_ascii_case(b"authorization:"))
413}
414
415fn decode_basic_credentials(line: &str) -> Option<String> {
419 let (_, raw_value) = line.split_once(':')?;
420 let (scheme, encoded) = split_auth_scheme(raw_value.trim_start())?;
421 if !scheme.eq_ignore_ascii_case("basic") {
422 return None;
423 }
424 let bytes = BASE64.decode(encoded.trim()).ok()?;
425 String::from_utf8(bytes).ok()
426}
427
428fn split_auth_scheme(header_value: &str) -> Option<(&str, &str)> {
431 let split_at = header_value.find(char::is_whitespace)?;
432 let (scheme, rest) = header_value.split_at(split_at);
433 Some((scheme, rest.trim_start()))
434}
435
436fn basic_auth_decoded_contains(headers: &str, placeholder: &str) -> bool {
439 headers
440 .split("\r\n")
441 .filter(|line| is_authorization_header(line))
442 .filter_map(decode_basic_credentials)
443 .any(|decoded| decoded.contains(placeholder))
444}
445
446fn contains_bytes(haystack: &[u8], needle: &[u8]) -> bool {
448 if needle.is_empty() || haystack.len() < needle.len() {
449 return false;
450 }
451 haystack.windows(needle.len()).any(|w| w == needle)
452}
453
454fn replace_bytes(haystack: &[u8], needle: &[u8], replacement: &[u8]) -> Option<Vec<u8>> {
459 if !contains_bytes(haystack, needle) {
460 return None;
461 }
462
463 let mut result = Vec::with_capacity(haystack.len());
464 let mut cursor = 0;
465 while cursor < haystack.len() {
466 if haystack[cursor..].starts_with(needle) {
467 result.extend_from_slice(replacement);
468 cursor += needle.len();
469 } else {
470 result.push(haystack[cursor]);
471 cursor += 1;
472 }
473 }
474 Some(result)
475}
476
477fn url_decoded_contains(haystack: &[u8], needle: &[u8]) -> bool {
479 let decoded: Vec<u8> = percent_decode(haystack).collect();
480 contains_bytes(&decoded, needle)
481}
482
483fn json_escaped_contains(haystack: &[u8], needle: &[u8]) -> bool {
487 let mut decoded = Vec::with_capacity(haystack.len());
488 let mut i = 0;
489 while i < haystack.len() {
490 if haystack[i] == b'\\'
491 && i + 5 < haystack.len()
492 && haystack[i + 1] == b'u'
493 && let (Some(a), Some(b), Some(c), Some(d)) = (
494 hex_digit(haystack[i + 2]),
495 hex_digit(haystack[i + 3]),
496 hex_digit(haystack[i + 4]),
497 hex_digit(haystack[i + 5]),
498 )
499 {
500 let cp = ((a as u32) << 12) | ((b as u32) << 8) | ((c as u32) << 4) | (d as u32);
501 if let Some(ch) = char::from_u32(cp) {
502 let mut buf = [0u8; 4];
503 decoded.extend_from_slice(ch.encode_utf8(&mut buf).as_bytes());
504 }
505 i += 6;
506 continue;
507 }
508 decoded.push(haystack[i]);
509 i += 1;
510 }
511 contains_bytes(&decoded, needle)
512}
513
514fn hex_digit(b: u8) -> Option<u8> {
515 (b as char).to_digit(16).map(|d| d as u8)
516}
517
518fn update_content_length(headers: &str, new_len: usize) -> String {
523 let mut result = String::with_capacity(headers.len());
524 for (i, line) in headers.split("\r\n").enumerate() {
525 if i > 0 {
526 result.push_str("\r\n");
527 }
528 if line
529 .as_bytes()
530 .get(..15)
531 .is_some_and(|b| b.eq_ignore_ascii_case(b"content-length:"))
532 {
533 result.push_str(&format!("Content-Length: {new_len}"));
534 } else {
535 result.push_str(line);
536 }
537 }
538 result
539}
540
541fn find_header_boundary(data: &[u8]) -> Option<usize> {
543 data.windows(4)
544 .position(|w| w == b"\r\n\r\n")
545 .map(|pos| pos + 4)
546}
547
548fn strictest_violation_action(
551 current: Option<BlockingAction>,
552 candidate: BlockingAction,
553) -> BlockingAction {
554 match (current, candidate) {
555 (Some(BlockingAction::BlockAndTerminate), _) | (_, BlockingAction::BlockAndTerminate) => {
556 BlockingAction::BlockAndTerminate
557 }
558 (Some(BlockingAction::BlockAndLog), _) | (_, BlockingAction::BlockAndLog) => {
559 BlockingAction::BlockAndLog
560 }
561 (Some(BlockingAction::Block), _) | (None, BlockingAction::Block) => BlockingAction::Block,
562 }
563}
564
565#[cfg(test)]
570mod tests {
571 use super::*;
572 use crate::secrets::config::*;
573
574 fn make_config(secrets: Vec<SecretEntry>) -> SecretsConfig {
575 SecretsConfig {
576 secrets,
577 on_violation: ViolationAction::Block,
578 }
579 }
580
581 fn make_secret(placeholder: &str, value: &str, host: &str) -> SecretEntry {
582 SecretEntry {
583 env_var: "TEST_KEY".into(),
584 value: value.into(),
585 placeholder: placeholder.into(),
586 allowed_hosts: vec![HostPattern::Exact(host.into())],
587 injection: SecretInjection::default(),
588 on_violation: None,
589 require_tls_identity: true,
590 }
591 }
592
593 fn basic_auth_only() -> SecretInjection {
594 SecretInjection {
595 headers: false,
596 basic_auth: true,
597 query_params: false,
598 body: false,
599 }
600 }
601
602 #[test]
603 fn substitute_in_headers() {
604 let config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
605 let mut handler = SecretsHandler::new(&config, "api.openai.com", true);
606
607 let input = b"GET / HTTP/1.1\r\nAuthorization: Bearer $KEY\r\n\r\n";
608 let output = handler.substitute(input).unwrap();
609 assert_eq!(
610 String::from_utf8(output.into_owned()).unwrap(),
611 "GET / HTTP/1.1\r\nAuthorization: Bearer real-secret\r\n\r\n"
612 );
613 }
614
615 #[test]
616 fn no_substitute_for_wrong_host() {
617 let config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
618 let mut handler = SecretsHandler::new(&config, "evil.com", true);
619
620 let input = b"GET / HTTP/1.1\r\nAuthorization: Bearer $KEY\r\n\r\n";
621 assert_eq!(
622 handler.substitute(input).unwrap_err(),
623 ViolationAction::Block
624 );
625 }
626
627 #[test]
628 fn allowed_placeholder_substitutes_when_another_secret_is_ineligible() {
629 let allowed = make_secret("$ALLOWED", "allowed-secret", "api.openai.com");
630 let blocked = make_secret("$BLOCKED", "blocked-secret", "api.github.com");
631 let config = make_config(vec![allowed, blocked]);
632 let mut handler = SecretsHandler::new(&config, "api.openai.com", true);
633
634 let input = b"GET / HTTP/1.1\r\nAuthorization: Bearer $ALLOWED\r\n\r\n";
635 let output = handler.substitute(input).unwrap();
636
637 assert_eq!(
638 String::from_utf8(output.into_owned()).unwrap(),
639 "GET / HTTP/1.1\r\nAuthorization: Bearer allowed-secret\r\n\r\n"
640 );
641 }
642
643 #[test]
644 fn global_passthrough_host_forwards_placeholder_unchanged() {
645 let mut config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
646 config.on_violation =
647 ViolationAction::Passthrough(vec![HostPattern::Exact("api.anthropic.com".into())]);
648 let mut handler = SecretsHandler::new(&config, "api.anthropic.com", true);
649
650 let input = b"GET / HTTP/1.1\r\nAuthorization: Bearer $KEY\r\n\r\n";
651 let output = handler.substitute(input).unwrap();
652 assert_eq!(&*output, input);
653 }
654
655 #[test]
656 fn per_secret_passthrough_host_forwards_placeholder_unchanged() {
657 let mut secret = make_secret("$KEY", "real-secret", "api.openai.com");
658 secret.on_violation = Some(ViolationAction::Passthrough(vec![HostPattern::Exact(
659 "api.anthropic.com".into(),
660 )]));
661 let config = make_config(vec![secret]);
662 let mut handler = SecretsHandler::new(&config, "api.anthropic.com", true);
663
664 let input = b"GET / HTTP/1.1\r\nAuthorization: Bearer $KEY\r\n\r\n";
665 let output = handler.substitute(input).unwrap();
666 assert_eq!(&*output, input);
667 }
668
669 #[test]
670 fn global_passthrough_action_forwards_disallowed_placeholder_unchanged() {
671 let mut config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
672 config.on_violation = ViolationAction::Passthrough(vec![HostPattern::Any]);
673 let mut handler = SecretsHandler::new(&config, "evil.com", true);
674
675 let input = b"GET / HTTP/1.1\r\nAuthorization: Bearer $KEY\r\n\r\n";
676 let output = handler.substitute(input).unwrap();
677 assert_eq!(&*output, input);
678 }
679
680 #[test]
681 fn passthrough_only_connection_has_no_handler_work() {
682 let mut config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
683 config.on_violation = ViolationAction::Passthrough(vec![HostPattern::Any]);
684 let handler = SecretsHandler::new(&config, "evil.com", true);
685
686 assert!(handler.is_empty());
687 }
688
689 #[test]
690 fn passthrough_host_does_not_allow_other_disallowed_placeholders() {
691 let mut passthrough = make_secret("$PASSTHROUGH", "real-secret-a", "api.openai.com");
692 passthrough.on_violation = Some(ViolationAction::Passthrough(vec![HostPattern::Exact(
693 "api.anthropic.com".into(),
694 )]));
695 let blocked = make_secret("$BLOCKED", "real-secret-b", "api.github.com");
696 let config = make_config(vec![passthrough, blocked]);
697 let mut handler = SecretsHandler::new(&config, "api.anthropic.com", true);
698
699 let input = b"GET / HTTP/1.1\r\nX-A: $PASSTHROUGH\r\nX-B: $BLOCKED\r\n\r\n";
700 assert_eq!(
701 handler.substitute(input).unwrap_err(),
702 ViolationAction::Block
703 );
704 }
705
706 #[test]
707 fn per_secret_passthrough_blocks_for_non_matching_host() {
708 let mut secret = make_secret("$KEY", "real-secret", "api.openai.com");
709 secret.on_violation = Some(ViolationAction::Passthrough(vec![HostPattern::Exact(
710 "api.anthropic.com".into(),
711 )]));
712 let config = make_config(vec![secret]);
713 let mut handler = SecretsHandler::new(&config, "evil.com", true);
714
715 let input = b"GET / HTTP/1.1\r\nAuthorization: Bearer $KEY\r\n\r\n";
716 assert_eq!(
717 handler.substitute(input).unwrap_err(),
718 ViolationAction::BlockAndLog
719 );
720 }
721
722 #[test]
723 fn global_passthrough_blocks_for_non_matching_host() {
724 let mut config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
725 config.on_violation =
726 ViolationAction::Passthrough(vec![HostPattern::Exact("api.anthropic.com".into())]);
727 let mut handler = SecretsHandler::new(&config, "evil.com", true);
728
729 let input = b"GET / HTTP/1.1\r\nAuthorization: Bearer $KEY\r\n\r\n";
730 assert_eq!(
731 handler.substitute(input).unwrap_err(),
732 ViolationAction::BlockAndLog
733 );
734 }
735
736 #[test]
737 fn global_block_and_terminate_marks_violation_as_terminating() {
738 let mut config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
739 config.on_violation = ViolationAction::BlockAndTerminate;
740 let mut handler = SecretsHandler::new(&config, "evil.com", true);
741
742 let input = b"GET / HTTP/1.1\r\nAuthorization: Bearer $KEY\r\n\r\n";
743 assert_eq!(
744 handler.substitute(input).unwrap_err(),
745 ViolationAction::BlockAndTerminate
746 );
747 }
748
749 #[test]
750 fn per_secret_block_and_terminate_marks_violation_as_terminating() {
751 let mut secret = make_secret("$KEY", "real-secret", "api.openai.com");
752 secret.on_violation = Some(ViolationAction::BlockAndTerminate);
753 let config = make_config(vec![secret]);
754 let mut handler = SecretsHandler::new(&config, "evil.com", true);
755
756 let input = b"GET / HTTP/1.1\r\nAuthorization: Bearer $KEY\r\n\r\n";
757 assert_eq!(
758 handler.substitute(input).unwrap_err(),
759 ViolationAction::BlockAndTerminate
760 );
761 }
762
763 #[test]
764 fn body_injection_disabled_by_default() {
765 let config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
766 let mut handler = SecretsHandler::new(&config, "api.openai.com", true);
767
768 let input = b"POST / HTTP/1.1\r\n\r\n{\"key\": \"$KEY\"}";
769 let output = handler.substitute(input).unwrap();
770 assert!(
771 String::from_utf8(output.into_owned())
772 .unwrap()
773 .contains("$KEY")
774 );
775 }
776
777 #[test]
778 fn body_injection_when_enabled() {
779 let mut secret = make_secret("$KEY", "real-secret", "api.openai.com");
780 secret.injection.body = true;
781 let config = make_config(vec![secret]);
782 let mut handler = SecretsHandler::new(&config, "api.openai.com", true);
783
784 let input = b"POST / HTTP/1.1\r\n\r\n{\"key\": \"$KEY\"}";
785 let output = handler.substitute(input).unwrap();
786 assert_eq!(
787 String::from_utf8(output.into_owned()).unwrap(),
788 "POST / HTTP/1.1\r\n\r\n{\"key\": \"real-secret\"}"
789 );
790 }
791
792 #[test]
793 fn body_injection_updates_content_length() {
794 let mut secret = make_secret("$KEY", "a]longer]secret]value", "api.openai.com");
795 secret.injection.body = true;
796 let config = make_config(vec![secret]);
797 let mut handler = SecretsHandler::new(&config, "api.openai.com", true);
798
799 let body = "{\"key\": \"$KEY\"}";
800 let input = format!(
801 "POST / HTTP/1.1\r\nContent-Length: {}\r\n\r\n{}",
802 body.len(),
803 body
804 );
805 let output = handler.substitute(input.as_bytes()).unwrap();
806 let result = String::from_utf8(output.into_owned()).unwrap();
807
808 let expected_body = "{\"key\": \"a]longer]secret]value\"}";
809 assert!(result.contains(expected_body));
810 assert!(result.contains(&format!("Content-Length: {}", expected_body.len())));
811 }
812
813 #[test]
814 fn body_injection_no_content_length_header() {
815 let mut secret = make_secret("$KEY", "longer-secret", "api.openai.com");
816 secret.injection.body = true;
817 let config = make_config(vec![secret]);
818 let mut handler = SecretsHandler::new(&config, "api.openai.com", true);
819
820 let input = b"POST / HTTP/1.1\r\nTransfer-Encoding: chunked\r\n\r\n{\"key\": \"$KEY\"}";
822 let output = handler.substitute(input).unwrap();
823 let result = String::from_utf8(output.into_owned()).unwrap();
824 assert!(result.contains("longer-secret"));
825 assert!(!result.contains("Content-Length"));
826 }
827
828 #[test]
829 fn header_only_substitution_preserves_content_length() {
830 let config = make_config(vec![make_secret("$KEY", "longer-value", "api.openai.com")]);
831 let mut handler = SecretsHandler::new(&config, "api.openai.com", true);
832
833 let input =
834 b"GET / HTTP/1.1\r\nAuthorization: Bearer $KEY\r\nContent-Length: 5\r\n\r\nhello";
835 let output = handler.substitute(input).unwrap();
836 let result = String::from_utf8(output.into_owned()).unwrap();
837 assert!(result.contains("Content-Length: 5"));
839 assert!(result.ends_with("hello"));
840 }
841
842 #[test]
843 fn eligible_secret_preserves_binary_body_without_placeholder() {
844 let config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
845 let mut handler = SecretsHandler::new(&config, "api.openai.com", true);
846
847 let body = vec![0x1f, 0x8b, 0x08, 0x00, 0xff, 0x00, 0x80, 0xfe];
848 let mut input = format!(
849 "POST /git-upload-pack HTTP/1.1\r\nContent-Encoding: gzip\r\nContent-Length: {}\r\n\r\n",
850 body.len()
851 )
852 .into_bytes();
853 input.extend_from_slice(&body);
854
855 let output = handler.substitute(&input).unwrap();
856 assert_eq!(&*output, input.as_slice());
857 }
858
859 #[test]
860 fn eligible_secret_preserves_binary_chunk_without_placeholder() {
861 let config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
862 let mut handler = SecretsHandler::new(&config, "api.openai.com", true);
863
864 let input = [0x1f, 0x8b, 0x08, 0x00, 0xff, 0x00, 0x80, 0xfe];
865 let output = handler.substitute(&input).unwrap();
866 assert_eq!(&*output, input.as_slice());
867 }
868
869 #[test]
870 fn body_injection_preserves_non_utf8_bytes() {
871 let mut secret = make_secret("$KEY", "real-secret", "api.openai.com");
872 secret.injection.body = true;
873 let config = make_config(vec![secret]);
874 let mut handler = SecretsHandler::new(&config, "api.openai.com", true);
875
876 let body = [0xff, b'$', b'K', b'E', b'Y', 0xfe];
877 let mut input =
878 format!("POST / HTTP/1.1\r\nContent-Length: {}\r\n\r\n", body.len()).into_bytes();
879 input.extend_from_slice(&body);
880
881 let output = handler.substitute(&input).unwrap().into_owned();
882 let expected_body = [b"\xffreal-secret".as_slice(), &[0xfe]].concat();
883 let expected = [
884 format!(
885 "POST / HTTP/1.1\r\nContent-Length: {}\r\n\r\n",
886 expected_body.len()
887 )
888 .as_bytes(),
889 expected_body.as_slice(),
890 ]
891 .concat();
892
893 assert_eq!(output, expected);
894 }
895
896 #[test]
897 fn no_secrets_passthrough() {
898 let config = make_config(vec![]);
899 let mut handler = SecretsHandler::new(&config, "anything.com", true);
900
901 let input = b"GET / HTTP/1.1\r\n\r\n";
902 let output = handler.substitute(input).unwrap();
903 assert_eq!(&*output, input);
904 }
905
906 #[test]
907 fn require_tls_identity_blocks_on_non_intercepted() {
908 let config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
909 let mut handler = SecretsHandler::new(&config, "api.openai.com", false);
911
912 let input = b"GET / HTTP/1.1\r\nAuthorization: Bearer $KEY\r\n\r\n";
913 let output = handler.substitute(input).unwrap();
914 assert!(
916 String::from_utf8(output.into_owned())
917 .unwrap()
918 .contains("$KEY")
919 );
920 }
921
922 #[test]
923 fn basic_auth_only_does_not_substitute_other_schemes() {
924 let mut secret = make_secret("$KEY", "real-secret", "api.openai.com");
925 secret.injection = basic_auth_only();
926 let config = make_config(vec![secret]);
927 let mut handler = SecretsHandler::new(&config, "api.openai.com", true);
928
929 let input = b"GET / HTTP/1.1\r\nAuthorization: Bearer $KEY\r\nX-Custom: $KEY\r\n\r\n";
931 let output = handler.substitute(input).unwrap();
932 let result = String::from_utf8(output.into_owned()).unwrap();
933 assert!(result.contains("Authorization: Bearer $KEY"));
934 assert!(result.contains("X-Custom: $KEY"));
935 }
936
937 #[test]
938 fn basic_auth_decodes_substitutes_and_reencodes_credentials() {
939 let mut user = make_secret("$MSB_USER", "alice", "api.openai.com");
940 user.env_var = "USER".into();
941 user.injection = basic_auth_only();
942 let mut password = make_secret("$MSB_PASSWORD", "s3cr3t", "api.openai.com");
943 password.env_var = "PASSWORD".into();
944 password.injection = basic_auth_only();
945 let config = make_config(vec![user, password]);
946 let mut handler = SecretsHandler::new(&config, "api.openai.com", true);
947
948 let encoded = BASE64.encode(b"$MSB_USER:$MSB_PASSWORD");
949 let input = format!("GET / HTTP/1.1\r\nAuthorization: Basic {encoded}\r\n\r\n");
950 let output = handler.substitute(input.as_bytes()).unwrap();
951 let result = String::from_utf8(output.into_owned()).unwrap();
952
953 assert!(result.contains(&format!(
954 "Authorization: Basic {}",
955 BASE64.encode(b"alice:s3cr3t")
956 )));
957 assert!(!result.contains("$MSB_USER"));
958 assert!(!result.contains("$MSB_PASSWORD"));
959 }
960
961 #[test]
962 fn basic_auth_encoded_placeholder_is_blocked_for_wrong_host() {
963 let mut secret = make_secret("$MSB_PASSWORD", "s3cr3t", "api.openai.com");
964 secret.injection = basic_auth_only();
965 let config = make_config(vec![secret]);
966 let mut handler = SecretsHandler::new(&config, "evil.com", true);
967
968 let encoded = BASE64.encode(b"user:$MSB_PASSWORD");
969 let input = format!("GET / HTTP/1.1\r\nAuthorization: Basic {encoded}\r\n\r\n");
970
971 assert_eq!(
972 handler.substitute(input.as_bytes()).unwrap_err(),
973 ViolationAction::Block
974 );
975 }
976
977 #[test]
978 fn basic_auth_encoded_placeholder_is_not_replaced_when_scope_disabled() {
979 let mut secret = make_secret("$MSB_PASSWORD", "s3cr3t", "api.openai.com");
980 secret.injection = SecretInjection {
981 headers: false,
982 basic_auth: false,
983 query_params: false,
984 body: false,
985 };
986 let config = make_config(vec![secret]);
987 let mut handler = SecretsHandler::new(&config, "api.openai.com", true);
988
989 let encoded = BASE64.encode(b"user:$MSB_PASSWORD");
990 let input = format!("GET / HTTP/1.1\r\nAuthorization: Basic {encoded}\r\n\r\n");
991 let output = handler.substitute(input.as_bytes()).unwrap();
992
993 assert_eq!(String::from_utf8(output.into_owned()).unwrap(), input);
994 }
995
996 #[test]
997 fn query_params_substitution() {
998 let mut secret = make_secret("$KEY", "real-secret", "api.openai.com");
999 secret.injection = SecretInjection {
1000 headers: false,
1001 basic_auth: false,
1002 query_params: true,
1003 body: false,
1004 };
1005 let config = make_config(vec![secret]);
1006 let mut handler = SecretsHandler::new(&config, "api.openai.com", true);
1007
1008 let input = b"GET /api?key=$KEY HTTP/1.1\r\nHost: api.openai.com\r\n\r\n";
1009 let output = handler.substitute(input).unwrap();
1010 let result = String::from_utf8(output.into_owned()).unwrap();
1011 assert!(result.contains("GET /api?key=real-secret HTTP/1.1"));
1013 }
1015
1016 #[test]
1017 fn url_encoded_placeholder_in_query_blocks_for_wrong_host() {
1018 let config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
1019 let mut handler = SecretsHandler::new(&config, "evil.com", true);
1020
1021 let input = b"GET /api?token=%24KEY HTTP/1.1\r\nHost: evil.com\r\n\r\n";
1023 assert_eq!(
1024 handler.substitute(input).unwrap_err(),
1025 ViolationAction::Block
1026 );
1027 }
1028
1029 #[test]
1030 fn url_encoded_placeholder_in_body_blocks_for_wrong_host() {
1031 let config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
1032 let mut handler = SecretsHandler::new(&config, "evil.com", true);
1033
1034 let input = b"POST / HTTP/1.1\r\nContent-Length: 13\r\n\r\nkey=%24KEY&x=1";
1035 assert_eq!(
1036 handler.substitute(input).unwrap_err(),
1037 ViolationAction::Block
1038 );
1039 }
1040
1041 #[test]
1042 fn json_escaped_placeholder_in_body_blocks_for_wrong_host() {
1043 let config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
1044 let mut handler = SecretsHandler::new(&config, "evil.com", true);
1045
1046 let input =
1048 b"POST / HTTP/1.1\r\nContent-Type: application/json\r\n\r\n{\"k\":\"\\u0024KEY\"}";
1049 assert_eq!(
1050 handler.substitute(input).unwrap_err(),
1051 ViolationAction::Block
1052 );
1053 }
1054
1055 #[test]
1056 fn placeholder_split_across_writes_blocks_for_wrong_host() {
1057 let config = make_config(vec![make_secret("$KEY", "real-secret", "api.openai.com")]);
1058 let mut handler = SecretsHandler::new(&config, "evil.com", true);
1059
1060 let first = b"GET / HTTP/1.1\r\nX-Token: $K";
1062 let second = b"EY\r\nHost: evil.com\r\n\r\n";
1063
1064 assert!(handler.substitute(first).is_ok());
1066 assert_eq!(
1068 handler.substitute(second).unwrap_err(),
1069 ViolationAction::Block
1070 );
1071 }
1072
1073 #[test]
1074 fn url_decoded_contains_basic() {
1075 assert!(url_decoded_contains(b"foo%24KEYbar", b"$KEY"));
1076 assert!(!url_decoded_contains(b"fooKEYbar", b"$KEY"));
1077 assert!(url_decoded_contains(b"%2", b"%2"));
1079 }
1080
1081 #[test]
1082 fn json_escaped_contains_basic() {
1083 assert!(json_escaped_contains(b"\"\\u0024KEY\"", b"$KEY"));
1084 assert!(json_escaped_contains(
1085 b"\\u0024\\u004B\\u0045\\u0059",
1086 b"$KEY"
1087 ));
1088 assert!(!json_escaped_contains(b"KEY", b"$KEY"));
1089 }
1090}