wraith/navigation/
module.rs

1//! High-level module abstraction
2
3use crate::error::{Result, WraithError};
4use crate::structures::pe::{
5    DataDirectoryType, DosHeader, ExportDirectory, NtHeaders, NtHeaders32, NtHeaders64,
6};
7use crate::structures::{LdrDataTableEntry, ListEntry};
8use crate::structures::pe::nt_headers::{NT_SIGNATURE, PE32_MAGIC};
9use core::ptr::NonNull;
10
11// max length for export/import names to prevent unbounded reads
12const MAX_NAME_LENGTH: usize = 512;
13// max reasonable number of exports to prevent DoS
14const MAX_EXPORT_COUNT: usize = 0x10000;
15
16/// immutable reference to a loaded module
17pub struct Module<'a> {
18    entry: &'a LdrDataTableEntry,
19}
20
21impl<'a> Module<'a> {
22    /// create module from LDR_DATA_TABLE_ENTRY reference
23    pub(crate) fn from_entry(entry: &'a LdrDataTableEntry) -> Self {
24        Self { entry }
25    }
26
27    /// get module base address
28    pub fn base(&self) -> usize {
29        self.entry.base()
30    }
31
32    /// get module size in bytes
33    pub fn size(&self) -> usize {
34        self.entry.size()
35    }
36
37    /// get entry point address (may be 0 if no entry point)
38    pub fn entry_point(&self) -> usize {
39        self.entry.entry_point()
40    }
41
42    /// get full path to module
43    pub fn full_path(&self) -> String {
44        // SAFETY: full_dll_name is valid for loaded modules
45        unsafe { self.entry.full_name() }
46    }
47
48    /// get module filename only (e.g., "ntdll.dll")
49    pub fn name(&self) -> String {
50        // SAFETY: base_dll_name is valid for loaded modules
51        unsafe { self.entry.base_name() }
52    }
53
54    /// get name as lowercase for comparison
55    pub fn name_lowercase(&self) -> String {
56        self.name().to_lowercase()
57    }
58
59    /// check if address falls within this module
60    pub fn contains(&self, address: usize) -> bool {
61        self.entry.contains_address(address)
62    }
63
64    /// check if module name matches (case-insensitive)
65    pub fn matches_name(&self, name: &str) -> bool {
66        // SAFETY: name buffers are valid for loaded modules
67        unsafe { self.entry.matches_name(name) }
68    }
69
70    /// get raw LDR entry reference
71    pub fn as_ldr_entry(&self) -> &LdrDataTableEntry {
72        self.entry
73    }
74
75    /// get DOS header
76    pub fn dos_header(&self) -> Result<&DosHeader> {
77        let base = self.base();
78        if base == 0 {
79            return Err(WraithError::NullPointer {
80                context: "module base",
81            });
82        }
83
84        // SAFETY: base points to valid PE image for loaded modules
85        let dos = unsafe { &*(base as *const DosHeader) };
86
87        if !dos.is_valid() {
88            return Err(WraithError::InvalidPeFormat {
89                reason: "invalid DOS signature".into(),
90            });
91        }
92
93        Ok(dos)
94    }
95
96    /// get NT headers
97    pub fn nt_headers(&self) -> Result<NtHeaders> {
98        let dos = self.dos_header()?;
99
100        // validate e_lfanew is reasonable
101        if !dos.is_nt_offset_valid() {
102            let lfanew = dos.e_lfanew; // copy from packed struct
103            return Err(WraithError::InvalidPeFormat {
104                reason: format!("invalid e_lfanew: {:#x}", lfanew),
105            });
106        }
107
108        let nt_offset = dos.nt_headers_offset();
109
110        // validate nt_offset is within module bounds (need space for NT headers)
111        const MIN_NT_HEADERS_SIZE: usize = 256; // enough for signature + file header + optional header
112        if nt_offset + MIN_NT_HEADERS_SIZE > self.size() {
113            return Err(WraithError::InvalidPeFormat {
114                reason: "e_lfanew points outside module bounds".into(),
115            });
116        }
117
118        let nt_addr = self.base() + nt_offset;
119
120        // SAFETY: validated that nt_addr is within module bounds
121        let signature = unsafe { *(nt_addr as *const u32) };
122        if signature != NT_SIGNATURE {
123            return Err(WraithError::InvalidPeFormat {
124                reason: "invalid NT signature".into(),
125            });
126        }
127
128        // check if 32 or 64 bit
129        let magic_offset = nt_addr + 4 + 20; // after signature + file header
130        let magic = unsafe { *(magic_offset as *const u16) };
131
132        if magic == PE32_MAGIC {
133            let headers = unsafe { &*(nt_addr as *const NtHeaders32) };
134            Ok(NtHeaders::Headers32(*headers))
135        } else {
136            let headers = unsafe { &*(nt_addr as *const NtHeaders64) };
137            Ok(NtHeaders::Headers64(*headers))
138        }
139    }
140
141    /// convert RVA to absolute address
142    pub fn rva_to_va(&self, rva: u32) -> usize {
143        self.base() + rva as usize
144    }
145
146    /// convert absolute address to RVA
147    pub fn va_to_rva(&self, va: usize) -> Option<u32> {
148        if va >= self.base() && va < self.base() + self.size() {
149            Some((va - self.base()) as u32)
150        } else {
151            None
152        }
153    }
154
155    /// get export by name
156    pub fn get_export(&self, name: &str) -> Result<usize> {
157        let nt = self.nt_headers()?;
158        let export_dir = nt
159            .data_directory(DataDirectoryType::Export.index())
160            .ok_or_else(|| WraithError::InvalidPeFormat {
161                reason: "no export directory".into(),
162            })?;
163
164        if !export_dir.is_present() {
165            return Err(WraithError::InvalidPeFormat {
166                reason: "export directory not present".into(),
167            });
168        }
169
170        // validate export directory RVA is within module
171        if !self.is_rva_valid(export_dir.virtual_address, core::mem::size_of::<ExportDirectory>())
172        {
173            return Err(WraithError::InvalidPeFormat {
174                reason: "export directory RVA outside module bounds".into(),
175            });
176        }
177
178        let export_va = self.rva_to_va(export_dir.virtual_address);
179        // SAFETY: validated RVA is within bounds
180        let exports = unsafe { &*(export_va as *const ExportDirectory) };
181
182        let num_names = exports.number_of_names as usize;
183        let num_functions = exports.number_of_functions as usize;
184
185        // sanity check export counts
186        if num_names > MAX_EXPORT_COUNT || num_functions > MAX_EXPORT_COUNT {
187            return Err(WraithError::InvalidPeFormat {
188                reason: format!("unreasonable export count: {num_names} names, {num_functions} functions"),
189            });
190        }
191
192        // validate array RVAs
193        let names_size = num_names.saturating_mul(4);
194        let ordinals_size = num_names.saturating_mul(2);
195        let functions_size = num_functions.saturating_mul(4);
196
197        if !self.is_rva_valid(exports.address_of_names, names_size)
198            || !self.is_rva_valid(exports.address_of_name_ordinals, ordinals_size)
199            || !self.is_rva_valid(exports.address_of_functions, functions_size)
200        {
201            return Err(WraithError::InvalidPeFormat {
202                reason: "export table array RVA outside module bounds".into(),
203            });
204        }
205
206        let names_va = self.rva_to_va(exports.address_of_names);
207        let ordinals_va = self.rva_to_va(exports.address_of_name_ordinals);
208        let functions_va = self.rva_to_va(exports.address_of_functions);
209
210        // search for name
211        for i in 0..num_names {
212            // SAFETY: validated arrays are within bounds
213            let name_rva = unsafe { *((names_va + i * 4) as *const u32) };
214
215            // validate name RVA
216            if !self.is_rva_valid(name_rva, 1) {
217                continue; // skip invalid entries
218            }
219
220            let name_va = self.rva_to_va(name_rva);
221            let export_name = match self.read_string_at(name_va) {
222                Some(s) => s,
223                None => continue, // skip unreadable names
224            };
225
226            if export_name == name {
227                let ordinal = unsafe { *((ordinals_va + i * 2) as *const u16) } as usize;
228
229                // validate ordinal is within functions array
230                if ordinal >= num_functions {
231                    return Err(WraithError::InvalidPeFormat {
232                        reason: format!("ordinal {ordinal} exceeds function count {num_functions}"),
233                    });
234                }
235
236                let func_rva = unsafe { *((functions_va + ordinal * 4) as *const u32) };
237
238                // check for forwarded export
239                if func_rva >= export_dir.virtual_address
240                    && func_rva < export_dir.virtual_address + export_dir.size
241                {
242                    // forwarder string is at the func_rva location
243                    let forwarder_va = self.rva_to_va(func_rva);
244                    let forwarder = self.read_string_at(forwarder_va)
245                        .unwrap_or("unknown")
246                        .to_string();
247                    return Err(WraithError::ForwardedExport { forwarder });
248                }
249
250                let func_va = self.rva_to_va(func_rva);
251                return Ok(func_va);
252            }
253        }
254
255        Err(WraithError::ModuleNotFound {
256            name: format!("export {name} not found"),
257        })
258    }
259
260    /// check if RVA + size is within module bounds
261    fn is_rva_valid(&self, rva: u32, size: usize) -> bool {
262        let rva = rva as usize;
263        rva < self.size() && size <= self.size() - rva
264    }
265
266    /// safely read a null-terminated string at address within module
267    fn read_string_at(&self, addr: usize) -> Option<&str> {
268        let base = self.base();
269        let end = base + self.size();
270
271        if addr < base || addr >= end {
272            return None;
273        }
274
275        let max_len = (end - addr).min(MAX_NAME_LENGTH);
276        let ptr = addr as *const u8;
277
278        // find null terminator within bounds
279        let mut len = 0;
280        while len < max_len {
281            // SAFETY: addr is within module bounds, iterating up to max_len
282            let byte = unsafe { *ptr.add(len) };
283            if byte == 0 {
284                break;
285            }
286            len += 1;
287        }
288
289        if len == 0 || len >= max_len {
290            return None; // empty or no null terminator found
291        }
292
293        // SAFETY: we've verified bounds and found null terminator
294        let bytes = unsafe { core::slice::from_raw_parts(ptr, len) };
295        core::str::from_utf8(bytes).ok()
296    }
297
298    /// get export by ordinal
299    pub fn get_export_by_ordinal(&self, ordinal: u16) -> Result<usize> {
300        let nt = self.nt_headers()?;
301        let export_dir = nt
302            .data_directory(DataDirectoryType::Export.index())
303            .ok_or_else(|| WraithError::InvalidPeFormat {
304                reason: "no export directory".into(),
305            })?;
306
307        if !export_dir.is_present() {
308            return Err(WraithError::InvalidPeFormat {
309                reason: "export directory not present".into(),
310            });
311        }
312
313        let export_va = self.rva_to_va(export_dir.virtual_address);
314        // SAFETY: export directory is present and valid
315        let exports = unsafe { &*(export_va as *const ExportDirectory) };
316
317        let index = ordinal as usize - exports.base as usize;
318        if index >= exports.number_of_functions as usize {
319            return Err(WraithError::InvalidPeFormat {
320                reason: "ordinal out of range".into(),
321            });
322        }
323
324        let functions_va = self.rva_to_va(exports.address_of_functions);
325        let func_rva = unsafe { *((functions_va + index * 4) as *const u32) };
326
327        Ok(self.rva_to_va(func_rva))
328    }
329}
330
331/// owned handle to a module (allows modifications)
332pub struct ModuleHandle {
333    entry: NonNull<LdrDataTableEntry>,
334}
335
336impl ModuleHandle {
337    /// create handle from raw LDR entry pointer
338    ///
339    /// # Safety
340    /// pointer must be valid LDR_DATA_TABLE_ENTRY
341    pub unsafe fn from_raw(ptr: *mut LdrDataTableEntry) -> Option<Self> {
342        NonNull::new(ptr).map(|entry| Self { entry })
343    }
344
345    /// get raw pointer
346    pub fn as_ptr(&self) -> *mut LdrDataTableEntry {
347        self.entry.as_ptr()
348    }
349
350    /// borrow as immutable Module
351    pub fn as_module(&self) -> Module<'_> {
352        // SAFETY: entry pointer is valid for lifetime of handle
353        Module::from_entry(unsafe { self.entry.as_ref() })
354    }
355
356    /// get mutable access to LDR entry (for unlinking)
357    ///
358    /// # Safety
359    /// caller must ensure no other references exist
360    pub unsafe fn as_entry_mut(&mut self) -> &mut LdrDataTableEntry {
361        unsafe { self.entry.as_mut() }
362    }
363
364    /// get pointers to all three list links
365    pub fn get_link_pointers(&self) -> ModuleLinkPointers {
366        // SAFETY: entry is valid
367        let entry = unsafe { self.entry.as_ref() };
368        ModuleLinkPointers {
369            in_load_order: &entry.in_load_order_links as *const _ as *mut ListEntry,
370            in_memory_order: &entry.in_memory_order_links as *const _ as *mut ListEntry,
371            in_initialization_order: &entry.in_initialization_order_links as *const _
372                as *mut ListEntry,
373        }
374    }
375}
376
377// ModuleHandle doesn't impl Drop - it's a reference, not ownership
378
379// SAFETY: ModuleHandle is a pointer to process-wide structure
380unsafe impl Send for ModuleHandle {}
381unsafe impl Sync for ModuleHandle {}
382
383/// pointers to module's list entry links
384pub struct ModuleLinkPointers {
385    pub in_load_order: *mut ListEntry,
386    pub in_memory_order: *mut ListEntry,
387    pub in_initialization_order: *mut ListEntry,
388}
wraith/navigation/module.rs

wraith/navigation/
module.rs
