yara_x/scanner/

mod.rs

/*! This module implements the YARA scanner.

The scanner takes the rules produces by the compiler and scans data with them.
*/
use std::collections::{BTreeMap, HashMap};
use std::fmt::{Debug, Formatter};
use std::fs;
use std::io::Read;
use std::mem::transmute;
use std::ops::Range;
use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::slice::Iter;
use std::sync::Once;
use std::sync::atomic::AtomicU64;
use std::time::Duration;

use bitvec::prelude::*;
#[cfg(unix)]
use memmap2::Advice;
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::{
    ModuleContext, ModuleError, RegisteredModule, module_by_name,
};
pub(crate) use crate::scanner::context::RuntimeObject;
pub(crate) use crate::scanner::context::RuntimeObjectHandle;
pub(crate) use crate::scanner::context::ScanContext;
pub(crate) use crate::scanner::context::ScanState;
use crate::scanner::context::create_wasm_store_and_ctx;
pub(crate) use crate::scanner::matches::Match;
use crate::types::{Struct, TypeValue};
use crate::variables::VariableError;
use crate::wasm::MATCHING_RULES_BITMAP_BASE;
use crate::wasm::runtime::Store;

mod context;
mod matches;

pub mod blocks;

#[cfg(test)]
mod tests;

/// Error returned when a scan operation fails.
#[derive(Error, Debug)]
#[non_exhaustive]
pub enum ScanError {
    /// The scan was aborted after the timeout period.
    #[error("timeout")]
    Timeout,
    /// Could not open the scanned file.
    #[error("can not open `{path}`: {err}")]
    OpenError {
        /// Path of the file being scanned.
        path: PathBuf,
        /// Error that occurred.
        err: std::io::Error,
    },
    /// Could not map the scanned file into memory.
    #[error("can not map `{path}`: {err}")]
    MapError {
        /// Path of the file being scanned.
        path: PathBuf,
        /// Error that occurred.
        err: std::io::Error,
    },
    /// Could not deserialize the protobuf message for some YARA module.
    #[error(
        "can not deserialize protobuf message for YARA module `{module}`: {err}"
    )]
    ProtoError {
        /// Module name.
        module: String,
        /// Error that occurred.
        err: protobuf::Error,
    },
    /// The module is unknown.
    #[error("unknown module `{module}`")]
    UnknownModule {
        /// Module name.
        module: String,
    },
    /// Some module produced an error when it was invoked.
    #[error("error in module `{module}`: {err}")]
    ModuleError {
        /// Module name.
        module: String,
        /// Error that occurred.
        err: ModuleError,
    },
}

/// Global counter that gets incremented every 1 second by a dedicated thread.
///
/// This counter is used for determining when a scan operation has timed out.
static HEARTBEAT_COUNTER: AtomicU64 = AtomicU64::new(0);

/// Used for spawning the thread that increments `HEARTBEAT_COUNTER`.
static INIT_HEARTBEAT: Once = Once::new();

/// Represents the data being scanned.
///
/// The scanned data can be backed by a slice owned by someone else, or a
/// vector or memory-mapped file owned by `ScannedData` itself.
pub enum ScannedData<'d> {
    Slice(&'d [u8]),
    Vec(Vec<u8>),
    Mmap { mmap: Mmap, len: usize },
}

impl AsRef<[u8]> for ScannedData<'_> {
    fn as_ref(&self) -> &[u8] {
        match self {
            ScannedData::Slice(s) => s,
            ScannedData::Vec(v) => v.as_ref(),
            ScannedData::Mmap { mmap, len } => &mmap.as_ref()[..*len],
        }
    }
}

impl ScannedData<'_> {
    #[inline]
    fn len(&self) -> usize {
        self.as_ref().len()
    }
}

impl<'d> TryInto<ScannedData<'d>> for &'d [u8] {
    type Error = ScanError;
    fn try_into(self) -> Result<ScannedData<'d>, Self::Error> {
        Ok(ScannedData::Slice(self))
    }
}

impl<'d, const N: usize> TryInto<ScannedData<'d>> for &'d [u8; N] {
    type Error = ScanError;
    fn try_into(self) -> Result<ScannedData<'d>, Self::Error> {
        Ok(ScannedData::Slice(self))
    }
}

/// Contains information about the time spent on a rule.
#[cfg(feature = "rules-profiling")]
pub struct ProfilingData<'r> {
    /// Rule namespace.
    pub namespace: &'r str,
    /// Rule name.
    pub rule: &'r str,
    /// Time spent executing the rule's condition.
    pub condition_exec_time: Duration,
    /// Time spent matching the rule's patterns.
    pub pattern_matching_time: Duration,
}

