agent_first_data/

redaction.rs

mod url;

use self::url::{redact_url_field_value, redact_url_in_str, redact_urls_in_text_with_context};
use serde_json::Value;
use std::collections::HashSet;

// ═══════════════════════════════════════════
// Public API: Output Formatters
// ═══════════════════════════════════════════

/// Which fields a [`Redactor`] scrubs. The default is [`RedactionPolicy::All`].
///
/// The policy selects a *scope* inside a structured value. A command line
/// ([`Redactor::argv`]) and a bare URL string ([`Redactor::url`]) have no
/// `result`/`trace` split to scope to, so on those two paths `TraceOnly`
/// redacts in full like `All`, and only [`RedactionPolicy::Off`] turns
/// redaction off.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum RedactionPolicy {
    /// Redact every secret field anywhere in the value (the default).
    #[default]
    All,
    /// Redact only inside the top-level `trace` object.
    TraceOnly,
    /// Do not redact anything.
    Off,
}

impl RedactionPolicy {
    /// Whether this policy redacts an input that carries no scope of its own —
    /// a command line or a single URL string, as opposed to a JSON value with a
    /// `trace` object to narrow to.
    ///
    /// `TraceOnly` narrows *where* redaction applies within a value; it does not
    /// weaken redaction. A standalone argv or URL has no non-`trace` half to
    /// leave alone — and both are diagnostic material by construction, the very
    /// thing `TraceOnly` scrubs — so `TraceOnly` redacts them in full, exactly
    /// like `All`. Only `Off`, the caller explicitly asking for raw output,
    /// disables redaction, and it does so on all three paths alike.
    fn redacts_unscoped_input(self) -> bool {
        !matches!(self, RedactionPolicy::Off)
    }
}

/// Rendering style for plain (logfmt) output only. JSON and YAML are always
/// structure-preserving and ignore this.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum PlainStyle {
    /// Human-readable AFDATA rendering: strip suffixes and format values.
    #[default]
    Readable,
    /// Schema-preserving rendering: keep keys and values unchanged after redaction.
    Raw,
}

/// Configurable redaction builder for secrets and legacy field names.
///
/// `Redactor` encapsulates redaction policy, custom secret field names, and
/// exact legacy URL field names.
/// Build with [`Redactor::new()`], configure via builder methods, then pass to
/// redaction functions like [`redacted_value`] or [`redact_url_secrets`].
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Redactor {
    policy: RedactionPolicy,
    secret_names: Vec<String>,
    url_names: Vec<String>,
    reveal_names: Vec<String>,
}

impl Redactor {
    /// Create a new default redactor (full redaction, no custom secret names).
    pub fn new() -> Self {
        Self::default()
    }

    /// Set custom field names to treat as secrets in addition to `_secret` suffixes.
    ///
    /// Matching is exact field-name equality at any nesting level. The same
    /// list also matches URL query-parameter names inside `_url` fields.
    /// Builder style: returns `self`.
    pub fn secret_names<I: IntoIterator<Item = S>, S: Into<String>>(mut self, names: I) -> Self {
        self.secret_names = names.into_iter().map(|s| s.into()).collect();
        self
    }

    /// Set exact legacy field names whose values should receive URL-aware
    /// redaction in addition to `_url`/`_URL` suffixed fields.
    ///
    /// A matching string is handled like an `_url` value. Arrays and nested
    /// collections recursively apply that treatment to every string leaf.
    /// Matching is exact and does not normalize case, whitespace, or spelling.
    #[must_use]
    pub fn url_names<I: IntoIterator<Item = S>, S: Into<String>>(mut self, names: I) -> Self {
        self.url_names = names.into_iter().map(Into::into).collect();
        self
    }

    /// Reveal exactly these structured field names at any depth. Other secrets
    /// remain masked; URL parameters and argv never use this exception.
    #[must_use]
    pub fn reveal_names<I: IntoIterator<Item = S>, S: Into<String>>(mut self, names: I) -> Self {
        self.reveal_names = names.into_iter().map(Into::into).collect();
        self
    }

    /// Set the redaction policy (default: full redaction).
    /// Builder style: returns `self`.
    pub fn policy(mut self, policy: RedactionPolicy) -> Self {
        self.policy = policy;
        self
    }

    /// Redact a JSON value copy using this redactor's policy and secret names.
    ///
    /// Clones `value` first; for a large payload you already own and can
    /// mutate, prefer [`Redactor::redact_in_place`] to avoid the copy.
    pub fn value(&self, value: &Value) -> Value {
        let mut v = value.clone();
        self.redact_in_place(&mut v);
        v
    }

