1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
/// ## purpose
/// to match multiple byte patterns against a byte slice in parallel.
/// we should get all valid matches at the end.
/// does not have to support scanning across the byte slice, only anchored at
/// the start. need support for single character wild cards (`.`).
///
/// implemented via [RegexSet](https://docs.rs/regex/1.3.9/regex/struct.RegexSet.html)
///
/// > Match multiple (possibly overlapping) regular expressions in a single
/// scan. >
/// > A regex set corresponds to the union of two or more regular expressions.
/// > That is, a regex set will match text where at least one of its constituent
/// > regular expressions matches. A regex set as its formulated here provides a
/// touch more power: >  it will also report which regular expressions in the
/// set match. > Indeed, this is the key difference between regex sets and a
/// single Regex with many alternates, > since only one alternate can match at a
/// time.
use anyhow::Result;
use nom::{
    branch::alt,
    bytes::complete::{tag, take_while_m_n},
    combinator::{map, map_res},
    multi::many1,
    IResult,
};

// u16 because we need 257 possible values, all unsigned.
#[derive(Copy, Clone, Hash, Eq, PartialEq)]
pub struct Symbol(pub u16);

// impl note: value 256 is WILDCARD.
pub const WILDCARD: Symbol = Symbol(0x100);

// byte values map directly into their Symbol indices.
impl std::convert::From<u8> for Symbol {
    fn from(v: u8) -> Self {
        Symbol(v as u16)
    }
}

impl std::fmt::Display for Symbol {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        if self.0 == WILDCARD.0 {
            write!(f, ".")
        } else {
            write!(f, r"\x{:02X}", self.0)
        }
    }
}

// a pattern is just a sequence of symbols.
#[derive(Hash, PartialEq, Eq, Clone)]
pub struct Pattern(pub Vec<Symbol>);

impl std::fmt::Display for Pattern {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        let parts: Vec<String> = self.0.iter().map(|s| format!("{}", s)).collect();
        write!(
            f,
            r"(?x)    # whitespace allowed
              (?-u)   # disable unicode mode, so we can match raw bytes
              ^       # match only from start of data
              ({})    # capture the byte sequence
            ",
            parts.join("")
        )
    }
}

fn is_hex_digit(c: char) -> bool {
    c.is_digit(16)
}

fn from_hex(input: &str) -> Result<u8, std::num::ParseIntError> {
    u8::from_str_radix(input, 16)
}

/// parse a single hex byte, like `AB`
fn hex(input: &str) -> IResult<&str, u8> {
    map_res(take_while_m_n(2, 2, is_hex_digit), from_hex)(input)
}

/// parse a single byte signature element, which is either a hex byte or a
/// wildcard.
fn sig_element(input: &str) -> IResult<&str, Symbol> {
    alt((map(hex, Symbol::from), map(tag(".."), |_| WILDCARD)))(input)
}

/// parse byte signature elements, hex or wildcard.
fn byte_signature(input: &str) -> IResult<&str, Pattern> {
    let (input, elems) = many1(sig_element)(input)?;
    Ok((input, Pattern(elems)))
}

/// parse a pattern from a string like `AABB..DD`.
impl std::convert::From<&str> for Pattern {
    fn from(v: &str) -> Self {
        byte_signature(v).expect("failed to parse pattern").1
    }
}

pub struct PatternSet {
    patterns: Vec<Pattern>,
    re:       regex::bytes::RegexSet,
}

impl std::fmt::Debug for PatternSet {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        for pattern in self.patterns.iter() {
            writeln!(f, "  - {}", pattern)?;
        }
        Ok(())
    }
}

impl PatternSet {
    pub fn r#match(&self, buf: &[u8]) -> Vec<&Pattern> {
        self.re.matches(buf).into_iter().map(|i| &self.patterns[i]).collect()
    }

    pub fn builder() -> PatternSetBuilder {
        PatternSetBuilder { patterns: vec![] }
    }

    pub fn from_patterns(patterns: Vec<Pattern>) -> PatternSet {
        PatternSetBuilder { patterns }.build()
    }
}

pub struct PatternSetBuilder {
    patterns: Vec<Pattern>,
}

impl PatternSetBuilder {
    pub fn add_pattern(&mut self, pattern: Pattern) {
        self.patterns.push(pattern)
    }

