specter/memory/info/
code_cave.rs

1//! Code cave finder and manager
2
3use crate::memory::info::{image, scan};
4#[cfg(feature = "dev_release")]
5use crate::utils::logger;
6use once_cell::sync::Lazy;
7use parking_lot::Mutex;
8use std::collections::HashMap;
9use thiserror::Error;
10
11/// ARM64 NOP instruction encoding
12const ARM64_NOP: u32 = 0x1F2003D5;
13
14#[derive(Error, Debug)]
15/// Errors that can occur during code cave operations
16pub enum CodeCaveError {
17    /// No suitable cave was found for the requested size
18    #[error("No suitable cave found")]
19    NoCaveFound,
20    /// The specified cave is already allocated
21    #[error("Cave already allocated at {0:#x}")]
22    AlreadyAllocated(usize),
23    /// The specified cave was not found in the registry
24    #[error("Cave not found at {0:#x}")]
25    NotFound(usize),
26    /// The requested size is invalid (e.g., 0)
27    #[error("Invalid size: {0}")]
28    InvalidSize(usize),
29    /// Image lookup failed
30    #[error("Image not found: {0}")]
31    ImageNotFound(#[from] image::ImageError),
32    /// Memory scan failed
33    #[error("Scan error: {0}")]
34    ScanError(#[from] scan::ScanError),
35    /// Custom error message
36    #[error("{0}")]
37    Custom(String),
38}
39
40/// Represents a code cave (unused memory region)
41#[derive(Debug, Clone)]
42pub struct CodeCave {
43    /// The start address of the cave
44    pub address: usize,
45    /// The size of the cave in bytes
46    pub size: usize,
47    /// Whether the cave is currently in use
48    pub allocated: bool,
49    /// Optional description of what the cave is used for or how it was found
50    pub description: Option<String>,
51}
52
53impl CodeCave {
54    /// Creates a new code cave
55    ///
56    /// # Arguments
57    /// * `address` - The start address
58    /// * `size` - The size in bytes
59    pub fn new(address: usize, size: usize) -> Self {
60        Self {
61            address,
62            size,
63            allocated: false,
64            description: None,
65        }
66    }
67
68    /// Creates a new code cave with a description
69    ///
70    /// # Arguments
71    /// * `address` - The start address
72    /// * `size` - The size in bytes
73    /// * `description` - A description of the cave
74    pub fn with_description(address: usize, size: usize, description: String) -> Self {
75        Self {
76            address,
77            size,
78            allocated: false,
79            description: Some(description),
80        }
81    }
82
83    /// Marks the cave as allocated
84    pub fn allocate(&mut self) {
85        self.allocated = true;
86    }
87
88    /// Marks the cave as free
89    pub fn free(&mut self) {
90        self.allocated = false;
91    }
92}
93
94/// Registry to track allocated code caves
95struct CaveRegistry {
96    caves: HashMap<usize, CodeCave>,
97}
98
99impl CaveRegistry {
100    fn new() -> Self {
101        Self {
102            caves: HashMap::new(),
103        }
104    }
105
106    fn register(&mut self, cave: CodeCave) -> Result<(), CodeCaveError> {
107        if self.caves.contains_key(&cave.address) {
108            return Err(CodeCaveError::AlreadyAllocated(cave.address));
109        }
110        self.caves.insert(cave.address, cave);
111        Ok(())
112    }
113
114    fn unregister(&mut self, address: usize) -> Result<CodeCave, CodeCaveError> {
115        self.caves
116            .remove(&address)
117            .ok_or(CodeCaveError::NotFound(address))
118    }
119
120    fn get(&self, address: usize) -> Option<&CodeCave> {
121        self.caves.get(&address)
122    }
123
124    fn list_all(&self) -> Vec<CodeCave> {
125        self.caves.values().cloned().collect()
126    }
127
128    fn clear(&mut self) {
129        self.caves.clear();
130    }
131}
132
133static REGISTRY: Lazy<Mutex<CaveRegistry>> = Lazy::new(|| Mutex::new(CaveRegistry::new()));
134
135/// Finds sequences of NOP instructions in a memory range
136///
137/// # Arguments
138/// * `start` - The start address of the range
139/// * `size` - The size of the range
140/// * `min_count` - The minimum number of NOP instructions to consider a cave
141///
142/// # Returns
143/// * `Result<Vec<CodeCave>, CodeCaveError>` - A list of found caves or an error
144pub fn find_nop_sequences(
145    start: usize,
146    size: usize,
147    min_count: usize,
148) -> Result<Vec<CodeCave>, CodeCaveError> {
149    if min_count == 0 {
150        return Err(CodeCaveError::InvalidSize(min_count));
151    }
152
153    let min_size = min_count * 4;
154    let mut caves = Vec::new();
155
156    unsafe {
157        let mut current_addr = start;
158        let end_addr = start + size;
159
160        while current_addr < end_addr {
161            if let Ok(instr) = crate::memory::rw::read::<u32>(current_addr) {
162                if instr == ARM64_NOP {
163                    let cave_start = current_addr;
164                    let mut cave_size = 0;
165                    let mut temp_addr = current_addr;
166
167                    while temp_addr < end_addr {
168                        if let Ok(i) = crate::memory::rw::read::<u32>(temp_addr) {
169                            if i == ARM64_NOP {
170                                cave_size += 4;
171                                temp_addr += 4;
172                            } else {
173                                break;
174                            }
175                        } else {
176                            break;
177                        }
178                    }
179
180                    if cave_size >= min_size {
181                        caves.push(CodeCave::with_description(
182                            cave_start,
183                            cave_size,
184                            format!("NOP sequence ({} instructions)", cave_size / 4),
185                        ));
186                    }
187
188                    current_addr = temp_addr;
189                } else {
190                    current_addr += 4;
191                }
192            } else {
193                current_addr += 4;
194            }
195        }
196    }
197
198    Ok(caves)
199}
200
201/// Finds alignment padding (zero bytes) in a memory range
202///
203/// # Arguments
204/// * `start` - The start address of the range
205/// * `size` - The size of the range
206///
207/// # Returns
208/// * `Result<Vec<CodeCave>, CodeCaveError>` - A list of found caves or an error
209pub fn find_alignment_padding(start: usize, size: usize) -> Result<Vec<CodeCave>, CodeCaveError> {
210    let mut caves = Vec::new();
211
212    unsafe {
213        let mut current_addr = start;
214        let end_addr = start + size;
215
216        while current_addr < end_addr {
217            if let Ok(byte) = crate::memory::rw::read::<u8>(current_addr) {
218                if byte == 0x00 {
219                    let cave_start = current_addr;
220                    let mut cave_size = 0;
221                    let mut temp_addr = current_addr;
222
223                    while temp_addr < end_addr {
224                        if let Ok(b) = crate::memory::rw::read::<u8>(temp_addr) {
225                            if b == 0x00 {
226                                cave_size += 1;
227                                temp_addr += 1;
228                            } else {
229                                break;
230                            }
231                        } else {
232                            break;
233                        }
234                    }
235
236                    if cave_size >= 16 {
237                        caves.push(CodeCave::with_description(
238                            cave_start,
239                            cave_size,
240                            format!("Padding ({} bytes)", cave_size),
241                        ));
242                    }
243
244                    current_addr = temp_addr;
245                } else {
246                    current_addr += 1;
247                }
248            } else {
249                current_addr += 1;
250            }
251        }
252    }
253
254    Ok(caves)
255}
256
257/// Finds all code caves (NOPs and padding) in a memory range
258///
259/// # Arguments
260/// * `start` - The start address of the range
261/// * `size` - The size of the range
262/// * `min_size` - The minimum size in bytes
263///
264/// # Returns
265/// * `Result<Vec<CodeCave>, CodeCaveError>` - A list of found caves or an error
266pub fn find_caves(
267    start: usize,
268    size: usize,
269    min_size: usize,
270) -> Result<Vec<CodeCave>, CodeCaveError> {
271    let mut all_caves = Vec::new();
272
273    let min_nops = min_size.div_ceil(4);
274    if let Ok(nop_caves) = find_nop_sequences(start, size, min_nops) {
275        all_caves.extend(nop_caves);
276    }
277
278    if let Ok(padding_caves) = find_alignment_padding(start, size) {
279        all_caves.extend(padding_caves.into_iter().filter(|c| c.size >= min_size));
280    }
281
282    all_caves.sort_by_key(|c| c.address);
283
284    Ok(all_caves)
285}
286
287/// Finds code caves in an entire image
288///
289/// # Arguments
290/// * `image_name` - The name of the image to scan
291/// * `min_size` - The minimum size in bytes
292///
293/// # Returns
294/// * `Result<Vec<CodeCave>, CodeCaveError>` - A list of found caves or an error
295pub fn find_caves_in_image(
296    image_name: &str,
297    min_size: usize,
298) -> Result<Vec<CodeCave>, CodeCaveError> {
299    let base = image::get_image_base(image_name)?;
300
301    let scan_size = 32 * 1024 * 1024; // 32MB
302
303    #[cfg(feature = "dev_release")]
304    logger::info(&format!(
305        "Scanning image '{}' at {:#x} for code caves (min size: {} bytes)",
306        image_name, base, min_size
307    ));
308
309    find_caves(base, scan_size, min_size)
310}
311
312/// Allocates a code cave of the requested size
313///
314/// This function scans the target image for a suitable cave, marks it as allocated, and returns it.
315///
316/// # Arguments
317/// * `size` - The required size in bytes
318///
319/// # Returns
320/// * `Result<CodeCave, CodeCaveError>` - The allocated cave or an error
321pub fn allocate_cave(size: usize) -> Result<CodeCave, CodeCaveError> {
322    if size == 0 {
323        return Err(CodeCaveError::InvalidSize(size));
324    }
325
326    let image_name = crate::config::get_target_image_name()
327        .ok_or_else(|| image::ImageError::NotFound("call mem_init first".to_string()))?;
328    let caves = find_caves_in_image(&image_name, size)?;
329
330    for mut cave in caves {
331        let registry = REGISTRY.lock();
332        if registry.get(cave.address).is_some() {
333            continue;
334        }
335        drop(registry);
336
337        if cave.size >= size {
338            cave.allocate();
339            REGISTRY.lock().register(cave.clone())?;
340            #[cfg(feature = "dev_release")]
341            logger::info(&format!(
342                "Allocated code cave at {:#x} (size: {} bytes)",
343                cave.address, cave.size
344            ));
345            return Ok(cave);
346        }
347    }
348
349    Err(CodeCaveError::NoCaveFound)
350}
351
352/// Allocates a code cave near a target address (within branch range)
353///
354/// Useful for creating trampolines that need to be within ±128MB of the target.
355///
356/// # Arguments
357/// * `target` - The target address
358/// * `size` - The required size in bytes
359///
360/// # Returns
361/// * `Result<CodeCave, CodeCaveError>` - The allocated cave or an error
362pub fn allocate_cave_near(target: usize, size: usize) -> Result<CodeCave, CodeCaveError> {
363    if size == 0 {
364        return Err(CodeCaveError::InvalidSize(size));
365    }
366
367    const BRANCH_RANGE: usize = 128 * 1024 * 1024; // ±128MB for ARM64 B instruction
368
369    let image_name = crate::config::get_target_image_name()
370        .ok_or_else(|| image::ImageError::NotFound("call mem_init first".to_string()))?;
371    let caves = find_caves_in_image(&image_name, size)?;
372
373    for mut cave in caves {
374        let distance = cave.address.abs_diff(target);
375
376        if distance <= BRANCH_RANGE && cave.size >= size {
377            let registry = REGISTRY.lock();
378            if registry.get(cave.address).is_some() {
379                continue;
380            }
381            drop(registry);
382
383            cave.allocate();
384            REGISTRY.lock().register(cave.clone())?;
385            #[cfg(feature = "dev_release")]
386            logger::info(&format!(
387                "Allocated code cave near {:#x} at {:#x} (size: {} bytes)",
388                target, cave.address, cave.size
389            ));
390            return Ok(cave);
391        }
392    }
393
394    Err(CodeCaveError::NoCaveFound)
395}
396
397/// Frees a previously allocated code cave
398///
399/// # Arguments
400/// * `address` - The address of the cave to free
401///
402/// # Returns
403/// * `Result<(), CodeCaveError>` - Result indicating success or failure
404pub fn free_cave(address: usize) -> Result<(), CodeCaveError> {
405    let mut cave = REGISTRY.lock().unregister(address)?;
406    cave.free();
407    #[cfg(feature = "dev_release")]
408    logger::info(&format!("Freed code cave at {:#x}", address));
409    Ok(())
410}
411
412/// Checks if a cave at the given address is available (not modified)
413///
414/// # Arguments
415/// * `address` - The address to check
416/// * `size` - The size to check
417///
418/// # Returns
419/// * `bool` - `true` if the memory contains valid cave content (NOPs or zeros)
420pub fn is_cave_available(address: usize, size: usize) -> bool {
421    let registry = REGISTRY.lock();
422
423    if registry.get(address).is_some() {
424        return false;
425    }
426
427    unsafe {
428        for offset in (0..size).step_by(4) {
429            if let Ok(instr) = crate::memory::rw::read::<u32>(address + offset) {
430                if instr != ARM64_NOP && instr != 0 {
431                    return false;
432                }
433            } else {
434                return false;
435            }
436        }
437    }
438
439    true
440}
441
442/// Lists all currently allocated code caves
443///
444/// # Returns
445/// * `Vec<CodeCave>` - A list of allocated caves
446pub fn list_allocated_caves() -> Vec<CodeCave> {
447    REGISTRY.lock().list_all()
448}
449
450/// Returns statistics about code cave usage
451///
452/// # Returns
453/// * `(usize, usize)` - A tuple containing (count, total_size_in_bytes)
454pub fn get_cave_stats() -> (usize, usize) {
455    let registry = REGISTRY.lock();
456    let caves = registry.list_all();
457    let total_size: usize = caves.iter().map(|c| c.size).sum();
458    (caves.len(), total_size)
459}
460
461/// Clears all allocated code caves from the registry
462pub fn clear_all_caves() {
463    REGISTRY.lock().clear();
464    #[cfg(feature = "dev_release")]
465    logger::info("Cleared all allocated code caves");
466}
specter/memory/info/code_cave.rs

specter/memory/info/
code_cave.rs
