[pull] main from VirusTotal:main

lib/src/compiler/mod.rs

@@ -5,7 +5,6 @@ module implements the YARA compiler.
 */
 
 use std::cell::RefCell;
-use std::collections::hash_map::Entry;
 use std::collections::{HashMap, HashSet};
 use std::io::Write;
 use std::path::{Path, PathBuf};
@@ -814,7 +813,7 @@ impl<'a> Compiler<'a> {
             wasm_mod,
             compiled_wasm_mod: Some(compiled_wasm_mod),
             relaxed_re_syntax: self.relaxed_re_syntax,
-            ac: None,
+            ac: AhoCorasick::default(),
             num_patterns: self.next_pattern_id.0 as usize,
             ident_pool: self.ident_pool,
             regex_pool: self.regex_pool,
@@ -1757,27 +1756,24 @@ impl Compiler<'_> {
                 num_private_patterns += 1;
             }
 
-            // Check if this pattern has been declared before, in this rule or
-            // in some other rule. In such cases the pattern ID is re-used, and
-            // we don't need to process (i.e: extract atoms and add them to
-            // Aho-Corasick automaton) the pattern again. Two patterns are
-            // considered equal if they are exactly the same, including any
-            // modifiers associated to the pattern, both are non-anchored
-            // or anchored at the same file offset, and if they have the same
-            // file size bounds.
             let pattern_id =
-                match self.patterns.entry(pattern.pattern().clone()) {
-                    // The pattern already exists, return the existing ID.
-                    Entry::Occupied(entry) => *entry.get(),
-                    // The pattern didn't exist.
-                    Entry::Vacant(entry) => {
-                        let pattern_id = self.next_pattern_id;
-                        self.next_pattern_id.incr(1);
-                        self.fast_scan_patterns.push(true);
-                        pending_patterns.insert(pattern_id);
-                        entry.insert(pattern_id);
-                        pattern_id
-                    }
+                // Check if this pattern has been declared before, in this rule or
+                // in some other rule. In such cases the pattern ID is re-used, and
+                // we don't need to process (i.e: extract atoms and add them to
+                // Aho-Corasick automaton) the pattern again. Two patterns are
+                // considered equal if they are exactly the same, including any
+                // modifiers associated to the pattern, both are non-anchored
+                // or anchored at the same file offset, and if they have the same
+                // file size bounds.
+                if let Some(pattern_id) = self.patterns.get(pattern.pattern()) {
+                    *pattern_id
+                } else {
+                    let pattern_id = self.next_pattern_id;
+                    self.next_pattern_id.incr(1);
+                    self.fast_scan_patterns.push(true);
+                    pending_patterns.insert(pattern_id);
+                    self.patterns.insert(pattern.pattern().clone(), pattern_id);
+                    pattern_id
                 };
 
             if !pattern.fast_scan_allowed() {

lib/src/compiler/rules.rs

@@ -33,7 +33,7 @@ const MAGIC: &[u8] = b"YARA-X\0\0";
 ///
 /// This version is incremented every time a change is made to the binary
 /// format in a way that breaks backwards compatibility.
-const SERIALIZATION_VERSION: u32 = 4;
+const SERIALIZATION_VERSION: u32 = 5;
 
 /// Aho-Corasick automaton bundled with an optional Teddy scanner if the
 /// number of patterns is low enough. If the Teddy scanner is present, and
@@ -45,6 +45,75 @@ pub(crate) struct AhoCorasick {
     pub(crate) teddy: Option<teddy::Searcher>,
 }
 
+impl Serialize for AhoCorasick {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        let bytes = self.daachorse.serialize();
+        bytes.serialize(serializer)
+    }
+}
+
+impl<'de> Deserialize<'de> for AhoCorasick {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let bytes = <&'de [u8]>::deserialize(deserializer)?;
+        let (daachorse, _) = DoubleArrayAhoCorasick::deserialize(bytes)
+            .map_err(|e| serde::de::Error::custom(format!("{}", e)))?;
+        Ok(Self { daachorse, teddy: None })
+    }
+}
+
+impl Default for AhoCorasick {
+    fn default() -> Self {
+        let daachorse =
+            DoubleArrayAhoCorasick::new(std::iter::empty::<&[u8]>()).unwrap();
+        Self { daachorse, teddy: None }
+    }
+}
+
+impl AhoCorasick {
+    pub(crate) fn new<'a, I>(patterns: I) -> Self
+    where
+        I: Iterator<Item = &'a [u8]> + ExactSizeIterator + Clone,
+    {
+        let daachorse = DoubleArrayAhoCorasick::new(patterns.clone())
+            .expect("failed to build Aho-Corasick automaton");
+
+        let mut ac = Self { daachorse, teddy: None };
+        ac.rebuild_teddy(patterns);
+        ac
+    }
+
+    pub(crate) fn rebuild_teddy<'a, I>(&mut self, patterns: I)
+    where
+        I: Iterator<Item = &'a [u8]> + ExactSizeIterator,
+    {
+        // Teddy only works from 1 to 64 patterns.
+        let len = patterns.len();
+        if len == 0 || len > 64 {
+            self.teddy = None;
+            return;
+        }
+
+        let mut teddy_builder = teddy::Builder::new();
+
+        for pattern in patterns {
+            // Teddy doesn't admit empty patterns.
+            if pattern.is_empty() {
+                self.teddy = None;
+                return;
+            }
+            teddy_builder.add(pattern);
+        }
+
+        self.teddy = teddy_builder.build();
+    }
+}
+
 /// A set of YARA rules in compiled form.
 ///
 /// This is the result from [`crate::Compiler::build`].
