scour_secrets/processor/

yaml_proc.rs

//! YAML structured processor.
//!
//! The CLI uses [`process_to_edits`](YamlProcessor::process_to_edits): it drives
//! a span-aware event parser (`saphyr-parser`) and replaces each matched scalar
//! at its exact source span, preserving comments, anchors, key order, and quote
//! style byte-for-byte (escaped/quoted scalars are hit as written, so they never
//! leak). `process` is the re-serializing fallback (which normalizes some
//! whitespace).
//!
//! Key paths use the same dot-separated convention as the JSON processor.

use crate::error::{Result, SanitizeError};
use crate::processor::limits::{DEFAULT_DEPTH, YAML_INPUT_SIZE, YAML_NODE_COUNT};
use crate::processor::{
    build_path, edit_token, walk_tree, FileTypeProfile, Processor, Replacement, TreeNode,
};
use crate::store::MappingStore;
use saphyr_parser::{Event, Parser, ScalarStyle};
use serde_yaml_ng::Value;

/// Byte length of a leading quoted scalar (`"..."` or `'...'`) within `src` —
/// the index just past its closing quote — or `None` for non-quoted styles or
/// when no closing quote is found (then the caller keeps saphyr's span).
///
/// saphyr's reported span for a quoted scalar can extend to end-of-line
/// (including a trailing inline comment), so we re-scan to the actual closing
/// quote to avoid clobbering comments and trailing formatting.
fn quoted_scalar_end(src: &[u8], style: ScalarStyle) -> Option<usize> {
    let quote = match style {
        ScalarStyle::DoubleQuoted => b'"',
        ScalarStyle::SingleQuoted => b'\'',
        _ => return None,
    };
    if src.first() != Some(&quote) {
        return None;
    }
    let mut i = 1;
    while i < src.len() {
        let b = src[i];
        // Double-quoted YAML uses backslash escapes; skip the escaped byte.
        if style == ScalarStyle::DoubleQuoted && b == b'\\' {
            i += 2;
            continue;
        }
        if b == quote {
            // Single-quoted YAML escapes a quote by doubling it (`''`).
            if style == ScalarStyle::SingleQuoted && src.get(i + 1) == Some(&quote) {
                i += 2;
                continue;
            }
            return Some(i + 1);
        }
        i += 1;
    }
    None
}

/// A container frame for the YAML event-stream walk (span-based editing).
enum YamlFrame {
    /// Inside a mapping. `path`/`key` locate the mapping itself; `current_key`
    /// is the key whose value is being read; `expecting_key` alternates.
    Mapping {
        path: String,
        expecting_key: bool,
        current_key: Option<String>,
    },
    /// Inside a sequence. Items are path-transparent (keep the parent key/path).
    Sequence { path: String, key: String },
}

/// The (dot-path, bare-key) of the value about to be read, given the frame stack.
fn yaml_value_position(frames: &[YamlFrame]) -> (String, String) {
    match frames.last() {
        Some(YamlFrame::Mapping {
            path, current_key, ..
        }) => {
            let key = current_key.clone().unwrap_or_default();
            (build_path(path, &key), key)
        }
        Some(YamlFrame::Sequence { path, key }) => (path.clone(), key.clone()),
        None => (String::new(), String::new()),
    }
}

/// After a value (scalar, or a container that just ended) is consumed, a parent
/// mapping should expect the next key.
fn yaml_note_value_consumed(frames: &mut [YamlFrame]) {
    if let Some(YamlFrame::Mapping { expecting_key, .. }) = frames.last_mut() {
        *expecting_key = true;
    }
}

/// Replacement text for a matched YAML scalar, given its source bytes (`span_src`)
/// and scalar style.
///
/// Flow scalars (plain, single/double-quoted) become a double-quoted token.
/// **Block** scalars (`|` literal, `>` folded) must stay block-valid: their value
/// spans the indented content lines, so replacing it with an inline `"token"`
/// would collapse the block and absorb the following keys. Instead, emit a single
/// indented line — `<indent>token` plus the trailing newline the span consumed —
/// keeping the block structure intact.
fn yaml_scalar_replacement(token: &str, style: ScalarStyle, span_src: &[u8]) -> String {
    match style {
        ScalarStyle::Literal | ScalarStyle::Folded => {
            let indent: String = span_src
                .iter()
                .take_while(|&&b| b == b' ' || b == b'\t')
                .map(|&b| b as char)
                .collect();
            let trailing_nl = if span_src.last() == Some(&b'\n') {
                "\n"
            } else {
                ""
            };
            format!("{indent}{token}{trailing_nl}")
        }
        _ => format!("\"{token}\""),
    }
}