/// Optional information for the scan operation.
#[derive(Debug, Default)]
pub struct ScanOptions<'a> {
    module_metadata: HashMap<&'a str, &'a [u8]>,
}

impl<'a> ScanOptions<'a> {
    /// Creates a new instance of `ScanOptions` with no additional information
    /// for the scan operation.
    ///
    /// Use other methods to add additional information.
    pub fn new() -> Self {
        Self { module_metadata: Default::default() }
    }

    /// Adds metadata for a YARA module.
    pub fn set_module_metadata(
        mut self,
        module_name: &'a str,
        metadata: &'a [u8],
    ) -> Self {
        self.module_metadata.insert(module_name, metadata);
        self
    }
}

/// Scans data with already compiled YARA rules.
///
/// The scanner receives a set of compiled [`Rules`] and scans data with those
/// rules. The same scanner can be used for scanning multiple files or
/// in-memory data sequentially, but you need multiple scanners for scanning in
/// parallel.
pub struct Scanner<'r> {
    _rules: &'r Rules,
    wasm_store: Pin<Box<Store<ScanContext<'static, 'static>>>>,
    use_mmap: bool,
    max_scan_size: Option<usize>,
}

impl<'r> Scanner<'r> {
    /// Creates a new scanner.
    pub fn new(rules: &'r Rules) -> Self {
        let wasm_store = create_wasm_store_and_ctx(rules);
        Self { _rules: rules, wasm_store, use_mmap: true, max_scan_size: None }
    }

    /// Sets a timeout for scan operations.
    ///
    /// The scan functions will return an [ScanError::Timeout] once the
    /// provided timeout duration has elapsed. The scanner will make every
    /// effort to stop promptly after the designated timeout duration. However,
    /// in some cases, particularly with rules containing only a few patterns,
    /// the scanner could potentially continue running for a longer period than
    /// the specified timeout.
    pub fn set_timeout(&mut self, timeout: Duration) -> &mut Self {
        self.scan_context_mut().set_timeout(timeout);
        self
    }

    /// Sets the maximum number of matches per pattern.
    ///
    /// When some pattern reaches the maximum number of patterns it won't
    /// produce more matches.
    pub fn max_matches_per_pattern(&mut self, n: usize) -> &mut Self {
        self.scan_context_mut()
            .tracker
            .pattern_matches
            .max_matches_per_pattern(n);
        self
    }

    /// Enables or disables fast scan mode.
    ///
    /// During rule compilation, the compiler analyzes rule conditions to
    /// identify patterns that are only ever used in simple boolean existence
    /// checks (e.g., `$a` in YARA). If a pattern is never queried for its match
    /// count (`#a`), specific match offset (`@a`), match length (`!a`), or
    /// evaluated inside a loop, it is classified as a fast-scan pattern.
    ///
    /// In fast scan mode, the scanner optimizes scans by stopping the search
    /// and match tracking for these fast-scan patterns once their **very first
    /// match** is found. Subsequent occurrences in the input data are ignored,
    /// preventing redundant Aho-Corasick scans, regex evaluations, and match
    /// memory allocations.
    ///
    /// Note that using fast scan mode implies that not all matches will be
    /// reported. For instance, when iterating matches using [`ScanResults`],
    /// you won't get all occurrences of the pattern in the file, only the first
    /// one.
    ///
    /// ### Example
    ///
    /// ```
    /// # use yara_x::{Compiler, Scanner};
    /// let mut compiler = Compiler::new();
    /// compiler.add_source(r#"
    ///     rule test {
    ///         strings:
    ///             $a = "abc"
    ///         condition:
    ///             $a
    ///     }
    /// "#).unwrap();
    ///
    /// let rules = compiler.build();
    /// let mut scanner = Scanner::new(&rules);
    ///
    /// // Enable fast scan mode.
    /// scanner.fast_scan(true);
    ///
    /// // The haystack contains two matches of "abc".
    /// let results = scanner.scan(b"abc...abc").unwrap();
    ///
    /// // Find the matching rule.
    /// let matching_rule = results.matching_rules().next().unwrap();
    ///
    /// // Only a single match is returned for pattern $a.
    /// let pattern = matching_rule.patterns().next().unwrap();
    /// let mut matches = pattern.matches();
    /// assert_eq!(matches.next().unwrap().range().start, 0); // The first match
    /// assert!(matches.next().is_none()); // No other matches are returned
    /// ```
    pub fn fast_scan(&mut self, yes: bool) -> &mut Self {
        self.scan_context_mut().tracker.fast_scan = yes;
        self
    }

