diff --git a/lib/src/compiler/atoms/mod.rs b/lib/src/compiler/atoms/mod.rs index 904668f0..43fa43a9 100644 --- a/lib/src/compiler/atoms/mod.rs +++ b/lib/src/compiler/atoms/mod.rs @@ -54,6 +54,8 @@ and `"efg"` will be used because `"a"` and `"h"` are too short. mod mask; mod quality; +#[cfg(test)] +mod tests; use std::collections::Bound; use std::iter; @@ -170,29 +172,29 @@ impl Atom { /// # Example /// /// ```ignore - /// let atom = Atom::from_slice_range(&[0x00, 0x01, 0x02, 0x03], 1..=2); + /// let atom = Atom::from_slice_range(&[0x00, 0x01, 0x02, 0x03], 1..=2)?; /// assert_eq!(atom.as_ref(), &[0x01, 0x02]); /// assert_eq!(atom.backtrack, 1) /// assert(!atom.is_exact); /// ``` /// - pub fn from_slice_range(s: &[u8], range: R) -> Self + pub fn from_slice_range(s: &[u8], range: R) -> Option where R: RangeBounds + SliceIndex<[u8], Output = [u8]>, { let backtrack = match range.start_bound() { - Bound::Included(b) => *b as u16, - Bound::Excluded(b) => (*b + 1) as u16, + Bound::Included(b) => u16::try_from(*b).ok()?, + Bound::Excluded(b) => u16::try_from(*b + 1).ok()?, Bound::Unbounded => 0, }; - let atom: &[u8] = &s[range]; + let atom: &[u8] = s.get(range)?; - Self { + Some(Self { bytes: atom.to_smallvec(), backtrack, exact: atom.len() == s.len(), - } + }) } #[inline] @@ -264,25 +266,68 @@ impl Atom { /// Returns a [`CaseCombinations`] iterator that produces all possible case /// combinations of this atom. - pub fn case_combinations(self) -> CaseCombinations { + pub fn case_combinations(&self) -> CaseCombinations { CaseCombinations::new(self) } +} - /// Returns a [`MaskCombinations`] iterator which produces all possible - /// combinations that result from applying a given mask to the atom. +/// An association of an [`Atom`] and a mask of exactly the same length. +/// +/// This type is used to generate all possible atoms that match a masked +/// pattern (e.g. `1? 2?` -> `10 20`, `10 21`, ..., `1f 2f`). +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct MaskedAtom { + pub(crate) atom: Atom, + pub(crate) mask: SmallVec<[u8; DESIRED_ATOM_SIZE]>, +} + +impl MaskedAtom { + /// Creates a [`MaskedAtom`] from a slice range of a pattern and its + /// corresponding mask. /// - /// The mask's length must be identical to the atom's length. In scenarios - /// where the mask is shorter than the atom, the iterator will yield atoms - /// matching the mask's length. If the mask is longer than the atom, it - /// gets truncated to the atom's length. + /// The pattern slice `s` and the `mask` slice must have the same length. + /// + /// The backtrack value of the resulting atom will be equal to the start + /// offset of the range within the slices. + pub fn from_slice_range(s: &[u8], mask: &[u8], range: R) -> Option + where + R: RangeBounds + SliceIndex<[u8], Output = [u8]> + Clone, + { + assert_eq!(s.len(), mask.len()); + let atom = Atom::from_slice_range(s, range.clone())?; + let mask_slice = mask.get(range)?; + Some(Self { atom, mask: SmallVec::from_slice(mask_slice) }) + } + + /// Returns a [`MaskCombinations`] iterator which produces all possible + /// combinations that result from applying the mask to the atom. /// /// The mask's binary 1 bits are constant, adopting the corresponding bits /// from the atom. Conversely, binary 0 bits within the mask are variable; /// consequently, the resultant atoms encompass all feasible permutations /// for these variable bits. - #[allow(dead_code)] - pub fn mask_combinations(self, mask: &[u8]) -> MaskCombinations { - MaskCombinations::new(self, mask) + pub fn mask_combinations(&self) -> MaskCombinations { + MaskCombinations::new(self) + } + + /// Reduces the length of the atom by 1 byte, removing the last byte. + /// + /// If the resulting atom ends with a full mask `??` (which is represented + /// by a mask byte of `0x00`), that byte is also removed. This process continues + /// until the last byte of the atom is not fully masked with `??` or the atom + /// becomes empty. + pub fn trim(&mut self) { + if self.atom.bytes.is_empty() { + return; + } + self.atom.bytes.pop(); + self.mask.pop(); + self.atom.set_exact(false); + + while self.mask.last() == Some(&0x00) { + self.atom.bytes.pop(); + self.mask.pop(); + } } } @@ -305,7 +350,7 @@ pub(crate) fn extract_atoms( } if flags.contains(SubPatternFlags::Nocase) { - Box::new(CaseCombinations::new(best_atom)) + Box::new(CaseCombinations::new(&best_atom)) } else { Box::new(iter::once(best_atom)) } @@ -336,15 +381,21 @@ pub(crate) struct MaskCombinations { cartesian_product: MultiProduct, backtrack: u16, exact: bool, + total_combinations: usize, } impl MaskCombinations { - fn new(atom: Atom, mask: &[u8]) -> Self { + fn new(masked_atom: &MaskedAtom) -> Self { Self { - exact: atom.exact, - backtrack: atom.backtrack, - cartesian_product: zip(atom.bytes, mask) - .map(|(byte, mask)| ByteMaskCombinator::new(byte, *mask)) + exact: masked_atom.atom.exact, + backtrack: masked_atom.atom.backtrack, + total_combinations: masked_atom + .mask + .iter() + .map(|m| 1 << m.count_zeros()) + .product(), + cartesian_product: zip(&masked_atom.atom.bytes, &masked_atom.mask) + .map(|(byte, mask)| ByteMaskCombinator::new(*byte, *mask)) .multi_cartesian_product(), } } @@ -361,6 +412,12 @@ impl Iterator for MaskCombinations { } } +impl ExactSizeIterator for MaskCombinations { + fn len(&self) -> usize { + self.total_combinations + } +} + /// Iterator that returns a sequence of atoms that covers all the possible case /// combinations for the ASCII characters in the original atom. The original /// atom is always included in the sequence, and non-alphabetic characters are @@ -376,7 +433,7 @@ pub(crate) struct CaseCombinations { } impl CaseCombinations { - fn new(atom: Atom) -> Self { + fn new(atom: &Atom) -> Self { Self { exact: atom.exact, backtrack: atom.backtrack, @@ -442,76 +499,3 @@ impl Iterator for XorCombinations { Some(atom) } } - -#[cfg(test)] -mod test { - use pretty_assertions::assert_eq; - - use crate::compiler::atoms::Atom; - - #[test] - fn mask_combinations() { - let atom = Atom::exact([0x11, 0x22, 0x33, 0x44]); - let mut c = atom.mask_combinations(&[0xff, 0xf0, 0xff, 0xff]); - - assert_eq!(c.next(), Some(Atom::exact([0x11, 0x20, 0x33, 0x44]))); - assert_eq!(c.next(), Some(Atom::exact([0x11, 0x21, 0x33, 0x44]))); - assert_eq!(c.next(), Some(Atom::exact([0x11, 0x22, 0x33, 0x44]))); - - let mut c = c.skip(10); - - assert_eq!(c.next(), Some(Atom::exact([0x11, 0x2d, 0x33, 0x44]))); - assert_eq!(c.next(), Some(Atom::exact([0x11, 0x2e, 0x33, 0x44]))); - assert_eq!(c.next(), Some(Atom::exact([0x11, 0x2f, 0x33, 0x44]))); - assert_eq!(c.next(), None); - } - - #[test] - fn case_combinations() { - let atom = Atom::exact(b"a1B2c"); - let mut c = atom.case_combinations(); - - assert_eq!(c.next(), Some(Atom::exact(b"a1b2c"))); - assert_eq!(c.next(), Some(Atom::exact(b"a1b2C"))); - assert_eq!(c.next(), Some(Atom::exact(b"a1B2c"))); - assert_eq!(c.next(), Some(Atom::exact(b"a1B2C"))); - assert_eq!(c.next(), Some(Atom::exact(b"A1b2c"))); - assert_eq!(c.next(), Some(Atom::exact(b"A1b2C"))); - assert_eq!(c.next(), Some(Atom::exact(b"A1B2c"))); - assert_eq!(c.next(), Some(Atom::exact(b"A1B2C"))); - assert_eq!(c.next(), None); - - let mut atom = Atom::exact([0x00_u8, 0x01, 0x02]); - atom.set_backtrack(2); - - let mut c = atom.clone().case_combinations(); - - assert_eq!(c.next(), Some(atom)); - assert_eq!(c.next(), None); - } - - #[test] - fn xor_combinations() { - let atom = Atom::exact([0x00_u8, 0x01, 0x02]); - let mut c = atom.xor_combinations(0..=1); - - assert_eq!(c.next(), Some(Atom::inexact([0x00, 0x01, 0x02]))); - assert_eq!(c.next(), Some(Atom::inexact([0x01, 0x00, 0x03]))); - assert_eq!(c.next(), None); - } - - #[test] - fn make_wide() { - let mut atom = Atom::exact([0x01_u8, 0x02, 0x03]); - atom.set_backtrack(2); - - let atom = atom.make_wide(); - - assert_eq!( - atom.bytes.as_slice(), - &[0x01, 0x00, 0x02, 0x00, 0x03, 0x00] - ); - - assert_eq!(atom.backtrack, 4); - } -} diff --git a/lib/src/compiler/atoms/quality.rs b/lib/src/compiler/atoms/quality.rs index 664284d3..fff44857 100644 --- a/lib/src/compiler/atoms/quality.rs +++ b/lib/src/compiler/atoms/quality.rs @@ -391,18 +391,16 @@ impl Ord for AtomsQuality { /// Returns the range for the best possible atom that can be extracted from /// the slice and its quality. -pub(crate) fn best_range_in_bytes( - bytes: &[u8], -) -> (Option>, i32) { +pub(crate) fn best_range_in_bytes(bytes: &[u8]) -> (Range, i32) { let mut best_quality = i32::MIN; - let mut best_range = None; + let mut best_range = 0..0; for i in 0..=bytes.len().saturating_sub(DESIRED_ATOM_SIZE) { let range = i..min(bytes.len(), i + DESIRED_ATOM_SIZE); let quality = atom_quality(&bytes[range.clone()]); if quality > best_quality { best_quality = quality; - best_range = Some(range); + best_range = range; } } @@ -411,11 +409,10 @@ pub(crate) fn best_range_in_bytes( /// Returns the range for the best possible atom that can be extracted from /// the masked slice. -#[allow(dead_code)] pub(crate) fn best_range_in_masked_bytes( bytes: &[u8], mask: &[u8], -) -> (Option>, i32) { +) -> (Range, i32) { let mut finders: Vec = (1..=DESIRED_ATOM_SIZE).map(BestAtomFinder::new).collect(); @@ -425,14 +422,15 @@ pub(crate) fn best_range_in_masked_bytes( } } - let mut best_result = (None, i32::MIN); + let mut best_result = (0..0, i32::MIN); for finder in finders { - let (range, quality) = finder.finalize(); - if range.is_none() { - continue; - } - if best_result.0.is_none() { + let (range, quality) = match finder.finalize() { + (Some(range), quality) => (range, quality), + _ => continue, + }; + + if best_result.0.is_empty() { best_result = (range, quality); continue; } @@ -440,8 +438,8 @@ pub(crate) fn best_range_in_masked_bytes( let is_better = match quality.cmp(&best_result.1) { Ordering::Greater => true, Ordering::Equal => { - let len = range.as_ref().unwrap().len(); - let best_len = best_result.0.as_ref().unwrap().len(); + let len = range.len(); + let best_len = best_result.0.len(); len < best_len } Ordering::Less => false, @@ -465,21 +463,8 @@ pub(crate) fn best_range_in_masked_bytes( /// correspond to the start of the slice in the data. pub(crate) fn best_atom_in_bytes(bytes: &[u8]) -> Atom { let (range, _) = best_range_in_bytes(bytes); - Atom::from_slice_range(bytes, range.unwrap()) -} - -/// Computes the quality of a masked atom. -#[cfg(test)] -pub fn masked_atom_quality<'a, B, M>(bytes: B, masks: M) -> i32 -where - B: IntoIterator, - M: IntoIterator, -{ - let mut finder = BestAtomFinder::new(DESIRED_ATOM_SIZE); - for (byte, mask) in zip(bytes, masks) { - finder.feed(*byte, *mask); - } - finder.finalize().1 + Atom::from_slice_range(bytes, range) + .expect("best_range_in_bytes returned invalid range") } /// Compute the quality of an atom. @@ -498,11 +483,24 @@ where #[cfg(test)] mod test { use super::{BestAtomFinder, atom_quality}; - use crate::compiler::atoms::quality::masked_atom_quality; - use crate::compiler::{AtomsQuality, atoms}; + use crate::compiler::{AtomsQuality, DESIRED_ATOM_SIZE, atoms}; use itertools::Itertools; use regex_syntax::hir::literal::Literal; use regex_syntax::hir::literal::Seq; + use std::iter::zip; + + /// Computes the quality of a masked atom. + fn masked_atom_quality<'a, B, M>(bytes: B, masks: M) -> i32 + where + B: IntoIterator, + M: IntoIterator, + { + let mut finder = BestAtomFinder::new(DESIRED_ATOM_SIZE); + for (byte, mask) in zip(bytes, masks) { + finder.feed(*byte, *mask); + } + finder.finalize().1 + } #[rustfmt::skip] #[allow(non_snake_case)] @@ -705,7 +703,7 @@ mod test { fn best_range_in_masked_bytes() { assert_eq!( atoms::best_range_in_masked_bytes(&[], &[],), - (None, i32::MIN), + (0..0, i32::MIN), ); assert_eq!( @@ -713,7 +711,7 @@ mod test { &[0x01, 0x02, 0x03, 0x04, 0x05], &[0xFF, 0xFF, 0xFF, 0xFF, 0xFF], ), - (Some(0..4), 88), + (0..4, 88), ); assert_eq!( @@ -721,7 +719,7 @@ mod test { &[0x01, 0x02, 0x00, 0x00], &[0xFF, 0xFF, 0x00, 0x00], ), - (Some(0..2), 44), + (0..2, 44), ); assert_eq!( @@ -729,7 +727,7 @@ mod test { &[0x01, 0x02, 0x00, 0x00, 0x05], &[0xFF, 0xFF, 0x00, 0x00, 0xFF], ), - (Some(0..2), 44), + (0..2, 44), ); assert_eq!( @@ -737,7 +735,7 @@ mod test { &[0x01, 0x02, 0x03, 0x04], &[0xFF, 0xFF, 0x0F, 0xFF], ), - (Some(0..4), 70), + (0..4, 70), ); assert_eq!( @@ -745,7 +743,7 @@ mod test { &[0x01, 0x02, 0x00, 0x04], &[0xFF, 0xFF, 0x00, 0xFF] ), - (Some(0..4), 51), + (0..4, 51), ); assert_eq!( @@ -753,7 +751,7 @@ mod test { &[0x01, 0x02, 0x00, 0x04, 0x05, 0x06, 0x07], &[0xFF, 0xFF, 0x00, 0xFF, 0xFF, 0xFF, 0xFF], ), - (Some(3..7), 88), + (3..7, 88), ); assert_eq!( @@ -761,7 +759,23 @@ mod test { &[0x01, 0x00, 0x00, 0x00, 0x00, 0xFF], &[0xFF, 0x00, 0x00, 0xFF, 0xFF, 0xFF], ), - (Some(3..6), 28), + (3..6, 28), + ); + + assert_eq!( + atoms::best_range_in_masked_bytes( + &[0x01, 0x02, 0x03, 0x04, 0x05, 0x06], + &[0x0F, 0x0F, 0x0F, 0x0F, 0x0F, 0x0F], + ), + (0..4, 16), + ); + + assert_eq!( + atoms::best_range_in_masked_bytes( + &[0x00, 0xA1, 0x81, 0x52, 0x00, 0x41, 0xA0, 0x72], + &[0x00, 0xFF, 0xFF, 0xFF, 0x00, 0xFF, 0xFF, 0xFF], + ), + (1..4, 64), ); } @@ -822,5 +836,22 @@ mod test { assert_eq!(finder_2.finalize(), (Some(3..5), 44)); assert_eq!(finder_3.finalize(), (Some(2..5), 29)); assert_eq!(finder_4.finalize(), (Some(0..4), 14)); + + let bytes = &[0x00, 0xA1, 0x81, 0x52, 0x00, 0x41, 0xA0, 0x72]; + let masks = &[0x00, 0xFF, 0xFF, 0xFF, 0x00, 0xFF, 0xFF, 0xFF]; + + let mut finder_2 = BestAtomFinder::new(2); + let mut finder_3 = BestAtomFinder::new(3); + let mut finder_4 = BestAtomFinder::new(4); + + for (byte, mask) in zip(bytes, masks) { + finder_2.feed(*byte, *mask); + finder_3.feed(*byte, *mask); + finder_4.feed(*byte, *mask); + } + + assert_eq!(finder_2.finalize(), (Some(1..3), 44)); + assert_eq!(finder_3.finalize(), (Some(1..4), 64)); + assert_eq!(finder_4.finalize(), (Some(0..4), 49)); } } diff --git a/lib/src/compiler/atoms/tests/mod.rs b/lib/src/compiler/atoms/tests/mod.rs new file mode 100644 index 00000000..e09ab57c --- /dev/null +++ b/lib/src/compiler/atoms/tests/mod.rs @@ -0,0 +1,256 @@ +use pretty_assertions::assert_eq; + +use crate::compiler::atoms::{Atom, MaskedAtom}; + +#[test] +fn mask_combinations() { + let masked_atom = MaskedAtom::from_slice_range( + &[0x11, 0x22, 0x33, 0x44], + &[0xff, 0xf0, 0xff, 0xff], + .., + ) + .unwrap(); + let mut c = masked_atom.mask_combinations(); + + assert_eq!(c.len(), 16); + + assert_eq!(c.next(), Some(Atom::exact([0x11, 0x20, 0x33, 0x44]))); + assert_eq!(c.next(), Some(Atom::exact([0x11, 0x21, 0x33, 0x44]))); + assert_eq!(c.next(), Some(Atom::exact([0x11, 0x22, 0x33, 0x44]))); + + let mut c = c.skip(10); + + assert_eq!(c.next(), Some(Atom::exact([0x11, 0x2d, 0x33, 0x44]))); + assert_eq!(c.next(), Some(Atom::exact([0x11, 0x2e, 0x33, 0x44]))); + assert_eq!(c.next(), Some(Atom::exact([0x11, 0x2f, 0x33, 0x44]))); + assert_eq!(c.next(), None); +} + +#[test] +fn case_combinations() { + let atom = Atom::exact(b"a1B2c"); + let mut c = atom.case_combinations(); + + assert_eq!(c.next(), Some(Atom::exact(b"a1b2c"))); + assert_eq!(c.next(), Some(Atom::exact(b"a1b2C"))); + assert_eq!(c.next(), Some(Atom::exact(b"a1B2c"))); + assert_eq!(c.next(), Some(Atom::exact(b"a1B2C"))); + assert_eq!(c.next(), Some(Atom::exact(b"A1b2c"))); + assert_eq!(c.next(), Some(Atom::exact(b"A1b2C"))); + assert_eq!(c.next(), Some(Atom::exact(b"A1B2c"))); + assert_eq!(c.next(), Some(Atom::exact(b"A1B2C"))); + assert_eq!(c.next(), None); + + let mut atom = Atom::exact([0x00_u8, 0x01, 0x02]); + atom.set_backtrack(2); + + let mut c = atom.clone().case_combinations(); + + assert_eq!(c.next(), Some(atom)); + assert_eq!(c.next(), None); +} + +#[test] +fn xor_combinations() { + let atom = Atom::exact([0x00_u8, 0x01, 0x02]); + let mut c = atom.xor_combinations(0..=1); + + assert_eq!(c.next(), Some(Atom::inexact([0x00, 0x01, 0x02]))); + assert_eq!(c.next(), Some(Atom::inexact([0x01, 0x00, 0x03]))); + assert_eq!(c.next(), None); +} + +#[test] +fn make_wide() { + let mut atom = Atom::exact([0x01_u8, 0x02, 0x03]); + atom.set_backtrack(2); + + let atom = atom.make_wide(); + + assert_eq!(atom.bytes.as_slice(), &[0x01, 0x00, 0x02, 0x00, 0x03, 0x00]); + assert_eq!(atom.backtrack, 4); +} + +fn check_atoms(pattern_src: &str, expected_atoms: I) +where + I: IntoIterator, + T: AsRef<[u8]>, +{ + let rule_src = format!( + r#" + rule test {{ + strings: + $a = {} + condition: + $a + }} + "#, + pattern_src + ); + + let rules = crate::compiler::compile(rule_src.as_str()).unwrap(); + let actual_atoms = rules.atoms(); + + let mut expected_iter = expected_atoms.into_iter(); + for actual in actual_atoms { + let expected = + expected_iter.next().expect("actual has more atoms than expected"); + assert_eq!(actual.as_slice(), expected.as_ref()); + } + + assert!( + expected_iter.next().is_none(), + "expected has more atoms than actual" + ); +} + +#[test] +fn atoms() { + check_atoms("{ 11 ?? ?? 22 }", [&[0x11]]); + check_atoms("{ 11 22 33 }", [&[0x11, 0x22, 0x33]]); + check_atoms("{ 1? 22 }", (0x10..=0x1f).map(|b| [b, 0x22])); + + check_atoms("{ 1? 2? 3? 44 55 66 77 8? 9? }", [&[0x44, 0x55, 0x66, 0x77]]); + + check_atoms( + "{ 11 22 33 44 [2-5] 55 66 77 88 }", + [&[0x11, 0x22, 0x33, 0x44]], + ); + + check_atoms( + "{ 11 22 33 [2-5] 44 55 66 77 88 }", + [&[0x55, 0x66, 0x77, 0x88]], + ); + + check_atoms( + "{ ( 11 22 33 44 | 55 66 77 88 ) }", + [&[0x11, 0x22, 0x33, 0x44], &[0x55, 0x66, 0x77, 0x88]], + ); + + check_atoms( + "{ ( 11 22 | 33 44 55 | 66 77 88 99 ) }", + [ + &[0x11_u8, 0x22].as_slice(), + &[0x33_u8, 0x44, 0x55].as_slice(), + &[0x66_u8, 0x77, 0x88, 0x99].as_slice(), + ], + ); + + check_atoms( + "{ ( 1? | 2? ) 33 }", + (0x10..=0x1f) + .map(|b| [b, 0x33]) + .chain((0x20..=0x2f).map(|b| [b, 0x33])), + ); + + check_atoms( + "{ 2? F? ?8 6? ?? 0? ?B ?B 3? ?? B? 2? ?7 5? }", + itertools::iproduct!( + 0x20..=0x2f, + 0xf0..=0xff, + (0x08..=0xf8).step_by(0x10) + ) + .map(|(b1, b2, b3)| [b1, b2, b3]), + ); + + check_atoms( + "{ 11 2? 33 4? 55 }", + itertools::iproduct!(0x20..=0x2f, 0x40..=0x4f) + .map(|(b2, b4)| [0x11, b2, 0x33, b4]), + ); + + check_atoms( + "{ 1? 2? 3? 44 }", + itertools::iproduct!(0x10..=0x1f, 0x20..=0x2f, 0x30..=0x3f) + .map(|(b1, b2, b3)| [b1, b2, b3, 0x44]), + ); + + check_atoms( + "{ 11 ?? 1? 11 }", + itertools::iproduct!(0x00..=0xff, 0x10..=0x1f) + .map(|(b2, b3)| [0x11, b2, b3, 0x11]), + ); +} + +#[test] +fn masked_atom_trim() { + // Normal trim: removes last byte + let mut masked_atom = MaskedAtom::from_slice_range( + &[0x11, 0x22, 0x33, 0x44], + &[0xff, 0xff, 0xff, 0xff], + .., + ) + .unwrap(); + masked_atom.trim(); + assert_eq!(masked_atom.atom.as_ref(), &[0x11, 0x22, 0x33]); + assert_eq!(masked_atom.mask.as_slice(), &[0xff, 0xff, 0xff]); + assert!(!masked_atom.atom.is_exact()); + + // Trim with ?? at the end of resulting atom: removes last byte and ?? + let mut masked_atom = MaskedAtom::from_slice_range( + &[0x11, 0x22, 0x33, 0x44], + &[0xff, 0xff, 0x00, 0xff], + .., + ) + .unwrap(); + masked_atom.trim(); + // After trimming 0x44, it becomes [0x11, 0x22, 0x33] with mask [0xff, 0xff, 0x00]. + // Since the last byte (0x33) has mask 0x00, it's also trimmed. + assert_eq!(masked_atom.atom.as_ref(), &[0x11, 0x22]); + assert_eq!(masked_atom.mask.as_slice(), &[0xff, 0xff]); + assert!(!masked_atom.atom.is_exact()); + + // Trim with multiple ?? at the end + let mut masked_atom = MaskedAtom::from_slice_range( + &[0x11, 0x22, 0x33, 0x44], + &[0xff, 0x00, 0x00, 0xff], + .., + ) + .unwrap(); + masked_atom.trim(); + // After trimming 0x44, the two 0x00s are also trimmed. + assert_eq!(masked_atom.atom.as_ref(), &[0x11]); + assert_eq!(masked_atom.mask.as_slice(), &[0xff]); + + // Trim to empty + let mut masked_atom = MaskedAtom::from_slice_range( + &[0x11, 0x22, 0x33], + &[0x00, 0x00, 0xff], + .., + ) + .unwrap(); + masked_atom.trim(); + // Trim 0x33, then both 0x00s are trimmed, resulting in empty. + assert!(masked_atom.atom.as_ref().is_empty()); + assert!(masked_atom.mask.is_empty()); + + // Trim on empty + let mut masked_atom = MaskedAtom::from_slice_range(&[], &[], ..).unwrap(); + masked_atom.trim(); + assert!(masked_atom.atom.as_ref().is_empty()); + assert!(masked_atom.mask.is_empty()); +} + +#[test] +fn from_slice_range_limits() { + let bytes = [0x11, 0x22, 0x33, 0x44]; + let mask = [0xff, 0xff, 0xff, 0xff]; + + // Out of bounds range + assert!(Atom::from_slice_range(&bytes, 2..10).is_none()); + assert!(MaskedAtom::from_slice_range(&bytes, &mask, 2..10).is_none()); + + // Backtrack overflow (start offset > u16::MAX) + let large_bytes = vec![0; 70000]; + let large_mask = vec![0; 70000]; + assert!(Atom::from_slice_range(&large_bytes, 65536..65538).is_none()); + assert!( + MaskedAtom::from_slice_range(&large_bytes, &large_mask, 65536..65538) + .is_none() + ); + + // Valid bounds and backtrack + let atom = Atom::from_slice_range(&large_bytes, 65535..65537); + assert!(atom.is_some()); + assert_eq!(atom.unwrap().backtrack(), 65535); +} diff --git a/lib/src/compiler/mod.rs b/lib/src/compiler/mod.rs index 155490ca..acc9ca21 100644 --- a/lib/src/compiler/mod.rs +++ b/lib/src/compiler/mod.rs @@ -2228,7 +2228,7 @@ impl Compiler<'_> { } // If this point is reached, this is a pattern that can't be split into - // multiple chained patterns, and is neither a literal or alternation + // multiple chained patterns, and is neither a literal nor alternation // of literals. Most patterns fall in this category. let mut flags = SubPatternFlags::empty(); @@ -2245,6 +2245,26 @@ impl Compiler<'_> { flags.insert(SubPatternFlags::GreedyRegexp); } + // If the pattern doesn't have the `nocase` or `wide` modifier, and it + // is a simple literal pattern with masks (e.g., `{ 01 02 ?? 04 }` or + // `{ 1? 2? 3? }`), we can treat it as a `LiteralWithMask` sub-pattern. + // This is much more efficient than executing it as a regular expression. + if !pattern.flags.contains(PatternFlags::Nocase) + && !pattern.flags.contains(PatternFlags::Wide) + && let Some((pattern, mask, atoms)) = + head.try_extract_literal_with_mask() + { + let pattern = self.intern_literal(pattern.as_slice(), false); + let mask = self.intern_literal(mask.as_slice(), false); + self.add_sub_pattern( + pattern_id, + SubPattern::LiteralWithMask { pattern, mask, flags }, + atoms, + SubPatternAtom::from_atom, + ); + return Ok(()); + }; + let (atoms, is_fast_regexp) = self.c_regexp(&head, span)?; if is_fast_regexp { @@ -2999,6 +3019,12 @@ pub(crate) enum SubPattern { flags: SubPatternFlags, }, + LiteralWithMask { + pattern: LiteralId, + mask: LiteralId, + flags: SubPatternFlags, + }, + LiteralChainHead { pattern: LiteralId, flags: SubPatternFlags, diff --git a/lib/src/compiler/rules.rs b/lib/src/compiler/rules.rs index e8d186de..d0d9c3f4 100644 --- a/lib/src/compiler/rules.rs +++ b/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 = 2; +const SERIALIZATION_VERSION: u32 = 3; /// Aho-Corasick automaton bundled with an optional Teddy scanner if the /// number of patterns is low enough. If the Teddy scanner is present, and diff --git a/lib/src/compiler/tests/mod.rs b/lib/src/compiler/tests/mod.rs index 6e85924c..d875546c 100644 --- a/lib/src/compiler/tests/mod.rs +++ b/lib/src/compiler/tests/mod.rs @@ -28,7 +28,7 @@ fn serialization() { // `DecodeError`. let mut data = Vec::new(); data.extend(b"YARA-X\0\0"); - data.extend(2u32.to_le_bytes()); + data.extend(3u32.to_le_bytes()); data.extend(b"foo"); assert!(matches!( diff --git a/lib/src/re/fast/compiler.rs b/lib/src/re/fast/compiler.rs index 7cea14d6..f6afc7fd 100644 --- a/lib/src/re/fast/compiler.rs +++ b/lib/src/re/fast/compiler.rs @@ -5,7 +5,9 @@ use std::mem::size_of; use bstr::ByteSlice; use regex_syntax::hir::{Class, Hir, HirKind, Visitor, visit}; -use crate::compiler::{Atom, best_range_in_bytes, best_range_in_masked_bytes}; +use crate::compiler::{ + Atom, MaskedAtom, best_range_in_bytes, best_range_in_masked_bytes, +}; use crate::re; use crate::re::fast::instr::Instr; use crate::re::{BckCodeLoc, Error, FwdCodeLoc, MAX_ALTERNATIVES, RegexpAtom}; @@ -99,7 +101,7 @@ impl Compiler { } } PatternPiece::JumpExact(..) | PatternPiece::Jump(..) => { - piece_atoms.push((None, None, None, i32::MIN)) + piece_atoms.push((None, None, 0..0, i32::MIN)) } }; @@ -163,10 +165,12 @@ impl Compiler { }; for (best_bytes, best_mask, best_range, _) in best_atoms { - match (best_bytes, best_mask, best_range) { - (Some(bytes), Some(mask), Some(range)) => { - let atom = Atom::from_slice_range(bytes, range.clone()); - for atom in atom.mask_combinations(&mask[range]) { + match (best_bytes, best_mask) { + (Some(bytes), Some(mask)) => { + let masked_atom = + MaskedAtom::from_slice_range(bytes, mask, best_range) + .ok_or(Error::FastIncompatible)?; + for atom in masked_atom.mask_combinations() { atoms.push(RegexpAtom { atom, fwd_code: Some(FwdCodeLoc::from(fwd_code_start)), @@ -174,11 +178,13 @@ impl Compiler { }) } } - (Some(bytes), None, Some(range)) => atoms.push(RegexpAtom { - atom: Atom::from_slice_range(bytes, range), + (Some(bytes), None) => atoms.push(RegexpAtom { + atom: Atom::from_slice_range(bytes, best_range) + .ok_or(Error::FastIncompatible)?, fwd_code: Some(FwdCodeLoc::from(fwd_code_start)), bck_code: bck_code_start.map(BckCodeLoc::from), }), + (None, None) => {} _ => unreachable!(), } } diff --git a/lib/src/re/hir.rs b/lib/src/re/hir.rs index b7fef4cf..a1ab1339 100644 --- a/lib/src/re/hir.rs +++ b/lib/src/re/hir.rs @@ -1,4 +1,5 @@ use std::hash::{Hash, Hasher}; +use std::iter::repeat_n; use std::mem; use std::ops::{RangeFrom, RangeInclusive}; @@ -15,6 +16,10 @@ use serde::{Deserialize, Serialize}; use yara_x_parser::ast; +use crate::compiler::{ + MAX_ATOMS_PER_REGEXP, MaskCombinations, MaskedAtom, + best_range_in_masked_bytes, +}; use crate::utils::cast; #[derive(Clone, Copy, Debug, PartialEq)] @@ -212,6 +217,96 @@ impl Hir { (chain.remove(0).hir, chain) } + /// Tries to extract a literal pattern with its corresponding mask and a + /// list of generated [`Atom`]s from the regular expression's [`Hir`]. + /// + /// This method succeeds if the regular expression represents a sequence of + /// masked bytes (e.g., `{ 01 02 ?? 04 }` or `{ 1? 2? 3? }`). If successful, + /// it returns: + /// + /// - The raw pattern bytes (`Vec`). + /// - The mask bytes (`Vec`). + /// - The list of combinations of atoms to search for (`Vec`), where + /// the atoms are trimmed if necessary to avoid exceeding + /// `MAX_ATOMS_PER_REGEXP` combinations. + /// + /// If the regular expression is not a simple masked literal sequence, it + /// returns `None`. + pub(crate) fn try_extract_literal_with_mask( + &self, + ) -> Option<(Vec, Vec, MaskCombinations)> { + use regex_syntax::hir::{HirKind, Repetition}; + + let mut pattern = Vec::new(); + let mut mask = Vec::new(); + let mut stack = vec![&self.inner]; + + fn is_any_byte(hir: ®ex_syntax::hir::Hir) -> bool { + match hir.kind() { + HirKind::Class(Class::Bytes(class)) => { + class.ranges().len() == 1 + && class.ranges()[0].start() == 0 + && class.ranges()[0].end() == 255 + } + _ => false, + } + } + + while let Some(hir) = stack.pop() { + match hir.kind() { + HirKind::Literal(lit) => { + let bytes = lit.0.as_ref(); + pattern.extend_from_slice(bytes); + mask.extend(repeat_n(0xff, bytes.len())); + } + HirKind::Repetition(Repetition { min, max, sub, .. }) => { + if *max == Some(*min) && is_any_byte(sub.as_ref()) { + let count = *min as usize; + pattern.extend(repeat_n(0x00, count)); + mask.extend(repeat_n(0x00, count)); + } else { + return None; + } + } + HirKind::Concat(sub_hirs) => { + // Push in reverse order so they are processed left-to-right. + for sub in sub_hirs.iter().rev() { + stack.push(sub); + } + } + HirKind::Class(Class::Bytes(class)) => { + if let Some(hex_byte) = class_to_masked_byte(class) { + pattern.push(hex_byte.value); + mask.push(hex_byte.mask); + } else { + return None; + } + } + HirKind::Capture(group) => { + stack.push(&group.sub); + } + _ => return None, + } + } + + debug_assert_eq!(pattern.len(), mask.len()); + + let (best_range, _) = best_range_in_masked_bytes(&pattern, &mask); + + let mut masked_atom = + MaskedAtom::from_slice_range(&pattern, &mask, best_range)?; + + while masked_atom.mask_combinations().len() > MAX_ATOMS_PER_REGEXP { + masked_atom.trim(); + } + + if masked_atom.atom.as_ref().is_empty() { + return None; + } + + Some((pattern, mask, masked_atom.mask_combinations())) + } + pub fn set_greedy(mut self, greediness: Option) -> Self { self.greedy = greediness; self diff --git a/lib/src/re/thompson/compiler.rs b/lib/src/re/thompson/compiler.rs index 14e400e5..266cf981 100644 --- a/lib/src/re/thompson/compiler.rs +++ b/lib/src/re/thompson/compiler.rs @@ -220,6 +220,10 @@ impl Compiler { } impl Compiler { + /// Repetition threshold. Repetitions with count(s) below or equal to this + /// threshold are unrolled (i.e., the repeated expression's code is + /// duplicated). Repetitions with count(s) above this threshold are + /// compiled using the `repeat` instruction to avoid bloated bytecode. const REPEAT_INSTR_THRESHOLD: u32 = 10; pub(super) fn compile_internal( @@ -679,7 +683,7 @@ impl Compiler { } // e{0,max} (not inside repetition_start/repetition_end yet) // - // l1: split_a l4 ( split_a for the non-greedy e{0,max}? ) + // l1: split_a l4 ( split_b for the non-greedy e{0,max}? ) // l2 ... code for e ... // l3: repeat l2, 0, max // l4: diff --git a/lib/src/re/thompson/instr.rs b/lib/src/re/thompson/instr.rs index eb079616..abf8fe5f 100644 --- a/lib/src/re/thompson/instr.rs +++ b/lib/src/re/thompson/instr.rs @@ -190,7 +190,7 @@ pub enum Instr<'a> { /// due to its more compact representation. ClassBitmap(ClassBitmap<'a>), - /// Matches a byte class. The class is represented 1 or more byte ranges + /// Matches a byte class. The class represents one or more byte ranges /// The first `u8` after the opcode indicates the number of ranges, then /// follows one pair `[u8, u8]` per range, indicating starting and ending /// bytes for the range, both inclusive. With 16 ranges this instruction diff --git a/lib/src/scanner/context.rs b/lib/src/scanner/context.rs index 74aaf864..0daa65ae 100644 --- a/lib/src/scanner/context.rs +++ b/lib/src/scanner/context.rs @@ -43,7 +43,7 @@ use crate::wasm::runtime::{ AsContext, AsContextMut, Global, GlobalType, Instance, MemoryType, Mutability, Store, TypedFunc, Val, ValType, }; -use crate::{Variable, wasm}; +use crate::{Variable, teddy, wasm}; /// Represents the states in which a scanner can be. pub(crate) enum ScanState<'a> { @@ -873,6 +873,7 @@ impl ScanContext<'_, '_> { SubPattern::Literal { flags, .. } | SubPattern::LiteralChainHead { flags, .. } | SubPattern::LiteralChainTail { flags, .. } + | SubPattern::LiteralWithMask { flags, .. } | SubPattern::Regexp { flags, .. } | SubPattern::RegexpChainHead { flags, .. } | SubPattern::RegexpChainTail { flags, .. } => flags, @@ -905,7 +906,7 @@ impl ScanContext<'_, '_> { .get_bytes(*pattern) .unwrap(); - if verify_literal_match(pattern, data, atom_pos, *flags) { + if verify_literal(pattern, data, atom_pos, *flags) { handle_sub_pattern_match( &mut self.tracker, &mut self.wasm, @@ -917,16 +918,28 @@ impl ScanContext<'_, '_> { ); } } - SubPattern::Regexp { flags, .. } - | SubPattern::RegexpChainHead { flags, .. } - | SubPattern::RegexpChainTail { flags, .. } => { - verify_regexp_match( - &mut self.vm, - data, - atom_pos, - atom, - *flags, - |match_range| { + SubPattern::LiteralWithMask { pattern, mask, flags } => { + let pattern = self + .compiled_rules + .lit_pool() + .get_bytes(*pattern) + .unwrap(); + + let mask = + self.compiled_rules.lit_pool().get_bytes(*mask).unwrap(); + + debug_assert_eq!(pattern.len(), mask.len()); + + if let Some(data) = + data.get(atom_pos..atom_pos + pattern.len()) + && verify_literal_with_mask(data, pattern, mask) + { + let match_range = atom_pos..atom_pos + pattern.len(); + if !flags.intersects( + SubPatternFlags::FullwordLeft + | SubPatternFlags::FullwordRight, + ) || verify_full_word(data, &match_range, *flags, None) + { handle_sub_pattern_match( &mut self.tracker, &mut self.wasm, @@ -935,9 +948,28 @@ impl ScanContext<'_, '_> { *pattern_id, Match::new(match_range).rebase(base), ); - }, - ) + } + } } + SubPattern::Regexp { flags, .. } + | SubPattern::RegexpChainHead { flags, .. } + | SubPattern::RegexpChainTail { flags, .. } => verify_regexp( + &mut self.vm, + data, + atom_pos, + atom, + *flags, + |match_range| { + handle_sub_pattern_match( + &mut self.tracker, + &mut self.wasm, + sub_pattern_id, + sub_pattern, + *pattern_id, + Match::new(match_range).rebase(base), + ); + }, + ), SubPattern::Xor { pattern, flags } => { let pattern = self @@ -947,7 +979,7 @@ impl ScanContext<'_, '_> { .unwrap(); if let Some(key) = - verify_xor_match(pattern, data, atom_pos, atom, *flags) + verify_xor(pattern, data, atom_pos, atom, *flags) { handle_sub_pattern_match( &mut self.tracker, @@ -964,7 +996,7 @@ impl ScanContext<'_, '_> { SubPattern::Base64 { pattern, padding } | SubPattern::Base64Wide { pattern, padding } => { - if let Some(match_range) = verify_base64_match( + if let Some(match_range) = verify_base64( self.compiled_rules .lit_pool() .get_bytes(*pattern) @@ -1007,7 +1039,7 @@ impl ScanContext<'_, '_> { }, ); - if let Some(match_range) = verify_base64_match( + if let Some(match_range) = verify_base64( self.compiled_rules .lit_pool() .get_bytes(*pattern) @@ -1170,8 +1202,7 @@ impl ScanContext<'_, '_> { .get_bytes(*pattern) .unwrap(); - if verify_literal_match(pattern, data, offset, *flags) - { + if verify_literal(pattern, data, offset, *flags) { handle_sub_pattern_match( &mut self.tracker, &mut self.wasm, @@ -1191,7 +1222,7 @@ impl ScanContext<'_, '_> { } /// Verifies if a literal `pattern` matches at `match_start` in `scanned_data`. -fn verify_literal_match( +fn verify_literal( pattern: &[u8], scanned_data: &[u8], match_start: usize, @@ -1223,6 +1254,104 @@ fn verify_literal_match( } } +#[cfg(any( + all(target_arch = "x86_64", target_feature = "sse2"), + all( + target_arch = "aarch64", + target_feature = "neon", + target_endian = "little" + ) +))] +#[inline(always)] +unsafe fn verify_literal_with_mask_simd( + data: &[u8], + pattern: &[u8], + mask: &[u8], +) -> bool { + let len = data.len(); + debug_assert!(len <= u16::MAX as usize); + + unsafe { + if len <= 16 { + let mut d_bytes = [0_u8; 16]; + let mut m_bytes = [0_u8; 16]; + let mut p_bytes = [0_u8; 16]; + + d_bytes[..len].copy_from_slice(data); + p_bytes[..len].copy_from_slice(&pattern[..len]); + m_bytes[..len].copy_from_slice(&mask[..len]); + + let data = V::load_unaligned(d_bytes.as_ptr()); + let pattern = V::load_unaligned(p_bytes.as_ptr()); + let mask = V::load_unaligned(m_bytes.as_ptr()); + let eq = data.and(mask).cmpeq(pattern); + + eq.cmpeq(V::splat(0x00)).is_zero() + } else { + let verify_chunk = |offset: usize| -> bool { + let data = V::load_unaligned(data[offset..].as_ptr()); + let pattern = V::load_unaligned(pattern[offset..].as_ptr()); + let mask = V::load_unaligned(mask[offset..].as_ptr()); + let eq = data.and(mask).cmpeq(pattern); + + eq.cmpeq(V::splat(0x00)).is_zero() + }; + + let mut offset = 0; + + while offset + 16 <= len { + if !verify_chunk(offset) { + return false; + } + offset += 16; + } + + if offset < len && !verify_chunk(len - 16) { + return false; + } + + true + } + } +} + +fn verify_literal_with_mask(data: &[u8], pattern: &[u8], mask: &[u8]) -> bool { + #[cfg(all(target_arch = "x86_64", target_feature = "sse2"))] + unsafe { + verify_literal_with_mask_simd::( + data, pattern, mask, + ) + } + + #[cfg(all( + target_arch = "aarch64", + target_feature = "neon", + target_endian = "little" + ))] + unsafe { + verify_literal_with_mask_simd::( + data, pattern, mask, + ) + } + + #[cfg(not(any( + all(target_arch = "x86_64", target_feature = "sse2"), + all( + target_arch = "aarch64", + target_feature = "neon", + target_endian = "little" + ) + )))] + { + for ((&d, &m), &t) in data.iter().zip(mask).zip(pattern) { + if (d & m) != t { + return false; + } + } + true + } +} + /// Returns true if the match delimited by `match_range` is a full word match. /// This means that the bytes before the range's start and after the range's /// end are both non-alphanumeric. @@ -1365,7 +1494,7 @@ fn verify_chain_of_matches( /// /// This function can produce multiple matches, `f` is called for every /// match found. -fn verify_regexp_match( +fn verify_regexp( vm: &mut VM, scanned_data: &[u8], match_start: usize, @@ -1459,7 +1588,7 @@ fn verify_regexp_match( /// within `scanned_data`. /// /// Returns the XOR key if the match was confirmed, or [`None`] if otherwise. -fn verify_xor_match( +fn verify_xor( pattern: &[u8], scanned_data: &[u8], match_start: usize, @@ -1545,7 +1674,7 @@ fn xor_slices_eq(pattern: &[u8], candidate: &[u8], key: u8) -> bool { /// within `scanned_data`. /// /// Returns the range where the match was found or [`None`] if otherwise. -fn verify_base64_match( +fn verify_base64( pattern: &[u8], scanned_data: &[u8], padding: usize, @@ -1696,6 +1825,7 @@ fn handle_sub_pattern_match( ) { match sub_pattern { SubPattern::Literal { .. } + | SubPattern::LiteralWithMask { .. } | SubPattern::Xor { .. } | SubPattern::Base64 { .. } | SubPattern::Base64Wide { .. } diff --git a/lib/src/teddy/mod.rs b/lib/src/teddy/mod.rs index 2eaf4899..0c83dd7c 100644 --- a/lib/src/teddy/mod.rs +++ b/lib/src/teddy/mod.rs @@ -13,7 +13,7 @@ use core::fmt::Debug; use std::sync::Arc; mod generic; -mod vector; +pub(crate) mod vector; pub(crate) use self::generic::{Match, Patterns};