[pull] main from VirusTotal:main

Cargo.toml

@@ -112,6 +112,7 @@ yara-x-macros = { path = "macros", version = "1.18.0" }
 yara-x-parser = { path = "parser", version = "1.18.0" }
 yara-x-proto = { path = "proto", version = "1.18.0"}
 zip = { version = "8.2.0", default-features = false }
+tinyzip = "0.4.0"
 simd-adler32 = "0.3.9"
 simd_cesu8 = "1.1.1"
 assert-call = "0.1.2"

docs/ModuleDeveloperGuide.md

@@ -313,7 +313,7 @@ So, let's create our `lib/src/modules/text.rs` file:
 use crate::mods::prelude::*;
 use crate::modules::protos::text::*;
 
-fn main(data: &[u8], _meta: Option<&[u8]>) -> Result<Text, ModuleError> {
+fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result<Text, ModuleError> {
     let mut text_proto = Text::new();
 
     // TODO: parse the data and populate text_proto.
@@ -353,16 +353,16 @@ will be `crate::modules::protos::foobar`
 Next comes the module's main function and the module registration:
 
 ```rust
-fn main(data: &[u8], _meta: Option<&[u8]>) -> Result<Text, ModuleError> {
+fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result<Text, ModuleError> {
     ...
 }
 
 register_module!("text", Text, main);
 ```
 
 The module's main function is called for every file scanned by YARA. This
-function receives a byte slice with the content of the file being scanned and an
-optional byte slice with per-scan metadata, and it returns a `Result` containing the
+function receives a mutable reference to a `ModuleContext` and a byte slice with
+the content of the file being scanned. It returns a `Result` containing the
 `Text` structure that was generated from the `text.proto` file (or a `ModuleError`).
 
 Registering the module is as simple as calling the `register_module!` macro.
@@ -379,7 +379,7 @@ use crate::modules::protos::text::*;
 use std::io;
 use std::io::BufRead;
 
-fn main(data: &[u8], _meta: Option<&[u8]>) -> Result<Text, ModuleError> {
+fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result<Text, ModuleError> {
     // Create an empty instance of the Text protobuf.
     let mut text_proto = Text::new();
 
@@ -1093,11 +1093,11 @@ contract as in built-in modules—it receives the scanned data, populates the
 protobuf, and returns it:
 
 ```rust
-use yara_x::errors::ModuleError;
+use yara_x::mods::prelude::*;
 
 fn foobar_main(
+    _ctx: &mut ModuleContext,
     data: &[u8],
-    _meta: Option<&[u8]>,
 ) -> Result<Foobar, ModuleError> {
     let mut out = Foobar::new();
     out.count = Some(data.len() as u64);

examples/custom-module/src/lib.rs

@@ -19,8 +19,8 @@ pub mod proto {
 pub use proto::foobar::Foobar;
 
 fn foobar_main(
+    _ctx: &mut ModuleContext,
     data: &[u8],
-    _meta: Option<&[u8]>,
 ) -> Result<Foobar, ModuleError> {
     let mut out = Foobar::new();
     out.count = Some(data.len() as u64);

lib/Cargo.toml

@@ -234,6 +234,9 @@ vt-module = [
     "dep:psl",
 ]
 
+# The `zip` module provides container file extraction capabilities.
+zip-module = ["dep:tinyzip"]
+
 # Enables all the default modules.
 default-modules = [
     "console-module",
@@ -252,6 +255,7 @@ default-modules = [
     "test_proto2-module",
     "test_proto3-module",
     "vt-module",
+    "zip-module",
 ]
 
 # Features that are enabled by default.
@@ -317,6 +321,7 @@ x509-parser = { workspace = true, optional = true }
 yara-x-macros = { workspace = true }
 yara-x-parser = { workspace = true, features = ["serde"] }
 zip = { workspace = true, optional = true, features = ["deflate"] }
+tinyzip = { workspace = true, optional = true }
 simd-adler32 = { workspace = true, optional = true }
 simd_cesu8 = { workspace = true, optional = true }
 

lib/src/compiler/context.rs

@@ -127,8 +127,7 @@ impl<'src> CompileContext<'_, 'src> {
             // If the current symbol table is `None` it means that the
             // identifier is not a field or method of some structure.
             return if symbol_table.is_none() {
-                let module = crate::modules::registered_modules()
-                    .find(|module| module.name() == ident.name);
+                let module = crate::modules::module_by_name(ident.name);
                 // Build the error for the unknown identifier.
                 let mut err = UnknownIdentifier::build(
                     self.report_builder,

lib/src/compiler/mod.rs

@@ -1913,8 +1913,7 @@ impl Compiler<'_> {
 
     fn c_import(&mut self, import: &Import) -> Result<(), CompileError> {
         let module_name = import.module_name;
-        let module = crate::modules::registered_modules()
-            .find(|m| m.name() == module_name);
+        let module = crate::modules::module_by_name(module_name);
 
         // Does a module with the given name actually exist? ...
         if module.is_none() {

lib/src/modules/console.rs

@@ -3,7 +3,7 @@ use std::borrow::Cow;
 use crate::mods::prelude::*;
 use crate::modules::protos::console::*;
 
-fn main(_data: &[u8], _meta: Option<&[u8]>) -> Result<Console, ModuleError> {
+fn main(_ctx: &mut ModuleContext, _data: &[u8]) -> Result<Console, ModuleError> {
     // Nothing to do, but we have to return our protobuf
     Ok(Console::new())
 }

lib/src/modules/crx/mod.rs

@@ -18,7 +18,7 @@ thread_local!(
         const { RefCell::new(None) };
 );
 
-fn main(data: &[u8], _meta: Option<&[u8]>) -> Result<Crx, ModuleError> {
+fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result<Crx, ModuleError> {
     PERMHASH_CACHE.with(|cache| *cache.borrow_mut() = None);
     match parser::Crx::parse(data) {
         Ok(crx) => Ok(crx.into()),

lib/src/modules/cuckoo/mod.rs

@@ -22,8 +22,11 @@ fn set_local(value: schema::CuckooJson) {
     });
 }
 
-fn main(_data: &[u8], meta: Option<&[u8]>) -> Result<Cuckoo, ModuleError> {
-    let meta = match meta {
+fn main(
+    _ctx: &mut ModuleContext,
+    _data: &[u8],
+) -> Result<Cuckoo, ModuleError> {
+    let meta = match _ctx.get_module_metadata("cuckoo") {
         None | Some([]) => {
             set_local(schema::CuckooJson::default());
             return Ok(Cuckoo::new());

lib/src/modules/dex/mod.rs

@@ -20,7 +20,7 @@ thread_local!(
         const { RefCell::new(None) };
 );
 
-fn main(data: &[u8], _meta: Option<&[u8]>) -> Result<Dex, ModuleError> {
+fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result<Dex, ModuleError> {
     CHECKSUM_CACHE.with(|cache| *cache.borrow_mut() = None);
     SIGNATURE_CACHE.with(|cache| *cache.borrow_mut() = None);
 

lib/src/modules/dotnet/mod.rs

@@ -7,7 +7,7 @@ use crate::modules::protos::dotnet::*;
 
 pub mod parser;
 
-fn main(data: &[u8], _meta: Option<&[u8]>) -> Result<Dotnet, ModuleError> {
+fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result<Dotnet, ModuleError> {
     match parser::Dotnet::parse(data) {
         Ok(dotnet) => Ok(dotnet.into()),
         Err(_) => {

lib/src/modules/elf/mod.rs

@@ -26,7 +26,7 @@ thread_local!(
     static TLSH_CACHE: RefCell<Option<String>> = const { RefCell::new(None) };
 );
 
-fn main(data: &[u8], _meta: Option<&[u8]>) -> Result<ELF, ModuleError> {
+fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result<ELF, ModuleError> {
     IMPORT_MD5_CACHE.with(|cache| *cache.borrow_mut() = None);
     TLSH_CACHE.with(|cache| *cache.borrow_mut() = None);
 

lib/src/modules/hash/mod.rs

@@ -29,7 +29,7 @@ thread_local!(
         RefCell::new(FxHashMap::default());
 );
 
-fn main(_data: &[u8], _meta: Option<&[u8]>) -> Result<Hash, ModuleError> {
+fn main(_ctx: &mut ModuleContext, _data: &[u8]) -> Result<Hash, ModuleError> {
     // With every scanned file the cache must be cleared.
     SHA256_CACHE.with(|cache| cache.borrow_mut().clear());
     SHA1_CACHE.with(|cache| cache.borrow_mut().clear());

lib/src/modules/lnk/mod.rs

@@ -15,7 +15,7 @@ use crate::mods::prelude::*;
 use crate::modules::protos::lnk::*;
 pub mod parser;
 
-fn main(data: &[u8], _meta: Option<&[u8]>) -> Result<Lnk, ModuleError> {
+fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result<Lnk, ModuleError> {
     match parser::LnkParser::new().parse(data) {
         Ok(lnk) => Ok(lnk),
         Err(_) => {

lib/src/modules/macho/mod.rs

@@ -593,7 +593,7 @@ fn symhash(ctx: &mut ScanContext) -> Option<Lowercase<FixedLenString<32>>> {
     Some(Lowercase::<FixedLenString<32>>::new(digest))
 }
 
-fn main(data: &[u8], _meta: Option<&[u8]>) -> Result<Macho, ModuleError> {
+fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result<Macho, ModuleError> {
     DYLIB_MD5_CACHE.with(|cache| *cache.borrow_mut() = None);
     ENTITLEMENT_MD5_CACHE.with(|cache| *cache.borrow_mut() = None);
     EXPORT_MD5_CACHE.with(|cache| *cache.borrow_mut() = None);

lib/src/modules/magic/mod.rs

@@ -27,7 +27,7 @@ thread_local! {
     static MIME_TYPE_CACHE: RefCell<Option<String>> = const { RefCell::new(None) };
 }
 
-fn main(_data: &[u8], _meta: Option<&[u8]>) -> Result<Magic, ModuleError> {
+fn main(_ctx: &mut ModuleContext, _data: &[u8]) -> Result<Magic, ModuleError> {
     // With every scanned file the cache must be cleared.
     TYPE_CACHE.set(None);
     MIME_TYPE_CACHE.set(None);

lib/src/modules/math.rs

@@ -19,7 +19,7 @@ thread_local!(
     > = RefCell::new(FxHashMap::default());
 );
 
-fn main(_data: &[u8], _meta: Option<&[u8]>) -> Result<Math, ModuleError> {
+fn main(_ctx: &mut ModuleContext, _data: &[u8]) -> Result<Math, ModuleError> {
     DISTRIBUTION_CACHE.with(|cache| cache.borrow_mut().clear());
 
     // Nothing to do, but we have to return our protobuf

lib/src/modules/mod.rs

@@ -1,6 +1,6 @@
 use protobuf::MessageDyn;
 use protobuf::reflect::MessageDescriptor;
-
+use rustc_hash::FxHashMap;
 use thiserror::Error;
 
 pub mod protos {
@@ -36,6 +36,29 @@ pub enum ModuleError {
     },
 }
 
+/// Context passed to the main function of YARA modules.
+#[derive(Default)]
+pub struct ModuleContext<'a> {
+    module_metadata: FxHashMap<&'static str, &'a [u8]>,
+}
+
+impl<'a> ModuleContext<'a> {
+    /// Set the metadata associated to the module with the given name.
+    pub fn set_module_metadata(
+        &mut self,
+        module_name: &'static str,
+        metadata: &'a [u8],
+    ) {
+        self.module_metadata.insert(module_name, metadata);
+    }
+
+    /// Returns the metadata explicitly provided to the module with the given
+    /// name, if any.
+    pub fn get_module_metadata(&self, module_name: &str) -> Option<&[u8]> {
+        self.module_metadata.get(module_name).copied()
+    }
+}
+
 /// The trait implemented by all registered modules.
 pub trait RegisteredModule: Send + Sync {
     /// Name used for the module in `import` statements (e.g. `"my_module"`).
@@ -47,10 +70,10 @@ pub trait RegisteredModule: Send + Sync {
 
     /// Main function called every time YARA scans some data, before
     /// evaluating the rules. Set to `None` for data-only modules.
-    fn main_fn(
+    fn main_fn<'a>(
         &self,
-        data: &[u8],
-        meta: Option<&[u8]>,
+        ctx: &mut ModuleContext<'a>,
+        data: &'a [u8],
     ) -> Option<Result<Box<dyn MessageDyn>, ModuleError>>;
 
     /// Rust module path of the submodule inside the external crate that
@@ -62,8 +85,8 @@ pub trait RegisteredModule: Send + Sync {
     fn rust_module_name(&self) -> Option<&'static str>;
 }
 
-/// Main function in a YARA module.
-pub type ModuleMainFn<T> = fn(&[u8], Option<&[u8]>) -> Result<T, ModuleError>;
+pub type ModuleMainFn<T> =
+    for<'a> fn(&mut ModuleContext<'a>, &'a [u8]) -> Result<T, ModuleError>;
 
 /// Description of a YARA module, generic over the type `T` returned by the
 /// main function.
@@ -93,13 +116,13 @@ where
         T::descriptor()
     }
 
-    fn main_fn(
+    fn main_fn<'a>(
         &self,
-        data: &[u8],
-        meta: Option<&[u8]>,
+        ctx: &mut ModuleContext<'a>,
+        data: &'a [u8],
     ) -> Option<Result<Box<dyn MessageDyn>, ModuleError>> {
         self.main_fn.map(|f| {
-            f(data, meta).map(|ok| Box::new(ok) as Box<dyn MessageDyn>)
+            f(ctx, data).map(|ok| Box::new(ok) as Box<dyn MessageDyn>)
         })
     }
 
@@ -154,6 +177,14 @@ pub(crate) fn registered_modules()
     inventory::iter::<&'static dyn RegisteredModule>().copied()
 }
 
+/// Returns a registered module given its name.
+#[inline]
+pub(crate) fn module_by_name(
+    name: &str,
+) -> Option<&'static dyn RegisteredModule> {
+    registered_modules().find(|m| m.name() == name)
+}
+
 pub mod mods {
     /*! Utility functions and structures that allow invoking YARA modules directly.
 
@@ -297,7 +328,13 @@ pub mod mods {
         let module = super::registered_modules()
             .find(|m| m.root_descriptor().full_name() == proto_name)?;
 
-        module.main_fn(data, meta)?.ok()
+        let mut ctx = super::ModuleContext::default();
+
+        if let Some(m) = meta {
+            ctx.module_metadata.insert(module.name(), m);
+        }
+
+        module.main_fn(&mut ctx, data)?.ok()
     }
 
     /// Invokes all YARA modules and returns the data produced by them.
@@ -333,8 +370,7 @@ pub mod mods {
     /// Returns the definition of the module with the given name.
     pub fn module_definition(name: &str) -> Option<reflect::Struct> {
         use std::rc::Rc;
-        super::registered_modules()
-            .find(|m| m.name() == name)
+        super::module_by_name(name)
             .map(|m| reflect::Struct::new(Rc::<crate::types::Struct>::from(m)))
     }
 
@@ -343,6 +379,7 @@ pub mod mods {
     #[allow(missing_docs)]
     pub mod prelude {
         pub use crate::modules::Module;
+        pub use crate::modules::ModuleContext;
         pub use crate::modules::ModuleError;
         pub use crate::modules::RegisteredModule;
         pub use crate::register_module;

lib/src/modules/modules.rs

@@ -32,4 +32,6 @@ mod test_proto3;
 #[cfg(feature = "time-module")]
 mod time;
 #[cfg(feature = "vt-module")]
-mod vt;
+mod vt;
+#[cfg(feature = "zip-module")]
+mod zip;

lib/src/modules/pe/mod.rs

@@ -35,7 +35,7 @@ thread_local!(
     static CHECKSUM_CACHE: RefCell<Option<i64>> = const { RefCell::new(None) };
 );
 
-fn main(data: &[u8], _meta: Option<&[u8]>) -> Result<PE, ModuleError> {
+fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result<PE, ModuleError> {
     IMPHASH_CACHE.with(|cache| *cache.borrow_mut() = None);
     CHECKSUM_CACHE.with(|cache| *cache.borrow_mut() = None);
 

lib/src/modules/protos/generated/mod.rs

@@ -28,3 +28,4 @@ pub mod time;
 pub mod titan;
 pub mod vtnet;
 pub mod yara;
+pub mod zip;