    /// Specifies whether [`Scanner::scan_file`] and [`Scanner::scan_file_with_options`]
    /// may use memory-mapped files to read input.
    ///
    /// By default, the scanner uses memory mapping for very large files, as this
    /// is typically faster than copying file contents into memory. However, this
    /// approach has a drawback: if another process truncates the file during
    /// scanning, a `SIGBUS` signal may occur.
    ///
    /// Setting this option disables memory mapping and forces the scanner to
    /// always read files into an in-memory buffer instead. This method is slower,
    /// but safer.
    pub fn use_mmap(&mut self, yes: bool) -> &mut Self {
        self.use_mmap = yes;
        self
    }

    /// Sets the maximum size of the data that will be scanned.
    ///
    /// If the scanned data (either a file or an in-memory buffer) is larger
    /// than this value, it will be truncated to the given size.
    ///
    /// The value returned by `filesize` will be also limited to the given
    /// size.
    ///
    /// Also notice that some modules (pe, elf, macho, etc) may be unable
    /// to properly parse truncated files.
    pub fn max_scan_size(&mut self, size: usize) -> &mut Self {
        self.max_scan_size = Some(size);
        self
    }

    /// Sets a callback that is invoked every time a YARA rule calls the
    /// `console` module.
    ///
    /// The `callback` function is invoked with a string representing the
    /// message being logged. The function can print the message to stdout,
    /// append it to a file, etc. If no callback is set these messages are
    /// ignored.
    pub fn console_log<F>(&mut self, callback: F) -> &mut Self
    where
        F: FnMut(String) + 'r,
    {
        self.scan_context_mut().console_log = Some(Box::new(callback));
        self
    }

    /// Sets the context size for matches.
    ///
    /// This specifies how many bytes at the left and right of each match will
    /// be reported by [`crate::Match::data_with_context`]. By default, the
    /// match context size is 0, which means that [`crate::Match::data_with_context`]
    /// will return exactly the same data as [`crate::Match::data`].
    pub fn match_context_size(&mut self, size: usize) -> &mut Self {
        self.scan_context_mut().match_context_size = size;
        self
    }

