rust_sanitize/processor/
yaml_proc.rs

1//! YAML structured processor.
2//!
3//! Parses YAML input, walks the value tree, replaces matched field
4//! values, and serializes back. Structure is preserved but minor
5//! formatting differences are possible (serde_yaml normalizes some
6//! whitespace).
7//!
8//! Key paths use the same dot-separated convention as the JSON processor.
9
10use crate::error::{Result, SanitizeError};
11use crate::processor::limits::{DEFAULT_DEPTH, YAML_INPUT_SIZE, YAML_NODE_COUNT};
12use crate::processor::{walk_tree, FileTypeProfile, Processor, TreeNode};
13use crate::store::MappingStore;
14use serde_yaml_ng::Value;
15
16/// Structured processor for YAML files.
17pub struct YamlProcessor;
18
19impl Processor for YamlProcessor {
20    fn name(&self) -> &'static str {
21        "yaml"
22    }
23
24    fn can_handle(&self, content: &[u8], profile: &FileTypeProfile) -> bool {
25        if profile.processor == "yaml" {
26            return true;
27        }
28        // Heuristic: starts with `---` or a YAML-ish key: value.
29        let text = String::from_utf8_lossy(content);
30        let trimmed = text.trim_start();
31        trimmed.starts_with("---")
32            || trimmed.starts_with("- ")
33            || trimmed.starts_with('{')
34            || trimmed.contains(": ")
35    }
36
37    fn process(
38        &self,
39        content: &[u8],
40        profile: &FileTypeProfile,
41        store: &MappingStore,
42    ) -> Result<Vec<u8>> {
43        // Guard against alias bombs: reject inputs above YAML_INPUT_SIZE.
44        let text = crate::processor::check_size_and_decode(content, "YAML", YAML_INPUT_SIZE)?;
45
46        let mut value: Value =
47            serde_yaml_ng::from_str(text).map_err(|e| SanitizeError::ParseError {
48                format: "YAML".into(),
49                message: format!("YAML parse error: {}", e),
50            })?;
51
52        // F-06 fix: count total nodes in the deserialized tree to detect
53        // alias bombs. After expansion, aliased subtrees become
54        // independent copies in memory, so the node count reflects the
55        // true memory footprint.
56        let node_count = count_yaml_nodes(&value);
57        if node_count > YAML_NODE_COUNT {
58            return Err(SanitizeError::InputTooLarge {
59                size: node_count,
60                limit: YAML_NODE_COUNT,
61            });
62        }
63
64        walk_yaml(&mut value, "", profile, store, 0)?;
65
66        let output = serde_yaml_ng::to_string(&value).map_err(|e| {
67            SanitizeError::IoError(std::io::Error::other(format!("YAML serialize error: {e}")))
68        })?;
69
70        Ok(output.into_bytes())
71    }
72}
73
74/// Count the total number of nodes in a YAML value tree (F-06 fix).
75/// Used to detect alias bombs that produce a small source document
76/// but expand to millions of nodes after alias resolution.
77fn count_yaml_nodes(value: &Value) -> usize {
78    count_yaml_nodes_inner(value, 0)
79}
80
81/// Inner recursive counter with depth guard to prevent stack overflow
82/// on deeply nested YAML before `walk_yaml`'s depth check is reached.
83fn count_yaml_nodes_inner(value: &Value, depth: usize) -> usize {
84    if depth > DEFAULT_DEPTH {
85        return 1; // Stop counting deeper; walk_yaml will catch depth violations
86    }
87    match value {
88        Value::Mapping(map) => {
89            1 + map
90                .iter()
91                .map(|(k, v)| {
92                    count_yaml_nodes_inner(k, depth + 1) + count_yaml_nodes_inner(v, depth + 1)
93                })
94                .sum::<usize>()
95        }
96        Value::Sequence(seq) => {
97            1 + seq
98                .iter()
99                .map(|v| count_yaml_nodes_inner(v, depth + 1))
100                .sum::<usize>()
101        }
102        Value::Tagged(tagged) => 1 + count_yaml_nodes_inner(&tagged.value, depth + 1),
103        _ => 1, // Null, Bool, Number, String
104    }
105}
106
107impl TreeNode for Value {
108    fn for_each_map_entry<F>(&mut self, mut f: F) -> Result<()>
109    where
110        F: FnMut(&str, &mut Self) -> Result<()>,
111    {
112        if let Self::Mapping(map) = self {
113            let keys: Vec<Self> = map.keys().cloned().collect();
114            for key in keys {
115                let key_str = yaml_key_to_string(&key);
116                if let Some(v) = map.get_mut(&key) {
117                    f(&key_str, v)?;
118                }
119            }
120        }
121        Ok(())
122    }
123
124    fn for_each_seq_item<F>(&mut self, mut f: F) -> Result<()>
125    where
126        F: FnMut(&mut Self) -> Result<()>,
127    {
128        if let Self::Sequence(seq) = self {
129            for item in seq.iter_mut() {
130                f(item)?;
131            }
132        }
133        Ok(())
134    }
135
136    fn as_str_mut(&mut self) -> Option<&mut String> {
137        if let Self::String(s) = self {
138            Some(s)
139        } else {
140            None
141        }
142    }
143
144    fn is_scalar(&self) -> bool {
145        matches!(self, Self::Number(_) | Self::Bool(_))
146    }
147
148    fn scalar_to_string(&self) -> String {
149        yaml_scalar_to_string(self)
150    }
151
152    fn set_string(&mut self, s: String) {
153        *self = Self::String(s);
154    }
155}
156
157/// Recursively walk a YAML value tree, replacing matched field values.
158fn walk_yaml(
159    value: &mut Value,
160    prefix: &str,
161    profile: &FileTypeProfile,
162    store: &MappingStore,
163    depth: usize,
164) -> Result<()> {
165    walk_tree(value, prefix, profile, store, depth, "YAML")
166}
167
168fn yaml_key_to_string(key: &Value) -> String {
169    match key {
170        Value::String(s) => s.clone(),
171        Value::Number(n) => n.to_string(),
172        Value::Bool(b) => b.to_string(),
173        _ => format!("{:?}", key),
174    }
175}
176
177fn yaml_scalar_to_string(v: &Value) -> String {
178    match v {
179        Value::String(s) => s.clone(),
180        Value::Number(n) => n.to_string(),
181        Value::Bool(b) => b.to_string(),
182        _ => String::new(),
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189    use crate::category::Category;
190    use crate::generator::HmacGenerator;
191    use crate::processor::profile::FieldRule;
192    use std::sync::Arc;
193
194    fn make_store() -> MappingStore {
195        let gen = Arc::new(HmacGenerator::new([42u8; 32]));
196        MappingStore::new(gen, None)
197    }
198
199    #[test]
200    fn basic_yaml_replacement() {
201        let store = make_store();
202        let proc = YamlProcessor;
203
204        let content = b"database:\n  host: db.corp.com\n  password: s3cret\nport: 5432\n";
205        let profile = FileTypeProfile::new(
206            "yaml",
207            vec![
208                FieldRule::new("database.password").with_category(Category::Custom("pw".into())),
209                FieldRule::new("database.host").with_category(Category::Hostname),
210            ],
211        );
212
213        let result = proc.process(content, &profile, &store).unwrap();
214        let out = String::from_utf8(result).unwrap();
215
216        assert!(!out.contains("s3cret"));
217        assert!(!out.contains("db.corp.com"));
218        // port should be preserved
219        assert!(out.contains("5432"));
220    }
221
222    #[test]
223    fn can_handle_by_profile_name() {
224        let proc = YamlProcessor;
225        let profile = FileTypeProfile::new("yaml", vec![]).with_extension(".yaml");
226        assert!(proc.can_handle(b"anything", &profile));
227    }
228
229    #[test]
230    fn can_handle_detects_document_marker() {
231        let proc = YamlProcessor;
232        let profile = FileTypeProfile::new("json", vec![]).with_extension(".json");
233        assert!(proc.can_handle(b"---\nkey: value\n", &profile));
234    }
235
236    #[test]
237    fn can_handle_detects_key_value_heuristic() {
238        let proc = YamlProcessor;
239        let profile = FileTypeProfile::new("other", vec![]).with_extension(".conf");
240        assert!(proc.can_handle(b"host: localhost\nport: 5432\n", &profile));
241    }
242
243    #[test]
244    fn can_handle_detects_sequence_heuristic() {
245        let proc = YamlProcessor;
246        let profile = FileTypeProfile::new("other", vec![]).with_extension(".txt");
247        assert!(proc.can_handle(b"- item1\n- item2\n", &profile));
248    }
249
250    #[test]
251    fn can_handle_rejects_plaintext() {
252        let proc = YamlProcessor;
253        let profile = FileTypeProfile::new("json", vec![]).with_extension(".json");
254        assert!(!proc.can_handle(b"just plain text with no yaml markers", &profile));
255    }
256
257    #[test]
258    fn non_string_scalars_not_targeted_pass_through() {
259        let store = make_store();
260        let proc = YamlProcessor;
261        // Only target the 'secret' field; booleans and numbers are untouched.
262        let content = b"enabled: true\ncount: 42\nsecret: hunter2\n";
263        let profile = FileTypeProfile::new(
264            "yaml",
265            vec![FieldRule::new("secret").with_category(Category::Custom("pw".into()))],
266        );
267        let result = proc.process(content, &profile, &store).unwrap();
268        let out = String::from_utf8(result).unwrap();
269        assert!(!out.contains("hunter2"), "secret must be replaced");
270        assert!(out.contains("42"), "integer must be preserved");
271    }
272
273    #[test]
274    fn deeply_nested_yaml_replaced() {
275        let store = make_store();
276        let proc = YamlProcessor;
277        let content = b"a:\n  b:\n    c:\n      secret: hunter2\n";
278        let profile = FileTypeProfile::new(
279            "yaml",
280            vec![FieldRule::new("a.b.c.secret").with_category(Category::Custom("pw".into()))],
281        );
282        let result = proc.process(content, &profile, &store).unwrap();
283        let out = String::from_utf8(result).unwrap();
284        assert!(!out.contains("hunter2"));
285    }
286
287    #[test]
288    fn invalid_utf8_returns_parse_error() {
289        let store = make_store();
290        let proc = YamlProcessor;
291        let bad = b"\xff\xfe invalid";
292        let profile = FileTypeProfile::new("yaml", vec![]);
293        let err = proc.process(bad, &profile, &store).unwrap_err();
294        assert!(matches!(
295            err,
296            crate::error::SanitizeError::ParseError { .. }
297        ));
298    }
299
300    #[test]
301    fn invalid_yaml_returns_parse_error() {
302        let store = make_store();
303        let proc = YamlProcessor;
304        let bad = b"key: [unclosed";
305        let profile = FileTypeProfile::new("yaml", vec![]);
306        let err = proc.process(bad, &profile, &store).unwrap_err();
307        assert!(matches!(
308            err,
309            crate::error::SanitizeError::ParseError { .. }
310        ));
311    }
312
313    #[test]
314    fn yaml_sequence_traversal() {
315        let store = make_store();
316        let proc = YamlProcessor;
317
318        let content = b"users:\n  - email: a@b.com\n  - email: c@d.com\n";
319        let profile = FileTypeProfile::new(
320            "yaml",
321            vec![FieldRule::new("users.email").with_category(Category::Email)],
322        );
323
324        let result = proc.process(content, &profile, &store).unwrap();
325        let out = String::from_utf8(result).unwrap();
326
327        assert!(!out.contains("a@b.com"));
328        assert!(!out.contains("c@d.com"));
329    }
330}
rust_sanitize/processor/yaml_proc.rs

rust_sanitize/processor/
yaml_proc.rs
