robs/
scanner.rs

1use crate::signature::Signature;
2use rayon::prelude::*;
3use std::usize;
4
5pub trait AobScanner {
6    fn find_signature(&self, signature: &Signature) -> Option<usize>;
7}
8
9impl AobScanner for &[u8] {
10    fn find_signature(&self, signature: &Signature) -> Option<usize> {
11        crate::scanner::find_signature(self, signature)
12    }
13}
14
15impl AobScanner for [u8] {
16    fn find_signature(&self, signature: &Signature) -> Option<usize> {
17        crate::scanner::find_signature(self, signature)
18    }
19}
20
21impl AobScanner for Vec<u8> {
22    fn find_signature(&self, signature: &Signature) -> Option<usize> {
23        crate::scanner::find_signature(self, signature)
24    }
25}
26
27pub fn find_signature(search_region: &[u8], signature: &Signature) -> Option<usize> {
28    let first_index = signature.first_byte?;
29    let first_item = signature.pattern[first_index];
30    let mask_len = signature.mask.len();
31
32    search_region
33        .par_iter()
34        .positions(|&item| item == first_item)
35        .find_any(|&index| {
36            check_mask(
37                &search_region[(index - first_index)..(index + mask_len)],
38                signature,
39            )
40        })
41        .map(|index| index + signature.offset)
42}
43
44fn check_mask(search_region: &[u8], signature: &Signature) -> bool {
45    signature
46        .matching_indices
47        .iter()
48        .all(|&i| search_region[i] == signature.pattern[i])
49}
50
51#[cfg(test)]
52mod tests {
53    use crate::{
54        scanner::{check_mask, find_signature},
55        signature::Signature,
56    };
57
58    #[test]
59    fn test_check_mask() {
60        let search_region = &[0x0B, 0x0C, 0x0D, 0x0E, 0x0F];
61        let signature = Signature::new("0B ?? 0D", 0);
62        assert!(signature.is_ok());
63        let signature = signature.unwrap();
64        let check = check_mask(search_region, &signature);
65        assert!(check);
66    }
67
68    #[test]
69    fn test_find_signature() {
70        let search_region = &[
71            0x00, 0x0B, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C,
72            0x0D, 0x0E, 0x0F,
73        ];
74        let signature = Signature::new("0B ?? 0D", 0);
75        assert!(signature.is_ok());
76        let signature = signature.unwrap();
77        let find = find_signature(search_region, &signature);
78        assert!(find.is_some());
79        let find = find.unwrap();
80        assert_eq!(find, 12);
81    }
82
83    #[test]
84    fn test_signature_not_found() {
85        let search_region = &[
86            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D,
87            0x0E, 0x0F,
88        ];
89        let signature = Signature::new("FF FF", 0);
90        assert!(signature.is_ok());
91        let signature = signature.unwrap();
92        let find = find_signature(search_region, &signature);
93        assert!(find.is_none());
94    }
95
96    #[test]
97    fn test_signature_exact() {
98        let search_region = &[
99            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D,
100            0x0E, 0x0F,
101        ];
102        let signature = Signature::new("00 01 02 03 04 05 06 07 08 09 0A 0B 0C 0D 0E 0F", 0);
103        assert!(signature.is_ok());
104        let signature = signature.unwrap();
105        let find = find_signature(search_region, &signature);
106        assert!(find.is_some());
107        let find = find.unwrap();
108        assert_eq!(find, 0);
109    }
110
111    #[test]
112    fn test_signature_exact_except_last() {
113        let search_region = &[
114            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D,
115            0x0E, 0x0F,
116        ];
117        let signature = Signature::new("00 01 02 03 04 05 06 07 08 09 0A 0B 0C 0D 0E FF", 0);
118        assert!(signature.is_ok());
119        let signature = signature.unwrap();
120        let find = find_signature(search_region, &signature);
121        assert!(find.is_none());
122    }
123
124    #[test]
125    fn test_signature_start_with_wildcard() {
126        let search_region = &[
127            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D,
128            0x0E, 0x0F, 0xFF, 0xCD, 0xAB, 0x89, 0x67, 0x45, 0x23, 0x01,
129        ];
130        let signature = Signature::new("?? 01 02 03 04 05 06 07 08 09 0A 0B 0C 0D 0E 0F", 0);
131        assert!(signature.is_ok());
132        let signature = signature.unwrap();
133        let find = find_signature(search_region, &signature);
134        assert!(find.is_some());
135        let find = find.unwrap();
136        assert_eq!(find, 1);
137    }
138}