/// Structured processor for YAML files.
pub struct YamlProcessor;

impl Processor for YamlProcessor {
    fn name(&self) -> &'static str {
        "yaml"
    }

    fn can_handle(&self, content: &[u8], profile: &FileTypeProfile) -> bool {
        if profile.processor == "yaml" {
            return true;
        }
        // Heuristic: starts with `---` or a YAML-ish key: value.
        let text = String::from_utf8_lossy(content);
        let trimmed = text.trim_start();
        trimmed.starts_with("---")
            || trimmed.starts_with("- ")
            || trimmed.starts_with('{')
            || trimmed.contains(": ")
    }

    fn process(
        &self,
        content: &[u8],
        profile: &FileTypeProfile,
        store: &MappingStore,
    ) -> Result<Vec<u8>> {
        // Guard against alias bombs: reject inputs above YAML_INPUT_SIZE.
        let text = crate::processor::check_size_and_decode(content, "YAML", YAML_INPUT_SIZE)?;

        let mut value: Value =
            serde_yaml_ng::from_str(text).map_err(|e| SanitizeError::ParseError {
                format: "YAML".into(),
                message: format!("YAML parse error: {}", e),
            })?;

        // F-06 fix: count total nodes in the deserialized tree to detect
        // alias bombs. After expansion, aliased subtrees become
        // independent copies in memory, so the node count reflects the
        // true memory footprint.
        let node_count = count_yaml_nodes(&value);
        if node_count > YAML_NODE_COUNT {
            return Err(SanitizeError::InputTooLarge {
                size: node_count,
                limit: YAML_NODE_COUNT,
            });
        }

        walk_yaml(&mut value, "", profile, store, 0)?;

        let output = serde_yaml_ng::to_string(&value).map_err(|e| {
            SanitizeError::IoError(std::io::Error::other(format!("YAML serialize error: {e}")))
        })?;

        Ok(output.into_bytes())
    }

    /// Span-based redaction: drive `saphyr-parser` (which yields each event with
    /// a byte `Span`) over the document, tracking the mapping/sequence path, and
    /// emit an edit replacing each matched value scalar's exact source span with
    /// a quoted token. Comments, anchors, key order, block/flow style, and the
    /// exact escaping of unrelated content are preserved; the value is hit in the
    /// source as written, so quoted/escaped scalars never leak.
    fn process_to_edits(
        &self,
        content: &[u8],
        profile: &FileTypeProfile,
        store: &MappingStore,
    ) -> Result<Option<Vec<Replacement>>> {
        let text = crate::processor::check_size_and_decode(content, "YAML", YAML_INPUT_SIZE)?;
        let mut edits = Vec::new();
        let mut frames: Vec<YamlFrame> = Vec::new();

        // saphyr reports span markers as CHARACTER counts, not byte offsets
        // (its `Marker::index()` "in bytes" doc is wrong — every char advances
        // the index by 1). For multi-byte UTF-8 we must translate char index →
        // byte offset before slicing `content`, or a scalar that follows
        // multi-byte content is sliced at the wrong position and the output is
        // corrupted. ASCII needs no map (char index == byte offset).
        let char_to_byte: Option<Vec<usize>> = if text.is_ascii() {
            None
        } else {
            Some(
                text.char_indices()
                    .map(|(b, _)| b)
                    .chain(std::iter::once(text.len()))
                    .collect(),
            )
        };
        let to_byte = |char_idx: usize| -> usize {
            char_to_byte
                .as_ref()
                .map_or(char_idx, |m| m.get(char_idx).copied().unwrap_or(text.len()))
        };

        for event in Parser::new_from_str(text) {
            let (event, span) = event.map_err(|e| SanitizeError::ParseError {
                format: "YAML".into(),
                message: format!("YAML parse error: {e}"),
            })?;
            match event {
                Event::Scalar(value, style, _aid, _tag) => {
                    let is_key = matches!(
                        frames.last(),
                        Some(YamlFrame::Mapping {
                            expecting_key: true,
                            ..
                        })
                    );
                    if is_key {
                        if let Some(YamlFrame::Mapping {
                            expecting_key,
                            current_key,
                            ..
                        }) = frames.last_mut()
                        {
                            *current_key = Some(value.into_owned());
                            *expecting_key = false;
                        }
                    } else {
                        let (path, key) = yaml_value_position(&frames);
                        if let Some(token) = edit_token(&key, &path, &value, profile, store)? {
                            let start = to_byte(span.start.index());
                            let mut end = to_byte(span.end.index());
                            // saphyr's span for a *quoted* scalar can run to the
                            // end of the line, swallowing trailing whitespace and
                            // an inline `# comment`. Clamp to the real closing
                            // quote so comments/formatting survive byte-for-byte.
                            if let Some(real) = quoted_scalar_end(&content[start..end], style) {
                                end = start + real;
                            }
                            let repl = yaml_scalar_replacement(&token, style, &content[start..end]);
                            edits.push(Replacement {
                                start,
                                end,
                                value: repl,
                            });
                        }
                        yaml_note_value_consumed(&mut frames);
                    }
                }
                Event::MappingStart(..) => {
                    let (path, _key) = yaml_value_position(&frames);
                    frames.push(YamlFrame::Mapping {
                        path,
                        expecting_key: true,
                        current_key: None,
                    });
                }
                Event::SequenceStart(..) => {
                    let (path, key) = yaml_value_position(&frames);
                    frames.push(YamlFrame::Sequence { path, key });
                }
                Event::MappingEnd | Event::SequenceEnd => {
                    frames.pop();
                    yaml_note_value_consumed(&mut frames);
                }
                Event::Alias(_) => {
                    // Alias references an anchored value (redacted at its
                    // definition); the alias node itself isn't a literal.
                    yaml_note_value_consumed(&mut frames);
                }
                _ => {}
            }
        }
        Ok(Some(edits))
    }
}

