feat: adding support for masked byte sequences

Author: WolverinDEV

bmatcher-core/src/atom.rs

@@ -18,6 +18,19 @@ pub enum Atom {
     /// Match a sequence of bytes from the sequence array.
     ByteSequence { seq_start: u16, seq_end: u16 },
 
+    /// Match a sequence of bytes from the sequence array
+    /// using a custom bitmask.
+    ByteSequenceMasked {
+        /// Start of the byte sequence in the sequence array.
+        seq_start: u16,
+
+        /// Start of the byte mask in the sequence array.
+        mask_start: u16,
+
+        /// Number of bytes to match
+        len: u16,
+    },
+
     /// Skip a fixed number of bytes.
     WildcardFixed(u16),
     /// Skip a variable number of bytes.

bmatcher-core/src/compiler/lexer.rs

@@ -11,6 +11,9 @@ pub enum Token<'a> {
     /// A one byte wildcard token: "?"
     Whildcard,
 
+    /// A mask token: "&"
+    Mask,
+
     /// A range open token: "["
     RangeOpen,
     /// A range token: "-"
@@ -73,6 +76,7 @@ impl<'a> Lexer<'a> {
     fn char_to_token(&self, value: char) -> Option<Token<'a>> {
         Some(match value {
             '?' => Token::Whildcard,
+            '&' => Token::Mask,
 
             '{' => Token::BlockOpen,
             '}' => Token::BlockClose,

bmatcher-core/src/compiler/parser.rs

@@ -18,6 +18,8 @@ pub enum ParseError {
     UnexpectedToken,
     UnexpectedEnd,
 
+    MaskByteLenMismatch,
+
     HexValueInvalid,
     HexValueIncomplete,
 
@@ -114,7 +116,7 @@ impl<'a> PatternParser<'a> {
         Ok(false)
     }
 
-    fn parse_byte_sequence(&mut self) -> Result<(), PositionedError<ParseError>> {
+    fn parse_bytes(&mut self) -> Result<(u16, u16), PositionedError<ParseError>> {
         let Token::Text(value) = self.pop_token()? else {
             return Err(PositionedError::new(
                 self.lexer.token_range(),
@@ -164,10 +166,38 @@ impl<'a> PatternParser<'a> {
             ));
         }
 
-        self.atoms.push(Atom::ByteSequence {
-            seq_start: bytes_start as u16,
-            seq_end: bytes_end as u16,
-        });
+        Ok((bytes_start as u16, bytes_end as u16))
+    }
+
+    fn parse_byte_sequence(&mut self) -> Result<(), PositionedError<ParseError>> {
+        let (bytes_start, bytes_end) = self.parse_bytes()?;
+
+        if let Some(Token::Mask) = self.peek_token() {
+            /* masked byte sequence */
+            let _ = self.pop_token()?;
+
+            let (mask_start, mask_end) = self.parse_bytes()?;
+            let mask_length = mask_end - mask_start;
+            let bytes_length = bytes_end - bytes_start;
+            if mask_length != bytes_length {
+                return Err(PositionedError::new(
+                    self.lexer.token_range(),
+                    ParseError::MaskByteLenMismatch,
+                ));
+            }
+
+            self.atoms.push(Atom::ByteSequenceMasked {
+                seq_start: bytes_start,
+                mask_start,
+                len: bytes_length,
+            });
+        } else {
+            /* normal byte sequence */
+            self.atoms.push(Atom::ByteSequence {
+                seq_start: bytes_start as u16,
+                seq_end: bytes_end as u16,
+            });
+        }
 
         Ok(())
     }
@@ -468,6 +498,46 @@ mod test {
         }
     }
 
+    #[test]
+    fn test_byte_sequence_mask() {
+        {
+            let parser = PatternParser::new("FF & F0");
+            let result = parser.parse().unwrap();
+            assert_eq!(
+                result.atoms(),
+                &[Atom::ByteSequenceMasked {
+                    seq_start: 0,
+                    mask_start: 1,
+                    len: 1
+                }]
+            );
+            assert_eq!(result.byte_sequence(), &[0xFF, 0xF0]);
+        }
+
+        {
+            let parser = PatternParser::new("FFEE & F0F0");
+            let result = parser.parse().unwrap();
+            assert_eq!(
+                result.atoms(),
+                &[Atom::ByteSequenceMasked {
+                    seq_start: 0,
+                    mask_start: 2,
+                    len: 2
+                }]
+            );
+            assert_eq!(result.byte_sequence(), &[0xFF, 0xEE, 0xF0, 0xF0]);
+        }
+
+        {
+            let parser = PatternParser::new("FFFF & FF");
+            let result = parser.parse().unwrap_err();
+            assert_eq!(
+                &result,
+                &PositionedError::new(7..9, ParseError::MaskByteLenMismatch)
+            );
+        }
+    }
+
     #[test]
     fn test_jump() {
         {

bmatcher-core/src/matcher.rs

@@ -100,6 +100,38 @@ impl<'a, S: Stack<u32>, C: Stack<usize>> BinaryMatcher<'a, S, C> {
                     atom_cursor += 1;
                     data_cursor += expected_bytes.len();
                 }
+                Atom::ByteSequenceMasked {
+                    seq_start,
+                    mask_start,
+                    len,
+                } => {
+                    let target_bytes = &self.pattern_byte_sequence
+                        [seq_start as usize..seq_start as usize + len as usize];
+
+                    let target_mask = &self.pattern_byte_sequence
+                        [mask_start as usize..mask_start as usize + len as usize];
+
+                    let actual_bytes = self.target.subrange(data_cursor, target_mask.len())?;
+
+                    let target_bytes = target_bytes
+                        .iter()
+                        .zip(target_mask)
+                        .map(|(value, mask)| *value & *mask);
+                    let actual_bytes = actual_bytes
+                        .iter()
+                        .zip(target_mask)
+                        .map(|(value, mask)| *value & *mask);
+
+                    if target_bytes
+                        .zip(actual_bytes)
+                        .any(|(expected, data)| expected != data)
+                    {
+                        return None;
+                    }
+
+                    atom_cursor += 1;
+                    data_cursor += len as usize;
+                }
                 Atom::WildcardFixed(length) => {
                     atom_cursor += 1;
                     data_cursor += length as usize;
@@ -295,6 +327,11 @@ mod test {
         test_single("B7 69 3D", DATA, None);
     }
 
+    #[test]
+    fn test_binary_mask() {
+        test_single("B7682D & FFFEFF", DATA, Some(&[0x41]));
+    }
+
     #[test]
     fn test_range() {
         test_single("B7 69 2D [0-3] 85 2C '", DATA, Some(&[0x41, 0x48]));

bmatcher-proc/src/macro_pattern.rs

@@ -31,6 +31,18 @@ fn emit_atom(output: &mut dyn Write, atom: &Atom) -> io::Result<()> {
             )?;
         }
 
+        Atom::ByteSequenceMasked {
+            seq_start,
+            mask_start,
+            len,
+        } => {
+            write!(
+                output,
+                "bmatcher::Atom::ByteSequenceMasked{{ seq_start: 0x{:X}, mask_start: 0x{:X}, len: 0x{:X} }}",
+                seq_start, mask_start, len
+            )?;
+        }
+
         Atom::WildcardFixed(value) => {
             write!(output, "bmatcher::Atom::WildcardFixed(0x{:X})", value)?;
         }