rust_sanitize/processor/
toml_proc.rs1use crate::error::{Result, SanitizeError};
23use crate::processor::limits::DEFAULT_INPUT_SIZE;
24use crate::processor::{walk_tree, FileTypeProfile, Processor, TreeNode};
25use crate::store::MappingStore;
26use toml::Value;
27
28pub struct TomlProcessor;
30
31impl Processor for TomlProcessor {
32 fn name(&self) -> &'static str {
33 "toml"
34 }
35
36 fn can_handle(&self, _content: &[u8], profile: &FileTypeProfile) -> bool {
37 profile.processor == "toml"
38 }
39
40 fn process(
41 &self,
42 content: &[u8],
43 profile: &FileTypeProfile,
44 store: &MappingStore,
45 ) -> Result<Vec<u8>> {
46 let text = crate::processor::check_size_and_decode(content, "TOML", DEFAULT_INPUT_SIZE)?;
47
48 let mut value: Value = toml::from_str(text).map_err(|e| SanitizeError::ParseError {
49 format: "TOML".into(),
50 message: format!("TOML parse error: {}", e),
51 })?;
52
53 walk_toml(&mut value, "", profile, store, 0)?;
54
55 let output = toml::to_string_pretty(&value).map_err(|e| {
56 SanitizeError::IoError(std::io::Error::other(format!("TOML serialize error: {e}")))
57 })?;
58
59 Ok(output.into_bytes())
60 }
61}
62
63impl TreeNode for Value {
64 fn for_each_map_entry<F>(&mut self, mut f: F) -> Result<()>
65 where
66 F: FnMut(&str, &mut Self) -> Result<()>,
67 {
68 if let Self::Table(map) = self {
69 let keys: Vec<String> = map.keys().cloned().collect();
70 for key in keys {
71 if let Some(v) = map.get_mut(&key) {
72 f(&key, v)?;
73 }
74 }
75 }
76 Ok(())
77 }
78
79 fn for_each_seq_item<F>(&mut self, mut f: F) -> Result<()>
80 where
81 F: FnMut(&mut Self) -> Result<()>,
82 {
83 if let Self::Array(arr) = self {
84 for item in arr.iter_mut() {
85 f(item)?;
86 }
87 }
88 Ok(())
89 }
90
91 fn as_str_mut(&mut self) -> Option<&mut String> {
92 if let Self::String(s) = self {
93 Some(s)
94 } else {
95 None
96 }
97 }
98
99 fn is_scalar(&self) -> bool {
100 matches!(
103 self,
104 Self::Integer(_) | Self::Float(_) | Self::Boolean(_) | Self::Datetime(_)
105 )
106 }
107
108 fn scalar_to_string(&self) -> String {
109 self.to_string()
110 }
111
112 fn set_string(&mut self, s: String) {
113 *self = Self::String(s);
114 }
115}
116
117fn walk_toml(
119 value: &mut Value,
120 prefix: &str,
121 profile: &FileTypeProfile,
122 store: &MappingStore,
123 depth: usize,
124) -> Result<()> {
125 walk_tree(value, prefix, profile, store, depth, "TOML")
126}
127
128#[cfg(test)]
129mod tests {
130 use super::*;
131 use crate::category::Category;
132 use crate::generator::HmacGenerator;
133 use crate::processor::profile::FieldRule;
134 use std::sync::Arc;
135
136 fn make_store() -> MappingStore {
137 let gen = Arc::new(HmacGenerator::new([42u8; 32]));
138 MappingStore::new(gen, None)
139 }
140
141 #[test]
142 fn basic_toml_replacement() {
143 let store = make_store();
144 let proc = TomlProcessor;
145 let content = br#"[database]
146host = "db.corp.com"
147password = "s3cret"
148port = 5432
149
150[smtp]
151user = "admin@corp.com"
152"#;
153 let profile = FileTypeProfile::new(
154 "toml",
155 vec![
156 FieldRule::new("database.password"),
157 FieldRule::new("smtp.user").with_category(Category::Email),
158 ],
159 );
160 let output = proc.process(content, &profile, &store).unwrap();
161 let text = String::from_utf8(output).unwrap();
162 assert!(!text.contains("s3cret"));
164 assert!(text.contains("db.corp.com"));
165 assert!(text.contains("5432"));
166 assert!(!text.contains("admin@corp.com"));
168 }
169
170 #[test]
171 fn wildcard_replaces_all_strings() {
172 let store = make_store();
173 let proc = TomlProcessor;
174 let content = b"api_key = \"secret\"\ndb_url = \"postgres://user:pass@host/db\"\n";
175 let profile = FileTypeProfile::new("toml", vec![FieldRule::new("*")]);
176 let output = proc.process(content, &profile, &store).unwrap();
177 let text = String::from_utf8(output).unwrap();
178 assert!(!text.contains("secret"));
179 assert!(!text.contains("postgres://user:pass@host/db"));
180 }
181
182 #[test]
183 fn invalid_toml_returns_parse_error() {
184 let store = make_store();
185 let proc = TomlProcessor;
186 let content = b"this is not valid toml [[[";
187 let profile = FileTypeProfile::new("toml", vec![FieldRule::new("*")]);
188 let result = proc.process(content, &profile, &store);
189 assert!(result.is_err());
190 }
191
192 #[test]
193 fn deeply_nested_toml() {
194 let store = make_store();
195 let proc = TomlProcessor;
196 let content = b"[a.b.c]\nkey = \"value\"\n";
197 let profile = FileTypeProfile::new("toml", vec![FieldRule::new("a.b.c.key")]);
198 let output = proc.process(content, &profile, &store).unwrap();
199 let text = String::from_utf8(output).unwrap();
200 assert!(!text.contains("value"));
201 }
202}