wraith/navigation/
memory_regions.rs

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

wraith/navigation/
memory_regions.rs