    /// Scans in-memory data.
    pub fn scan<'a>(
        &'a mut self,
        data: &'a [u8],
    ) -> Result<ScanResults<'a, 'r>, ScanError> {
        let mut data = data;
        if let Some(max) = self.max_scan_size
            && data.len() > max
        {
            data = &data[..max];
        }
        self.scan_impl(ScannedData::Slice(data), None)
    }

    /// Scans a file.
    pub fn scan_file<'a, P>(
        &'a mut self,
        target: P,
    ) -> Result<ScanResults<'a, 'r>, ScanError>
    where
        P: AsRef<Path>,
    {
        self.scan_impl(self.load_file(target.as_ref())?, None)
    }

    /// Like [`Scanner::scan`], but allows to specify additional scan options.
    pub fn scan_with_options<'a, 'opts>(
        &'a mut self,
        data: &'a [u8],
        options: ScanOptions<'opts>,
    ) -> Result<ScanResults<'a, 'r>, ScanError> {
        let mut data = data;
        if let Some(max) = self.max_scan_size
            && data.len() > max
        {
            data = &data[..max];
        }
        self.scan_impl(ScannedData::Slice(data), Some(options))
    }

    /// Like [`Scanner::scan_file`], but allows to specify additional scan
    /// options.
    pub fn scan_file_with_options<'opts, P>(
        &mut self,
        target: P,
        options: ScanOptions<'opts>,
    ) -> Result<ScanResults<'_, 'r>, ScanError>
    where
        P: AsRef<Path>,
    {
        self.scan_impl(self.load_file(target.as_ref())?, Some(options))
    }

    /// Sets the value of a global variable.
    ///
    /// The variable must has been previously defined by calling
    /// [`crate::Compiler::define_global`], and the type it has during the
    /// definition must match the type of the new value (`T`).
    ///
    /// The variable will retain the new value in subsequent scans, unless this
    /// function is called again for setting a new value.
    pub fn set_global<T: TryInto<Variable>>(
        &mut self,
        ident: &str,
        value: T,
    ) -> Result<&mut Self, VariableError>
    where
        VariableError: From<<T as TryInto<Variable>>::Error>,
    {
        self.scan_context_mut().set_global(ident, value)?;
        Ok(self)
    }

    /// Sets the output data for a YARA module.
    ///
    /// Each YARA module generates an output consisting of a data structure that
    /// contains information about the scanned file. This data structure is
    /// represented by a Protocol Buffer message. Typically, you won't need to
    /// provide this data yourself, as the YARA module automatically generates
    /// different outputs for each file it scans.
    ///
    /// However, there are two scenarios in which you may want to provide the
    /// output for a module yourself:
    ///
    /// 1) When the module does not produce any output on its own.
    /// 2) When you already know the output of the module for the upcoming file
    ///    to be scanned, and you prefer to reuse this data instead of generating
    ///    it again.
    ///
    /// Case 1) applies to certain modules lacking a main function, thus
    /// incapable of producing any output on their own. For such modules, you
    /// must set the output before scanning the associated data. Since the
    /// module's output typically varies with each scanned file, you need to
    /// call [`Scanner::set_module_output`] prior to each invocation of
    /// [`Scanner::scan`]. Once [`Scanner::scan`] is executed, the module's
    /// output is consumed and will be empty unless set again before the
    /// subsequent call.
    ///
    /// Case 2) applies when you have previously stored the module's output for
    /// certain scanned data. In such cases, when rescanning the data, you can
    /// utilize this function to supply the module's output, thereby preventing
    /// redundant computation by the module. This optimization enhances
    /// performance by eliminating the need for the module to reparse the
    /// scanned data.
    ///
    /// <br>
    ///
    /// The `data` argument must be a Protocol Buffer message corresponding
    /// to any of the existing YARA modules.
    pub fn set_module_output(
        &mut self,
        data: Box<dyn MessageDyn>,
    ) -> Result<&mut Self, ScanError> {
        let descriptor = data.descriptor_dyn();
        let full_name = descriptor.full_name();

        // Check if the protobuf message passed to this function corresponds
        // with any of the existing modules.
        if !crate::modules::registered_modules()
            .any(|m| m.root_descriptor().full_name() == full_name)
        {
            return Err(ScanError::UnknownModule {
                module: full_name.to_string(),
            });
        }

        self.scan_context_mut()
            .user_provided_module_outputs
            .insert(full_name.to_string(), data);

        Ok(self)
    }

    /// Similar to [`Scanner::set_module_output`], but receives a module name
    /// and the protobuf message as raw data.
    ///
    /// `name` can be either the YARA module name (i.e: "pe", "elf", "dotnet",
    /// etc.) or the fully-qualified name for the protobuf message associated
    /// to the module (i.e: "pe.PE", "elf.ELF", "dotnet.Dotnet", etc.).
    pub fn set_module_output_raw(
        &mut self,
        name: &str,
        data: &[u8],
    ) -> Result<&mut Self, ScanError> {
        // Try to find the module by name first, if not found, then try
        // to find a module where the fully-qualified name for its protobuf
        // message matches the `name` arguments.
        let descriptor = module_by_name(name)
            .map(|module| module.root_descriptor())
            .or_else(|| {
                crate::modules::registered_modules()
                    .find(|module| {
                        module.root_descriptor().full_name() == name
                    })
                    .map(|module| module.root_descriptor())
            });

        if descriptor.is_none() {
            return Err(ScanError::UnknownModule { module: name.to_string() });
        }

        let mut is = CodedInputStream::from_bytes(data);

        // Default recursion limit is 100, that's not enough for some deeply
        // nested structures like the process tree in the `vt` module.
        is.set_recursion_limit(500);

        self.set_module_output(
            descriptor.unwrap().parse_from(&mut is).map_err(|err| {
                ScanError::ProtoError { module: name.to_string(), err }
            })?,
        )
    }

    /// Returns profiling data for the slowest N rules.
    ///
    /// The profiling data reflects the cumulative execution time of each rule
    /// across all scanned files. This information is useful for identifying
    /// performance bottlenecks. To reset the profiling data and start fresh
    /// for subsequent scans, use [`Scanner::clear_profiling_data`].
    #[cfg(feature = "rules-profiling")]
    pub fn slowest_rules(&self, n: usize) -> Vec<ProfilingData<'_>> {
        self.scan_context().slowest_rules(n)
    }

    /// Clears all accumulated profiling data.
    ///
    /// This method resets the profiling data collected during rule execution
    /// across scanned files. Use this to start a new profiling session, ensuring
    /// the results reflect only the data gathered after this method is called.
    #[cfg(feature = "rules-profiling")]
    pub fn clear_profiling_data(&mut self) {
        self.scan_context_mut().clear_profiling_data()
    }
}

