-
-
Notifications
You must be signed in to change notification settings - Fork 1.3k
perf: use Aho-Corasick for efficient redaction #7931
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 1 commit
3f13c99
ace6ea4
f1a1989
9287026
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -6,7 +6,6 @@ pub use settings::Settings; | |||||
| use std::collections::{BTreeMap, BTreeSet, HashMap}; | ||||||
| use std::fmt::{Debug, Formatter}; | ||||||
| use std::iter::once; | ||||||
| use std::ops::Deref; | ||||||
| use std::path::{Path, PathBuf}; | ||||||
| use std::sync::LazyLock as Lazy; | ||||||
| use std::sync::{Arc, Mutex, RwLock}; | ||||||
|
|
@@ -44,6 +43,7 @@ use crate::env_diff::EnvMap; | |||||
| use crate::hook_env::WatchFilePattern; | ||||||
| use crate::hooks::Hook; | ||||||
| use crate::plugins::PluginType; | ||||||
| use crate::redactions::Redactor; | ||||||
| use crate::tera::BASE_CONTEXT; | ||||||
| use crate::watch_files::WatchFile; | ||||||
| use crate::wildcard::Wildcard; | ||||||
|
|
@@ -79,7 +79,7 @@ pub struct Alias { | |||||
| } | ||||||
|
|
||||||
| static _CONFIG: RwLock<Option<Arc<Config>>> = RwLock::new(None); | ||||||
| static _REDACTIONS: Lazy<Mutex<Arc<IndexSet<String>>>> = Lazy::new(Default::default); | ||||||
| static _REDACTOR: Lazy<Mutex<Redactor>> = Lazy::new(Default::default); | ||||||
|
|
||||||
| pub fn is_loaded() -> bool { | ||||||
| _CONFIG.read().unwrap().is_some() | ||||||
|
|
@@ -688,60 +688,30 @@ impl Config { | |||||
| .collect() | ||||||
| } | ||||||
| pub fn add_redactions(&self, redactions: impl IntoIterator<Item = String>, env: &EnvMap) { | ||||||
| let mut r = _REDACTIONS.lock().unwrap(); | ||||||
| let redactions = redactions.into_iter().flat_map(|r| { | ||||||
| let matcher = Wildcard::new(vec![r]); | ||||||
| let mut r = _REDACTOR.lock().unwrap(); | ||||||
| let new_redactions = redactions.into_iter().flat_map(|pattern| { | ||||||
| let matcher = Wildcard::new(vec![pattern]); | ||||||
| env.iter() | ||||||
| .filter(|(k, _)| matcher.match_any(k)) | ||||||
| .map(|(_, v)| v.clone()) | ||||||
| .collect::<Vec<_>>() | ||||||
| }); | ||||||
| *r = Arc::new(r.iter().cloned().chain(redactions).collect()); | ||||||
| *r = r.with_additional(new_redactions); | ||||||
| } | ||||||
|
|
||||||
| /// Get the current redactor (with cached Aho-Corasick automaton). | ||||||
| pub fn redactor(&self) -> Redactor { | ||||||
| _REDACTOR.lock().unwrap().clone() | ||||||
| } | ||||||
|
|
||||||
| /// Get the current redaction patterns. | ||||||
| pub fn redactions(&self) -> Arc<IndexSet<String>> { | ||||||
| let r = _REDACTIONS.lock().unwrap(); | ||||||
| r.deref().clone() | ||||||
|
|
||||||
| // self.redactions.get_or_try_init(|| { | ||||||
| // let mut redactions = Redactions::default(); | ||||||
| // for cf in self.config_files.values() { | ||||||
| // let r = cf.redactions(); | ||||||
| // if !r.is_empty() { | ||||||
| // let mut r = r.clone(); | ||||||
| // let (tera, ctx) = self.tera(&cf.config_root()); | ||||||
| // r.render(&mut tera.clone(), &ctx)?; | ||||||
| // redactions.merge(r); | ||||||
| // } | ||||||
| // } | ||||||
| // if redactions.is_empty() { | ||||||
| // return Ok(Default::default()); | ||||||
| // } | ||||||
| // | ||||||
| // let ts = self.get_toolset()?; | ||||||
| // let env = ts.full_env()?; | ||||||
| // | ||||||
| // let env_matcher = Wildcard::new(redactions.env.clone()); | ||||||
| // let var_matcher = Wildcard::new(redactions.vars.clone()); | ||||||
| // | ||||||
| // let env_vals = env | ||||||
| // .into_iter() | ||||||
| // .filter(|(k, _)| env_matcher.match_any(k)) | ||||||
| // .map(|(_, v)| v); | ||||||
| // let var_vals = self | ||||||
| // .vars | ||||||
| // .iter() | ||||||
| // .filter(|(k, _)| var_matcher.match_any(k)) | ||||||
| // .map(|(_, v)| v.to_string()); | ||||||
| // Ok(env_vals.chain(var_vals).collect()) | ||||||
| // }) | ||||||
| } | ||||||
|
|
||||||
| pub fn redact(&self, mut input: String) -> String { | ||||||
| for redaction in self.redactions().deref() { | ||||||
| input = input.replace(redaction, "[redacted]"); | ||||||
| } | ||||||
| input | ||||||
| Arc::new(_REDACTOR.lock().unwrap().patterns().clone()) | ||||||
|
||||||
| Arc::new(_REDACTOR.lock().unwrap().patterns().clone()) | |
| Arc::clone(_REDACTOR.lock().unwrap().patterns()) |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -23,18 +23,33 @@ impl log::Log for Logger { | |
| } | ||
|
|
||
| fn log(&self, record: &Record) { | ||
| if record.level() <= self.file_level | ||
| && let Some(log_file) = &self.log_file | ||
| { | ||
| let mut log_file = log_file.lock().unwrap(); | ||
| let out = self.render(record, self.file_level); | ||
| if !out.is_empty() { | ||
| let _ = writeln!(log_file, "{}", console::strip_ansi_codes(&out)); | ||
| let term_level = *self.term_level.lock().unwrap(); | ||
| let will_log_file = record.level() <= self.file_level && self.log_file.is_some(); | ||
| let will_log_term = record.level() <= term_level; | ||
|
|
||
| if !will_log_file && !will_log_term { | ||
| return; | ||
| } | ||
|
|
||
| // Redact once for all outputs (Aho-Corasick makes this efficient) | ||
| let args = record.args().to_string(); | ||
| let args = if config::is_loaded() { | ||
| Config::get_().redact(&args) | ||
| } else { | ||
| args | ||
| }; | ||
|
Comment on lines
+36
to
+40
|
||
|
|
||
| if will_log_file { | ||
| if let Some(log_file) = &self.log_file { | ||
| let mut log_file = log_file.lock().unwrap(); | ||
| let out = self.render(record, self.file_level, &args); | ||
| if !out.is_empty() { | ||
| let _ = writeln!(log_file, "{}", console::strip_ansi_codes(&out)); | ||
| } | ||
| } | ||
| } | ||
| let term_level = *self.term_level.lock().unwrap(); | ||
| if record.level() <= term_level { | ||
| let out = self.render(record, term_level); | ||
| if will_log_term { | ||
| let out = self.render(record, term_level, &args); | ||
| if !out.is_empty() { | ||
| ui::multi_progress_report::MultiProgressReport::suspend_if_active(|| { | ||
| eprintln!("{out}"); | ||
|
|
@@ -66,12 +81,7 @@ impl Logger { | |
| logger | ||
| } | ||
|
|
||
| fn render(&self, record: &Record, level: LevelFilter) -> String { | ||
| let mut args = record.args().to_string(); | ||
| if config::is_loaded() { | ||
| let config = Config::get_(); | ||
| args = config.redact(args); | ||
| } | ||
| fn render(&self, record: &Record, level: LevelFilter, args: &str) -> String { | ||
| match level { | ||
| LevelFilter::Off => "".to_string(), | ||
| LevelFilter::Trace => { | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,4 +1,6 @@ | ||
| use aho_corasick::AhoCorasick; | ||
| use indexmap::IndexSet; | ||
| use std::sync::Arc; | ||
|
|
||
| #[derive(Default, Clone, Debug, serde::Deserialize)] | ||
| pub struct Redactions(pub IndexSet<String>); | ||
|
|
@@ -19,3 +21,140 @@ impl Redactions { | |
| self.0.is_empty() | ||
| } | ||
| } | ||
|
|
||
| /// A redactor that uses Aho-Corasick for efficient multi-pattern string replacement. | ||
| /// | ||
| /// This is more efficient than iterating through patterns and calling `str::replace()` | ||
| /// for each one, especially when there are many patterns. Aho-Corasick finds all | ||
| /// matches in a single pass through the text - O(n + z) vs O(n * m). | ||
| #[derive(Clone)] | ||
| pub struct Redactor { | ||
| patterns: Arc<IndexSet<String>>, | ||
| automaton: Option<Arc<AhoCorasick>>, | ||
| } | ||
|
|
||
| impl Default for Redactor { | ||
| fn default() -> Self { | ||
| Self { | ||
| patterns: Arc::new(IndexSet::new()), | ||
| automaton: None, | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl Redactor { | ||
| /// Create a new redactor from a set of patterns to redact. | ||
| pub fn new(patterns: impl IntoIterator<Item = String>) -> Self { | ||
| let patterns: IndexSet<String> = patterns.into_iter().filter(|p| !p.is_empty()).collect(); | ||
| let automaton = if patterns.is_empty() { | ||
| None | ||
| } else { | ||
| // Build the Aho-Corasick automaton - O(m) where m is total pattern length | ||
| AhoCorasick::new(patterns.iter()).ok().map(Arc::new) | ||
|
cursor[bot] marked this conversation as resolved.
|
||
| }; | ||
| Self { | ||
| patterns: Arc::new(patterns), | ||
| automaton, | ||
| } | ||
| } | ||
|
|
||
| /// Create a new redactor by adding more patterns to an existing one. | ||
| pub fn with_additional(&self, additional: impl IntoIterator<Item = String>) -> Self { | ||
| let mut patterns = (*self.patterns).clone(); | ||
| for p in additional { | ||
| if !p.is_empty() { | ||
| patterns.insert(p); | ||
| } | ||
| } | ||
| Self::new(patterns) | ||
| } | ||
|
|
||
| /// Returns true if there are no patterns to redact. | ||
| pub fn is_empty(&self) -> bool { | ||
| self.patterns.is_empty() | ||
| } | ||
|
cursor[bot] marked this conversation as resolved.
Outdated
|
||
|
|
||
| /// Returns the patterns being redacted. | ||
| pub fn patterns(&self) -> &IndexSet<String> { | ||
| &self.patterns | ||
| } | ||
|
|
||
| /// Redact all matching patterns in the input string, replacing them with `[redacted]`. | ||
| /// | ||
| /// This is O(n + z) where n is the input length and z is the number of matches, | ||
| /// compared to O(n * m) for the naive approach of iterating through m patterns. | ||
| pub fn redact(&self, input: &str) -> String { | ||
| match &self.automaton { | ||
| Some(ac) => { | ||
| // Each pattern needs its own replacement string | ||
| let replacements: Vec<&str> = vec!["[redacted]"; self.patterns.len()]; | ||
| ac.replace_all(input, &replacements) | ||
| } | ||
| None => input.to_string(), | ||
| } | ||
| } | ||
|
|
||
| /// Redact in place, returning the same String if no changes were made. | ||
| pub fn redact_string(&self, input: String) -> String { | ||
| if self.automaton.is_none() { | ||
| return input; | ||
| } | ||
| self.redact(&input) | ||
| } | ||
|
||
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod tests { | ||
| use super::*; | ||
|
|
||
| #[test] | ||
| fn test_empty_redactor() { | ||
| let r = Redactor::default(); | ||
| assert_eq!(r.redact("hello world"), "hello world"); | ||
| } | ||
|
|
||
| #[test] | ||
| fn test_single_pattern() { | ||
| let r = Redactor::new(["secret".to_string()]); | ||
| assert_eq!(r.redact("my secret value"), "my [redacted] value"); | ||
| } | ||
|
|
||
| #[test] | ||
| fn test_multiple_patterns() { | ||
| let r = Redactor::new(["secret".to_string(), "password".to_string()]); | ||
| assert_eq!( | ||
| r.redact("secret and password here"), | ||
| "[redacted] and [redacted] here" | ||
| ); | ||
| } | ||
|
|
||
| #[test] | ||
| fn test_overlapping_patterns() { | ||
| let r = Redactor::new(["abc".to_string(), "bc".to_string()]); | ||
| let result = r.redact("abcd"); | ||
| // Should replace "abc" first, leaving "d" | ||
| assert_eq!(result, "[redacted]d"); | ||
| } | ||
|
|
||
| #[test] | ||
| fn test_multiple_occurrences() { | ||
| let r = Redactor::new(["token".to_string()]); | ||
| assert_eq!(r.redact("token1 and token2"), "[redacted]1 and [redacted]2"); | ||
| } | ||
|
|
||
| #[test] | ||
| fn test_with_additional() { | ||
| let r1 = Redactor::new(["secret".to_string()]); | ||
| let r2 = r1.with_additional(["password".to_string()]); | ||
|
|
||
| assert_eq!(r1.redact("secret password"), "[redacted] password"); | ||
| assert_eq!(r2.redact("secret password"), "[redacted] [redacted]"); | ||
| } | ||
|
|
||
| #[test] | ||
| fn test_empty_patterns_filtered() { | ||
| let r = Redactor::new(["".to_string(), "secret".to_string(), "".to_string()]); | ||
| assert_eq!(r.patterns().len(), 1); | ||
| assert_eq!(r.redact("my secret"), "my [redacted]"); | ||
| } | ||
| } | ||
Uh oh!
There was an error while loading. Please reload this page.