diff --git a/lib/src/scanner/context.rs b/lib/src/scanner/context.rs index 37f0fee92..84de00f6a 100644 --- a/lib/src/scanner/context.rs +++ b/lib/src/scanner/context.rs @@ -278,13 +278,17 @@ impl ScanContext<'_> { /// increased by the time elapsed since `rule_execution_start_time`. #[cfg(feature = "rules-profiling")] pub(crate) fn update_time_spent_in_rule(&mut self, rule_id: RuleId) { - self.time_spent_in_rule - .get_mut::(rule_id.into()) - .unwrap() - .add_assign(self.clock.delta_as_nanos( + // The RuleId is not guaranteed to be a valid one. It may be larger + // than the last RuleId, so we can't assume that the `get_mut` will + // be successful. + if let Some(time_spend_in_rule) = + self.time_spent_in_rule.get_mut::(rule_id.into()) + { + time_spend_in_rule.add_assign(self.clock.delta_as_nanos( self.rule_execution_start_time, self.clock.raw(), )); + } } /// Called during the scan process when a rule didn't match. diff --git a/lib/src/scanner/mod.rs b/lib/src/scanner/mod.rs index 458046e0d..b08c4d693 100644 --- a/lib/src/scanner/mod.rs +++ b/lib/src/scanner/mod.rs @@ -734,13 +734,21 @@ impl<'r> Scanner<'r> { #[cfg(feature = "rules-profiling")] if func_result.is_err() { let ctx = self.wasm_store.data_mut(); - // When a timeout occurs, neither `ctx.track_rule_no_match` nor - // `ctx.track_rule_match` are invoked for the rule that was being - // executed at that moment. This implies that the time spent in - // that rule has not being updated yet, and we must do it here. - // The ID of the rule that was being executed is the one that - // comes after the last executed one. This assumes that rules are - // executed in strict ID increasing order. + // If a timeout occurs, the methods `ctx.track_rule_no_match` or + // `ctx.track_rule_match` may not be invoked for the currently + // executing rule. This means that the time spent within that rule + // has not been recorded yet, so we need to update it here. + // + // The ID of the rule that was running during the timeout can be + // determined as the one immediately following the last executed + // rule, based on the assumption that rules are processed in a + // strictly ascending ID order. + // + // Additionally, if the timeout happens after `ctx.last_executed_rule` + // has been updated with the last rule ID, we might end up calling + // `update_time_spent_in_rule` with an ID that is off by one. + // However, this function is designed to handle such cases + // gracefully. ctx.update_time_spent_in_rule( ctx.last_executed_rule .map_or(RuleId::from(0), |rule_id| rule_id.next()),