impl<'r> Scanner<'r> {
    #[cfg(feature = "rules-profiling")]
    #[inline]
    fn scan_context<'a>(&self) -> &ScanContext<'r, 'a> {
        unsafe {
            transmute::<&ScanContext<'static, 'static>, &ScanContext<'r, '_>>(
                self.wasm_store.data(),
            )
        }
    }
    #[inline]
    fn scan_context_mut<'a>(&mut self) -> &mut ScanContext<'r, 'a> {
        unsafe {
            transmute::<
                &mut ScanContext<'static, 'static>,
                &mut ScanContext<'r, '_>,
            >(self.wasm_store.data_mut())
        }
    }

    fn load_file<'a>(
        &self,
        path: &Path,
    ) -> Result<ScannedData<'a>, ScanError> {
        let file = fs::File::open(path).map_err(|err| {
            ScanError::OpenError { path: path.to_path_buf(), err }
        })?;

        let mut size = file.metadata().map(|m| m.len()).unwrap_or(0) as usize;

        if let Some(max_scan_size) = self.max_scan_size {
            size = std::cmp::min(size, max_scan_size);
        }

        // For files smaller than ~500MB reading the whole file is faster than
        // using a memory-mapped file.
        let data = if self.use_mmap && size > 500_000_000 {
            let mapped_file = unsafe {
                MmapOptions::new().map_copy_read_only(&file).map_err(|err| {
                    ScanError::MapError { path: path.to_path_buf(), err }
                })
            }?;
            #[cfg(unix)]
            mapped_file.advise(Advice::Sequential).map_err(|err| {
                ScanError::MapError { path: path.to_path_buf(), err }
            })?;
            ScannedData::Mmap { mmap: mapped_file, len: size }
        } else {
            let mut buffered_file = Vec::with_capacity(size);
            (&file)
                .take(size as u64)
                .read_to_end(&mut buffered_file)
                .map_err(|err| ScanError::OpenError {
                    path: path.to_path_buf(),
                    err,
                })?;
            ScannedData::Vec(buffered_file)
        };

        Ok(data)
    }

    fn scan_impl<'a, 'opts>(
        &'a mut self,
        data: ScannedData<'a>,
        options: Option<ScanOptions<'opts>>,
    ) -> Result<ScanResults<'a, 'r>, ScanError> {
        let ctx = self.scan_context_mut();

        // Clear information about matches found in a previous scan, if any.
        ctx.reset();

        // Set the global variable `filesize` to the size of the scanned data.
        ctx.set_filesize(data.len() as i64);

        // Create the context that will be passed to the main function of each
        // module.
        let mut mod_ctx = ModuleContext::default();

        // For each item in options.module_metadata, check if the module name
        // is actually a registered module, and then add the corresponding
        // metadata to mod_ctx.
        for (module, meta) in options
            .map(|options| options.module_metadata)
            .into_iter()
            .flatten()
            .filter_map(|(name, meta)| Some((module_by_name(name)?, meta)))
        {
            mod_ctx.set_module_metadata(module.name(), meta);
        }

        // Indicate that the scanner is currently scanning the given data.
        ctx.scan_state = ScanState::ScanningData(data);

        let data = match &ctx.scan_state {
            ScanState::ScanningData(data) => data.as_ref(),
            _ => unreachable!(),
        };

        for module_name in ctx.compiled_rules.imports() {
            // Look up the module in the module registry.
            let module = module_by_name(module_name)
                .unwrap_or_else(|| panic!("module `{module_name}` not found"));

            let module_root_descriptor = module.root_descriptor();
            let root_struct_name = module_root_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 {
                if let Some(main_res) = module.main_fn(&mut mod_ctx, data) {
                    module_output = Some(main_res.map_err(|err| {
                        ScanError::ModuleError {
                            module: module_name.to_string(),
                            err,
                        }
                    })?);
                } else {
                    module_output = None;
                }
            }

            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(),
                    root_struct_name,
                    "main function of module `{}` must return `{}`, but returned `{}`",
                    module_name,
                    root_struct_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,
                    root_struct_name
                );
            }

            // 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_descriptor,
                module_output.as_deref(),
                generate_fields_for_enums,
                false,
            );

            if let Some(module_output) = module_output {
                ctx.module_outputs
                    .insert(root_struct_name.to_string(), module_output);
            }

            // 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 data = match ctx.scan_state.take() {
            ScanState::ScanningData(data) => data,
            _ => unreachable!(),
        };

        ctx.scan_state = ScanState::Finished(DataSnippets::SingleBlock(data));

        Ok(ScanResults::new(ctx))
    }
}