    /// Redact secret components of a URL string using this redactor's settings.
    ///
    /// A query parameter is redacted iff its (form-decoded) name ends in
    /// `_secret`/`_SECRET` or matches an exact entry in `secret_names`. A
    /// fragment written in the same `k=v&k=v` shape — how an OAuth implicit-flow
    /// response carries its token — is redacted by that same rule; any other
    /// fragment passes through byte-for-byte. The userinfo password
    /// (`scheme://user:pass@host`) is always redacted as a structural rule.
    /// Only the secret spans are replaced with `***`; every other byte is
    /// preserved. A string that is not a single, whitespace-free,
    /// scheme-prefixed URL (including a URL embedded in surrounding prose) is
    /// returned unchanged.
    ///
    /// A URL is unscoped input: `RedactionPolicy::Off` returns it unchanged,
    /// every other policy redacts it in full.
    pub fn url(&self, url: &str) -> String {
        if !self.policy.redacts_unscoped_input() {
            return url.to_string();
        }
        let context = RedactionContext::from_redactor(self);
        redact_url_in_str(url, &context).unwrap_or_else(|| url.to_string())
    }

    /// Redact complete scheme-prefixed URLs embedded in prose.
    ///
    /// This is deliberately explicit and independent from structured field
    /// redaction. It recognizes URL spans only; it never scans surrounding text
    /// for secret-looking names or values.
    pub fn urls_in_text(&self, text: &str) -> String {
        if !self.policy.redacts_unscoped_input() {
            return text.to_string();
        }
        let context = RedactionContext::from_redactor(self);
        redact_urls_in_text_with_context(text, &context)
    }

    /// Redact `value` in place, using this redactor's policy and secret names.
    ///
    /// The zero-copy counterpart of [`Redactor::value`] — use it on a large
    /// payload you already own to avoid cloning.
    pub fn redact_in_place(&self, value: &mut Value) {
        let context = RedactionContext::from_redactor(self);
        apply_redaction_policy_with_context(value, self.policy, &context);
    }

    /// Redact secret *values* out of a command line, using this redactor's
    /// policy and secret names.
    ///
    /// A long flag whose name is secret by AFDATA naming (`--api-key-secret`,
    /// or an exact `secret_names` entry) has its value replaced by `***`, in
    /// both `--flag=value` and `--flag value` spellings. Everything else is
    /// preserved byte-for-byte.
    ///
    /// Free text is deliberately never scanned: a bare `api_key_secret=sk-live`
    /// positional, or a secret-looking token after a non-secret flag, is left
    /// alone. AFDATA decides sensitivity from the *field name*, and argv is no
    /// exception — rename the flag rather than pattern-matching values. A flag
    /// with no value (end of argv, or followed by another flag) is likewise
    /// left inspectable.
    ///
    /// Only long (`--`) flags are recognized, matching the convention's
    /// long-flags-only rule.
    ///
    /// A command line is unscoped input: `RedactionPolicy::Off` returns `args`
    /// unchanged, every other policy redacts in full.
    pub fn argv<S: AsRef<str>>(&self, args: &[S]) -> Vec<String> {
        if !self.policy.redacts_unscoped_input() {
            return args.iter().map(|arg| arg.as_ref().to_string()).collect();
        }
        let context = RedactionContext::from_redactor(self);
        let mut out = Vec::with_capacity(args.len());
        let mut redact_next = false;
        for arg in args {
            let arg = arg.as_ref();
            if redact_next {
                redact_next = false;
                if !arg.starts_with('-') {
                    out.push(REDACTED_MARKER.to_string());
                    continue;
                }
            }
            if let Some(rest) = arg.strip_prefix("--") {
                if let Some((name, _)) = rest.split_once('=') {
                    if is_secret_flag_name(name, &context) {
                        out.push(format!("--{name}={REDACTED_MARKER}"));
                        continue;
                    }
                } else if is_secret_flag_name(rest, &context) {
                    redact_next = true;
                }
            }
            out.push(arg.to_string());
        }
        out
    }

    /// True when `name` would be treated as a secret field name by this
    /// redactor: an exact `_secret`/`_SECRET` suffix, or an exact match
    /// against a configured `secret_names` entry.
    ///
    /// Exposed for callers that must gate on a single *targeted* field name
    /// (for example a CLI dot-path leaf) rather than redact a whole value —
    /// [`Redactor::value`] only rewrites fields it finds while walking an
    /// object, so a bare scalar pulled out from under its field name needs
    /// this explicit check instead.
    pub fn is_secret_name(&self, name: &str) -> bool {
        RedactionContext::from_redactor(self).is_secret_key(name)
    }