    pub fn build(self) -> PatternSet {
        // should not be possible to generate invalid regex from a pattern
        // otherwise, programming error.
        // must reject invalid patterns when deserializing from pat/sig.

        let mut patterns = vec![];
        for pattern in self.patterns.iter() {
            patterns.push(format!("{}", pattern));
        }

        let re = regex::bytes::RegexSet::new(patterns).expect("invalid regex");

        PatternSet {
            patterns: self.patterns,
            re,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_empty_build() {
        PatternSet::builder().build();
    }

    // patterns:
    //   - pat0: aabbccdd
    //
    // transition table:
    //
    //  aa  bb  cc  dd  ..
    //  0:  1                   alive: pat0
    //  1:      2               alive: pat0
    //  2:          3           alive: pat0
    //  3:              4       alive: pat0
    //  4:                      matches: pat0
    #[test]
    fn test_add_one_pattern() {
        let mut b = PatternSet::builder();
        b.add_pattern(Pattern::from("AABBCCDD"));

        println!("{:?}", b.build());
    }

    // patterns:
    //   - pat0: aabbccdd
    //   - pat1: aabbcccc
    //
    // transition table:
    //
    //  aa  bb  cc  dd  ..
    //  0:  1                   alive: pat0 pat1
    //  1:      2               alive: pat0 pat1
    //  2:          3           alive: pat0 pat1
    //  3:          5   4       alive: pat0 pat1
    //  4:                      matches: pat0
    //  5:                      matches: pat1
    #[test]
    fn test_add_two_patterns() {
        let mut b = PatternSet::builder();
        b.add_pattern(Pattern::from("AABBCCDD"));
        b.add_pattern(Pattern::from("AABBCCCC"));

        println!("{:?}", b.build());
    }

    // patterns:
    //   - pat0: aabbccdd
    //   - pat1: aabbcc..
    //
    // transition table:
    //       aa  bb  cc  dd  ..
    //    0:  1                   alive: pat0 pat1
    //    1:      2               alive: pat0 pat1
    //    2:          3           alive: pat0 pat1
    //    3:              4   5   alive: pat0 pat1
    //    4:                      matches: pat0 pat1
    //    5:                      matches: pat1
    #[test]
    fn test_add_one_wildcard() {
        let mut b = PatternSet::builder();
        b.add_pattern(Pattern::from("AABBCCDD"));
        b.add_pattern(Pattern::from("AABBCC.."));

        println!("{:?}", b.build());
    }

    // we don't match when we don't have any patterns.
    #[test]
    fn test_match_empty() {
        let pattern_set = PatternSet::builder().build();
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\xCC\xDD").len(), 0);
    }

    // we match things we want to, and don't match other data.
    #[test]
    fn test_match_one() {
        let mut b = PatternSet::builder();
        b.add_pattern(Pattern::from("AABBCCDD"));
        let pattern_set = b.build();

        // true positive
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\xCC\xDD").len(), 1);
        // true negative
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\xCC\xEE").len(), 0);
    }

    // we match from the beginning of the buffer onwards,
    // ignoring trailing bytes beyond the length of the pattern.
    #[test]
    fn test_match_long() {
        let mut b = PatternSet::builder();
        b.add_pattern(Pattern::from("AABBCCDD"));
        let pattern_set = b.build();

        assert_eq!(pattern_set.r#match(b"\xAA\xBB\xCC\xDD\x00").len(), 1);
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\xCC\xDD\x11").len(), 1);
    }

    // we can match when there are single character wildcards present,
    // and order of the pattern declarations should not matter.
    #[test]
    fn test_match_one_tail_wildcard() {
        let mut b = PatternSet::builder();
        b.add_pattern(Pattern::from("AABBCC.."));
        b.add_pattern(Pattern::from("AABBCCDD"));
        let pattern_set = b.build();

        assert_eq!(pattern_set.r#match(b"\xAA\xBB\xCC\xDD").len(), 2);
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\xCC\xEE").len(), 1);
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\x00\x00").len(), 0);

        // order of patterns should not matter
        let mut b = PatternSet::builder();
        b.add_pattern(Pattern::from("AABBCCDD"));
        b.add_pattern(Pattern::from("AABBCC.."));
        let pattern_set = b.build();

        assert_eq!(pattern_set.r#match(b"\xAA\xBB\xCC\xDD").len(), 2);
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\xCC\xEE").len(), 1);
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\x00\x00").len(), 0);
    }

    // wildcards can be found in the middle of patterns, too.
    #[test]
    fn test_match_one_middle_wildcard() {
        let pattern_set = PatternSet::from_patterns(vec![Pattern::from("AABB..DD"), Pattern::from("AABBCCDD")]);

        assert_eq!(pattern_set.r#match(b"\xAA\xBB\xCC\xDD").len(), 2);
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\xEE\xDD").len(), 1);
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\x00\x00").len(), 0);

        // order of patterns should not matter
        let pattern_set = PatternSet::from_patterns(vec![Pattern::from("AABBCCDD"), Pattern::from("AABB..DD")]);

        assert_eq!(pattern_set.r#match(b"\xAA\xBB\xCC\xDD").len(), 2);
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\xEE\xDD").len(), 1);
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\x00\x00").len(), 0);
    }

    // we can have an arbitrary mix of wildcards and literals.
    #[test]
    fn test_match_many() {
        let pattern_set = PatternSet::from_patterns(vec![
            Pattern::from("AABB..DD"),
            Pattern::from("AABBCCDD"),
            Pattern::from("........"),
            Pattern::from("....CCDD"),
        ]);
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\xCC\xDD").len(), 4);
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\x00\xDD").len(), 2);
        assert_eq!(pattern_set.r#match(b"\xAA\xBB\x00\x00").len(), 1);
        assert_eq!(pattern_set.r#match(b"\x00\x00\xCC\xDD").len(), 2);
        assert_eq!(pattern_set.r#match(b"\x00\x00\x00\x00").len(), 1);
    }

    #[test]
    fn test_match_pathological_case() {
        let pattern_set = PatternSet::from_patterns(vec![
            Pattern::from("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"),
            Pattern::from("................................................................"),
        ]);
        assert_eq!(pattern_set.r#match(b"\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA\xAA").len(), 2);
    }
}