/// Helper type that exposes the data matched during a scan operation.
///
/// Matching data can be accessed through the [`Match::data`] method. Normally,
/// this data can be retrieved by slicing directly into the scanned input.
/// However, that requires the original input to remain valid until the scan
/// results are processed. This works fine for a single contiguous block of
/// memory, but is impractical when scanning multiple blocks, since holding
/// onto all of them until the end would consume excessive memory.
///
/// To handle this, two strategies are used:
///
/// - **Single-block scans**: Data is accessed directly from the input slice.
/// - **Multi-block scans**: Matching fragments are copied and retained in a
///   BTreeMap until the results are processed. The keys in the btree are
///   the offsets where the snippets start and the values are vectors with
///   the snippet's data.
///
/// Each strategy corresponds to a variant in this enum.
pub(crate) enum DataSnippets<'d> {
    SingleBlock(ScannedData<'d>),
    MultiBlock(BTreeMap<usize, Vec<u8>>),
}

impl DataSnippets<'_> {
    pub(crate) fn get(&self, range: Range<usize>) -> Option<&[u8]> {
        self.get_with_context(range, 0).map(|(data, _)| data)
    }

    /// Gets the data for the given `range`, but adding `context_size` additional
    /// bytes to the left and right.
    ///
    /// Returns a tuple where the first item is the data slice with context,
    /// and the second item is a range relative to the slice indicating where
    /// the `range` part is located.
    ///
    /// The result will be `None` only if the data for `range` can't be found.
    /// The additional bytes at the left and right will be added if possible,
    /// but otherwise won't affect the result.
    pub(crate) fn get_with_context(
        &self,
        range: Range<usize>,
        context_size: usize,
    ) -> Option<(&[u8], Range<usize>)> {
        match self {
            Self::SingleBlock(data) => {
                let start = range.start.saturating_sub(context_size);
                let end = range.end.saturating_add(context_size);
                let end = std::cmp::min(end, data.len());

                let slice = data.as_ref().get(start..end)?;
                let rel_start = range.start - start;
                let rel_end = range.end - start;

                Some((slice, rel_start..rel_end))
            }
            Self::MultiBlock(btree) => {
                for (snippet_offset, snippet_data) in
                    btree.range(..=range.start).rev()
                {
                    // Calculate the start and end of the slice within the snippet.
                    let start = range.start.saturating_sub(*snippet_offset);
                    let end = range.end.saturating_sub(*snippet_offset);

                    if end > snippet_data.len() {
                        continue;
                    }

                    let start = start.saturating_sub(context_size);
                    let end = end.saturating_add(context_size);
                    let end = std::cmp::min(end, snippet_data.len());

                    match snippet_data.get(start..end) {
                        Some(data) if !data.is_empty() => {
                            let rel_start =
                                range.start - (*snippet_offset + start);
                            let rel_end =
                                range.end - (*snippet_offset + start);
                            return Some((data, rel_start..rel_end));
                        }
                        _ => continue,
                    }
                }

                None
            }
        }
    }
}

/// Results of a scan operation.
///
/// Allows iterating over both the matching and non-matching rules.
pub struct ScanResults<'a, 'r> {
    ctx: &'a ScanContext<'r, 'a>,
}

impl Debug for ScanResults<'_, '_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.write_str("ScanResults")
    }
}

impl<'a, 'r> ScanResults<'a, 'r> {
    fn new(ctx: &'a ScanContext<'r, 'a>) -> Self {
        Self { ctx }
    }

    /// Returns an iterator that yields the matching rules in arbitrary order.
    pub fn matching_rules(&self) -> MatchingRules<'_, 'r> {
        MatchingRules::new(self.ctx)
    }

    /// Returns an iterator that yields the non-matching rules in arbitrary
    /// order.
    pub fn non_matching_rules(&self) -> NonMatchingRules<'_, 'r> {
        NonMatchingRules::new(self.ctx)
    }

    /// Returns the protobuf produced by a YARA module after processing the
    /// data.
    ///
    /// The result will be `None` if the module doesn't exist or didn't
    /// produce any output.
    pub fn module_output(
        &self,
        module_name: &str,
    ) -> Option<&'a dyn MessageDyn> {
        let module_descriptor =
            module_by_name(module_name).map(|m| m.root_descriptor())?;
        let module_output = self
            .ctx
            .module_outputs
            .get(module_descriptor.full_name())?
            .as_ref();
        Some(module_output)
    }

    /// 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.
    pub fn module_outputs(&self) -> ModuleOutputs<'a, 'r> {
        ModuleOutputs::new(self.ctx)
    }
}

