Skip to main content

microsandbox_network/secrets/
handler.rs

1//! Secret substitution handler for the TLS proxy.
2//!
3//! Scans decrypted plaintext for placeholder strings and replaces them
4//! with real secret values, but only when the destination host is allowed.
5
6use 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
13//--------------------------------------------------------------------------------------------------
14// Types
15//--------------------------------------------------------------------------------------------------
16
17/// Handles secret placeholder substitution in TLS-intercepted plaintext.
18///
19/// Created from [`SecretsConfig`] and the destination SNI. Determines which
20/// secrets are eligible for this connection based on host matching.
21pub struct SecretsHandler {
22    /// Secrets eligible for substitution on this connection.
23    eligible_for_substitution: Vec<EligibleSecret>,
24    /// Secret placeholders that should trigger an effective blocking action.
25    ineligible_for_substitution: Vec<IneligibleSecret>,
26    /// Whether this connection is TLS-intercepted (not bypass).
27    tls_intercepted: bool,
28    /// Longest placeholder length. Sizes the sliding-window tail.
29    max_placeholder_len: usize,
30    /// Trailing bytes carried over from the previous `substitute` call so a
31    /// placeholder split across TCP writes still trips the violation check.
32    /// Capped at `max_placeholder_len - 1` bytes.
33    prev_tail: Vec<u8>,
34}
35
36/// A secret that passed host matching for this connection.
37struct 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
47/// A secret that did not pass substitution or passthrough host matching.
48struct IneligibleSecret {
49    placeholder: String,
50    action: BlockingAction,
51}
52
53/// Blocking action to take when an ineligible placeholder is detected.
54#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
55enum BlockingAction {
56    Block,
57    #[default]
58    BlockAndLog,
59    BlockAndTerminate,
60}
61
62//--------------------------------------------------------------------------------------------------
63// Methods
64//--------------------------------------------------------------------------------------------------
65
66impl EligibleSecret {
67    /// Returns true if any of the header-side injection scopes is enabled
68    /// (`headers`, `basic_auth`, or `query_params`).
69    fn wants_header_injection(&self) -> bool {
70        self.inject_headers || self.inject_basic_auth || self.inject_query_params
71    }
72
73    /// Returns true when the current header bytes contain this secret's
74    /// placeholder in a header-substitution scope.
75    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        // Search decoded Basic auth credentials, not the raw header value.
86        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    /// Substitute this secret's placeholder in the headers portion, scoped by
97    /// the secret's `headers` / `basic_auth` / `query_params` flags.
98    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    /// Substitute this secret's placeholder in a single header line. Returns
113    /// `None` if the line is not in scope for any of the requested injection
114    /// modes.
115    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    /// Decode `Basic <base64>` credentials, substitute the placeholder in the
132    /// decoded `user:password`, and return the re-encoded line. Returns `None`
133    /// if the line isn't `Basic` scheme or the decoded credentials don't
134    /// contain the placeholder. Non-Basic schemes (e.g. `Bearer`) are handled
135    /// by `inject_headers` instead.
136    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    /// Create a handler for a specific connection.
171    ///
172    /// Filters secrets by host matching against the SNI. Only secrets
173    /// whose `allowed_hosts` match `sni` will be substituted.
174    /// `tls_intercepted` indicates whether this is a MITM connection
175    /// (true) or a bypass/plain connection (false).
176    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 the SNI matches an allowed host for this secret, add it to the
188            // eligible list for substitution, and skip violation checks for this secret.
189            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            // Passthrough means the placeholder can be forwarded unchanged to this SNI.
206            if let ViolationAction::Passthrough(hosts) = action
207                && hosts.iter().any(|p| p.matches(sni))
208            {
209                continue;
210            }
211
212            // Non-matching passthrough policies fall back to the default blocking action.
213            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    /// Substitute secrets in plaintext data (guest → server direction).
229    ///
230    /// Splits the HTTP message on `\r\n\r\n` to scope substitution:
231    /// - `headers`: substitutes in the header portion (before boundary)
232    /// - `basic_auth`: substitutes in Authorization headers specifically
233    /// - `query_params`: substitutes in the request line (first line, query portion)
234    /// - `body`: substitutes in the body portion (after boundary)
235    ///
236    /// Returns the violation action if a placeholder is detected going to a
237    /// disallowed host.
238    pub fn substitute<'a>(&mut self, data: &'a [u8]) -> Result<Cow<'a, [u8]>, ViolationAction> {
239        // Split raw bytes at the header boundary BEFORE converting to owned strings.
240        // This avoids position shifts from from_utf8_lossy replacement chars.
241        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        // Check for disallowed placeholders before forwarding or substituting data.
248        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            // No substitution needed. Return borrowed slice (zero-copy).
269            return Ok(Cow::Borrowed(data));
270        }
271
272        // Start with borrowed bytes; allocate only when a substitution is needed.
273        let mut header_str = None;
274        let mut body = None;
275
276        for secret in &self.eligible_for_substitution {
277            // Skip secrets that require TLS identity on non-intercepted connections.
278            if secret.require_tls_identity && !self.tls_intercepted {
279                continue;
280            }
281
282            // Header substitution still uses string helpers after a scoped match.
283            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            // Body substitution works on bytes so encoded payloads stay valid.
290            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        // No header or body replacement was produced. Return original bytes.
308        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        // Update Content-Length only when body substitution changed the size.
320        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    /// Returns true if this connection needs no secret substitution or violation detection.
340    pub fn is_empty(&self) -> bool {
341        self.eligible_for_substitution.is_empty() && self.ineligible_for_substitution.is_empty()
342    }
343
344    /// Returns the strongest blocking action for any placeholder appearing in data
345    /// for a host that isn't allowed to receive either the real secret or the placeholder.
346    ///
347    /// Scans the raw bytes (stitched with the previous call's tail for
348    /// cross-write detection), plus URL- and JSON-decoded variants for
349    /// encoded-placeholder bypass attempts, plus base64-decoded Basic auth
350    /// credentials.
351    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    /// Update the sliding-window tail with the trailing bytes of `data`, so
382    /// the next `substitute` call can detect placeholders split across the
383    /// boundary.
384    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
403//--------------------------------------------------------------------------------------------------
404// Functions
405//--------------------------------------------------------------------------------------------------
406
407/// Returns true if `line` starts with the `Authorization:` header name
408/// (case-insensitive).
409fn 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
415/// Decode the credentials of a `Basic` `Authorization` header line. Returns
416/// `None` if the line is not `Basic`-scheme or the payload is not valid
417/// base64 / UTF-8.
418fn 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
428/// Split an `Authorization` header value into `(scheme, rest)` at the first
429/// whitespace. Returns `None` if no whitespace separator is found.
430fn 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
436/// Returns true if any `Authorization: Basic` line in `headers` decodes to
437/// credentials containing `placeholder`.
438fn 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
446/// Byte-slice substring check.
447fn 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
454/// Replace all occurrences of `needle` in `haystack`.
455///
456/// Returns `None` when no replacement is needed so callers can preserve the
457/// original byte slice without rebuilding arbitrary binary payloads.
458fn 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
477/// Returns true if `haystack`, after URL percent-decoding, contains `needle`.
478fn url_decoded_contains(haystack: &[u8], needle: &[u8]) -> bool {
479    let decoded: Vec<u8> = percent_decode(haystack).collect();
480    contains_bytes(&decoded, needle)
481}
482
483/// Returns true if `haystack`, after JSON `\uXXXX` decoding, contains `needle`.
484/// Only `\uXXXX` escapes are expanded (sufficient to detect ASCII placeholders
485/// hidden via unicode escapes); other JSON escapes pass through.
486fn 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
518/// Update the Content-Length header value in `headers` to `new_len`.
519///
520/// Performs a case-insensitive line scan. If no Content-Length header exists
521/// (e.g. chunked transfer encoding), the headers are returned unchanged.
522fn 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
541/// Find the `\r\n\r\n` boundary between HTTP headers and body.
542fn 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
548/// Returns the stricter of two blocking actions, where
549/// `BlockAndTerminate` > `BlockAndLog` > `Block`.
550fn 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//--------------------------------------------------------------------------------------------------
566// Tests
567//--------------------------------------------------------------------------------------------------
568
569#[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        // No Content-Length header (e.g. chunked).
821        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        // Body unchanged, Content-Length should stay 5.
838        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        // tls_intercepted = false — secret requires TLS identity
910        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        // Placeholder should NOT be substituted.
915        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        // basic_auth only handles Basic credentials; Bearer needs inject_headers.
930        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        // Request line should be substituted.
1012        assert!(result.contains("GET /api?key=real-secret HTTP/1.1"));
1013        // Other headers should NOT be substituted.
1014    }
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        // `%24KEY` is the URL-encoded form of `$KEY`.
1022        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        // `$KEY` is the JSON unicode-escape form of `$KEY`.
1047        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        // Send the placeholder bytes across two separate substitute() calls.
1061        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        // The first chunk doesn't contain the full placeholder, so it forwards.
1065        assert!(handler.substitute(first).is_ok());
1066        // The second chunk completes the placeholder when stitched with the tail.
1067        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        // Invalid escapes pass through unchanged.
1078        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}