perf(scanner): lazily materialize imported modules

Author: king-tero

lib/src/modules/mod.rs

@@ -148,6 +148,30 @@ pub(crate) static BUILTIN_MODULES: LazyLock<FxHashMap<&'static str, Module>> =
         modules
     });
 
+pub(crate) fn module_name_from_root_struct(
+    root_struct_name: &str,
+) -> Option<&'static str> {
+    BUILTIN_MODULES.iter().find_map(|(module_name, module)| {
+        (module.root_struct_descriptor.full_name() == root_struct_name)
+            .then_some(*module_name)
+    })
+}
+
+pub(crate) fn module_name_from_rust_module_path(
+    rust_module_path: &str,
+) -> Option<&'static str> {
+    let module_path = rust_module_path.strip_prefix("yara_x::modules::")?;
+    BUILTIN_MODULES.iter().find_map(|(module_name, module)| {
+        module.rust_module_name.and_then(|name| {
+            (module_path == name
+                || module_path
+                    .strip_prefix(name)
+                    .is_some_and(|suffix| suffix.starts_with("::")))
+            .then_some(*module_name)
+        })
+    })
+}
+
 pub mod mods {
     /*! Utility functions and structures that allow invoking YARA modules directly.
 

lib/src/scanner/context.rs

@@ -27,6 +27,7 @@ use crate::compiler::{
     SubPatternAtom, SubPatternFlags, SubPatternId,
 };
 use crate::errors::VariableError;
+use crate::modules::{self, BUILTIN_MODULES};
 use crate::re::Action;
 use crate::re::fast::FastVM;
 use crate::re::hir::ChainedPatternGap;
@@ -120,6 +121,14 @@ pub(crate) struct ScanContext<'r, 'd> {
     /// operation. Keys are the fully qualified protobuf message names, and
     /// values are the protobuf messages set with [`Scanner::set_module_output`].
     pub user_provided_module_outputs: FxHashMap<String, Box<dyn MessageDyn>>,
+    /// Metadata passed to module main functions for the current scan.
+    pub module_metadata: FxHashMap<String, Vec<u8>>,
+    /// Modules that have been materialized during the current scan.
+    pub materialized_modules: FxHashSet<String>,
+    /// Whether imported modules should be materialized on first use.
+    pub lazy_modules: bool,
+    /// Error captured while materializing a module from a field lookup.
+    pub module_materialization_error: Option<ScanError>,
     /// Hash map that tracks the matches occurred during a scan. The keys
     /// are the PatternId of the matching pattern, and values are a list
     /// of matches.
@@ -290,6 +299,98 @@ impl ScanContext<'_, '_> {
     /// use crate::modules::protos::my_module::MyModuleProto;
     /// let module_data: MyModuleProto = ctx.module_data::<MyModuleProto>()
     /// ```
+    pub(crate) fn materialize_module_for_root_field(
+        &mut self,
+        field_index: usize,
+    ) {
+        if self.module_materialization_error.is_some() {
+            return;
+        }
+
+        let module_name =
+            self.root_struct.field_by_index(field_index).and_then(|field| {
+                if let TypeValue::Struct(structure) = &field.type_value {
+                    structure
+                        .protobuf_type_name()
+                        .and_then(modules::module_name_from_root_struct)
+                } else {
+                    None
+                }
+            });
+
+        if let Some(module_name) = module_name
+            && let Err(err) = self.materialize_module(module_name)
+        {
+            self.module_materialization_error = Some(err);
+        }
+    }
+
+    pub(crate) fn materialize_module(
+        &mut self,
+        module_name: &str,
+    ) -> Result<(), ScanError> {
+        if !self.materialized_modules.insert(module_name.to_string()) {
+            return Ok(());
+        }
+
+        let module = BUILTIN_MODULES
+            .get(module_name)
+            .unwrap_or_else(|| panic!("module `{module_name}` not found"));
+
+        let root_struct_name = module.root_struct_descriptor.full_name();
+        let root_struct_name = root_struct_name.to_string();
+
+        let module_output = if let Some(output) =
+            self.user_provided_module_outputs.remove(root_struct_name.as_str())
+        {
+            Some(output)
+        } else if let Some(main_fn) = module.main_fn
+            && let Some(data) = self.scanned_data()
+        {
+            let meta =
+                self.module_metadata.get(module_name).map(Vec::as_slice);
+            Some(main_fn(data, meta).map_err(|err| {
+                ScanError::ModuleError { module: module_name.to_string(), err }
+            })?)
+        } else {
+            None
+        };
+
+        if let Some(module_output) = &module_output {
+            debug_assert_eq!(
+                module_output.descriptor_dyn().full_name(),
+                module.root_struct_descriptor.full_name(),
+                "main function of module `{}` must return `{}`, but returned `{}`",
+                module_name,
+                module.root_struct_descriptor.full_name(),
+                module_output.descriptor_dyn().full_name(),
+            );
+
+            debug_assert!(
+                module_output.is_initialized_dyn(),
+                "module `{}` returned a protobuf `{}` where some required fields are not initialized ",
+                module_name,
+                module.root_struct_descriptor.full_name()
+            );
+        }
+
+        let generate_fields_for_enums = !cfg!(feature = "constant-folding");
+        let module_struct = Struct::from_proto_descriptor_and_msg(
+            &module.root_struct_descriptor,
+            module_output.as_deref(),
+            generate_fields_for_enums,
+        );
+
+        if let Some(module_output) = module_output {
+            self.module_outputs.insert(root_struct_name, module_output);
+        }
+
+        self.root_struct
+            .add_field(module_name, TypeValue::Struct(module_struct));
+
+        Ok(())
+    }
+
     pub(crate) fn module_output<T: MessageFull>(&self) -> Option<&T> {
         let m = self.module_outputs.get(T::descriptor().full_name())?.as_ref();
         <dyn MessageDyn>::downcast_ref(m)
@@ -472,10 +573,15 @@ impl ScanContext<'_, '_> {
         // is running. In that case, the function returns `Ok(0)` but the
         // scan state is updated to `ScanState::Timeout`.
         match eval_result {
-            Ok(0) => match self.scan_state {
-                ScanState::Timeout => Err(ScanError::Timeout),
-                _ => Ok(()),
-            },
+            Ok(0) => {
+                if let Some(err) = self.module_materialization_error.take() {
+                    return Err(err);
+                }
+                match self.scan_state {
+                    ScanState::Timeout => Err(ScanError::Timeout),
+                    _ => Ok(()),
+                }
+            }
             Ok(v) => panic!("WASM main returned: {v}"),
             Err(err) if err.is::<ScanError>() => {
                 Err(err.downcast::<ScanError>().unwrap())
@@ -514,6 +620,9 @@ impl ScanContext<'_, '_> {
 
         // Clear module outputs from previous scans.
         self.module_outputs.clear();
+        self.module_metadata.clear();
+        self.materialized_modules.clear();
+        self.module_materialization_error = None;
 
         // Move the matching rules to the `matching_rules` vector, leaving the
         // `matching_rules_per_ns` map empty.
@@ -1804,6 +1913,10 @@ pub fn create_wasm_store_and_ctx<'r>(
         wasm_pattern_search_done: None,
         module_outputs: FxHashMap::default(),
         user_provided_module_outputs: FxHashMap::default(),
+        module_metadata: FxHashMap::default(),
+        materialized_modules: FxHashSet::default(),
+        lazy_modules: false,
+        module_materialization_error: None,
         pattern_matches: PatternMatches::new(),
         unconfirmed_matches: FxHashMap::default(),
         deadline: 0,

lib/src/scanner/mod.rs

@@ -20,15 +20,14 @@ use memmap2::{Mmap, MmapOptions};
 use protobuf::{CodedInputStream, MessageDyn};
 use thiserror::Error;
 
+use crate::Variable;
 use crate::compiler::{RuleId, Rules};
 use crate::models::Rule;
 use crate::modules::{BUILTIN_MODULES, Module, ModuleError};
 use crate::scanner::context::create_wasm_store_and_ctx;
-use crate::types::{Struct, TypeValue};
 use crate::variables::VariableError;
 use crate::wasm::MATCHING_RULES_BITMAP_BASE;
 use crate::wasm::runtime::Store;
-use crate::{Variable, modules};
 
 pub(crate) use crate::scanner::context::RuntimeObject;
 pub(crate) use crate::scanner::context::RuntimeObjectHandle;
@@ -152,6 +151,7 @@ pub struct ProfilingData<'r> {
 #[derive(Debug, Default)]
 pub struct ScanOptions<'a> {
     module_metadata: HashMap<&'a str, &'a [u8]>,
+    lazy_modules: bool,
 }
 
 impl<'a> ScanOptions<'a> {
@@ -160,7 +160,7 @@ impl<'a> ScanOptions<'a> {
     ///
     /// Use other methods to add additional information.
     pub fn new() -> Self {
-        Self { module_metadata: Default::default() }
+        Self { module_metadata: Default::default(), lazy_modules: false }
     }
 
     /// Adds metadata for a YARA module.
@@ -172,6 +172,18 @@ impl<'a> ScanOptions<'a> {
         self.module_metadata.insert(module_name, metadata);
         self
     }
+
+    /// Enables or disables lazy module execution for this scan.
+    ///
+    /// When enabled, imported modules are executed only if rule condition
+    /// evaluation actually accesses one of their fields or functions. This can
+    /// avoid expensive parsers for rules that short-circuit on strings, but
+    /// modules skipped this way won't appear in [`ScanResults::module_output`]
+    /// or [`ScanResults::module_outputs`].
+    pub fn lazy_modules(mut self, yes: bool) -> Self {
+        self.lazy_modules = yes;
+        self
+    }
 }
 
 /// Scans data with already compiled YARA rules.
@@ -502,111 +514,33 @@ impl<'r> Scanner<'r> {
         // Indicate that the scanner is currently scanning the given data.
         ctx.scan_state = ScanState::ScanningData(data);
 
-        for module_name in ctx.compiled_rules.imports() {
-            // Lookup the module in the list of built-in modules.
-            let module = modules::BUILTIN_MODULES
-                .get(module_name)
-                .unwrap_or_else(|| panic!("module `{module_name}` not found"));
-
-            let root_struct_name = module.root_struct_descriptor.full_name();
-
-            let module_output;
-            // If the user already provided some output for the module by
-            // calling `Scanner::set_module_output`, use that output. If not,
-            // call the module's main function (if the module has a main
-            // function) for getting its output.
-            if let Some(output) =
-                ctx.user_provided_module_outputs.remove(root_struct_name)
-            {
-                module_output = Some(output);
-            } else {
-                let meta: Option<&'opts [u8]> =
-                    options.as_ref().and_then(|options| {
-                        options.module_metadata.get(module_name).copied()
-                    });
-
-                if let Some(main_fn) = module.main_fn {
-                    module_output = Some(
-                        main_fn(ctx.scanned_data().unwrap(), meta).map_err(
-                            |err| ScanError::ModuleError {
-                                module: module_name.to_string(),
-                                err,
-                            },
-                        )?,
-                    );
-                } else {
-                    module_output = None;
-                }
-            }
+        ctx.lazy_modules =
+            options.as_ref().is_some_and(|options| options.lazy_modules);
 
-            if let Some(module_output) = &module_output {
-                // Make sure that the module is returning a protobuf message of
-                // the expected type.
-                debug_assert_eq!(
-                    module_output.descriptor_dyn().full_name(),
-                    module.root_struct_descriptor.full_name(),
-                    "main function of module `{}` must return `{}`, but returned `{}`",
-                    module_name,
-                    module.root_struct_descriptor.full_name(),
-                    module_output.descriptor_dyn().full_name(),
-                );
-
-                // Make sure that the module is returning a protobuf message
-                // where all required fields are initialized. This only applies
-                // to proto2, proto3 doesn't have "required" fields, all fields
-                // are optional.
-                debug_assert!(
-                    module_output.is_initialized_dyn(),
-                    "module `{}` returned a protobuf `{}` where some required fields are not initialized ",
-                    module_name,
-                    module.root_struct_descriptor.full_name()
-                );
+        if let Some(options) = options.as_ref() {
+            for (module_name, metadata) in &options.module_metadata {
+                ctx.module_metadata
+                    .insert((*module_name).to_string(), (*metadata).to_vec());
             }
+        }
 
-            // When constant folding is enabled we don't need to generate
-            // structure fields for enums. This is because during the
-            // optimization process symbols like MyEnum.ENUM_ITEM are resolved
-            // to their constant values at compile time. In other words, the
-            // compiler determines that MyEnum.ENUM_ITEM is equal to some value
-            // X, and uses that value in the generated code.
-            //
-            // However, without constant folding, enums are treated as any
-            // other field in a struct, and their values are determined at scan
-            // time. For that reason these fields must be generated for enums
-            // when constant folding is disabled.
-            let generate_fields_for_enums =
-                !cfg!(feature = "constant-folding");
-
-            let module_struct = Struct::from_proto_descriptor_and_msg(
-                &module.root_struct_descriptor,
-                module_output.as_deref(),
-                generate_fields_for_enums,
-            );
-
-            if let Some(module_output) = module_output {
-                ctx.module_outputs
-                    .insert(root_struct_name.to_string(), module_output);
+        if !ctx.lazy_modules {
+            let imported_modules: Vec<String> =
+                ctx.compiled_rules.imports().map(str::to_owned).collect();
+            for module_name in imported_modules {
+                ctx.materialize_module(module_name.as_str())?;
             }
-
-            // The data structure obtained from the module is added to the
-            // root structure. Any data from previous scans will be replaced
-            // with the new data structure.
-            ctx.root_struct
-                .add_field(module_name, TypeValue::Struct(module_struct));
         }
 
-        // The user provided module outputs are not needed anymore. Let's
-        // clear any remaining entry in the hash map (which can happen if
-        // the user has set outputs for modules that are not even imported
-        // by the rules.
-        ctx.user_provided_module_outputs.clear();
-
         // Clear the flag that indicates that the search phase was done.
         ctx.set_pattern_search_done(false);
 
         // Evaluate the conditions of every rule, this will call
         // `ScanContext::search_for_patterns` if necessary.
-        ctx.eval_conditions()?;
+        let eval_result = ctx.eval_conditions();
+        ctx.module_metadata.clear();
+        ctx.user_provided_module_outputs.clear();
+        eval_result?;
 
         let data = match ctx.scan_state.take() {
             ScanState::ScanningData(data) => data,
@@ -698,7 +632,9 @@ impl<'a, 'r> ScanResults<'a, 'r> {
     /// data.
     ///
     /// The result will be `None` if the module doesn't exist or didn't
-    /// produce any output.
+    /// produce any output. When [`ScanOptions::lazy_modules`] is enabled,
+    /// imported modules that were never accessed during condition evaluation
+    /// are omitted as well.
     pub fn module_output(
         &self,
         module_name: &str,
@@ -715,7 +651,9 @@ impl<'a, 'r> ScanResults<'a, 'r> {
     /// Returns an iterator that yields tuples composed of a YARA module name
     /// and the protobuf produced by that module.
     ///
-    /// Only returns the modules that produced some output.
+    /// Only returns the modules that produced some output. When
+    /// [`ScanOptions::lazy_modules`] is enabled, imported modules that were
+    /// never accessed during condition evaluation are not included.
     pub fn module_outputs(&self) -> ModuleOutputs<'a, 'r> {
         ModuleOutputs::new(self.ctx)
     }