    /// True when `name` receives URL-aware structured-field redaction: an
    /// exact `_url`/`_URL` suffix or an exact configured `url_names` entry.
    ///
    /// The counterpart to [`Redactor::is_secret_name`], for the same targeted
    /// single-field case. Kept even with no in-tree caller: a consumer that
    /// gates on one name has to ask both questions, and an API that answers
    /// only half of a symmetric pair sends the caller back to reimplementing
    /// the suffix rule — which is exactly how the rule drifts.
    pub fn is_url_name(&self, name: &str) -> bool {
        RedactionContext::from_redactor(self).is_url_key(name)
    }
}

impl From<RedactionPolicy> for Redactor {
    fn from(policy: RedactionPolicy) -> Self {
        Self {
            policy,
            secret_names: Vec::new(),
            url_names: Vec::new(),
            reveal_names: Vec::new(),
        }
    }
}

/// Output options combining redaction and rendering style.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct OutputOptions {
    /// Redactor applied before rendering.
    pub redaction: Redactor,
    /// Rendering style for plain output only.
    pub style: PlainStyle,
}

impl From<RedactionPolicy> for OutputOptions {
    fn from(policy: RedactionPolicy) -> Self {
        Self {
            redaction: Redactor::from(policy),
            style: PlainStyle::default(),
        }
    }
}

// ═══════════════════════════════════════════
// Public API: Redaction & Utility
// ═══════════════════════════════════════════

/// Return a JSON value copy with default `_secret` redaction applied.
pub fn redacted_value(value: &Value) -> Value {
    Redactor::new().value(value)
}

/// Redact secret values out of a command line, using default options.
///
/// Returns `args` with the value of every `_secret`-suffixed long flag replaced
/// by `***`, covering both `--flag=value` and `--flag value`. Use
/// [`Redactor::argv`] for custom `secret_names` or a non-default policy.
///
/// Intended for CLIs that record their own invocation — startup diagnostics,
/// audit trails, crash reports — where writing argv verbatim would put a
/// credential in the log.
pub fn redact_argv<S: AsRef<str>>(args: &[S]) -> Vec<String> {
    Redactor::new().argv(args)
}

/// Redact secret components of a single URL string, using default options.
///
/// Returns `url` with its userinfo password and any `_secret`-suffixed query
/// parameter values replaced by `***`.
pub fn redact_url_secrets(url: &str) -> String {
    Redactor::new().url(url)
}

/// Redact complete scheme-prefixed URLs embedded in prose, using default
/// secret-name rules.
///
/// This helper is opt-in. Ordinary structured redaction never scans prose.
pub fn redact_urls_in_text(text: &str) -> String {
    Redactor::new().urls_in_text(text)
}

// ═══════════════════════════════════════════
// Secret Redaction
// ═══════════════════════════════════════════

/// The scalar every redacted span, value, and subtree is replaced with. Also
/// the signal plain rendering reads to tell a hidden field from a live one.
pub(crate) const REDACTED_MARKER: &str = "***";

#[derive(Default)]
pub(crate) struct RedactionContext {
    secret_names: HashSet<String>,
    url_names: HashSet<String>,
    reveal_names: HashSet<String>,
}

impl RedactionContext {
    fn from_redactor(redactor: &Redactor) -> Self {
        let secret_names = redactor.secret_names.iter().cloned().collect();
        let url_names = redactor.url_names.iter().cloned().collect();
        Self {
            secret_names,
            url_names,
            reveal_names: redactor.reveal_names.iter().cloned().collect(),
        }
    }

    fn is_secret_key(&self, key: &str) -> bool {
        key_has_secret_suffix(key) || self.secret_names.contains(key)
    }

    fn is_url_key(&self, key: &str) -> bool {
        key_has_url_suffix(key) || self.url_names.contains(key)
    }
}

fn key_has_secret_suffix(key: &str) -> bool {
    key.ends_with("_secret") || key.ends_with("_SECRET")
}

fn key_has_url_suffix(key: &str) -> bool {
    key.ends_with("_url") || key.ends_with("_URL")
}

