wm_dispatch/
secret_scan.rs1use std::sync::Arc;
17use std::sync::atomic::{AtomicU64, Ordering};
18
19pub const DEFAULT_SAMPLE_EVERY: u64 = 100;
21
22pub const SAMPLE_EVERY_ENV: &str = "WM_SECRET_SCAN_EVERY";
25
26pub struct SecretSampler {
28 every: u64,
29 seen: AtomicU64,
30 sampled: AtomicU64,
31 hits: AtomicU64,
32}
33
34impl SecretSampler {
35 #[must_use]
37 pub const fn new(every: u64) -> Self {
38 Self {
39 every,
40 seen: AtomicU64::new(0),
41 sampled: AtomicU64::new(0),
42 hits: AtomicU64::new(0),
43 }
44 }
45
46 #[must_use]
48 pub fn from_env() -> Self {
49 Self::new(parse_every(std::env::var(SAMPLE_EVERY_ENV).ok().as_deref()))
50 }
51
52 #[must_use]
54 pub const fn every(&self) -> u64 {
55 self.every
56 }
57
58 #[must_use]
60 pub fn stats(&self) -> (u64, u64, u64) {
61 (
62 self.seen.load(Ordering::Relaxed),
63 self.sampled.load(Ordering::Relaxed),
64 self.hits.load(Ordering::Relaxed),
65 )
66 }
67
68 pub fn scan(&self, tool: &str, output: &serde_json::Value) -> Vec<&'static str> {
72 if self.every == 0 {
73 return Vec::new();
74 }
75 let n = self.seen.fetch_add(1, Ordering::Relaxed);
76 if n % self.every != 0 {
77 return Vec::new();
78 }
79 self.sampled.fetch_add(1, Ordering::Relaxed);
80 let Ok(text) = serde_json::to_string(output) else {
81 return Vec::new();
82 };
83 let kinds = wm_memory::credential_shaped_content(&text);
84 if !kinds.is_empty() {
85 self.hits.fetch_add(1, Ordering::Relaxed);
86 tracing::warn!(
87 tool,
88 kinds = ?kinds,
89 output_bytes = text.len(),
90 "secret-scan: successful output looks credential-bearing (warn-only; content withheld)"
91 );
92 }
93 kinds
94 }
95}
96
97#[must_use]
103pub fn parse_every(value: Option<&str>) -> u64 {
104 value
105 .and_then(|v| v.trim().parse::<u64>().ok())
106 .unwrap_or(DEFAULT_SAMPLE_EVERY)
107}
108
109pub type SharedSampler = Arc<SecretSampler>;
111
112#[cfg(test)]
113mod tests {
114 use super::*;
115
116 const EXAMPLE_AKIA: &str = "AKIAIOSFODNN7EXAMPLE";
118
119 fn key_output() -> serde_json::Value {
120 serde_json::json!({"data": format!("key={EXAMPLE_AKIA}")})
121 }
122
123 #[test]
124 fn disabled_sampler_scans_nothing() {
125 let s = SecretSampler::new(0);
126 assert!(s.scan("memory.read", &key_output()).is_empty());
127 assert_eq!(s.stats(), (0, 0, 0));
128 }
129
130 #[test]
131 fn cadence_is_counter_deterministic() {
132 let s = SecretSampler::new(3);
133 for _ in 0..9 {
134 s.scan("t", &serde_json::json!({"ok": true}));
135 }
136 assert_eq!(s.stats(), (9, 3, 0));
138 }
139
140 #[test]
141 fn hit_kinds_returned_and_counted() {
142 let s = SecretSampler::new(1);
143 let kinds = s.scan("memory.read", &key_output());
144 assert!(kinds.contains(&"aws_access_key_id"));
145 assert_eq!(s.stats(), (1, 1, 1));
146 }
147
148 #[test]
149 fn clean_output_sampled_without_hit() {
150 let s = SecretSampler::new(1);
151 let kinds = s.scan("memory.search", &serde_json::json!({"results": []}));
152 assert!(kinds.is_empty());
153 assert_eq!(s.stats(), (1, 1, 0));
154 }
155
156 #[test]
157 fn parse_every_defaults_and_parses() {
158 assert_eq!(parse_every(None), DEFAULT_SAMPLE_EVERY);
161 assert_eq!(parse_every(Some("7")), 7);
162 assert_eq!(parse_every(Some("0")), 0);
163 assert_eq!(parse_every(Some(" 25 ")), 25);
164 assert_eq!(parse_every(Some("garbage")), DEFAULT_SAMPLE_EVERY);
165 assert_eq!(parse_every(Some("")), DEFAULT_SAMPLE_EVERY);
166 }
167}