1use std::borrow::Cow;
14use std::sync::LazyLock;
15use std::sync::atomic::{AtomicBool, Ordering};
16
17pub const PREFIX: &str = "DRE_SECRET_";
18pub const MASK: &str = "*****";
19pub const MASKED_SOURCE_ERROR: &str =
20 "source error details omitted because the request contains a masked secret";
21const MIN_LEN: usize = 4;
22
23static ENABLED: AtomicBool = AtomicBool::new(true);
24
25fn eligible(value: &str) -> bool {
26 value.chars().count() >= MIN_LEN
27}
28
29static SECRETS: LazyLock<Vec<String>> = LazyLock::new(|| {
31 let mut v: Vec<String> = std::env::vars()
32 .filter(|(k, v)| k.starts_with(PREFIX) && eligible(v))
33 .map(|(_, v)| v)
34 .collect();
35 v.sort_by_key(|s| std::cmp::Reverse(s.len()));
36 v.dedup();
37 v
38});
39
40pub fn set_enabled(on: bool) {
41 ENABLED.store(on, Ordering::Relaxed);
42}
43
44pub(crate) fn enabled() -> bool {
45 ENABLED.load(Ordering::Relaxed)
46}
47
48fn mask_with<'a>(s: &'a str, secrets: &[String]) -> Cow<'a, str> {
49 if !secrets.iter().any(|x| s.contains(x.as_str())) {
50 return Cow::Borrowed(s);
51 }
52 let mut out = s.to_string();
53 for x in secrets {
54 out = out.replace(x.as_str(), MASK);
55 }
56 Cow::Owned(out)
57}
58
59pub fn mask(s: &str) -> Cow<'_, str> {
61 if !ENABLED.load(Ordering::Relaxed) {
62 return Cow::Borrowed(s);
63 }
64 mask_with(s, &SECRETS)
65}
66
67pub fn contains_secret(s: &str) -> bool {
69 ENABLED.load(Ordering::Relaxed) && SECRETS.iter().any(|secret| s.contains(secret))
70}
71
72fn mask_source_error_with<'a>(error: &'a str, sensitive: bool, secrets: &[String]) -> Cow<'a, str> {
73 if sensitive {
74 Cow::Borrowed(MASKED_SOURCE_ERROR)
75 } else {
76 mask_with(error, secrets)
77 }
78}
79
80pub fn mask_source_error(error: &str, sensitive: bool) -> Cow<'_, str> {
84 if !ENABLED.load(Ordering::Relaxed) {
85 return Cow::Borrowed(error);
86 }
87 mask_source_error_with(error, sensitive, &SECRETS)
88}
89
90fn unique_masked_names_with(names: &[String], secrets: &[String]) -> Vec<String> {
91 use std::collections::BTreeSet;
92
93 let masked: Vec<String> = names
94 .iter()
95 .map(|name| mask_with(name, secrets).into_owned())
96 .collect();
97 if names == masked {
98 return names.to_vec();
99 }
100 let mut used = BTreeSet::new();
101 let mut output = vec![String::new(); names.len()];
102 let mut pending = Vec::new();
103
104 for (index, (original, protected)) in names.iter().zip(&masked).enumerate() {
106 if original == protected {
107 used.insert(original.clone());
108 output[index].clone_from(original);
109 } else {
110 pending.push((protected.clone(), original.clone(), index));
111 }
112 }
113
114 pending.sort();
116 for (base, _, index) in pending {
117 let mut candidate = base.clone();
118 let mut suffix = 2;
119 while !used.insert(candidate.clone()) {
120 candidate = format!("{base}#{suffix}");
121 suffix += 1;
122 }
123 output[index] = candidate;
124 }
125 output
126}
127
128fn redact_json_with(v: &mut serde_json::Value, secrets: &[String]) {
129 use serde_json::Value;
130 match v {
131 Value::String(s) => {
132 if let Cow::Owned(m) = mask_with(s, secrets) {
133 *s = m;
134 }
135 }
136 Value::Array(a) => a.iter_mut().for_each(|v| redact_json_with(v, secrets)),
137 Value::Object(o) => {
138 let entries = std::mem::take(o);
139 let keys: Vec<String> = entries.keys().cloned().collect();
140 let protected_keys = unique_masked_names_with(&keys, secrets);
141 for ((_, mut val), protected_key) in entries.into_iter().zip(protected_keys) {
142 redact_json_with(&mut val, secrets);
143 o.insert(protected_key, val);
144 }
145 }
146 _ => {}
147 }
148}
149
150pub fn redact_json(v: &mut serde_json::Value) {
153 if ENABLED.load(Ordering::Relaxed) {
154 redact_json_with(v, &SECRETS);
155 }
156}
157
158fn redacted<T: serde::Serialize>(v: &T) -> serde_json::Result<serde_json::Value> {
159 let mut v = serde_json::to_value(v)?;
160 redact_json(&mut v);
161 Ok(v)
162}
163
164pub fn to_json_pretty<T: serde::Serialize>(v: &T) -> serde_json::Result<String> {
166 serde_json::to_string_pretty(&redacted(v)?)
167}
168
169pub fn to_json_line<T: serde::Serialize>(v: &T) -> serde_json::Result<String> {
171 serde_json::to_string(&redacted(v)?)
172}
173
174#[cfg(test)]
175mod tests {
176 use super::*;
177 use serde_json::json;
178
179 fn secrets(values: &[&str]) -> Vec<String> {
180 let mut values: Vec<String> = values.iter().map(|v| (*v).to_string()).collect();
181 values.sort_by_key(|s| std::cmp::Reverse(s.len()));
182 values
183 }
184
185 #[test]
186 fn redacts_nested_values_longest_first() {
187 let secrets = secrets(&["nested-secret", "nested-secret-and-more"]);
188 let mut value = json!({"items": ["nested-secret-and-more", {"value": "x nested-secret y"}]});
189 redact_json_with(&mut value, &secrets);
190 assert_eq!(value, json!({"items": ["*****", {"value": "x ***** y"}]}));
191 }
192
193 #[test]
194 fn preserves_entries_when_masked_keys_collide() {
195 let secrets = secrets(&["alpha-secret-key", "bravo-secret-key"]);
196 let mut value = json!({
197 "alpha-secret-key": 1,
198 "bravo-secret-key": 2,
199 "*****": 3,
200 "*****#2": 4,
201 });
202 redact_json_with(&mut value, &secrets);
203 assert_eq!(
204 value,
205 json!({
206 "*****": 3,
207 "*****#2": 4,
208 "*****#3": 1,
209 "*****#4": 2,
210 })
211 );
212 }
213
214 #[test]
215 fn masking_one_name_does_not_rename_duplicate_literal_names() {
216 let secrets = secrets(&["secret-column"]);
217 let names = ["a", "a", "secret-column"].map(str::to_string);
218 assert_eq!(unique_masked_names_with(&names, &secrets), ["a", "a", "*****"]);
219 }
220
221 #[test]
222 fn suppresses_transformed_source_errors_only_for_secret_requests() {
223 let secrets = secrets(&["long-secret-value-that-a-source-truncated"]);
224 let error = "LINE 1: select * from long-secret-value-that-a-source...";
225 assert_eq!(mask_source_error_with(error, true, &secrets), MASKED_SOURCE_ERROR);
226 assert_eq!(
227 mask_source_error_with("ordinary syntax error", false, &secrets),
228 "ordinary syntax error"
229 );
230 }
231
232 #[test]
233 fn minimum_length_counts_unicode_characters() {
234 assert!(!eligible("ééé"));
235 assert!(eligible("éééé"));
236 }
237}