/// Whether a long flag's name is secret, normalizing the kebab-case flag
/// spelling to the snake_case field spelling the convention is defined in.
///
/// Ungated: core argv redaction relies on it.
pub(crate) fn is_secret_flag_name(flag_name: &str, context: &RedactionContext) -> bool {
    let normalized = flag_name.replace('-', "_");
    context.is_secret_key(&normalized) || context.is_secret_key(flag_name)
}

const MAX_DEPTH: usize = 256;
const MAX_DEPTH_MARKER: &str = "<afdata:max-depth>";

fn redact_secrets_with_context(value: &mut Value, context: &RedactionContext) {
    redact_secrets_with_context_depth(value, context, 0);
}

fn redact_secrets_with_context_depth(value: &mut Value, context: &RedactionContext, depth: usize) {
    if depth >= MAX_DEPTH {
        *value = Value::String(MAX_DEPTH_MARKER.into());
        return;
    }
    match value {
        Value::Object(map) => {
            let keys: Vec<String> = map.keys().cloned().collect();
            for key in keys {
                if context.is_secret_key(&key) && !context.reveal_names.contains(&key) {
                    // A null secret is an *absent* secret. Masking it would
                    // manufacture the appearance of a configured credential:
                    // readers cannot tell `"***"`-because-set from
                    // `"***"`-because-null, so a tool showing its own config
                    // would report every unset secret as configured. Redaction
                    // hides a value that exists; it does not invent one.
                    if !map.get(&key).is_some_and(Value::is_null) {
                        map.insert(key, Value::String(REDACTED_MARKER.into()));
                    }
                } else if context.is_url_key(&key) {
                    if let Some(v) = map.get_mut(&key) {
                        redact_url_value_depth(v, context, depth + 1);
                    }
                } else if let Some(v) = map.get_mut(&key) {
                    redact_secrets_with_context_depth(v, context, depth + 1);
                }
            }
        }
        Value::Array(arr) => {
            for v in arr {
                redact_secrets_with_context_depth(v, context, depth + 1);
            }
        }
        _ => {}
    }
}

fn redact_url_value_depth(value: &mut Value, context: &RedactionContext, depth: usize) {
    if depth >= MAX_DEPTH {
        *value = Value::String(MAX_DEPTH_MARKER.into());
        return;
    }
    match value {
        Value::String(text) => {
            *text = redact_url_field_value(text, context);
        }
        Value::Array(values) => {
            for value in values {
                redact_url_value_depth(value, context, depth + 1);
            }
        }
        Value::Object(values) => {
            let keys = values.keys().cloned().collect::<Vec<_>>();
            for key in keys {
                if context.is_secret_key(&key) && !context.reveal_names.contains(&key) {
                    if !values.get(&key).is_some_and(Value::is_null) {
                        values.insert(key, Value::String(REDACTED_MARKER.into()));
                    }
                } else if let Some(value) = values.get_mut(&key) {
                    redact_url_value_depth(value, context, depth + 1);
                }
            }
        }
        _ => {}
    }
}

