forge
ExploreLog in

Tahinli/forge

public
0 watching
★ 0
7.9 KiB
1//! Secret masking for job logs. Replaces every secret value, and its standard base64 encoding,
2//! with `***`. [`MaskStream`] does the same over chunks, so a secret split across two chunks is
3//! still masked.
4
5use std::collections::BTreeSet;
6use std::fmt;
7
8use base64::Engine as _;
9use base64::engine::general_purpose::STANDARD;
10
11/// The text that replaces every match.
12pub const MASK: &str = "***";
13/// Values shorter than this (in chars) are not masked, since they would hide ordinary output.
14pub const MIN_SECRET_CHARS: usize = 4;
15
16/// Masks a fixed set of secret values.
17#[derive(Clone)]
18pub struct Masker {
19 /// Byte patterns (raw values and their base64), longest first, deduplicated.
20 patterns: Vec<Vec<u8>>,
21}
22
23impl fmt::Debug for Masker {
24 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25 f.debug_struct("Masker")
26 .field("patterns", &self.patterns.len())
27 .finish()
28 }
29}
30
31impl Masker {
32 /// Builds a masker. Values shorter than [`MIN_SECRET_CHARS`] are ignored.
33 pub fn new(secrets: impl IntoIterator<Item = String>) -> Self {
34 let mut set: BTreeSet<Vec<u8>> = BTreeSet::new();
35 for value in secrets {
36 if value.chars().count() < MIN_SECRET_CHARS {
37 continue;
38 }
39 let encoded = STANDARD.encode(value.as_bytes());
40 set.insert(value.into_bytes());
41 set.insert(encoded.into_bytes());
42 }
43 let mut patterns: Vec<Vec<u8>> = set.into_iter().collect();
44 patterns.sort_by_key(|p| std::cmp::Reverse(p.len()));
45 Self { patterns }
46 }
47
48 /// Masks every occurrence in `text`.
49 pub fn mask(&self, text: &str) -> String {
50 let (out, _) = scan(&self.patterns, text.as_bytes(), true);
51 bytes_to_string(out)
52 }
53
54 /// Starts a streaming masker that owns a copy of this masker.
55 pub fn stream(&self) -> MaskStream {
56 MaskStream::new(self.clone())
57 }
58}
59
60/// Masks a byte stream that arrives in chunks.
61#[derive(Debug, Clone)]
62pub struct MaskStream {
63 masker: Masker,
64 /// Bytes held back because they may be the start of a secret. Never more than the longest
65 /// pattern minus one.
66 held: Vec<u8>,
67}
68
69impl MaskStream {
70 pub fn new(masker: Masker) -> Self {
71 Self {
72 masker,
73 held: Vec::new(),
74 }
75 }
76
77 /// Feeds one chunk. Returns the masked bytes that are now safe to emit.
78 pub fn push(&mut self, chunk: &[u8]) -> Vec<u8> {
79 self.held.extend_from_slice(chunk);
80 let (out, rest) = scan(&self.masker.patterns, &self.held, false);
81 self.held = rest;
82 out
83 }
84
85 /// Ends the stream and returns whatever is still held back, masked.
86 pub fn finish(&mut self) -> Vec<u8> {
87 let held = std::mem::take(&mut self.held);
88 let (out, _) = scan(&self.masker.patterns, &held, true);
89 out
90 }
91}
92
93/// Walks `input` and writes masked bytes. A match is replaced by [`MASK`]. When `last` is false,
94/// the scan stops at the first position where a pattern could still complete in later input, and
95/// returns that unprocessed tail. When `last` is true, the whole input is processed.
96fn scan(patterns: &[Vec<u8>], input: &[u8], last: bool) -> (Vec<u8>, Vec<u8>) {
97 let mut out = Vec::with_capacity(input.len());
98 let mut i = 0;
99 while i < input.len() {
100 let tail = &input[i..];
101 if let Some(p) = patterns.iter().find(|p| tail.starts_with(p)) {
102 out.extend_from_slice(MASK.as_bytes());
103 i += p.len();
104 continue;
105 }
106 if !last
107 && patterns
108 .iter()
109 .any(|p| p.len() > tail.len() && p.starts_with(tail))
110 {
111 return (out, tail.to_vec());
112 }
113 out.push(input[i]);
114 i += 1;
115 }
116 (out, Vec::new())
117}
118
119fn bytes_to_string(bytes: Vec<u8>) -> String {
120 match String::from_utf8(bytes) {
121 Ok(s) => s,
122 Err(e) => String::from_utf8_lossy(e.as_bytes()).into_owned(),
123 }
124}
125
126#[cfg(test)]
127mod tests {
128 use super::*;
129
130 fn masker(values: &[&str]) -> Masker {
131 Masker::new(values.iter().map(|v| v.to_string()))
132 }
133
134 fn stream_all(m: &Masker, chunks: &[&[u8]]) -> String {
135 let mut s = m.stream();
136 let mut out = Vec::new();
137 for c in chunks {
138 out.extend(s.push(c));
139 }
140 out.extend(s.finish());
141 String::from_utf8(out).unwrap()
142 }
143
144 #[test]
145 fn masks_plain_values() {
146 let m = masker(&["hunter2-secret"]);
147 assert_eq!(m.mask("token=hunter2-secret ok"), "token=*** ok");
148 assert_eq!(
149 m.mask("two hunter2-secret and hunter2-secret"),
150 "two *** and ***"
151 );
152 assert_eq!(m.mask("nothing here"), "nothing here");
153 }
154
155 #[test]
156 fn masks_base64_encoding() {
157 let secret = "hunter2-secret";
158 let encoded = STANDARD.encode(secret.as_bytes());
159 let m = masker(&[secret]);
160 assert_eq!(m.mask(&format!("leak: {encoded}")), "leak: ***");
161 let other = masker(&["abcdefgh"]);
162 let enc = STANDARD.encode("abcdefgh");
163 assert_eq!(other.mask(&enc), "***");
164 }
165
166 #[test]
167 fn short_values_are_ignored() {
168 let m = masker(&["abc", "", "ab"]);
169 assert_eq!(m.mask("abc ab"), "abc ab");
170 // 4 chars is the minimum that is masked.
171 let m = masker(&["abcd"]);
172 assert_eq!(m.mask("xabcdx"), "x***x");
173 }
174
175 #[test]
176 fn longest_match_wins() {
177 let m = masker(&["secret-value", "secret"]);
178 assert_eq!(m.mask("a secret-value b"), "a *** b");
179 }
180
181 #[test]
182 fn empty_secret_set_is_a_no_op() {
183 let m = Masker::new(std::iter::empty());
184 assert_eq!(m.mask("anything"), "anything");
185 assert_eq!(stream_all(&m, &[b"anything"]), "anything");
186 }
187
188 #[test]
189 fn secret_split_across_two_chunks() {
190 let m = masker(&["hunter2-secret"]);
191 let out = stream_all(&m, &[b"pre hunter2-", b"secret post"]);
192 assert_eq!(out, "pre *** post");
193 }
194
195 #[test]
196 fn secret_split_one_byte_at_a_time() {
197 let m = masker(&["hunter2-secret"]);
198 let text = b"a hunter2-secret b hunter2-secret c";
199 let chunks: Vec<&[u8]> = text.chunks(1).collect();
200 assert_eq!(stream_all(&m, &chunks), "a *** b *** c");
201 }
202
203 #[test]
204 fn base64_split_across_chunks() {
205 let secret = "hunter2-secret";
206 let encoded = STANDARD.encode(secret);
207 let m = masker(&[secret]);
208 let (a, b) = encoded.split_at(5);
209 let out = stream_all(
210 &m,
211 &[format!("x{a}").as_bytes(), format!("{b}y").as_bytes()],
212 );
213 assert_eq!(out, "x***y");
214 }
215
216 #[test]
217 fn held_bytes_are_bounded() {
218 let m = masker(&["0123456789"]);
219 let mut s = m.stream();
220 let out = s.push(b"0123456");
221 assert!(out.is_empty());
222 assert!(s.held.len() <= 9);
223 let out = s.push(b"x");
224 assert_eq!(out, b"0123456x".to_vec());
225 }
226
227 #[test]
228 fn partial_prefix_that_does_not_complete_is_emitted() {
229 let m = masker(&["hunter2-secret"]);
230 let out = stream_all(&m, &[b"hunt", b"er is here"]);
231 assert_eq!(out, "hunter is here");
232 }
233
234 #[test]
235 fn finish_flushes_held_partial_match() {
236 let m = masker(&["hunter2-secret"]);
237 let mut s = m.stream();
238 let mut out = s.push(b"end hunter2-se");
239 out.extend(s.finish());
240 assert_eq!(String::from_utf8(out).unwrap(), "end hunter2-se");
241 }
242
243 #[test]
244 fn multibyte_text_survives() {
245 let m = masker(&["hunter2-secret"]);
246 assert_eq!(m.mask("héllo 🙂 hunter2-secret ✓"), "héllo 🙂 *** ✓");
247 let out = stream_all(&m, &["héllo 🙂 ".as_bytes(), "hunter2-secret ✓".as_bytes()]);
248 assert_eq!(out, "héllo 🙂 *** ✓");
249 }
250
251 #[test]
252 fn debug_does_not_print_values() {
253 let m = masker(&["hunter2-secret"]);
254 assert!(!format!("{m:?}").contains("hunter2"));
255 }
256}