/// Count the total number of nodes in a YAML value tree (F-06 fix).
/// Used to detect alias bombs that produce a small source document
/// but expand to millions of nodes after alias resolution.
fn count_yaml_nodes(value: &Value) -> usize {
    count_yaml_nodes_inner(value, 0)
}

/// Inner recursive counter with depth guard to prevent stack overflow
/// on deeply nested YAML before `walk_yaml`'s depth check is reached.
fn count_yaml_nodes_inner(value: &Value, depth: usize) -> usize {
    if depth > DEFAULT_DEPTH {
        return 1; // Stop counting deeper; walk_yaml will catch depth violations
    }
    match value {
        Value::Mapping(map) => {
            1 + map
                .iter()
                .map(|(k, v)| {
                    count_yaml_nodes_inner(k, depth + 1) + count_yaml_nodes_inner(v, depth + 1)
                })
                .sum::<usize>()
        }
        Value::Sequence(seq) => {
            1 + seq
                .iter()
                .map(|v| count_yaml_nodes_inner(v, depth + 1))
                .sum::<usize>()
        }
        Value::Tagged(tagged) => 1 + count_yaml_nodes_inner(&tagged.value, depth + 1),
        _ => 1, // Null, Bool, Number, String
    }
}

impl TreeNode for Value {
    fn for_each_map_entry<F>(&mut self, mut f: F) -> Result<()>
    where
        F: FnMut(&str, &mut Self) -> Result<()>,
    {
        if let Self::Mapping(map) = self {
            let keys: Vec<Self> = map.keys().cloned().collect();
            for key in keys {
                let key_str = yaml_key_to_string(&key);
                if let Some(v) = map.get_mut(&key) {
                    f(&key_str, v)?;
                }
            }
        }
        Ok(())
    }

    fn for_each_seq_item<F>(&mut self, mut f: F) -> Result<()>
    where
        F: FnMut(&mut Self) -> Result<()>,
    {
        if let Self::Sequence(seq) = self {
            for item in seq.iter_mut() {
                f(item)?;
            }
        }
        Ok(())
    }

    fn as_str_mut(&mut self) -> Option<&mut String> {
        if let Self::String(s) = self {
            Some(s)
        } else {
            None
        }
    }

    fn is_scalar(&self) -> bool {
        matches!(self, Self::Number(_) | Self::Bool(_))
    }

    fn scalar_to_string(&self) -> String {
        yaml_scalar_to_string(self)
    }

    fn set_string(&mut self, s: String) {
        *self = Self::String(s);
    }
}

/// Recursively walk a YAML value tree, replacing matched field values.
fn walk_yaml(
    value: &mut Value,
    prefix: &str,
    profile: &FileTypeProfile,
    store: &MappingStore,
    depth: usize,
) -> Result<()> {
    walk_tree(value, prefix, profile, store, depth, "YAML")
}

fn yaml_key_to_string(key: &Value) -> String {
    match key {
        Value::String(s) => s.clone(),
        Value::Number(n) => n.to_string(),
        Value::Bool(b) => b.to_string(),
        _ => format!("{:?}", key),
    }
}

