use crate::types::unwind::CompiledFunctionUnwindInfoReference;
use std::collections::HashMap;
use windows_sys::Win32::System::Diagnostics::Debug::{
IMAGE_RUNTIME_FUNCTION_ENTRY, RtlAddFunctionTable, RtlDeleteFunctionTable,
};
pub struct UnwindRegistry {
functions: HashMap<usize, Vec<IMAGE_RUNTIME_FUNCTION_ENTRY>>,
published: bool,
}
impl UnwindRegistry {
pub fn new() -> Self {
Self {
functions: HashMap::new(),
published: false,
}
}
pub fn register(
&mut self,
base_address: usize,
func_start: u32,
func_len: u32,
info: &CompiledFunctionUnwindInfoReference,
) -> Result<(), String> {
if self.published {
return Err("unwind registry has already been published".to_string());
}
match info {
CompiledFunctionUnwindInfoReference::WindowsX64(_) => {}
_ => return Err("unsupported unwind information".to_string()),
};
let mut entry: IMAGE_RUNTIME_FUNCTION_ENTRY = unsafe { std::mem::zeroed() };
entry.BeginAddress = func_start;
entry.EndAddress = func_start + func_len;
entry.Anonymous.UnwindInfoAddress = (entry.EndAddress + 3) & !3;
let entries = self.functions.entry(base_address).or_default();
entries.push(entry);
Ok(())
}
pub fn publish_eh_frame(&mut self, _eh_frame: Option<&[u8]>) -> Result<(), String> {
if self.published {
return Err("unwind registry has already been published".to_string());
}
self.published = true;
if !self.functions.is_empty() {
for (base_address, functions) in self.functions.iter_mut() {
assert_eq!(
(functions.as_mut_ptr() as u64) % 4,
0,
"function table allocation was not aligned"
);
unsafe {
if !RtlAddFunctionTable(
functions.as_mut_ptr(),
functions.len() as u32,
*base_address as u64,
) {
return Err("failed to register function tables".to_string());
}
}
}
}
Ok(())
}
}
impl Drop for UnwindRegistry {
fn drop(&mut self) {
if self.published {
unsafe {
for functions in self.functions.values_mut() {
RtlDeleteFunctionTable(functions.as_mut_ptr());
}
}
}
}
}