diff --git a/CHANGELOG.md b/CHANGELOG.md index df35aa4..a632df3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,23 @@ This project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html). `u16::from_be_bytes` received the chunk bytes in reverse order, turning every UTF-16BE string into non-ASCII garbage that was then dropped. +### Performance (closes [#18](https://github.com/marirs/capa-rs/issues/18)) + +- **Global features computed once per file instead of once per + instruction** `extract_global_features` paid a full + `goblin::Object::parse` for the OS feature on every call — and it was + called per instruction, per basic block and per function. The result + (OS/arch — constant per file) is now cached in a + `once_cell::sync::OnceCell` on the smda extractor (thread-safe for the + rayon per-function loop). Instruction-level extraction on + `data/mimikatz.exe_` (808 KiB, ~136k instructions) drops from ~8.4 s to + ~0.2 s (**~40×**); `data/Demo64.dll` from ~23 ms to ~3.4 ms. +- **Single hash lookup in `Feature::evaluate`** All 20 feature types used + `contains_key` + indexing, doing two hash-map lookups and cloning + `self` (including its `HashSet`) twice on every hit — the + hottest operation of the rule engine. All sites now do one + `HashMap::get`, matching the pattern `RangeStatement` already used. + ## [0.5.2] — xor-zero number(0), regex /i fast path, rule pre-pruning ### Fixed — feature extraction parity diff --git a/src/extractor/smda.rs b/src/extractor/smda.rs index 8edc668..a98f8ec 100644 --- a/src/extractor/smda.rs +++ b/src/extractor/smda.rs @@ -110,6 +110,14 @@ pub struct Extractor<'a> { report: DisassemblyReport<'a>, buf: &'a [u8], path: String, + /// Cache for `extract_global_features` (issue marirs/capa-rs#18): + /// OS/arch are constant per file, but computing the OS feature + /// costs a full `goblin::Object::parse` of the buffer — and the + /// feature was recomputed on every instruction and basic block. + /// `OnceCell` (not `lazy_static`) because the value is per-file; + /// the `sync` flavour because `find_capabilities` shares the + /// extractor across rayon worker threads. + global_features_cache: once_cell::sync::OnceCell>, } impl std::fmt::Debug for Extractor<'_> { @@ -175,22 +183,34 @@ impl<'data> super::Extractor for Extractor<'data> { } fn extract_global_features(&self) -> Result> { - Ok(vec![ - ( - crate::rules::features::Feature::Os(crate::rules::features::OsFeature::new( - &self.extract_os()?.to_string(), - "", - )?), - 0, - ), - ( - crate::rules::features::Feature::Arch(crate::rules::features::ArchFeature::new( - &self.extract_arch()?.to_string(), - "", - )?), - 0, - ), - ]) + // Issue marirs/capa-rs#18: computed once per extractor, then + // cloned — the uncached version paid a full goblin parse for + // the OS feature on every call (i.e. per instruction). + Ok(self + .global_features_cache + .get_or_try_init(|| -> Result> { + Ok(vec![ + ( + crate::rules::features::Feature::Os( + crate::rules::features::OsFeature::new( + &self.extract_os()?.to_string(), + "", + )?, + ), + 0, + ), + ( + crate::rules::features::Feature::Arch( + crate::rules::features::ArchFeature::new( + &self.extract_arch()?.to_string(), + "", + )?, + ), + 0, + ), + ]) + })? + .clone()) } fn extract_file_features(&self) -> Result> { @@ -455,6 +475,7 @@ impl<'data> Extractor<'data> { report, buf: data, path: path.to_string(), + global_features_cache: once_cell::sync::OnceCell::new(), }) } @@ -483,6 +504,7 @@ impl<'data> Extractor<'data> { // Synthetic path — no real file backs a buffer-mode // extractor. Kept stable so `Debug` output is uniform. path: "".to_string(), + global_features_cache: once_cell::sync::OnceCell::new(), }) } diff --git a/src/rules/features.rs b/src/rules/features.rs index 34bb1fc..99d5411 100755 --- a/src/rules/features.rs +++ b/src/rules/features.rs @@ -413,8 +413,8 @@ impl FunctionNameFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::FunctionName(self.clone())) { - return Ok((true, features[&Feature::FunctionName(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::FunctionName(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -458,8 +458,8 @@ impl SectionFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::Section(self.clone())) { - return Ok((true, features[&Feature::Section(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::Section(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -503,8 +503,8 @@ impl ImportFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::Import(self.clone())) { - return Ok((true, features[&Feature::Import(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::Import(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -548,8 +548,8 @@ impl ExportFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::Export(self.clone())) { - return Ok((true, features[&Feature::Export(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::Export(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -588,8 +588,8 @@ impl BasicBlockFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::BasicBlock(self.clone())) { - return Ok((true, features[&Feature::BasicBlock(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::BasicBlock(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -614,8 +614,8 @@ impl MnemonicFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::Mnemonic(self.clone())) { - return Ok((true, features[&Feature::Mnemonic(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::Mnemonic(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -661,8 +661,8 @@ impl OffsetFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::Offset(self.clone())) { - return Ok((true, features[&Feature::Offset(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::Offset(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -708,11 +708,8 @@ impl OperandOffsetFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::OperandOffset(self.clone())) { - return Ok(( - true, - features[&Feature::OperandOffset(self.clone())].clone(), - )); + if let Some(locations) = features.get(&Feature::OperandOffset(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -766,8 +763,8 @@ impl NumberFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::Number(self.clone())) { - return Ok((true, features[&Feature::Number(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::Number(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -813,11 +810,8 @@ impl OperandNumberFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::OperandNumber(self.clone())) { - return Ok(( - true, - features[&Feature::OperandNumber(self.clone())].clone(), - )); + if let Some(locations) = features.get(&Feature::OperandNumber(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -888,8 +882,8 @@ impl ApiFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::Api(self.clone())) { - return Ok((true, features[&Feature::Api(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::Api(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -936,8 +930,8 @@ impl PropertyFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::Property(self.clone())) { - return Ok((true, features[&Feature::Property(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::Property(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -989,8 +983,8 @@ impl MatchedRuleFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::MatchedRule(self.clone())) { - return Ok((true, features[&Feature::MatchedRule(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::MatchedRule(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -1068,11 +1062,8 @@ impl CharacteristicFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::Characteristic(self.clone())) { - return Ok(( - true, - features[&Feature::Characteristic(self.clone())].clone(), - )); + if let Some(locations) = features.get(&Feature::Characteristic(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -1124,8 +1115,8 @@ impl StringFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::String(self.clone())) { - return Ok((true, features[&Feature::String(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::String(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -1536,8 +1527,8 @@ impl ArchFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::Arch(self.clone())) { - return Ok((true, features[&Feature::Arch(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::Arch(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -1594,8 +1585,8 @@ impl NamespaceFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::Namespace(self.clone())) { - return Ok((true, features[&Feature::Namespace(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::Namespace(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -1645,8 +1636,8 @@ impl ClassFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::Class(self.clone())) { - return Ok((true, features[&Feature::Class(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::Class(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -1700,8 +1691,8 @@ impl OsFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::Os(self.clone())) { - return Ok((true, features[&Feature::Os(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::Os(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) } @@ -1762,8 +1753,8 @@ impl FormatFeature { &self, features: &std::collections::HashMap>, ) -> Result<(bool, Vec)> { - if features.contains_key(&Feature::Format(self.clone())) { - return Ok((true, features[&Feature::Format(self.clone())].clone())); + if let Some(locations) = features.get(&Feature::Format(self.clone())) { + return Ok((true, locations.clone())); } Ok((false, vec![])) }