specter/memory/allocation/
shellcode.rs

1//! Shellcode loader with symbol resolution and position-independent code support
2
3use crate::memory::{code_cave, info::symbol, rw};
4#[cfg(feature = "dev_release")]
5use crate::utils::logger;
6use std::arch::asm;
7use std::collections::HashMap;
8use std::ffi::c_void;
9use thiserror::Error;
10
11const CACHE_LINE_SIZE: usize = 64;
12
13#[derive(Error, Debug)]
14/// Errors that can occur during shellcode loading or execution
15pub enum LoaderError {
16    /// Failed to allocate memory for the shellcode (e.g., no suitable code cave found)
17    #[error("Code cave allocation failed: {0}")]
18    AllocationFailed(#[from] code_cave::CodeCaveError),
19    /// A required symbol for relocation was not found
20    #[error("Symbol not found: {0}")]
21    SymbolNotFound(String),
22    /// Failed to write shellcode to memory
23    #[error("Write failed: {0}")]
24    WriteFailed(#[from] rw::RwError),
25    /// The input bytes or allocation size is invalid
26    #[error("Invalid shellcode size: {0}")]
27    InvalidSize(usize),
28    /// Failed to apply a relocation at the specified offset
29    #[error("Relocation failed at offset {0}")]
30    RelocationFailed(usize),
31}
32
33/// Represents loaded executable shellcode in memory
34pub struct LoadedShellcode {
35    /// The base address of the loaded code
36    pub address: usize,
37    /// The size of the loaded code
38    pub size: usize,
39    /// Whether to automatically free the memory on drop
40    auto_free: bool,
41}
42
43impl LoadedShellcode {
44    /// Executes the shellcode as a function with no arguments, returning usize
45    ///
46    /// # Safety
47    /// Caller must ensure the shellcode is valid, matches the signature, and respects ABI conventions.
48    pub unsafe fn execute(&self) -> usize {
49        unsafe {
50            let func: extern "C" fn() -> usize = std::mem::transmute(self.address);
51            func()
52        }
53    }
54
55    /// Executes the shellcode with a custom function signature
56    ///
57    /// # Safety
58    /// Caller must ensure the signature `F` matches the shellcode's implementation.
59    pub unsafe fn execute_as<F, R>(&self, callback: impl FnOnce(F) -> R) -> R
60    where
61        F: Copy,
62    {
63        unsafe {
64            let func_ptr = self.address as *const ();
65            let func: F = *(&func_ptr as *const *const () as *const F);
66            callback(func)
67        }
68    }
69
70    /// Gets a function pointer to the shellcode
71    ///
72    /// # Safety
73    /// Caller must ensure the signature `F` matches the shellcode's implementation.
74    pub unsafe fn as_function<F>(&self) -> F
75    where
76        F: Copy,
77    {
78        unsafe {
79            let func_ptr = self.address as *const ();
80            *(&func_ptr as *const *const () as *const F)
81        }
82    }
83
84    /// Manually free the shellcode memory
85    pub fn free(self) {
86        if self.address != 0 {
87            let _ = code_cave::free_cave(self.address);
88        }
89    }
90}
91
92impl Drop for LoadedShellcode {
93    fn drop(&mut self) {
94        if self.auto_free && self.address != 0 {
95            let _ = code_cave::free_cave(self.address);
96        }
97    }
98}
99
100#[derive(Debug, Clone)]
101/// Represents a symbol that needs to be resolved and patched into the shellcode
102pub struct SymbolRelocation {
103    /// The offset within the shellcode where the address should be written
104    pub offset: usize,
105    /// The name of the symbol to resolve
106    pub symbol_name: String,
107}
108
109/// Builder for configuring and loading shellcode
110pub struct ShellcodeBuilder {
111    code: Vec<u8>,
112    relocations: Vec<SymbolRelocation>,
113    auto_free: bool,
114    target_address: Option<usize>,
115}
116
117impl ShellcodeBuilder {
118    /// Creates a new shellcode builder from raw bytes
119    ///
120    /// # Arguments
121    /// * `shellcode` - The raw machine code bytes
122    pub fn new(shellcode: &[u8]) -> Self {
123        Self {
124            code: shellcode.to_vec(),
125            relocations: Vec::new(),
126            auto_free: true,
127            target_address: None,
128        }
129    }
130
131    /// Creates a shellcode builder from ARM64 instructions (u32 array)
132    ///
133    /// # Arguments
134    /// * `instructions` - Slice of 32-bit ARM64 instructions
135    pub fn from_instructions(instructions: &[u32]) -> Self {
136        let bytes: Vec<u8> = instructions
137            .iter()
138            .flat_map(|&instr| instr.to_le_bytes())
139            .collect();
140        Self::new(&bytes)
141    }
142
143    /// Adds a symbol relocation for dynamic linking
144    ///
145    /// The loader will resolve the symbol address and write it at the specified offset.
146    ///
147    /// # Arguments
148    /// * `offset` - Byte offset in the shellcode
149    /// * `symbol_name` - Name of the symbol to resolve
150    pub fn with_symbol(mut self, offset: usize, symbol_name: &str) -> Self {
151        self.relocations.push(SymbolRelocation {
152            offset,
153            symbol_name: symbol_name.to_string(),
154        });
155        self
156    }
157
158    /// Disables automatic cleanup (shellcode won't be freed on drop)
159    ///
160    /// Use this if you want the shellcode to persist for the lifetime of the process.
161    pub fn no_auto_free(mut self) -> Self {
162        self.auto_free = false;
163        self
164    }
165
166    /// Tries to load shellcode near a specific address (within branch range)
167    ///
168    /// Use this if your shellcode contains relative branches to nearby code.
169    pub fn near_address(mut self, target: usize) -> Self {
170        self.target_address = Some(target);
171        self
172    }
173
174    /// Loads the shellcode into memory with all configured options
175    ///
176    /// This will:
177    /// 1. Allocate a code cave (optionally near a target)
178    /// 2. Resolve and apply symbol relocations
179    /// 3. Write the code to memory
180    /// 4. Flush the instruction cache
181    ///
182    /// # Returns
183    /// * `Result<LoadedShellcode, LoaderError>` - The loaded shellcode handle or an error
184    pub fn load(self) -> Result<LoadedShellcode, LoaderError> {
185        if self.code.is_empty() {
186            return Err(LoaderError::InvalidSize(0));
187        }
188
189        let mut cave = if let Some(target) = self.target_address {
190            code_cave::allocate_cave_near(target, self.code.len())?
191        } else {
192            code_cave::allocate_cave(self.code.len())?
193        };
194
195        let alignment = cave.address % 4;
196        if alignment != 0 {
197            let adjust = 4 - alignment;
198            cave.address += adjust;
199            cave.size = cave.size.saturating_sub(adjust);
200
201            if cave.size < self.code.len() {
202                return Err(LoaderError::InvalidSize(cave.size));
203            }
204        }
205
206        let mut code = self.code.clone();
207
208        let mut symbol_cache: HashMap<String, usize> = HashMap::new();
209        for reloc in &self.relocations {
210            let symbol_addr = if let Some(&cached) = symbol_cache.get(&reloc.symbol_name) {
211                cached
212            } else {
213                let addr = symbol::resolve_symbol(&reloc.symbol_name)
214                    .map_err(|_| LoaderError::SymbolNotFound(reloc.symbol_name.clone()))?;
215
216                symbol_cache.insert(reloc.symbol_name.clone(), addr);
217                addr
218            };
219
220            if reloc.offset + 8 > code.len() {
221                return Err(LoaderError::RelocationFailed(reloc.offset));
222            }
223
224            let addr_bytes = (symbol_addr as u64).to_le_bytes();
225            code[reloc.offset..reloc.offset + 8].copy_from_slice(&addr_bytes);
226        }
227
228        unsafe {
229            rw::write_bytes(cave.address, &code)?;
230        }
231
232        // Verify shellcode
233        unsafe {
234            for (i, &expected) in code.iter().enumerate() {
235                match rw::read::<u8>(cave.address + i) {
236                    Ok(actual) => {
237                        if actual != expected {
238                            return Err(LoaderError::WriteFailed(rw::RwError::NullPointer));
239                        }
240                    }
241                    Err(e) => {
242                        return Err(LoaderError::WriteFailed(e));
243                    }
244                }
245            }
246        }
247
248        if !crate::memory::protection::is_executable(cave.address) {
249            return Err(LoaderError::InvalidSize(0));
250        }
251
252        // Flush instruction cache
253        unsafe {
254            invalidate_icache(cave.address as *mut c_void, code.len());
255        }
256
257        #[cfg(feature = "dev_release")]
258        logger::info(&format!(
259            "Shellcode loaded at {:#x} ({} bytes)",
260            cave.address,
261            code.len()
262        ));
263
264        Ok(LoadedShellcode {
265            address: cave.address,
266            size: code.len(),
267            auto_free: self.auto_free,
268        })
269    }
270}
271
272/// Trait for types that can be converted to shellcode bytes
273pub trait ShellcodeSource {
274    fn into_bytes(self) -> Vec<u8>;
275}
276
277impl ShellcodeSource for &[u8] {
278    fn into_bytes(self) -> Vec<u8> {
279        self.to_vec()
280    }
281}
282
283impl ShellcodeSource for &[u32] {
284    fn into_bytes(self) -> Vec<u8> {
285        self.iter().flat_map(|&instr| instr.to_le_bytes()).collect()
286    }
287}
288
289impl<const N: usize> ShellcodeSource for &[u32; N] {
290    fn into_bytes(self) -> Vec<u8> {
291        self.iter().flat_map(|&instr| instr.to_le_bytes()).collect()
292    }
293}
294
295impl<const N: usize> ShellcodeSource for &[u8; N] {
296    fn into_bytes(self) -> Vec<u8> {
297        self.to_vec()
298    }
299}
300
301/// Unified helper to load shellcode from bytes or ARM64 instructions
302///
303/// # Arguments
304/// * `source` - The shellcode source (bytes, instruction array, etc.)
305///
306/// # Returns
307/// * `Result<LoadedShellcode, LoaderError>` - The loaded shellcode handle or an error
308pub fn load(source: impl ShellcodeSource) -> Result<LoadedShellcode, LoaderError> {
309    ShellcodeBuilder::new(&source.into_bytes()).load()
310}
311
312/// Invalidates the instruction cache for a memory range
313///
314/// Uses the full ARM64 cache maintenance sequence:
315/// dc cvau -> dsb ish -> ic ivau -> dsb ish -> isb
316#[inline]
317pub unsafe fn invalidate_icache(start: *mut c_void, len: usize) {
318    unsafe {
319        let start_addr = start as usize;
320        let end_addr = start_addr + len;
321        let mut addr = start_addr & !(CACHE_LINE_SIZE - 1);
322
323        // Clean data cache to Point of Unification
324        while addr < end_addr {
325            asm!("dc cvau, {x}", x = in(reg) addr, options(nostack, preserves_flags));
326            addr += CACHE_LINE_SIZE;
327        }
328        asm!("dsb ish", options(nostack, preserves_flags));
329
330        // Invalidate instruction cache
331        addr = start_addr & !(CACHE_LINE_SIZE - 1);
332        while addr < end_addr {
333            asm!("ic ivau, {x}", x = in(reg) addr, options(nostack, preserves_flags));
334            addr += CACHE_LINE_SIZE;
335        }
336        asm!("dsb ish", options(nostack, preserves_flags));
337        asm!("isb", options(nostack, preserves_flags));
338    }
339}
specter/memory/allocation/shellcode.rs

specter/memory/allocation/
shellcode.rs
