wraith/navigation/
memory_regions.rs

1//! Memory region enumeration using VirtualQuery
2
3#[cfg(all(not(feature = "std"), feature = "alloc"))]
4use alloc::vec::Vec;
5
6#[cfg(feature = "std")]
7use std::vec::Vec;
8
9use crate::error::{Result, WraithError};
10
11/// memory region information
12#[derive(Debug, Clone)]
13pub struct MemoryRegion {
14    pub base_address: usize,
15    pub allocation_base: usize,
16    pub allocation_protect: u32,
17    pub region_size: usize,
18    pub state: MemoryState,
19    pub protect: u32,
20    pub memory_type: MemoryType,
21}
22
23/// memory state
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum MemoryState {
26    Commit,
27    Reserve,
28    Free,
29}
30
31/// memory type
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum MemoryType {
34    Image,   // mapped executable image
35    Mapped,  // memory-mapped file
36    Private, // private memory
37    Unknown,
38}
39
40impl MemoryRegion {
41    /// check if region is executable
42    pub fn is_executable(&self) -> bool {
43        (self.protect & PAGE_EXECUTE) != 0
44            || (self.protect & PAGE_EXECUTE_READ) != 0
45            || (self.protect & PAGE_EXECUTE_READWRITE) != 0
46            || (self.protect & PAGE_EXECUTE_WRITECOPY) != 0
47    }
48
49    /// check if region is readable
50    pub fn is_readable(&self) -> bool {
51        (self.protect & PAGE_READONLY) != 0
52            || (self.protect & PAGE_READWRITE) != 0
53            || (self.protect & PAGE_WRITECOPY) != 0
54            || (self.protect & PAGE_EXECUTE_READ) != 0
55            || (self.protect & PAGE_EXECUTE_READWRITE) != 0
56            || (self.protect & PAGE_EXECUTE_WRITECOPY) != 0
57    }
58
59    /// check if region is writable
60    pub fn is_writable(&self) -> bool {
61        (self.protect & PAGE_READWRITE) != 0
62            || (self.protect & PAGE_WRITECOPY) != 0
63            || (self.protect & PAGE_EXECUTE_READWRITE) != 0
64            || (self.protect & PAGE_EXECUTE_WRITECOPY) != 0
65    }
66
67    /// check if region is committed (accessible)
68    pub fn is_committed(&self) -> bool {
69        self.state == MemoryState::Commit
70    }
71
72    /// check if this is part of an image
73    pub fn is_image(&self) -> bool {
74        self.memory_type == MemoryType::Image
75    }
76
77    /// get protection string (e.g., "RWX", "R--", etc.)
78    pub fn protection_string(&self) -> &'static str {
79        match self.protect {
80            PAGE_NOACCESS => "---",
81            PAGE_READONLY => "R--",
82            PAGE_READWRITE => "RW-",
83            PAGE_WRITECOPY => "RC-",
84            PAGE_EXECUTE => "--X",
85            PAGE_EXECUTE_READ => "R-X",
86            PAGE_EXECUTE_READWRITE => "RWX",
87            PAGE_EXECUTE_WRITECOPY => "RCX",
88            _ => "???",
89        }
90    }
91}
92
93/// iterator over memory regions in current process
94pub struct MemoryRegionIterator {
95    current_address: usize,
96    max_address: usize,
97}
98
99impl MemoryRegionIterator {
100    /// create new iterator starting from address 0
101    pub fn new() -> Self {
102        Self {
103            current_address: 0,
104            max_address: Self::max_user_address(),
105        }
106    }
107
108    /// create iterator starting from specific address
109    pub fn from_address(address: usize) -> Self {
110        Self {
111            current_address: address,
112            max_address: Self::max_user_address(),
113        }
114    }
115
116    fn max_user_address() -> usize {
117        #[cfg(target_arch = "x86_64")]
118        {
119            0x7FFFFFFFFFFF // typical x64 user space limit
120        }
121        #[cfg(target_arch = "x86")]
122        {
123            0x7FFFFFFF // typical x86 user space limit
124        }
125    }
126}
127
128impl Default for MemoryRegionIterator {
129    fn default() -> Self {
130        Self::new()
131    }
132}
133
134impl Iterator for MemoryRegionIterator {
135    type Item = MemoryRegion;
136
137    fn next(&mut self) -> Option<Self::Item> {
138        if self.current_address >= self.max_address {
139            return None;
140        }
141
142        let mut mbi = MemoryBasicInformation::default();
143        // SAFETY: VirtualQuery is safe to call with valid buffer
144        let result = unsafe {
145            VirtualQuery(
146                self.current_address as *const _,
147                &mut mbi,
148                core::mem::size_of::<MemoryBasicInformation>(),
149            )
150        };
151
152        if result == 0 {
153            return None;
154        }
155
156        // advance to next region
157        self.current_address = mbi.base_address + mbi.region_size;
158
159        let state = match mbi.state {
160            MEM_COMMIT => MemoryState::Commit,
161            MEM_RESERVE => MemoryState::Reserve,
162            MEM_FREE => MemoryState::Free,
163            _ => MemoryState::Free,
164        };
165
166        let memory_type = match mbi.memory_type {
167            MEM_IMAGE => MemoryType::Image,
168            MEM_MAPPED => MemoryType::Mapped,
169            MEM_PRIVATE => MemoryType::Private,
170            _ => MemoryType::Unknown,
171        };
172
173        Some(MemoryRegion {
174            base_address: mbi.base_address,
175            allocation_base: mbi.allocation_base,
176            allocation_protect: mbi.allocation_protect,
177            region_size: mbi.region_size,
178            state,
179            protect: mbi.protect,
180            memory_type,
181        })
182    }
183}
184
185/// find all executable memory regions
186pub fn find_executable_regions() -> Vec<MemoryRegion> {
187    MemoryRegionIterator::new()
188        .filter(|r| r.is_committed() && r.is_executable())
189        .collect()
190}
191
192/// find all image (module) regions
193pub fn find_image_regions() -> Vec<MemoryRegion> {
194    MemoryRegionIterator::new()
195        .filter(|r| r.is_committed() && r.is_image())
196        .collect()
197}
198
199/// find all private memory regions
200pub fn find_private_regions() -> Vec<MemoryRegion> {
201    MemoryRegionIterator::new()
202        .filter(|r| r.is_committed() && r.memory_type == MemoryType::Private)
203        .collect()
204}
205
206/// query single memory region at address
207pub fn query_region(address: usize) -> Result<MemoryRegion> {
208    let mut mbi = MemoryBasicInformation::default();
209    // SAFETY: VirtualQuery is safe to call with valid buffer
210    let result = unsafe {
211        VirtualQuery(
212            address as *const _,
213            &mut mbi,
214            core::mem::size_of::<MemoryBasicInformation>(),
215        )
216    };
217
218    if result == 0 {
219        return Err(WraithError::ReadFailed {
220            address: address as u64,
221            size: 0,
222        });
223    }
224
225    let state = match mbi.state {
226        MEM_COMMIT => MemoryState::Commit,
227        MEM_RESERVE => MemoryState::Reserve,
228        _ => MemoryState::Free,
229    };
230
231    let memory_type = match mbi.memory_type {
232        MEM_IMAGE => MemoryType::Image,
233        MEM_MAPPED => MemoryType::Mapped,
234        MEM_PRIVATE => MemoryType::Private,
235        _ => MemoryType::Unknown,
236    };
237
238    Ok(MemoryRegion {
239        base_address: mbi.base_address,
240        allocation_base: mbi.allocation_base,
241        allocation_protect: mbi.allocation_protect,
242        region_size: mbi.region_size,
243        state,
244        protect: mbi.protect,
245        memory_type,
246    })
247}
248
249// internal structures for VirtualQuery
250#[repr(C)]
251#[derive(Default)]
252struct MemoryBasicInformation {
253    base_address: usize,
254    allocation_base: usize,
255    allocation_protect: u32,
256    #[cfg(target_arch = "x86_64")]
257    partition_id: u16,
258    region_size: usize,
259    state: u32,
260    protect: u32,
261    memory_type: u32,
262}
263
264// memory state constants
265const MEM_COMMIT: u32 = 0x1000;
266const MEM_RESERVE: u32 = 0x2000;
267const MEM_FREE: u32 = 0x10000;
268
269// memory type constants
270const MEM_IMAGE: u32 = 0x1000000;
271const MEM_MAPPED: u32 = 0x40000;
272const MEM_PRIVATE: u32 = 0x20000;
273
274// page protection constants
275const PAGE_NOACCESS: u32 = 0x01;
276const PAGE_READONLY: u32 = 0x02;
277const PAGE_READWRITE: u32 = 0x04;
278const PAGE_WRITECOPY: u32 = 0x08;
279const PAGE_EXECUTE: u32 = 0x10;
280const PAGE_EXECUTE_READ: u32 = 0x20;
281const PAGE_EXECUTE_READWRITE: u32 = 0x40;
282const PAGE_EXECUTE_WRITECOPY: u32 = 0x80;
283
284#[link(name = "kernel32")]
285extern "system" {
286    fn VirtualQuery(
287        address: *const core::ffi::c_void,
288        buffer: *mut MemoryBasicInformation,
289        length: usize,
290    ) -> usize;
291}
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296
297    #[test]
298    fn test_memory_iterator() {
299        let regions: Vec<_> = MemoryRegionIterator::new().take(10).collect();
300        assert!(!regions.is_empty());
301    }
302
303    #[test]
304    fn test_find_executable() {
305        let exec_regions = find_executable_regions();
306        // should find at least our own code
307        assert!(!exec_regions.is_empty());
308    }
309
310    #[test]
311    fn test_query_region() {
312        // query our own code
313        let addr = test_query_region as usize;
314        let region = query_region(addr).expect("should query region");
315        assert!(region.is_executable());
316        assert!(region.is_committed());
317    }
318
319    #[test]
320    fn test_protection_string() {
321        let region = query_region(test_protection_string as usize).expect("should query");
322        let prot_str = region.protection_string();
323        // code should be executable
324        assert!(prot_str.contains('X'));
325    }
326}
wraith/navigation/memory_regions.rs

wraith/navigation/
memory_regions.rs
