verity_memory/ops/
write.rs

1use winapi::um::winnt::PAGE_EXECUTE_READWRITE;
2use winapi::{shared::minwindef::LPVOID, um::memoryapi::VirtualProtect};
3
4#[cfg(feature = "advanced-write")]
5use crate::macros::match_number::{FloatType, IntegerType, IntegralType, NumberType};
6#[cfg(feature = "advanced-write")]
7use crate::types::Instruction;
8use crate::{errors::WriteMemoryError, utils};
9#[cfg(feature = "advanced-write")]
10use crate::match_number;
11
12#[cfg(feature = "advanced-write")]
13use super::asm::{float_ret, get_instruction, integer_ret, integral_ret};
14
15/// Writes a value of type `T` to the specified memory location.
16///
17/// # Safety
18/// This function is unsafe because it directly manipulates raw pointers, which can cause undefined behavior
19/// if the pointer is invalid or points to memory that is not writable.
20///
21/// # Parameters
22/// - `dest_ptr`: A mutable pointer to the destination memory where the value will be written.
23/// - `value`: The value to write at the destination memory.
24///
25/// # Returns
26/// - `Ok(())` if the value was successfully written to memory.
27/// - `Err(WriteMemoryError)` if an error occurred, such as a null pointer or invalid alignment.
28///
29/// # Errors
30/// - `WriteMemoryError::NullPointer` if `dest_ptr` is null.
31/// - `WriteMemoryError::InvalidAlignment` if `dest_ptr` is not correctly aligned.
32/// - `WriteMemoryError::FailedToChangeProtection` if memory protection could not be modified.
33/// - `WriteMemoryError::FailedToRestoreProtection` if memory protection could not be restored.
34///
35/// # Example
36/// ```rust
37/// use verity_memory::ops::write;
38/// unsafe {
39///     let mut value: i32 = 42;
40///     let result = write::write_memory(&mut value as *mut i32, 100);
41///     assert!(result.is_ok());
42///     assert_eq!(value, 100);
43/// }
44/// ```
45pub unsafe fn write_memory<T: Copy>(dest_ptr: *mut T, value: T) -> Result<(), WriteMemoryError> {
46    if dest_ptr.is_null() {
47        return Err(WriteMemoryError::NullPointer);
48    }
49
50    if !utils::check_alignment(dest_ptr) {
51        return Err(WriteMemoryError::InvalidAlignment);
52    }
53
54    let mut old_protect = 0;
55    let size = std::mem::size_of::<T>();
56
57    let res = VirtualProtect(
58        dest_ptr as LPVOID,
59        size,
60        PAGE_EXECUTE_READWRITE,
61        &mut old_protect,
62    );
63    if res == 0 {
64        return Err(WriteMemoryError::FailedToChangeProtection);
65    }
66
67    *dest_ptr = value;
68
69    let res_restore = VirtualProtect(dest_ptr as LPVOID, size, old_protect, &mut old_protect);
70    if res_restore == 0 {
71        return Err(WriteMemoryError::FailedToRestoreProtection);
72    }
73
74    Ok(())
75}
76
77/// Replaces a specified number of instructions at a memory location with NOPs (0x90).
78///
79/// # Safety
80/// This function is unsafe because it directly modifies memory, which can corrupt the process
81/// if the memory is not writable or if the replaced instructions are critical.
82///
83/// # Parameters
84/// - `dest_ptr`: A mutable pointer to the memory location where instructions will be replaced.
85/// - `num_instructions`: The number of instructions to replace with NOPs.
86///
87/// # Returns
88/// - `Some(Vec<Instruction>)` containing the original instructions if successful.
89/// - `None` if it failed to read instructions or write memory.
90///
91/// # Example
92/// ```rust
93/// use verity_memory::ops::write;
94/// unsafe {
95///     let buffer = vec![0x55, 0x48, 0x89, 0xE5]; // Some sample machine code (push rbp; mov rbp, rsp)
96///     let original_instructions = write::nop_instructions(buffer.as_ptr() as *mut u8, 2);
97///     assert!(original_instructions.is_some());
98/// }
99/// ```
100#[cfg(feature = "advanced-write")]
101pub unsafe fn nop_instructions(dest_ptr: *mut u8, num_instructions: usize) -> Option<Vec<Instruction>> {
102    let mut instructions = Vec::new();
103    let mut current_ptr = dest_ptr;
104
105    if num_instructions <= 0 {
106        panic!("You must nop at least one instruction...");
107    }
108
109    for _ in 0..num_instructions {
110        if let Some(instr) = get_instruction(current_ptr, 16) {
111            instructions.push(instr.clone());
112            current_ptr = current_ptr.add(instr.size);
113        } else {
114            eprintln!("Failed to get instruction at memory address: {:?}", current_ptr);
115            return None;
116        }
117    }
118
119    let total_size: usize = instructions.iter().map(|instr| instr.size).sum();
120
121    let nops = vec![0x90; total_size];
122    let mut written_size = 0;
123    for i in 0..num_instructions {
124        let instruction = &instructions[i];
125        for j in 0..instruction.size {
126            let res = write_memory(dest_ptr.add(written_size + j), nops[written_size + j]);
127            if let Err(e) = res {
128                eprintln!("Failed to write memory at offset {}: {:?}", written_size + j, e);
129                return None;
130            }
131        }
132        written_size += instruction.size;
133    }
134
135    Some(instructions)
136}
137
138/// Replaces the return value of a function with a specified value or inserts a `RET` instruction.
139///
140/// # Safety
141/// This function is unsafe because it directly modifies memory, which can cause undefined behavior
142/// if the memory is not writable or if the return value type is not correctly handled.
143///
144/// # Parameters
145/// - `dest_ptr`: A mutable pointer to the function's first instruction.
146/// - `return_value`: An optional value to return. If `None`, a `RET` instruction is written instead.
147///
148/// # Returns
149/// - `Some(Instruction)` containing the original instruction if successful.
150/// - `None` if an error occurred during instruction writing.
151///
152/// # Example
153/// ```rust
154/// use verity_memory::ops::write;
155/// unsafe {
156///     let buffer = vec![0x55, 0x48, 0x89, 0xE5]; // Example machine code
157///     let result = write::replace_return_value::<i32>(buffer.as_ptr() as *mut u8, Some(123));
158///     assert!(result.is_some());
159/// }
160/// ```
161#[cfg(feature = "advanced-write")]
162pub unsafe fn replace_return_value<T: Copy + 'static>(
163    dest_ptr: *mut u8,
164    return_value: Option<T>,
165) -> Option<Instruction> {
166    let original_instruction = get_instruction(dest_ptr, 16)?;
167
168    let value = match return_value {
169        Some(val) => val,
170        None => {
171            if let Err(e) = write_memory(dest_ptr, 0xC3) {
172                eprintln!("Failed to write RET instruction: {:?}", e);
173                return None;
174            }
175            return Some(original_instruction);
176        }
177    };
178
179    let result = match_number!(value);
180    if result.is_none() {
181        return None;
182    }
183
184    let number_type = result.unwrap();
185    let instruction_bytes = match number_type {
186        NumberType::Float(float_type) => float_ret(float_type),
187        NumberType::Integer(integer_type) => integer_ret(integer_type),
188        NumberType::Integral(integral_type) => integral_ret(integral_type),
189        NumberType::Unknown => {
190            return None;
191        }
192    };
193
194    let mut current_ptr = dest_ptr;
195    for instruction_byte in instruction_bytes {
196        if write_memory(current_ptr, instruction_byte).is_err() {
197            return None;
198        }
199        current_ptr = current_ptr.add(1);
200    }
201
202    Some(original_instruction)
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208    use std::ptr;
209
210    fn mock_dest_ptr<T: Copy>(value: T) -> *mut T {
211        let mut boxed_value = Box::new(value);
212        let ptr = &mut *boxed_value as *mut T;
213        std::mem::forget(boxed_value);
214        ptr
215    }
216
217    #[test]
218    fn test_write_memory_success() {
219        let value: u32 = 42;
220        let dest_ptr = mock_dest_ptr(value);
221
222        let result = unsafe { write_memory(dest_ptr, 100_u32) };
223        assert!(result.is_ok());
224        assert_eq!(unsafe { *dest_ptr }, 100_u32);
225    }
226
227    #[test]
228    fn test_write_memory_null_pointer() {
229        let dest_ptr: *mut u32 = ptr::null_mut();
230        
231        let result = unsafe { write_memory(dest_ptr, 100_u32) };
232        assert!(matches!(result, Err(WriteMemoryError::NullPointer)));
233    }
234    
235    #[test]
236    #[cfg(feature = "advanced-write")]
237    fn test_nop_instructions_success() {
238        let data: Vec<u8> = vec![0x55, 0x48, 0x8B, 0xEC, 0x90];
239        let dest_ptr = data.as_ptr() as *mut u8;
240
241        unsafe {
242            if let Some(instructions) = nop_instructions(dest_ptr, 2) {
243                assert_eq!(instructions.len(), 2);
244            } else {
245                panic!("Failed to retrieve instructions");
246            }
247        }
248    }
249
250    #[test]
251    #[cfg(feature = "advanced-write")]
252    fn test_nop_instructions_failure() {
253        let dest_ptr: *mut u8 = ptr::null_mut();
254
255        unsafe {
256            let instructions = nop_instructions(dest_ptr, 1);
257            assert!(instructions.is_none());
258        }
259    }
260
261    #[test]
262    #[cfg(feature = "advanced-write")]
263    fn test_replace_return_value_integer() {
264        let data: Vec<u8> = vec![0x55, 0x48, 0x8B, 0xEC];
265        let dest_ptr = data.as_ptr() as *mut u8;
266
267        unsafe {
268            let result = replace_return_value(dest_ptr, Some(123_u32));
269            assert!(result.is_some());
270        }
271    }
272
273    #[test]
274    #[cfg(feature = "advanced-write")]
275    fn test_replace_return_value_float() {
276        let data: Vec<u8> = vec![0x55, 0x48, 0x8B, 0xEC];
277        let dest_ptr = data.as_ptr() as *mut u8;
278
279        unsafe {
280            let result = replace_return_value(dest_ptr, Some(123.45_f32));
281            assert!(result.is_some());
282        }
283    }
284
285    #[test]
286    #[cfg(feature = "advanced-write")]
287    fn test_replace_return_value_none() {
288        let data: Vec<u8> = vec![0x55, 0x48, 0x8B, 0xEC];
289        let dest_ptr = data.as_ptr() as *mut u8;
290
291        unsafe {
292            let result = replace_return_value::<i32>(dest_ptr, None);
293            assert!(result.is_some());
294        }
295    }
296}
verity_memory/ops/write.rs

verity_memory/ops/
write.rs
