use crate::{
Result,
error::MapleError,
memory_module::{DLL_PROCESS_ATTACH, DLL_PROCESS_DETACH, MemoryModule, MemoryModuleBuilder},
pe::*,
};
use std::collections::HashMap;
use std::ffi::{CString, OsStr};
use std::mem;
use std::os::windows::ffi::OsStrExt;
use std::ptr;
use winapi::shared::minwindef::{BOOL, DWORD, FARPROC, HMODULE, LPVOID};
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::libloaderapi::{GetProcAddress, LoadLibraryW};
use winapi::um::memoryapi::{VirtualAlloc, VirtualFree, VirtualProtect};
use winapi::um::winnt::{
IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE, MEM_COMMIT, MEM_RELEASE,
MEM_RESERVE, PAGE_EXECUTE, PAGE_EXECUTE_READ, PAGE_EXECUTE_READWRITE, PAGE_READONLY,
PAGE_READWRITE,
};
pub struct WindowsMemoryModule {
base_address: *mut u8,
size: usize,
entry_point: Option<unsafe extern "system" fn() -> BOOL>,
dll_entry: Option<unsafe extern "system" fn(HMODULE, DWORD, LPVOID) -> BOOL>,
is_dll: bool,
is_loaded: bool,
loaded_modules: HashMap<String, HMODULE>,
pe_data: Vec<u8>,
}
impl WindowsMemoryModule {
pub fn from_memory(data: &[u8]) -> Result<Self> {
let builder = MemoryModuleBuilder::default();
Self::from_memory_with_options(data, &builder)
}
pub fn from_memory_with_options(data: &[u8], options: &MemoryModuleBuilder) -> Result<Self> {
let pe = PEParser::new(data)?;
let image_size = pe.size_of_image() as usize;
let base_address = unsafe {
VirtualAlloc(
ptr::null_mut(),
image_size,
MEM_RESERVE | MEM_COMMIT,
PAGE_READWRITE,
)
};
if base_address.is_null() {
return Err(MapleError::MemoryAllocation(format!(
"Failed to allocate memory: {}",
unsafe { GetLastError() }
)));
}
let mut module = WindowsMemoryModule {
base_address: base_address as *mut u8,
size: image_size,
entry_point: None,
dll_entry: None,
is_dll: pe.is_dll(),
is_loaded: false,
loaded_modules: HashMap::new(),
pe_data: data.to_vec(),
};
module.copy_sections(data, &pe)?;
if options.process_relocations {
module.process_relocations(data, &pe)?;
}
if options.resolve_imports {
module.resolve_imports(data, &pe, options.ignore_missing_imports)?;
}
module.finalize_sections(&pe)?;
module.process_tls_callbacks(&pe)?;
let entry_point_rva = pe.entry_point();
if entry_point_rva != 0 {
let entry_point_va = unsafe { module.base_address.add(entry_point_rva as usize) };
if module.is_dll {
module.dll_entry = Some(unsafe {
mem::transmute::<
*mut u8,
unsafe extern "system" fn(HMODULE, DWORD, LPVOID) -> BOOL,
>(entry_point_va)
});
if options.call_dll_main && !options.is_application_dll {
let dll_main = module.dll_entry.unwrap();
let result = unsafe {
dll_main(
module.base_address as HMODULE,
DLL_PROCESS_ATTACH,
ptr::null_mut(),
)
};
if result == 0 {
return Err(MapleError::ExecutionFailed(
"DLL initialization failed".to_string(),
));
}
}
} else {
module.entry_point = Some(unsafe {
mem::transmute::<*mut u8, unsafe extern "system" fn() -> BOOL>(entry_point_va)
});
}
}
module.is_loaded = true;
Ok(module)
}
fn copy_sections(&mut self, data: &[u8], pe: &PEParser) -> Result<()> {
unsafe {
ptr::copy_nonoverlapping(
data.as_ptr(),
self.base_address,
pe.optional_header.size_of_headers as usize,
);
}
for section in pe.sections() {
if section.size_of_raw_data == 0 {
continue;
}
let dest = unsafe { self.base_address.add(section.virtual_address as usize) };
let src_offset = section.pointer_to_raw_data as usize;
let copy_size = std::cmp::min(
section.size_of_raw_data as usize,
section.virtual_size as usize,
);
if src_offset + copy_size > data.len() {
return Err(MapleError::InvalidPEFormat(
"Section data exceeds file size".to_string(),
));
}
unsafe {
ptr::copy_nonoverlapping(data[src_offset..].as_ptr(), dest, copy_size);
}
}
Ok(())
}
fn process_relocations(&mut self, _data: &[u8], pe: &PEParser) -> Result<()> {
let reloc_dir = match pe.get_base_relocation_directory() {
Some(dir) if dir.size > 0 => dir,
_ => return Ok(()),
};
let delta = self.base_address as i64 - pe.image_base() as i64;
if delta == 0 {
return Ok(());
}
let mut offset = 0u32;
while offset < reloc_dir.size {
if offset + mem::size_of::<ImageBaseRelocation>() as u32 > reloc_dir.size {
break;
}
let reloc_data = pe
.get_data_at_rva(
reloc_dir.virtual_address + offset,
mem::size_of::<ImageBaseRelocation>(),
)
.ok_or_else(|| {
MapleError::RelocationFailed("Failed to read relocation block".to_string())
})?;
let reloc_block = unsafe { &*(reloc_data.as_ptr() as *const ImageBaseRelocation) };
if reloc_block.size_of_block < mem::size_of::<ImageBaseRelocation>() as u32 {
break;
}
let entries_count =
(reloc_block.size_of_block as usize - mem::size_of::<ImageBaseRelocation>()) / 2;
let entries_data = pe
.get_data_at_rva(
reloc_dir.virtual_address
+ offset
+ mem::size_of::<ImageBaseRelocation>() as u32,
entries_count * 2,
)
.ok_or_else(|| {
MapleError::RelocationFailed("Failed to read relocation entries".to_string())
})?;
for i in 0..entries_count {
let entry_offset = i * 2;
if entry_offset + 2 > entries_data.len() {
break;
}
let entry = u16::from_le_bytes([
entries_data[entry_offset],
entries_data[entry_offset + 1],
]);
let reloc_type = entry >> 12;
let reloc_offset = entry & 0xFFF;
if reloc_type == IMAGE_REL_BASED_ABSOLUTE {
continue;
}
let reloc_va = reloc_block.virtual_address + reloc_offset as u32;
let reloc_ptr = unsafe { self.base_address.add(reloc_va as usize) };
match reloc_type {
IMAGE_REL_BASED_DIR64 => {
let old_value = unsafe { ptr::read(reloc_ptr as *const u64) };
let new_value = (old_value as i64 + delta) as u64;
unsafe { ptr::write(reloc_ptr as *mut u64, new_value) };
}
IMAGE_REL_BASED_HIGHLOW => {
let old_value = unsafe { ptr::read(reloc_ptr as *const u32) };
let new_value = (old_value as i64 + delta) as u32;
unsafe { ptr::write(reloc_ptr as *mut u32, new_value) };
}
_ => {
return Err(MapleError::RelocationFailed(format!(
"Unsupported relocation type: {}",
reloc_type
)));
}
}
}
offset += reloc_block.size_of_block;
}
Ok(())
}
fn resolve_imports(&mut self, data: &[u8], pe: &PEParser, ignore_missing: bool) -> Result<()> {
let import_dir = match pe.get_import_directory() {
Some(dir) if dir.size > 0 => dir,
_ => return Ok(()),
};
let mut offset = 0usize;
while offset + mem::size_of::<ImageImportDescriptor>() <= import_dir.size as usize {
let import_desc_data = pe
.get_data_at_rva(
import_dir.virtual_address + offset as u32,
mem::size_of::<ImageImportDescriptor>(),
)
.ok_or_else(|| {
MapleError::ImportResolution("Failed to read import descriptor".to_string())
})?;
let import_desc =
unsafe { &*(import_desc_data.as_ptr() as *const ImageImportDescriptor) };
if import_desc.name == 0 {
break;
}
let dll_name = self.read_string_at_rva(data, pe, import_desc.name)?;
let dll_handle = self.load_library(&dll_name)?;
let thunk_rva = if import_desc.original_first_thunk != 0 {
import_desc.original_first_thunk
} else {
import_desc.first_thunk
};
let mut thunk_offset = 0;
loop {
let thunk_va = thunk_rva + thunk_offset;
let thunk_data = pe.get_data_at_rva(thunk_va, 8);
if thunk_data.is_none() {
break;
}
let thunk_data = thunk_data.unwrap();
let thunk_value = u64::from_le_bytes([
thunk_data[0],
thunk_data[1],
thunk_data[2],
thunk_data[3],
thunk_data[4],
thunk_data[5],
thunk_data[6],
thunk_data[7],
]);
if thunk_value == 0 {
break;
}
let proc_address = if (thunk_value & IMAGE_ORDINAL_FLAG64) != 0 {
let ordinal = (thunk_value & 0xFFFF) as u16;
match self.get_proc_address_by_ordinal(dll_handle, ordinal) {
Ok(addr) => addr,
Err(e) => {
if ignore_missing {
eprintln!(
"Warning: Missing ordinal import {} from {}: {}",
ordinal, dll_name, e
);
std::ptr::null_mut()
} else {
return Err(e);
}
}
}
} else {
let import_name_rva = thunk_value as u32;
let func_name = self.read_import_name(data, pe, import_name_rva)?;
match self.get_proc_address_by_name(dll_handle, &func_name) {
Ok(addr) => addr,
Err(e) => {
if ignore_missing {
eprintln!(
"Warning: Missing named import {} from {}: {}",
func_name, dll_name, e
);
std::ptr::null_mut()
} else {
return Err(e);
}
}
}
};
let iat_va = import_desc.first_thunk + thunk_offset;
let iat_ptr = unsafe { self.base_address.add(iat_va as usize) as *mut u64 };
unsafe {
ptr::write(iat_ptr, proc_address as u64);
}
thunk_offset += 8;
}
offset += mem::size_of::<ImageImportDescriptor>();
}
Ok(())
}
fn finalize_sections(&mut self, pe: &PEParser) -> Result<()> {
for section in pe.sections() {
if section.virtual_size == 0 {
continue;
}
let section_ptr = unsafe { self.base_address.add(section.virtual_address as usize) };
let section_size = section.virtual_size as usize;
let mut protection = 0;
if (section.characteristics & IMAGE_SCN_MEM_READ) != 0 {
if (section.characteristics & IMAGE_SCN_MEM_WRITE) != 0 {
if (section.characteristics & IMAGE_SCN_MEM_EXECUTE) != 0 {
protection = PAGE_EXECUTE_READWRITE;
} else {
protection = PAGE_READWRITE;
}
} else if (section.characteristics & IMAGE_SCN_MEM_EXECUTE) != 0 {
protection = PAGE_EXECUTE_READ;
} else {
protection = PAGE_READONLY;
}
} else if (section.characteristics & IMAGE_SCN_MEM_EXECUTE) != 0 {
protection = PAGE_EXECUTE;
}
if protection != 0 {
let mut old_protection = 0;
let result = unsafe {
VirtualProtect(
section_ptr as LPVOID,
section_size,
protection,
&mut old_protection,
)
};
if result == 0 {
return Err(MapleError::MemoryAllocation(format!(
"Failed to set section protection: {}",
unsafe { GetLastError() }
)));
}
}
}
Ok(())
}
fn process_tls_callbacks(&mut self, pe: &PEParser) -> Result<()> {
let tls_dir = match pe.get_tls_directory() {
Some(dir) if dir.size > 0 => dir,
_ => return Ok(()),
};
let tls_data = pe
.get_data_at_rva(
tls_dir.virtual_address,
mem::size_of::<ImageTlsDirectory64>(),
)
.ok_or_else(|| {
MapleError::ExecutionFailed("Failed to read TLS directory".to_string())
})?;
let tls_directory = unsafe { &*(tls_data.as_ptr() as *const ImageTlsDirectory64) };
if tls_directory.address_of_callbacks == 0 {
return Ok(());
}
let delta = self.base_address as i64 - pe.image_base() as i64;
let callbacks_va = (tls_directory.address_of_callbacks as i64 + delta) as *const u64;
unsafe {
let mut callback_ptr = callbacks_va;
loop {
let callback_addr = ptr::read(callback_ptr);
if callback_addr == 0 {
break;
}
let callback: unsafe extern "system" fn(LPVOID, DWORD, LPVOID) =
mem::transmute(callback_addr);
callback(
self.base_address as LPVOID,
DLL_PROCESS_ATTACH,
ptr::null_mut(),
);
callback_ptr = callback_ptr.add(1);
}
}
Ok(())
}
fn load_library(&mut self, name: &str) -> Result<HMODULE> {
if let Some(&handle) = self.loaded_modules.get(name) {
return Ok(handle);
}
let wide_name: Vec<u16> = OsStr::new(name)
.encode_wide()
.chain(std::iter::once(0))
.collect();
let handle = unsafe { LoadLibraryW(wide_name.as_ptr()) };
if handle.is_null() {
return Err(MapleError::ImportResolution(format!(
"Failed to load library {}: {}",
name,
unsafe { GetLastError() }
)));
}
self.loaded_modules.insert(name.to_string(), handle);
Ok(handle)
}
fn get_proc_address_by_name(&self, module: HMODULE, name: &str) -> Result<FARPROC> {
let c_name = CString::new(name)
.map_err(|_| MapleError::ImportResolution("Invalid function name".to_string()))?;
let proc = unsafe { GetProcAddress(module, c_name.as_ptr()) };
if proc.is_null() {
return Err(MapleError::SymbolNotFound(format!(
"Function {} not found",
name
)));
}
Ok(proc)
}
fn get_proc_address_by_ordinal(&self, module: HMODULE, ordinal: u16) -> Result<FARPROC> {
let proc = unsafe { GetProcAddress(module, ordinal as usize as *const i8) };
if proc.is_null() {
return Err(MapleError::SymbolNotFound(format!(
"Function ordinal {} not found",
ordinal
)));
}
Ok(proc)
}
fn read_string_at_rva(&self, data: &[u8], pe: &PEParser, rva: u32) -> Result<String> {
let offset = pe
.rva_to_offset(rva)
.ok_or_else(|| MapleError::InvalidPEFormat("Invalid RVA for string".to_string()))?;
let mut end = offset;
while end < data.len() && data[end] != 0 {
end += 1;
}
String::from_utf8(data[offset..end].to_vec())
.map_err(|_| MapleError::InvalidPEFormat("Invalid UTF-8 string".to_string()))
}
fn read_import_name(&self, data: &[u8], pe: &PEParser, rva: u32) -> Result<String> {
let import_name_offset = pe.rva_to_offset(rva).ok_or_else(|| {
MapleError::InvalidPEFormat("Invalid RVA for import name".to_string())
})?;
if import_name_offset + 2 > data.len() {
return Err(MapleError::InvalidPEFormat(
"Import name out of bounds".to_string(),
));
}
let name_offset = import_name_offset + 2;
let mut end = name_offset;
while end < data.len() && data[end] != 0 {
end += 1;
}
String::from_utf8(data[name_offset..end].to_vec())
.map_err(|_| MapleError::InvalidPEFormat("Invalid UTF-8 import name".to_string()))
}
}
impl MemoryModule for WindowsMemoryModule {
fn get_proc_address(&self, name: &str) -> Result<*const u8> {
if !self.is_loaded {
return Err(MapleError::ExecutionFailed("Module not loaded".to_string()));
}
let pe = PEParser::new(&self.pe_data)?;
let export_dir = match pe.get_export_directory() {
Some(dir) if dir.size > 0 => dir,
_ => {
return Err(MapleError::SymbolNotFound(
"No export directory".to_string(),
));
}
};
let export_data = pe
.get_data_at_rva(
export_dir.virtual_address,
mem::size_of::<ImageExportDirectory>(),
)
.ok_or_else(|| MapleError::SymbolNotFound("Invalid export directory".to_string()))?;
let export_desc = unsafe { &*(export_data.as_ptr() as *const ImageExportDirectory) };
if export_desc.number_of_names == 0 {
return Err(MapleError::SymbolNotFound("No named exports".to_string()));
}
let names_rva = export_desc.address_of_names;
let ordinals_rva = export_desc.address_of_name_ordinals;
let functions_rva = export_desc.address_of_functions;
for i in 0..export_desc.number_of_names {
let name_ptr_rva_data = pe
.get_data_at_rva(names_rva + i * 4, 4)
.ok_or_else(|| MapleError::SymbolNotFound("Invalid name pointer".to_string()))?;
let name_ptr_rva = u32::from_le_bytes([
name_ptr_rva_data[0],
name_ptr_rva_data[1],
name_ptr_rva_data[2],
name_ptr_rva_data[3],
]);
let export_name = self.read_string_at_rva(&self.pe_data, &pe, name_ptr_rva)?;
if export_name == name {
let ordinal_data = pe
.get_data_at_rva(ordinals_rva + i * 2, 2)
.ok_or_else(|| MapleError::SymbolNotFound("Invalid ordinal".to_string()))?;
let ordinal = u16::from_le_bytes([ordinal_data[0], ordinal_data[1]]);
let function_rva_data = pe
.get_data_at_rva(functions_rva + ordinal as u32 * 4, 4)
.ok_or_else(|| {
MapleError::SymbolNotFound("Invalid function RVA".to_string())
})?;
let function_rva = u32::from_le_bytes([
function_rva_data[0],
function_rva_data[1],
function_rva_data[2],
function_rva_data[3],
]);
if function_rva == 0 {
return Err(MapleError::SymbolNotFound(
"Function not implemented".to_string(),
));
}
let function_va = unsafe { self.base_address.add(function_rva as usize) };
return Ok(function_va);
}
}
Err(MapleError::SymbolNotFound(format!(
"Function {} not found",
name
)))
}
fn get_proc_address_ordinal(&self, ordinal: u16) -> Result<*const u8> {
if !self.is_loaded {
return Err(MapleError::ExecutionFailed("Module not loaded".to_string()));
}
let pe = PEParser::new(&self.pe_data)?;
let export_dir = match pe.get_export_directory() {
Some(dir) if dir.size > 0 => dir,
_ => {
return Err(MapleError::SymbolNotFound(
"No export directory".to_string(),
));
}
};
let export_data = pe
.get_data_at_rva(
export_dir.virtual_address,
mem::size_of::<ImageExportDirectory>(),
)
.ok_or_else(|| MapleError::SymbolNotFound("Invalid export directory".to_string()))?;
let export_desc = unsafe { &*(export_data.as_ptr() as *const ImageExportDirectory) };
if ordinal < export_desc.base as u16
|| ordinal >= (export_desc.base + export_desc.number_of_functions) as u16
{
return Err(MapleError::SymbolNotFound("Invalid ordinal".to_string()));
}
let function_index = (ordinal - export_desc.base as u16) as u32;
let functions_rva = export_desc.address_of_functions;
let function_rva_data = pe
.get_data_at_rva(functions_rva + function_index * 4, 4)
.ok_or_else(|| MapleError::SymbolNotFound("Invalid function RVA".to_string()))?;
let function_rva = u32::from_le_bytes([
function_rva_data[0],
function_rva_data[1],
function_rva_data[2],
function_rva_data[3],
]);
if function_rva == 0 {
return Err(MapleError::SymbolNotFound(
"Function not implemented".to_string(),
));
}
let function_va = unsafe { self.base_address.add(function_rva as usize) };
Ok(function_va)
}
fn execute_entry_point(&self) -> Result<()> {
if !self.is_loaded {
return Err(MapleError::ExecutionFailed("Module not loaded".to_string()));
}
if self.is_dll {
return Err(MapleError::ExecutionFailed(
"Cannot execute DLL entry point".to_string(),
));
}
match self.entry_point {
Some(entry) => {
unsafe { entry() };
Ok(())
}
None => Err(MapleError::ExecutionFailed(
"No entry point found".to_string(),
)),
}
}
fn call_dll_entry_point(&self, reason: u32) -> Result<bool> {
if !self.is_loaded {
return Err(MapleError::ExecutionFailed("Module not loaded".to_string()));
}
if !self.is_dll {
return Err(MapleError::ExecutionFailed("Not a DLL".to_string()));
}
match self.dll_entry {
Some(dll_main) => {
let result =
unsafe { dll_main(self.base_address as HMODULE, reason, ptr::null_mut()) };
Ok(result != 0)
}
None => Ok(true),
}
}
fn execute_dll_application(&self) -> Result<()> {
if !self.is_loaded {
return Err(MapleError::ExecutionFailed("Module not loaded".to_string()));
}
if !self.is_dll {
return Err(MapleError::ExecutionFailed("Not a DLL".to_string()));
}
match self.dll_entry {
Some(dll_main) => {
let result = unsafe {
dll_main(
self.base_address as HMODULE,
DLL_PROCESS_ATTACH,
ptr::null_mut(),
)
};
if result == 0 {
return Err(MapleError::ExecutionFailed(
"DLL application initialization failed".to_string(),
));
}
Ok(())
}
None => Err(MapleError::ExecutionFailed(
"No DLL entry point found".to_string(),
)),
}
}
fn free(&mut self) -> Result<()> {
if !self.is_loaded {
return Ok(());
}
if self.is_dll {
let _ = self.call_dll_entry_point(DLL_PROCESS_DETACH);
}
let result = unsafe { VirtualFree(self.base_address as LPVOID, 0, MEM_RELEASE) };
if result == 0 {
return Err(MapleError::MemoryAllocation(format!(
"Failed to free memory: {}",
unsafe { GetLastError() }
)));
}
self.is_loaded = false;
Ok(())
}
fn is_loaded(&self) -> bool {
self.is_loaded
}
fn base_address(&self) -> *const u8 {
self.base_address
}
fn size(&self) -> usize {
self.size
}
fn find_resource(&self, name: Option<&str>, resource_type: Option<&str>) -> Result<*const u8> {
self.find_resource_ex(name, resource_type, 0)
}
fn find_resource_ex(
&self,
_name: Option<&str>,
_resource_type: Option<&str>,
_language: u16,
) -> Result<*const u8> {
if !self.is_loaded {
return Err(MapleError::ExecutionFailed("Module not loaded".to_string()));
}
let pe = PEParser::new(&self.pe_data)?;
let _resource_dir = match pe.get_resource_directory() {
Some(dir) if dir.size > 0 => dir,
_ => {
return Err(MapleError::SymbolNotFound(
"No resource directory".to_string(),
));
}
};
Err(MapleError::SymbolNotFound("Resource not found".to_string()))
}
fn sizeof_resource(&self, _resource: *const u8) -> Result<usize> {
if !self.is_loaded {
return Err(MapleError::ExecutionFailed("Module not loaded".to_string()));
}
Err(MapleError::SymbolNotFound(
"Resource size not available".to_string(),
))
}
fn load_resource(&self, _resource: *const u8) -> Result<*const u8> {
if !self.is_loaded {
return Err(MapleError::ExecutionFailed("Module not loaded".to_string()));
}
Err(MapleError::SymbolNotFound(
"Resource data not available".to_string(),
))
}
fn load_string(&self, id: u32, buffer: &mut [u16]) -> Result<usize> {
self.load_string_ex(id, buffer, 0)
}
fn load_string_ex(&self, _id: u32, _buffer: &mut [u16], _language: u16) -> Result<usize> {
if !self.is_loaded {
return Err(MapleError::ExecutionFailed("Module not loaded".to_string()));
}
Err(MapleError::SymbolNotFound(
"String resource not found".to_string(),
))
}
}
impl Drop for WindowsMemoryModule {
fn drop(&mut self) {
let _ = self.free();
}
}