/// Iterator that yields the rules that matched during a scan.
///
/// Private rules are not included by default, use
/// [`MatchingRules::include_private`] for changing this behaviour.
pub struct MatchingRules<'a, 'r> {
    ctx: &'a ScanContext<'r, 'a>,
    iterator: Iter<'a, RuleId>,
    len_non_private: usize,
    len_private: usize,
    include_private: bool,
}

impl<'a, 'r> MatchingRules<'a, 'r> {
    fn new(ctx: &'a ScanContext<'r, 'a>) -> Self {
        Self {
            ctx,
            iterator: ctx.matching_rules.iter(),
            include_private: false,
            len_non_private: ctx.matching_rules.len()
                - ctx.num_matching_private_rules,
            len_private: ctx.num_matching_private_rules,
        }
    }

    /// Specifies whether the iterator should yield private rules.
    ///
    /// This does not reset the iterator to its initial state, the iterator will
    /// continue from its current position.
    pub fn include_private(mut self, yes: bool) -> Self {
        self.include_private = yes;
        self
    }
}

impl<'a, 'r> Iterator for MatchingRules<'a, 'r> {
    type Item = Rule<'a, 'r>;

    fn next(&mut self) -> Option<Self::Item> {
        let rules = self.ctx.compiled_rules;
        loop {
            let rule_id = *self.iterator.next()?;
            let rule_info = rules.get(rule_id);
            if rule_info.is_private {
                self.len_private -= 1;
            } else {
                self.len_non_private -= 1;
            }
            if self.include_private || !rule_info.is_private {
                return Some(Rule { ctx: Some(self.ctx), rule_info, rules });
            }
        }
    }
}

impl ExactSizeIterator for MatchingRules<'_, '_> {
    #[inline]
    fn len(&self) -> usize {
        if self.include_private {
            self.len_non_private + self.len_private
        } else {
            self.len_non_private
        }
    }
}

/// Iterator that yields the rules that didn't match during a scan.
///
/// Private rules are not included by default, use
/// [`NonMatchingRules::include_private`] for changing this behaviour.
pub struct NonMatchingRules<'a, 'r> {
    ctx: &'a ScanContext<'r, 'a>,
    iterator: bitvec::slice::IterZeros<'a, u8, Lsb0>,
    include_private: bool,
    len_private: usize,
    len_non_private: usize,
}

impl<'a, 'r> NonMatchingRules<'a, 'r> {
    fn new(ctx: &'a ScanContext<'r, 'a>) -> Self {
        let num_rules = ctx.compiled_rules.num_rules();
        let main_memory = ctx
            .wasm
            .main_memory
            .unwrap()
            .data(unsafe { ctx.wasm.store.as_ref() });

        let base = MATCHING_RULES_BITMAP_BASE as usize;

        // Create a BitSlice that covers the region of main memory containing
        // the bitmap that tells which rules matched and which did not.
        let matching_rules_bitmap = BitSlice::<_, Lsb0>::from_slice(
            &main_memory[base..base + num_rules / 8 + 1],
        );

        // The BitSlice will cover more bits than necessary, for example, if
        // there are 3 rules the BitSlice will have 8 bits because it is
        // created from a u8 slice that has 1 byte. Here we make sure that
        // the BitSlice has exactly as many bits as existing rules.
        let matching_rules_bitmap = &matching_rules_bitmap[0..num_rules];

        Self {
            ctx,
            iterator: matching_rules_bitmap.iter_zeros(),
            include_private: false,
            len_non_private: ctx.compiled_rules.num_rules()
                - ctx.matching_rules.len()
                - ctx.num_non_matching_private_rules,
            len_private: ctx.num_non_matching_private_rules,
        }
    }

    /// Specifies whether the iterator should yield private rules.
    ///
    /// This does not reset the iterator to its initial state, the iterator will
    /// continue from its current position.
    pub fn include_private(mut self, yes: bool) -> Self {
        self.include_private = yes;
        self
    }
}

impl<'a, 'r> Iterator for NonMatchingRules<'a, 'r> {
    type Item = Rule<'a, 'r>;

    fn next(&mut self) -> Option<Self::Item> {
        let rules = self.ctx.compiled_rules;

        loop {
            let rule_id = RuleId::from(self.iterator.next()?);
            let rule_info = rules.get(rule_id);

            if rule_info.is_private {
                self.len_private -= 1;
            } else {
                self.len_non_private -= 1;
            }

            if self.include_private || !rule_info.is_private {
                return Some(Rule { ctx: Some(self.ctx), rule_info, rules });
            }
        }
    }
}

impl ExactSizeIterator for NonMatchingRules<'_, '_> {
    #[inline]
    fn len(&self) -> usize {
        if self.include_private {
            self.len_non_private + self.len_private
        } else {
            self.len_non_private
        }
    }
}

