Skip to main content

imping_hook/
scan.rs

1//! Finding code by a byte pattern.
2//!
3//! A pattern is hex bytes separated by spaces, with `?` or `??` for any byte:
4//! `"48 89 5C 24 ?? 57"`. It's searched for in a loaded module's executable
5//! sections and must match exactly once, so a pattern that no longer fits the
6//! program fails instead of landing somewhere else. [`find`] and [`find_in`]
7//! are the short way, [`Pattern`] can also step from the match to an address
8//! near it, and [`Hook::scan`] and [`Injection::scan`] patch what a pattern
9//! finds.
10//!
11//! [`Hook::scan`]: crate::Hook::scan
12//! [`Injection::scan`]: crate::Injection::scan
13//!
14//! # Examples
15//!
16//! ```no_run
17//! use imping_hook::scan::{self, Pattern};
18//!
19//! let address = scan::find("48 89 5C 24 ?? 57 48 83 EC 20")?;
20//! let address = scan::find_in("engine.dll", "48 89 5C 24 ?? 57")?;
21//! // The function a matched call goes to.
22//! let address = Pattern::new("E8 ?? ?? ?? ?? 48 8B D8").follow_call().find()?;
23//! # Ok::<(), imping_hook::Error>(())
24//! ```
25
26use crate::error::{Error, Result};
27use pelite::image::IMAGE_SCN_MEM_EXECUTE;
28use pelite::pattern::Atom;
29use pelite::pe::{Pe, PeView};
30
31/// Finds the one place in the main module's code that `pattern` matches.
32///
33/// # Errors
34///
35/// - [`Error::InvalidPattern`] if `pattern` can't be parsed.
36/// - [`Error::PatternNotFound`] or [`Error::PatternNotUnique`] unless it
37///   matches exactly once.
38pub fn find(pattern: &str) -> Result<*const ()> {
39    Pattern::new(pattern).find()
40}
41
42/// Finds the one place in a loaded module's code that `pattern` matches.
43///
44/// `module` is a module name such as `"engine.dll"`. It must stay loaded while
45/// it's searched.
46///
47/// # Errors
48///
49/// As for [`find`], plus [`Error::ModuleNotFound`] if no module of that name is
50/// loaded.
51pub fn find_in(module: &str, pattern: &str) -> Result<*const ()> {
52    Pattern::new(pattern).in_module(module).find()
53}
54
55/// Where a hook or injection goes: an address, or a pattern to find it by.
56pub(crate) enum Target {
57    Address(*const ()),
58    Pattern(Pattern),
59}
60
61impl Target {
62    /// Returns the address, searching for the pattern if there is one.
63    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/// A byte pattern to find, where to look for it, and how to get from the match
72/// to the address wanted.
73///
74/// See the [module docs](self) for the syntax.
75///
76/// # Examples
77///
78/// ```no_run
79/// use imping_hook::scan::Pattern;
80///
81/// // The match starts four bytes into the function.
82/// let function = Pattern::new("48 83 EC 20 33 DB").offset(-4).find()?;
83/// // The function called at the match, in another module.
84/// let callee = Pattern::new("E8 ?? ?? ?? ?? 48 8B D8")
85///     .in_module("engine.dll")
86///     .follow_call()
87///     .find()?;
88/// # Ok::<(), imping_hook::Error>(())
89/// ```
90#[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    /// Creates a pattern to find in the main module, such as
101    /// `"48 89 5C 24 ?? 57"`.
102    ///
103    /// It's parsed by [`find`](Self::find).
104    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    /// Searches a loaded module such as `"engine.dll"` instead of the main
114    /// module.
115    pub fn in_module(mut self, module: &str) -> Self {
116        self.module = Some(module.to_owned());
117        self
118    }
119
120    /// Moves the address this many bytes on from the start of the match, or
121    /// back for a negative count.
122    pub fn offset(mut self, bytes: isize) -> Self {
123        self.offset = bytes;
124        self
125    }
126
127    /// Follows the `call` or `jmp` at the address to where it goes.
128    ///
129    /// For a pattern that matches a call site rather than the function itself.
130    /// The pattern must spell out the instruction's `E8` (call) or `E9` (jmp)
131    /// opcode at the offset, and cover its 4-byte displacement, which may be
132    /// wildcards: `"E8 ?? ?? ?? ?? 48 8B D8"`.
133    pub fn follow_call(mut self) -> Self {
134        self.follow_call = true;
135        self
136    }
137
138    /// Finds the one place the pattern matches and returns the address it
139    /// leads to.
140    ///
141    /// The module must stay loaded while it's searched.
142    ///
143    /// # Errors
144    ///
145    /// - [`Error::InvalidPattern`] if the pattern can't be parsed, or doesn't
146    ///   spell out the call for [`follow_call`](Self::follow_call).
147    /// - [`Error::PatternNotFound`] or [`Error::PatternNotUnique`] unless it
148    ///   matches exactly once.
149    /// - [`Error::ModuleNotFound`] if the module isn't loaded.
150    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        // SAFETY: `base` is a module the loader has mapped, and the caller keeps
171        // it loaded while it's searched.
172        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            // SAFETY: `check_call` made sure the match covers the whole call, so
206            // its displacement is mapped code.
207            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    /// Checks the pattern spells out a `call` or `jmp` (`E8` or `E9`) at the
214    /// offset, with room for its 4-byte displacement.
215    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
234/// Parses a pattern into its bytes (`None` for a wildcard), each with the
235/// position of its token in `pattern`.
236fn 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            // from_str_radix alone would also take a sign, like "+1".
247            _ 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}