Skip to main content

agent_first_data/
redaction.rs

1use serde_json::Value;
2use std::collections::HashSet;
3
4// ═══════════════════════════════════════════
5// Public API: Output Formatters
6// ═══════════════════════════════════════════
7
8/// Redaction policy. Convert with `.into()` to [`Redactor`] or [`OutputOptions`].
9#[derive(Clone, Copy, Debug, PartialEq, Eq)]
10pub enum RedactionPolicy {
11    /// Redact only inside top-level `trace`.
12    RedactionTraceOnly,
13    /// Do not redact any fields.
14    RedactionNone,
15}
16
17/// Rendering style for YAML and plain output.
18#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
19pub enum OutputStyle {
20    /// Human-readable AFDATA rendering: strip suffixes and format values.
21    #[default]
22    Readable,
23    /// Schema-preserving rendering: keep keys and values unchanged after redaction.
24    Raw,
25}
26
27/// Configurable redaction builder for secrets and legacy field names.
28///
29/// `Redactor` encapsulates redaction policy and custom secret field names.
30/// Build with [`Redactor::new()`], configure via builder methods, then pass to
31/// redaction functions like [`redacted_value`] or [`redact_url_secrets`].
32#[derive(Clone, Debug, Default, PartialEq, Eq)]
33pub struct Redactor {
34    policy: Option<RedactionPolicy>,
35    secret_names: Vec<String>,
36}
37
38impl Redactor {
39    /// Create a new default redactor (full redaction, no custom secret names).
40    pub fn new() -> Self {
41        Self::default()
42    }
43
44    /// Set custom field names to treat as secrets in addition to `_secret` suffixes.
45    ///
46    /// Matching is exact field-name equality at any nesting level. The same
47    /// list also matches URL query-parameter names inside `_url` fields.
48    /// Builder style: returns `self`.
49    pub fn secret_names<I: IntoIterator<Item = S>, S: Into<String>>(mut self, names: I) -> Self {
50        self.secret_names = names.into_iter().map(|s| s.into()).collect();
51        self
52    }
53
54    /// Set the redaction policy (default: full redaction).
55    /// Builder style: returns `self`.
56    pub fn policy(mut self, policy: RedactionPolicy) -> Self {
57        self.policy = Some(policy);
58        self
59    }
60
61    /// Redact a JSON value copy using this redactor's policy and secret names.
62    pub fn value(&self, value: &Value) -> Value {
63        let mut v = value.clone();
64        self.redact_in_place(&mut v);
65        v
66    }
67
68    /// Redact secret components of a URL string using this redactor's settings.
69    ///
70    /// A query parameter is redacted iff its (form-decoded) name ends in
71    /// `_secret`/`_SECRET` or matches an exact entry in `secret_names`. The
72    /// userinfo password (`scheme://user:pass@host`) is always redacted as a
73    /// structural rule. Only the secret spans are replaced with `***`; every
74    /// other byte is preserved. A string that is not a single, whitespace-free,
75    /// scheme-prefixed URL (including a URL embedded in surrounding prose) is
76    /// returned unchanged.
77    pub fn url(&self, url: &str) -> String {
78        let context = RedactionContext::from_redactor(self);
79        redact_url_in_str(url, &context).unwrap_or_else(|| url.to_string())
80    }
81
82    pub(crate) fn redact_in_place(&self, value: &mut Value) {
83        let context = RedactionContext::from_redactor(self);
84        apply_redaction_policy_with_context(value, self.policy, &context);
85    }
86}
87
88impl From<RedactionPolicy> for Redactor {
89    fn from(policy: RedactionPolicy) -> Self {
90        Self {
91            policy: Some(policy),
92            secret_names: Vec::new(),
93        }
94    }
95}
96
97/// Output options combining redaction and rendering style.
98#[derive(Clone, Debug, Default, PartialEq, Eq)]
99pub struct OutputOptions {
100    /// Redactor applied before rendering.
101    pub redaction: Redactor,
102    /// Rendering style for YAML and plain output.
103    pub style: OutputStyle,
104}
105
106impl From<RedactionPolicy> for OutputOptions {
107    fn from(policy: RedactionPolicy) -> Self {
108        Self {
109            redaction: Redactor::from(policy),
110            style: OutputStyle::default(),
111        }
112    }
113}
114
115// ═══════════════════════════════════════════
116// Public API: Redaction & Utility
117// ═══════════════════════════════════════════
118
119/// Return a JSON value copy with default `_secret` redaction applied.
120pub fn redacted_value(value: &Value) -> Value {
121    Redactor::new().value(value)
122}
123
124/// Redact secret components of a single URL string, using default options.
125///
126/// Returns `url` with its userinfo password and any `_secret`-suffixed query
127/// parameter values replaced by `***`.
128pub fn redact_url_secrets(url: &str) -> String {
129    Redactor::new().url(url)
130}
131
132/// Parse a human-readable size string into bytes.
133///
134/// Accepts a number followed by an explicit unit. Decimal units are
135/// `B`, `kB`, `MB`, `GB`, `TB`; binary units are `KiB`, `MiB`, `GiB`, `TiB`.
136
137// ═══════════════════════════════════════════
138// Secret Redaction
139// ═══════════════════════════════════════════
140
141#[derive(Default)]
142pub(crate) struct RedactionContext {
143    secret_names: HashSet<String>,
144}
145
146impl RedactionContext {
147    fn from_redactor(redactor: &Redactor) -> Self {
148        let secret_names = redactor.secret_names.iter().cloned().collect();
149        Self { secret_names }
150    }
151
152    fn is_secret_key(&self, key: &str) -> bool {
153        key_has_secret_suffix(key) || self.secret_names.contains(key)
154    }
155}
156
157fn key_has_secret_suffix(key: &str) -> bool {
158    key.ends_with("_secret") || key.ends_with("_SECRET")
159}
160
161fn key_has_url_suffix(key: &str) -> bool {
162    key.ends_with("_url") || key.ends_with("_URL")
163}
164
165#[cfg(feature = "cli-help")]
166pub(crate) fn is_secret_flag_name(flag_name: &str, context: &RedactionContext) -> bool {
167    let normalized = flag_name.replace('-', "_");
168    context.is_secret_key(&normalized) || context.is_secret_key(flag_name)
169}
170
171const MAX_DEPTH: usize = 256;
172const MAX_DEPTH_MARKER: &str = "<afdata:max-depth>";
173
174fn redact_secrets_with_context(value: &mut Value, context: &RedactionContext) {
175    redact_secrets_with_context_depth(value, context, 0);
176}
177
178fn redact_secrets_with_context_depth(value: &mut Value, context: &RedactionContext, depth: usize) {
179    if depth >= MAX_DEPTH {
180        *value = Value::String(MAX_DEPTH_MARKER.into());
181        return;
182    }
183    match value {
184        Value::Object(map) => {
185            let keys: Vec<String> = map.keys().cloned().collect();
186            for key in keys {
187                if context.is_secret_key(&key) {
188                    map.insert(key, Value::String("***".into()));
189                } else if key_has_url_suffix(&key) {
190                    if let Some(Value::String(s)) = map.get_mut(&key) {
191                        *s = redact_url_field_value(s, context);
192                    } else if let Some(v) = map.get_mut(&key) {
193                        redact_secrets_with_context_depth(v, context, depth + 1);
194                    }
195                } else if let Some(v) = map.get_mut(&key) {
196                    redact_secrets_with_context_depth(v, context, depth + 1);
197                }
198            }
199        }
200        Value::Array(arr) => {
201            for v in arr {
202                redact_secrets_with_context_depth(v, context, depth + 1);
203            }
204        }
205        _ => {}
206    }
207}
208
209/// Redact secret components of a single URL string, returning `Some(redacted)`
210/// when `s` is a processable URL, or `None` when it is not (so callers can keep
211/// the original). Only secret spans change; all other bytes are preserved.
212fn redact_url_in_str(s: &str, context: &RedactionContext) -> Option<String> {
213    // Precondition (spec): a single, whitespace-free, scheme-prefixed URL.
214    // The gate is scheme + no-whitespace only — NOT "parses as a URL library
215    // object". Span location below is purely byte-wise, so we never re-serialize
216    // the URL; adding a `url::Url::parse` gate here would diverge across
217    // languages (e.g. ports > 65535 or empty hosts that one library rejects and
218    // another accepts) and silently leak secrets in the values it rejects.
219    if !s.contains("://") || !is_single_url(s) {
220        return None;
221    }
222    let scheme_sep = s.find("://")?;
223    let scheme = &s[..scheme_sep];
224    let rest = &s[scheme_sep + 3..];
225
226    // Authority runs from after "://" to the first '/', '?', or '#'.
227    let auth_end = rest.find(['/', '?', '#']).unwrap_or(rest.len());
228    let authority = &rest[..auth_end];
229    let remainder = &rest[auth_end..];
230
231    let new_authority = redact_userinfo_password(authority);
232
233    // Query runs from the first '?' to the first '#' (or end).
234    let new_remainder = match remainder.find('?') {
235        Some(q) => {
236            let (path, q_onwards) = remainder.split_at(q);
237            let query_body = &q_onwards[1..];
238            let (query, fragment) = match query_body.find('#') {
239                Some(h) => (&query_body[..h], &query_body[h..]),
240                None => (query_body, ""),
241            };
242            format!("{path}?{}{fragment}", redact_query(query, context))
243        }
244        None => remainder.to_string(),
245    };
246
247    Some(format!("{scheme}://{new_authority}{new_remainder}"))
248}
249
250fn redact_url_field_value(s: &str, context: &RedactionContext) -> String {
251    if let Some(redacted) = redact_url_in_str(s, context) {
252        return redacted;
253    }
254    let trimmed = s.trim();
255    if trimmed != s
256        && let Some(redacted) = redact_url_in_str(trimmed, context)
257    {
258        return redacted;
259    }
260    // Fail closed: a `_url` value we could not parse as a clean scheme-prefixed
261    // URL, yet which carries a credential sigil (`@` userinfo) or internal
262    // whitespace, is redacted wholesale rather than passed through. A schemeless
263    // connection string like `user:pass@host/db` has no scheme anchor for the
264    // surgical span logic above, so blanket redaction is the safe default.
265    if s.chars().any(char::is_whitespace) || s.contains('@') {
266        return "***".to_string();
267    }
268    s.to_string()
269}
270
271/// Replace the userinfo password (`user:pass@`) with `***`, preserving the
272/// username. Authority without `@`, or userinfo without `:`, is unchanged.
273fn redact_userinfo_password(authority: &str) -> String {
274    let Some(at) = authority.rfind('@') else {
275        return authority.to_string();
276    };
277    let userinfo = &authority[..at];
278    match userinfo.find(':') {
279        Some(colon) => format!("{}:***{}", &authority[..colon], &authority[at..]),
280        None => authority.to_string(),
281    }
282}
283
284/// Redact the values of secret-named query parameters, preserving raw bytes of
285/// every other segment (keys, benign values, encoding, ordering, separators).
286fn redact_query(query: &str, context: &RedactionContext) -> String {
287    query
288        .split('&')
289        .map(|segment| {
290            let Some(eq) = segment.find('=') else {
291                return segment.to_string();
292            };
293            let raw_key = &segment[..eq];
294            // Form-decode the name (`+` → space, percent-decode) for the check.
295            let name = url::form_urlencoded::parse(segment.as_bytes())
296                .next()
297                .map(|(k, _)| k.into_owned())
298                .unwrap_or_default();
299            if context.is_secret_key(&name) {
300                format!("{raw_key}=***")
301            } else {
302                segment.to_string()
303            }
304        })
305        .collect::<Vec<_>>()
306        .join("&")
307}
308
309/// True when `s` begins with a URL scheme (`ALPHA *(ALPHA / DIGIT / "+" / "-" /
310/// ".") "://"`) and contains no ASCII whitespace — i.e. a single bare URL, not
311/// a URL embedded in prose.
312fn is_single_url(s: &str) -> bool {
313    if s.bytes().any(|b| b.is_ascii_whitespace()) {
314        return false;
315    }
316    let bytes = s.as_bytes();
317    if !bytes.first().is_some_and(|b| b.is_ascii_alphabetic()) {
318        return false;
319    }
320    let mut i = 1;
321    while i < bytes.len() {
322        let c = bytes[i];
323        if c.is_ascii_alphanumeric() || matches!(c, b'+' | b'-' | b'.') {
324            i += 1;
325        } else {
326            break;
327        }
328    }
329    s[i..].starts_with("://")
330}
331
332fn apply_redaction_policy_with_context(
333    value: &mut Value,
334    redaction_policy: Option<RedactionPolicy>,
335    context: &RedactionContext,
336) {
337    match redaction_policy {
338        Some(RedactionPolicy::RedactionTraceOnly) => {
339            if let Value::Object(map) = value
340                && let Some(trace) = map.get_mut("trace")
341            {
342                redact_secrets_with_context(trace, context);
343            }
344        }
345        Some(RedactionPolicy::RedactionNone) => {}
346        None => redact_secrets_with_context(value, context),
347    }
348}