1use crate::error::{Error, Result};
27use pelite::image::IMAGE_SCN_MEM_EXECUTE;
28use pelite::pattern::Atom;
29use pelite::pe::{Pe, PeView};
30
31pub fn find(pattern: &str) -> Result<*const ()> {
39 Pattern::new(pattern).find()
40}
41
42pub fn find_in(module: &str, pattern: &str) -> Result<*const ()> {
52 Pattern::new(pattern).in_module(module).find()
53}
54
55pub(crate) enum Target {
57 Address(*const ()),
58 Pattern(Pattern),
59}
60
61impl Target {
62 pub(crate) fn resolve(self) -> Result<*const ()> {
64 match self {
65 Target::Address(address) => Ok(address),
66 Target::Pattern(pattern) => pattern.find(),
67 }
68 }
69}
70
71#[derive(Debug, Clone)]
91#[must_use = "a pattern does nothing until `find` is called"]
92pub struct Pattern {
93 pattern: String,
94 module: Option<String>,
95 offset: isize,
96 follow_call: bool,
97}
98
99impl Pattern {
100 pub fn new(pattern: &str) -> Self {
105 Pattern {
106 pattern: pattern.to_owned(),
107 module: None,
108 offset: 0,
109 follow_call: false,
110 }
111 }
112
113 pub fn in_module(mut self, module: &str) -> Self {
116 self.module = Some(module.to_owned());
117 self
118 }
119
120 pub fn offset(mut self, bytes: isize) -> Self {
123 self.offset = bytes;
124 self
125 }
126
127 pub fn follow_call(mut self) -> Self {
134 self.follow_call = true;
135 self
136 }
137
138 pub fn find(&self) -> Result<*const ()> {
151 let bytes = parse(&self.pattern)?;
152 if self.follow_call {
153 self.check_call(&bytes)?;
154 }
155 let base = crate::sys::module_base(self.module.as_deref()).ok_or_else(|| {
156 Error::ModuleNotFound {
157 name: self
158 .module
159 .clone()
160 .unwrap_or_else(|| "main module".to_owned()),
161 }
162 })?;
163
164 let mut atoms = vec![Atom::Save(0)];
165 atoms.extend(bytes.iter().map(|&(byte, _)| match byte {
166 Some(byte) => Atom::Byte(byte),
167 None => Atom::Skip(1),
168 }));
169
170 let view = unsafe { PeView::module(base) };
173 let scanner = view.scanner();
174 let (mut first, mut matches) = (None, 0);
175 for section in view.section_headers().iter() {
176 if section.Characteristics & IMAGE_SCN_MEM_EXECUTE == 0 {
177 continue;
178 }
179 let start = section.VirtualAddress;
180 let size = section.VirtualSize.max(section.SizeOfRawData);
181 let mut found = scanner.matches(&atoms, start..start + size);
182 let mut save = [0; 1];
183 while found.next(&mut save) {
184 first.get_or_insert(save[0]);
185 matches += 1;
186 }
187 }
188 let rva = match (first, matches) {
189 (Some(rva), 1) => rva,
190 (None, _) => {
191 return Err(Error::PatternNotFound {
192 pattern: self.pattern.clone(),
193 });
194 }
195 _ => {
196 return Err(Error::PatternNotUnique {
197 pattern: self.pattern.clone(),
198 matches,
199 });
200 }
201 };
202
203 let mut address = (base as usize + rva as usize).wrapping_add_signed(self.offset);
204 if self.follow_call {
205 let displacement = unsafe { ((address + 1) as *const i32).read_unaligned() };
208 address = (address + 5).wrapping_add_signed(displacement as isize);
209 }
210 Ok(address as *const ())
211 }
212
213 fn check_call(&self, bytes: &[(Option<u8>, usize)]) -> Result<()> {
216 let opcode = usize::try_from(self.offset)
217 .ok()
218 .filter(|&at| at + 5 <= bytes.len());
219 match opcode.map(|at| bytes[at]) {
220 Some((Some(0xE8 | 0xE9), _)) => Ok(()),
221 Some((_, position)) => Err(self.invalid_at(position)),
222 None => Err(self.invalid_at(self.pattern.len())),
223 }
224 }
225
226 fn invalid_at(&self, position: usize) -> Error {
227 Error::InvalidPattern {
228 pattern: self.pattern.clone(),
229 position,
230 }
231 }
232}
233
234fn parse(pattern: &str) -> Result<Vec<(Option<u8>, usize)>> {
237 let invalid = |position| Error::InvalidPattern {
238 pattern: pattern.to_owned(),
239 position,
240 };
241 let mut bytes = Vec::new();
242 for token in pattern.split_ascii_whitespace() {
243 let position = token.as_ptr() as usize - pattern.as_ptr() as usize;
244 let byte = match token {
245 "?" | "??" => None,
246 _ if token.len() == 2 && token.bytes().all(|c| c.is_ascii_hexdigit()) => {
248 Some(u8::from_str_radix(token, 16).map_err(|_| invalid(position))?)
249 }
250 _ => return Err(invalid(position)),
251 };
252 bytes.push((byte, position));
253 }
254 if bytes.iter().all(|(byte, _)| byte.is_none()) {
255 return Err(invalid(0));
256 }
257 Ok(bytes)
258}
259
260#[cfg(test)]
261mod tests {
262 use super::*;
263
264 fn position(pattern: &str) -> Option<usize> {
265 match parse(pattern) {
266 Err(Error::InvalidPattern { position, .. }) => Some(position),
267 _ => None,
268 }
269 }
270
271 #[test]
272 fn parses_bytes_and_wildcards() {
273 let bytes: Vec<_> = parse("48 ? 5c ??")
274 .unwrap()
275 .into_iter()
276 .map(|(b, _)| b)
277 .collect();
278 assert_eq!(bytes, [Some(0x48), None, Some(0x5c), None]);
279 }
280
281 #[test]
282 fn points_at_the_bad_token() {
283 assert_eq!(position("48 8B ZZ"), Some(6));
284 assert_eq!(position("48 8"), Some(3));
285 assert_eq!(position("48 +1"), Some(4));
286 assert_eq!(position("48 ???"), Some(3));
287 }
288
289 #[test]
290 fn rejects_empty_and_all_wildcard_patterns() {
291 assert_eq!(position(""), Some(0));
292 assert_eq!(position("?? ??"), Some(0));
293 }
294
295 #[test]
296 fn follow_call_needs_the_whole_call() {
297 let check = |pattern: &str, offset| Pattern::new(pattern).offset(offset).follow_call();
298 assert!(
299 check("E8 ?? ?? ?? ??", 0)
300 .check_call(&parse("E8 ?? ?? ?? ??").unwrap())
301 .is_ok()
302 );
303 assert!(
304 check("48 E9 ?? ?? ?? ??", 1)
305 .check_call(&parse("48 E9 ?? ?? ?? ??").unwrap())
306 .is_ok()
307 );
308 let short = check("E8 ?? ??", 0);
309 assert!(matches!(
310 short.check_call(&parse("E8 ?? ??").unwrap()),
311 Err(Error::InvalidPattern { position: 8, .. })
312 ));
313 let wrong = check("48 8B ?? ?? ?? ??", 1);
314 assert!(matches!(
315 wrong.check_call(&parse("48 8B ?? ?? ?? ??").unwrap()),
316 Err(Error::InvalidPattern { position: 3, .. })
317 ));
318 }
319}