/// Iterator that returns the outputs produced by YARA modules.
pub struct ModuleOutputs<'a, 'r> {
    ctx: &'a ScanContext<'r, 'a>,
    len: usize,
    iterator: Box<dyn Iterator<Item = &'static dyn RegisteredModule> + 'a>,
}

impl<'a, 'r> ModuleOutputs<'a, 'r> {
    fn new(ctx: &'a ScanContext<'r, 'a>) -> Self {
        Self {
            ctx,
            len: ctx.module_outputs.len(),
            iterator: Box::new(crate::modules::registered_modules()),
        }
    }
}

impl ExactSizeIterator for ModuleOutputs<'_, '_> {
    #[inline]
    fn len(&self) -> usize {
        self.len
    }
}

impl<'a> Iterator for ModuleOutputs<'a, '_> {
    type Item = (&'a str, &'a dyn MessageDyn);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let module = self.iterator.next()?;
            if let Some(module_output) = self
                .ctx
                .module_outputs
                .get(module.root_descriptor().full_name())
            {
                return Some((module.name(), module_output.as_ref()));
            }
        }
    }
}

#[cfg(test)]
mod snippet_tests {
    use super::DataSnippets;
    use std::collections::BTreeMap;

    #[test]
    fn snippets_multiblock() {
        let mut btree_map = BTreeMap::new();

        btree_map.insert(0, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
        btree_map.insert(50, vec![50, 51, 52, 53, 54]);
        btree_map.insert(52, vec![52, 53]);
        btree_map.insert(100, vec![100, 101, 102, 103]);

        let snippets = DataSnippets::MultiBlock(btree_map);

        assert_eq!(snippets.get(0..2), Some([0, 1].as_slice()));
        assert_eq!(snippets.get(1..3), Some([1, 2].as_slice()));
        assert_eq!(snippets.get(8..9), Some([8].as_slice()));
        assert_eq!(snippets.get(10..11), None);
        assert_eq!(snippets.get(50..51), Some([50].as_slice()));
        assert_eq!(snippets.get(51..53), Some([51, 52].as_slice()));
        assert_eq!(snippets.get(50..54), Some([50, 51, 52, 53].as_slice()));
        assert_eq!(snippets.get(52..54), Some([52, 53].as_slice()));
        assert_eq!(snippets.get(52..55), Some([52, 53, 54].as_slice()));
        assert_eq!(snippets.get(52..53), Some([52].as_slice()));
        assert_eq!(snippets.get(50..56), None);
        assert_eq!(snippets.get(100..101), Some([100].as_slice()));
        assert_eq!(snippets.get(101..103), Some([101, 102].as_slice()));

        assert_eq!(
            snippets.get_with_context(0..2, 1),
            Some(([0, 1, 2].as_slice(), 0..2))
        );

        assert_eq!(
            snippets.get_with_context(0..2, 2),
            Some(([0, 1, 2, 3].as_slice(), 0..2))
        );

        assert_eq!(
            snippets.get_with_context(2..4, 2),
            Some(([0, 1, 2, 3, 4, 5].as_slice(), 2..4))
        );

        assert_eq!(
            snippets.get_with_context(51..52, 3),
            Some(([50, 51, 52, 53, 54].as_slice(), 1..2))
        );

        assert_eq!(
            snippets.get_with_context(102..103, 3),
            Some(([100, 101, 102, 103].as_slice(), 2..3))
        );
    }

    #[test]
    fn snippets_singleblock() {
        let data = b"Lorem ipsum dolor sit amet".to_vec();
        let scanned_data = super::ScannedData::Vec(data);
        let snippets = DataSnippets::SingleBlock(scanned_data);

        // Test get
        assert_eq!(snippets.get(6..11), Some(b"ipsum".as_slice()));
        assert_eq!(snippets.get(0..5), Some(b"Lorem".as_slice()));
        assert_eq!(snippets.get(20..26), Some(b"t amet".as_slice()));
        assert_eq!(snippets.get(27..30), None);

        // Test get_with_context
        // context_size = 5
        assert_eq!(
            snippets.get_with_context(6..11, 5),
            Some((b"orem ipsum dolo".as_slice(), 5..10))
        );
        assert_eq!(
            snippets.get_with_context(0..5, 5),
            Some((b"Lorem ipsu".as_slice(), 0..5))
        );
        assert_eq!(
            snippets.get_with_context(20..26, 5),
            Some((b"or sit amet".as_slice(), 5..11))
        );
        assert_eq!(snippets.get_with_context(32..35, 5), None);
    }
}