fn yaml_scalar_to_string(v: &Value) -> String {
    match v {
        Value::String(s) => s.clone(),
        Value::Number(n) => n.to_string(),
        Value::Bool(b) => b.to_string(),
        _ => String::new(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::category::Category;
    use crate::generator::HmacGenerator;
    use crate::processor::profile::FieldRule;
    use std::sync::Arc;

    fn make_store() -> MappingStore {
        let gen = Arc::new(HmacGenerator::new([42u8; 32]));
        MappingStore::new(gen, None)
    }

    #[test]
    fn basic_yaml_replacement() {
        let store = make_store();
        let proc = YamlProcessor;

        let content = b"database:\n  host: db.corp.com\n  password: s3cret\nport: 5432\n";
        let profile = FileTypeProfile::new(
            "yaml",
            vec![
                FieldRule::new("database.password").with_category(Category::Custom("pw".into())),
                FieldRule::new("database.host").with_category(Category::Hostname),
            ],
        );

        let result = proc.process(content, &profile, &store).unwrap();
        let out = String::from_utf8(result).unwrap();

        assert!(!out.contains("s3cret"));
        assert!(!out.contains("db.corp.com"));
        // port should be preserved
        assert!(out.contains("5432"));
    }

    #[test]
    fn can_handle_by_profile_name() {
        let proc = YamlProcessor;
        let profile = FileTypeProfile::new("yaml", vec![]).with_extension(".yaml");
        assert!(proc.can_handle(b"anything", &profile));
    }

    #[test]
    fn can_handle_detects_document_marker() {
        let proc = YamlProcessor;
        let profile = FileTypeProfile::new("json", vec![]).with_extension(".json");
        assert!(proc.can_handle(b"---\nkey: value\n", &profile));
    }

    #[test]
    fn can_handle_detects_key_value_heuristic() {
        let proc = YamlProcessor;
        let profile = FileTypeProfile::new("other", vec![]).with_extension(".conf");
        assert!(proc.can_handle(b"host: localhost\nport: 5432\n", &profile));
    }

    #[test]
    fn can_handle_detects_sequence_heuristic() {
        let proc = YamlProcessor;
        let profile = FileTypeProfile::new("other", vec![]).with_extension(".txt");
        assert!(proc.can_handle(b"- item1\n- item2\n", &profile));
    }

    #[test]
    fn can_handle_rejects_plaintext() {
        let proc = YamlProcessor;
        let profile = FileTypeProfile::new("json", vec![]).with_extension(".json");
        assert!(!proc.can_handle(b"just plain text with no yaml markers", &profile));
    }

    #[test]
    fn non_string_scalars_not_targeted_pass_through() {
        let store = make_store();
        let proc = YamlProcessor;
        // Only target the 'secret' field; booleans and numbers are untouched.
        let content = b"enabled: true\ncount: 42\nsecret: hunter2\n";
        let profile = FileTypeProfile::new(
            "yaml",
            vec![FieldRule::new("secret").with_category(Category::Custom("pw".into()))],
        );
        let result = proc.process(content, &profile, &store).unwrap();
        let out = String::from_utf8(result).unwrap();
        assert!(!out.contains("hunter2"), "secret must be replaced");
        assert!(out.contains("42"), "integer must be preserved");
    }

    #[test]
    fn deeply_nested_yaml_replaced() {
        let store = make_store();
        let proc = YamlProcessor;
        let content = b"a:\n  b:\n    c:\n      secret: hunter2\n";
        let profile = FileTypeProfile::new(
            "yaml",
            vec![FieldRule::new("a.b.c.secret").with_category(Category::Custom("pw".into()))],
        );
        let result = proc.process(content, &profile, &store).unwrap();
        let out = String::from_utf8(result).unwrap();
        assert!(!out.contains("hunter2"));
        // Non-secret structure preserved: only the value changed, the nested
        // keys remain.
        assert!(out.contains("secret:"));
        let parsed: serde_yaml_ng::Value = serde_yaml_ng::from_str(&out).unwrap();
        assert!(parsed["a"]["b"]["c"]["secret"].as_str().is_some());
    }

    #[test]
    fn invalid_utf8_returns_parse_error() {
        let store = make_store();
        let proc = YamlProcessor;
        let bad = b"\xff\xfe invalid";
        let profile = FileTypeProfile::new("yaml", vec![]);
        let err = proc.process(bad, &profile, &store).unwrap_err();
        assert!(matches!(
            err,
            crate::error::SanitizeError::ParseError { .. }
        ));
    }

    #[test]
    fn invalid_yaml_returns_parse_error() {
        let store = make_store();
        let proc = YamlProcessor;
        let bad = b"key: [unclosed";
        let profile = FileTypeProfile::new("yaml", vec![]);
        let err = proc.process(bad, &profile, &store).unwrap_err();
        assert!(matches!(
            err,
            crate::error::SanitizeError::ParseError { .. }
        ));
    }

    #[test]
    fn yaml_sequence_traversal() {
        let store = make_store();
        let proc = YamlProcessor;

        let content = b"users:\n  - email: a@b.com\n  - email: c@d.com\n";
        let profile = FileTypeProfile::new(
            "yaml",
            vec![FieldRule::new("users.email").with_category(Category::Email)],
        );

        let result = proc.process(content, &profile, &store).unwrap();
        let out = String::from_utf8(result).unwrap();

        assert!(!out.contains("a@b.com"));
        assert!(!out.contains("c@d.com"));
        // Non-secret structure preserved: the key and both sequence items.
        assert!(out.contains("users:"));
        let parsed: serde_yaml_ng::Value = serde_yaml_ng::from_str(&out).unwrap();
        assert_eq!(parsed["users"].as_sequence().unwrap().len(), 2);
    }

    // ── process_to_edits (span-based, format-preserving) ─────────────────────

    /// Edit-mode alone must redact plain, double-quoted, single-quoted, and
    /// escaped scalars while preserving comments and unrelated values.
    #[test]
    fn edits_redact_all_scalar_styles_and_preserve_comments() {
        let store = make_store();
        let proc = YamlProcessor;
        let content = b"# top\ndb:\n  a: plain-SEC1   # inline\n  b: \"dq-SEC2\"\n  c: 'sq-SEC3'\n  d: \"x\\\"y-SEC4\"\n  host: keep.local\n";
        let profile = FileTypeProfile::new(
            "yaml",
            vec![
                FieldRule::new("db.a").with_category(Category::Custom("k".into())),
                FieldRule::new("db.b").with_category(Category::Custom("k".into())),
                FieldRule::new("db.c").with_category(Category::Custom("k".into())),
                FieldRule::new("db.d").with_category(Category::Custom("k".into())),
            ],
        );
        let edits = proc
            .process_to_edits(content, &profile, &store)
            .unwrap()
            .unwrap();
        let out = crate::processor::apply_edits(content, edits);
        let text = String::from_utf8(out).unwrap();
        for leak in ["SEC1", "SEC2", "SEC3", "SEC4"] {
            assert!(!text.contains(leak), "leaked {leak}: {text}");
        }
        assert!(text.contains("# top"), "top comment dropped: {text}");
        assert!(text.contains("# inline"), "inline comment dropped: {text}");
        assert!(
            text.contains("host: keep.local"),
            "non-secret changed: {text}"
        );
        // Output remains valid YAML.
        assert!(
            serde_yaml_ng::from_str::<serde_yaml_ng::Value>(&text).is_ok(),
            "invalid YAML: {text}"
        );
    }

    /// Regression: block scalars (`|`, `>`) must be redacted while staying
    /// block-valid — the inline-`"token"` replacement collapsed the block and
    /// absorbed following keys.
    #[test]
    fn edits_keep_block_scalars_valid() {
        let store = make_store();
        let proc = YamlProcessor;
        let content = b"lit: |\n  line1-SEC1\n  line2-SEC2\nfold: >\n  folded-SEC3\nnext: keep\n";
        let profile = FileTypeProfile::new(
            "yaml",
            vec![
                FieldRule::new("lit").with_category(Category::Custom("k".into())),
                FieldRule::new("fold").with_category(Category::Custom("k".into())),
            ],
        );
        let edits = proc
            .process_to_edits(content, &profile, &store)
            .unwrap()
            .unwrap();
        let out = crate::processor::apply_edits(content, edits);
        let text = String::from_utf8(out).unwrap();
        for leak in ["SEC1", "SEC2", "SEC3"] {
            assert!(!text.contains(leak), "leaked {leak}: {text}");
        }
        // Following keys are NOT absorbed into the block; output is valid YAML.
        assert!(text.contains("fold:"), "fold key absorbed: {text}");
        assert!(text.contains("next: keep"), "next key absorbed: {text}");
        let parsed: serde_yaml_ng::Value = serde_yaml_ng::from_str(&text).unwrap();
        assert_eq!(parsed["next"].as_str(), Some("keep"));
    }
}