wraith/navigation/
module.rs

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

wraith/navigation/
module.rs
