imping_runtime/module.rs
1//! Module identity and address resolution.
2
3use crate::error::{Error, Result, VtableMiss};
4use crate::sys;
5use core::fmt;
6use core::sync::atomic::{AtomicU32, AtomicUsize, Ordering};
7use std::sync::Mutex;
8
9/// A vtable looked up by RTTI type name and locator offset, and where it is.
10pub(crate) type FoundVtable = (&'static str, u32, core::result::Result<usize, VtableMiss>);
11
12/// One build of a module, identified by the timestamp and image size in its
13/// image header.
14///
15/// A packer or DRM wrapper can change the image size while keeping the
16/// timestamp, so the file a disassembler analysed and the one that runs may
17/// be two builds with the same code; a [`Module`] can accept both.
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
19pub struct Build {
20 /// The `TimeDateStamp` from the module's file header.
21 pub timestamp: u32,
22 /// The `SizeOfImage` from the module's optional header.
23 pub image_size: u32,
24}
25
26/// A module (executable or DLL) of the target process, as the bindings expect
27/// it: its file name and the builds they accept.
28///
29/// Declare one as a `static`, call [`init`](Self::init) once the module is
30/// loaded, then resolve addresses with [`address`](Self::address). Generated
31/// bindings declare one per IR module as `MODULE`, and their `init()` calls
32/// [`Module::init`].
33///
34/// # Examples
35///
36/// ```no_run
37/// use imping_runtime::{Build, Module};
38///
39/// static TARGET: Module = Module::new(
40/// "target.exe",
41/// &[Build { timestamp: 0x64b8f2a1, image_size: 0xa2b000 }],
42/// );
43///
44/// # fn main() -> imping_runtime::Result<()> {
45/// // SAFETY: the addresses used with `TARGET` belong to this build.
46/// unsafe { TARGET.init() }?;
47/// let update = TARGET.address(0x15040);
48/// # let _ = update;
49/// # Ok(())
50/// # }
51/// ```
52pub struct Module {
53 name: &'static str,
54 builds: &'static [Build],
55 /// Which labelled build each of `builds` is, as an index into `labels`;
56 /// empty when there is one unlabelled build.
57 build_of: &'static [usize],
58 /// The labels of the builds the bindings were generated from.
59 labels: &'static [&'static str],
60 /// The loaded build's index into `labels`, or `usize::MAX` if it is none
61 /// of `builds`; set before `base`.
62 build: AtomicUsize,
63 /// The loaded image's size; set before `base`.
64 image_size: AtomicU32,
65 /// The module's base address once `init` has succeeded, otherwise 0.
66 base: AtomicUsize,
67 /// Vtables looked up by RTTI type name and locator offset.
68 pub(crate) vtables: Mutex<Vec<FoundVtable>>,
69}
70
71impl Module {
72 /// Describes a module by its file name, such as `"target.exe"`, and the
73 /// builds it accepts.
74 ///
75 /// With no builds, [`init`](Self::init) always fails.
76 pub const fn new(name: &'static str, builds: &'static [Build]) -> Self {
77 Self::labelled(name, builds, &[], &[])
78 }
79
80 /// As [`new`](Self::new), for bindings generated from several labelled
81 /// builds: `build_of[i]` is the index into `labels` of `builds[i]`.
82 pub(crate) const fn labelled(
83 name: &'static str,
84 builds: &'static [Build],
85 build_of: &'static [usize],
86 labels: &'static [&'static str],
87 ) -> Self {
88 Self {
89 name,
90 builds,
91 build_of,
92 labels,
93 build: AtomicUsize::new(usize::MAX),
94 image_size: AtomicU32::new(0),
95 base: AtomicUsize::new(0),
96 vtables: Mutex::new(Vec::new()),
97 }
98 }
99
100 /// Finds the module in this process and checks its image header against
101 /// the accepted builds.
102 ///
103 /// Once this succeeds, further calls return `Ok` and do nothing. Failures
104 /// aren't remembered, so a DLL the target loads later can be initialized
105 /// once it's there.
106 ///
107 /// # Errors
108 ///
109 /// - [`Error::ModuleNotLoaded`] if no module of this name is loaded.
110 /// - [`Error::BuildMismatch`] if it is none of the accepted builds.
111 /// - [`Error::InvalidImage`] if its image header can't be read.
112 /// - [`Error::UnsupportedPlatform`] on platforms other than Windows.
113 ///
114 /// # Safety
115 ///
116 /// Everything later done through this module (addresses, function
117 /// signatures, type layouts) must be correct for every accepted build;
118 /// `init` only checks the identity. For generated bindings, the IR they
119 /// came from must describe each accepted build correctly.
120 pub unsafe fn init(&self) -> Result<()> {
121 if self.base().is_some() {
122 return Ok(());
123 }
124 let (base, identity) = sys::locate(self.name).map_err(|error| match error {
125 sys::LocateError::NotLoaded => Error::ModuleNotLoaded { module: self.name },
126 sys::LocateError::InvalidImage => Error::InvalidImage { module: self.name },
127 sys::LocateError::Unsupported => Error::UnsupportedPlatform,
128 })?;
129 let found = Build {
130 timestamp: identity.timestamp,
131 image_size: identity.image_size,
132 };
133 let Some(index) = self.builds.iter().position(|build| *build == found) else {
134 return Err(Error::BuildMismatch {
135 module: self.name,
136 found,
137 expected: self.builds,
138 });
139 };
140 self.build.store(self.build_of(index), Ordering::Release);
141 self.image_size.store(found.image_size, Ordering::Release);
142 self.base.store(base, Ordering::Release);
143 log::debug!("imping: `{}` found at {base:#x}", self.name);
144 Ok(())
145 }
146
147 /// Finds the module in this process and records its base whatever its
148 /// build, returning the build found.
149 ///
150 /// For trying a build the bindings don't list; [`init`](Self::init) is the
151 /// checked way in. An unlisted build is logged as a warning, and generated
152 /// functions then check their first bytes before they're called, falling
153 /// back to their byte pattern. If the module is already initialized,
154 /// nothing changes and the loaded build is returned.
155 ///
156 /// # Errors
157 ///
158 /// - [`Error::ModuleNotLoaded`] if no module of this name is loaded.
159 /// - [`Error::InvalidImage`] if its image header can't be read.
160 /// - [`Error::UnsupportedPlatform`] on platforms other than Windows.
161 ///
162 /// # Safety
163 ///
164 /// Everything later done through this module (addresses, function
165 /// signatures, type layouts) must be correct for the build actually
166 /// loaded, which nothing has checked. Generated functions still verify
167 /// their first bytes where the bindings know them; globals and raw
168 /// addresses are used unchecked.
169 pub unsafe fn init_any_build(&self) -> Result<Build> {
170 let (base, identity) = sys::locate(self.name).map_err(|error| match error {
171 sys::LocateError::NotLoaded => Error::ModuleNotLoaded { module: self.name },
172 sys::LocateError::InvalidImage => Error::InvalidImage { module: self.name },
173 sys::LocateError::Unsupported => Error::UnsupportedPlatform,
174 })?;
175 let found = Build {
176 timestamp: identity.timestamp,
177 image_size: identity.image_size,
178 };
179 if self.base().is_none() {
180 let index = self.builds.iter().position(|build| *build == found);
181 let build = index.map_or(usize::MAX, |index| self.build_of(index));
182 self.build.store(build, Ordering::Release);
183 if index.is_none() {
184 log::warn!(
185 "imping: `{}` (timestamp {:#x}, image size {:#x}) is not a build the \
186 bindings accept; initialized anyway",
187 self.name,
188 found.timestamp,
189 found.image_size
190 );
191 }
192 self.image_size.store(found.image_size, Ordering::Release);
193 self.base.store(base, Ordering::Release);
194 log::debug!("imping: `{}` found at {base:#x}", self.name);
195 }
196 Ok(found)
197 }
198
199 /// Returns the module's file name.
200 pub fn name(&self) -> &'static str {
201 self.name
202 }
203
204 /// Returns the module's base address, or `None` until [`init`](Self::init)
205 /// has succeeded.
206 pub fn base(&self) -> Option<usize> {
207 match self.base.load(Ordering::Acquire) {
208 0 => None,
209 base => Some(base),
210 }
211 }
212
213 /// The index into the labels of `builds[index]`: 0 without labels.
214 const fn build_of(&self, index: usize) -> usize {
215 if index < self.build_of.len() {
216 self.build_of[index]
217 } else {
218 0
219 }
220 }
221
222 /// The loaded build's index into the labels, or `None` before `init` or
223 /// for a build none of them is.
224 pub(crate) fn build_index(&self) -> Option<usize> {
225 self.base()?;
226 match self.build.load(Ordering::Acquire) {
227 usize::MAX => None,
228 index => Some(index),
229 }
230 }
231
232 /// The label of the build at `index`, or `""` without labels.
233 pub(crate) fn label(&self, index: usize) -> &'static str {
234 self.labels.get(index).copied().unwrap_or("")
235 }
236
237 /// The loaded image's size, or 0 before it's initialized.
238 pub(crate) fn image_size(&self) -> usize {
239 self.image_size.load(Ordering::Acquire) as usize
240 }
241
242 /// Returns the address of `rva`, an address relative to the module's base.
243 ///
244 /// # Panics
245 ///
246 /// If [`init`](Self::init) has not succeeded.
247 pub fn address(&self, rva: u32) -> usize {
248 match self.base() {
249 Some(base) => base + rva as usize,
250 None => panic!(
251 "imping: `{}` is not initialized; call its `init()` first",
252 self.name
253 ),
254 }
255 }
256}
257
258impl fmt::Debug for Module {
259 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
260 f.debug_struct("Module")
261 .field("name", &self.name)
262 .field("builds", &self.builds)
263 .field("base", &self.base().map(|base| base as *const ()))
264 .finish()
265 }
266}