@@ -157,12 +226,8 @@ pub struct Rules {
 
     /// Aho-Corasick automaton containing the atoms extracted from the patterns.
     /// This allows to search for all the atoms in the scanned data at the same
-    /// time in an efficient manner. The automaton is not serialized during when
-    /// [`Rules::serialize`] is called, it needs to be wrapped in [`Option`] so
-    /// that we can use `#[serde(skip)]` on it because [`AhoCorasick`] doesn't
-    /// implement the [`Default`] trait.
-    #[serde(skip)]
-    pub(in crate::compiler) ac: Option<AhoCorasick>,
+    /// time in an efficient manner.
+    pub(in crate::compiler) ac: AhoCorasick,
 
     /// Warnings that were produced while compiling these rules. These warnings
     /// are not serialized, rules that are obtained by deserializing previously
@@ -308,7 +373,7 @@ impl Rules {
             info!("WASM build time: {:?}", Instant::elapsed(&start));
         }
 
-        rules.build_ac_automaton();
+        rules.ac.rebuild_teddy(rules.atoms.iter().map(|x| x.atom.as_ref()));
 
         // Make sure that the maximum SubPatternId is within the boundaries
         // of sub_patterns array. This check is important because during
@@ -516,14 +581,10 @@ impl Rules {
     /// atoms.
     #[inline]
     pub(crate) fn ac_automaton(&self) -> &AhoCorasick {
-        self.ac.as_ref().expect("Aho-Corasick automaton not compiled")
+        &self.ac
     }
 
     pub(crate) fn build_ac_automaton(&mut self) {
-        if self.ac.is_some() {
-            return;
-        }
-
         #[cfg(feature = "logging")]
         let start = Instant::now();
 
@@ -554,26 +615,7 @@ impl Rules {
             }
         }
 
-        // The Teddy algorithm can't be used in all cases. It will be used if:
-        // - The number of atoms is between 1 and 64.
-        // - None of the atoms is empty.
-        let use_teddy = self.atoms.len() <= 64
-            && !self.atoms.is_empty()
-            && !self.atoms.iter().any(|x| x.atom.as_ref().is_empty());
-
-        let teddy_searcher = if use_teddy {
-            let mut teddy_builder = teddy::Builder::new();
-            self.atoms.iter().for_each(|x| teddy_builder.add(x.atom.as_ref()));
-            teddy_builder.build()
-        } else {
-            None
-        };
-
-        let atoms = self.atoms.iter().map(|x| x.atom.as_ref());
-        let ac = DoubleArrayAhoCorasick::new(atoms)
-            .expect("failed to build Aho-Corasick automaton");
-
-        self.ac = Some(AhoCorasick { daachorse: ac, teddy: teddy_searcher });
+        self.ac = AhoCorasick::new(self.atoms.iter().map(|x| x.atom.as_ref()));
 
         #[cfg(feature = "logging")]
         {

lib/src/compiler/tests/mod.rs

@@ -27,7 +27,7 @@ fn serialization() {
     // `DecodeError`.
     let mut data = Vec::new();
     data.extend(b"YARA-X\0\0");
-    data.extend(4u32.to_le_bytes());
+    data.extend(5u32.to_le_bytes());
     data.extend(b"foo");
 
     assert!(matches!(