Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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<Scope>`) 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
Expand Down
54 changes: 38 additions & 16 deletions src/extractor/smda.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Vec<(crate::rules::features::Feature, u64)>>,
}

impl std::fmt::Debug for Extractor<'_> {
Expand Down Expand Up @@ -175,22 +183,34 @@ impl<'data> super::Extractor for Extractor<'data> {
}

fn extract_global_features(&self) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
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<Vec<(crate::rules::features::Feature, u64)>> {
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<Vec<(crate::rules::features::Feature, u64)>> {
Expand Down Expand Up @@ -455,6 +475,7 @@ impl<'data> Extractor<'data> {
report,
buf: data,
path: path.to_string(),
global_features_cache: once_cell::sync::OnceCell::new(),
})
}

Expand Down Expand Up @@ -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: "<buffer>".to_string(),
global_features_cache: once_cell::sync::OnceCell::new(),
})
}

Expand Down
89 changes: 40 additions & 49 deletions src/rules/features.rs
Original file line number Diff line number Diff line change
Expand Up @@ -413,8 +413,8 @@ impl FunctionNameFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -458,8 +458,8 @@ impl SectionFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -503,8 +503,8 @@ impl ImportFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -548,8 +548,8 @@ impl ExportFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -588,8 +588,8 @@ impl BasicBlockFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand All @@ -614,8 +614,8 @@ impl MnemonicFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -661,8 +661,8 @@ impl OffsetFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -708,11 +708,8 @@ impl OperandOffsetFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -766,8 +763,8 @@ impl NumberFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -813,11 +810,8 @@ impl OperandNumberFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -888,8 +882,8 @@ impl ApiFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -936,8 +930,8 @@ impl PropertyFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -989,8 +983,8 @@ impl MatchedRuleFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -1068,11 +1062,8 @@ impl CharacteristicFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -1124,8 +1115,8 @@ impl StringFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -1536,8 +1527,8 @@ impl ArchFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -1594,8 +1585,8 @@ impl NamespaceFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -1645,8 +1636,8 @@ impl ClassFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -1700,8 +1691,8 @@ impl OsFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down Expand Up @@ -1762,8 +1753,8 @@ impl FormatFeature {
&self,
features: &std::collections::HashMap<Feature, Vec<u64>>,
) -> Result<(bool, Vec<u64>)> {
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![]))
}
Expand Down
Loading