/// Redact secret components of a single URL string, returning `Some(redacted)`
/// when `s` is a processable URL, or `None` when it is not (so callers can keep
/// the original). Only secret spans change; all other bytes are preserved.
fn apply_redaction_policy_with_context(
    value: &mut Value,
    redaction_policy: RedactionPolicy,
    context: &RedactionContext,
) {
    match redaction_policy {
        RedactionPolicy::All => redact_secrets_with_context(value, context),
        RedactionPolicy::TraceOnly => {
            if let Value::Object(map) = value
                && let Some(trace) = map.get_mut("trace")
            {
                redact_secrets_with_context(trace, context);
            }
        }
        RedactionPolicy::Off => {}
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    // ── URL fragments carry secrets too ──────────

    #[test]
    fn url_fragment_params_redacted_like_query_params() {
        assert_eq!(
            redact_url_secrets("https://h/p?token_secret=QUERYLEAK#token_secret=FRAGLEAK"),
            "https://h/p?token_secret=***#token_secret=***"
        );
    }

    #[test]
    fn url_fragment_params_redacted_without_a_query() {
        // The OAuth implicit-flow shape: the credential exists only after '#'.
        assert_eq!(
            redact_url_secrets("https://h/cb#access_token_secret=abc&state=xyz"),
            "https://h/cb#access_token_secret=***&state=xyz"
        );
    }

    #[test]
    fn url_fragment_without_params_is_preserved() {
        for url in [
            "https://h/p?a=1#section",
            "https://h/p#",
            "https://h/p#a/b?c",
        ] {
            assert_eq!(redact_url_secrets(url), url);
        }
    }

    #[test]
    fn url_fragment_honors_secret_names() {
        let redactor = Redactor::new().secret_names(vec!["token".to_string()]);
        assert_eq!(
            redactor.url("https://h/cb#token=abc&page=2"),
            "https://h/cb#token=***&page=2"
        );
    }

    #[test]
    fn url_names_are_exact_and_recurse_through_collections() {
        let redactor = Redactor::new().secret_names(["token"]).url_names([
            "url",
            "relays",
            "synapse_selector",
        ]);
        let input = json!({
            "url": "https://u:pw@h/?token=one",
            "URL": "https://u:visible@h/?token=two",
            "relays": [
                "wss://relay.example/?token=three",
                {"nested": "https://u:pw@nested.example/"}
            ],
            "synapse_selector": "user:pass@host/db",
            "other": "https://u:visible@h/?token=four"
        });

        assert_eq!(
            redactor.value(&input),
            json!({
                "url": "https://u:***@h/?token=***",
                "URL": "https://u:visible@h/?token=two",
                "relays": [
                    "wss://relay.example/?token=***",
                    {"nested": "https://u:***@nested.example/"}
                ],
                "synapse_selector": "***",
                "other": "https://u:visible@h/?token=four"
            })
        );
        assert!(redactor.is_url_name("url"));
        assert!(!redactor.is_url_name("URL"));
    }

    #[test]
    fn prose_url_redaction_is_explicit_and_preserves_punctuation() {
        let message = "connect (https://u:pw@h/?token_secret=one), then wss://h/?token_secret=two. bare user:pw@h stays";
        assert_eq!(
            redact_urls_in_text(message),
            "connect (https://u:***@h/?token_secret=***), then wss://h/?token_secret=***. bare user:pw@h stays"
        );
        assert_eq!(
            redacted_value(&json!({"message": message})),
            json!({"message": message})
        );
    }

    #[test]
    fn prose_url_redaction_honors_policy_and_secret_names() {
        let text = "see https://h/?token=abc";
        assert_eq!(
            Redactor::new().secret_names(["token"]).urls_in_text(text),
            "see https://h/?token=***"
        );
        assert_eq!(
            Redactor::new()
                .policy(RedactionPolicy::Off)
                .secret_names(["token"])
                .urls_in_text(text),
            text
        );
    }

    // ── One policy meaning across value, argv, url ──────────

    fn scoped_value() -> Value {
        json!({
            "result": {"api_key_secret": "sk-result"},
            "trace": {"api_key_secret": "sk-trace"}
        })
    }

    fn argv() -> Vec<String> {
        vec!["tool".to_string(), "--api-key-secret=sk-live".to_string()]
    }

    #[test]
    fn all_policy_redacts_every_path() {
        let redactor = Redactor::new().policy(RedactionPolicy::All);
        assert_eq!(
            redactor.value(&scoped_value()),
            json!({
                "result": {"api_key_secret": "***"},
                "trace": {"api_key_secret": "***"}
            })
        );
        assert_eq!(redactor.argv(&argv()), vec!["tool", "--api-key-secret=***"]);
        assert_eq!(
            redactor.url("https://u:pw@h/cb?token_secret=abc"),
            "https://u:***@h/cb?token_secret=***"
        );
    }

    #[test]
    fn trace_only_scopes_a_value_but_redacts_argv_and_url_in_full() {
        let redactor = Redactor::new().policy(RedactionPolicy::TraceOnly);
        // Scoped input: only the `trace` half is scrubbed.
        assert_eq!(
            redactor.value(&scoped_value()),
            json!({
                "result": {"api_key_secret": "sk-result"},
                "trace": {"api_key_secret": "***"}
            })
        );
        // Unscoped input: a command line and a bare URL have no non-`trace`
        // half to leave alone, so they are redacted like `All`.
        assert_eq!(redactor.argv(&argv()), vec!["tool", "--api-key-secret=***"]);
        assert_eq!(
            redactor.url("https://u:pw@h/cb?token_secret=abc"),
            "https://u:***@h/cb?token_secret=***"
        );
    }

    #[test]
    fn off_policy_disables_every_path() {
        let redactor = Redactor::new().policy(RedactionPolicy::Off);
        assert_eq!(redactor.value(&scoped_value()), scoped_value());
        assert_eq!(redactor.argv(&argv()), argv());
        assert_eq!(
            redactor.url("https://u:pw@h/cb?token_secret=abc"),
            "https://u:pw@h/cb?token_secret=abc"
        );
    }
}