Revision control
Copy as Markdown
Other Tools
use super::{
arch::ArchX86_64,
unwind_rule::{OffsetOrPop, UnwindRuleX86_64},
unwindregs::Reg,
};
use crate::arch::Arch;
use crate::pe::{PeSections, PeUnwinderError, PeUnwinding};
use crate::unwind_result::UnwindResult;
use core::ops::ControlFlow;
use alloc::vec::Vec;
use pe_unwind_info::x86_64::{
FunctionEpilogInstruction, FunctionTableEntries, Register, UnwindInfo, UnwindInfoTrailer,
UnwindOperation, UnwindState,
};
struct State<'a, F> {
regs: &'a mut <ArchX86_64 as Arch>::UnwindRegs,
read_stack: &'a mut F,
}
impl<F> UnwindState for State<'_, F>
where
F: FnMut(u64) -> Result<u64, ()>,
{
fn read_register(&mut self, register: Register) -> u64 {
self.regs.get(convert_pe_register(register))
}
fn read_stack(&mut self, addr: u64) -> Option<u64> {
(self.read_stack)(addr).ok()
}
fn write_register(&mut self, register: Register, value: u64) {
self.regs.set(convert_pe_register(register), value)
}
fn write_xmm_register(&mut self, _register: pe_unwind_info::x86_64::XmmRegister, _value: u128) {
// Ignore
}
}
fn convert_pe_register(r: Register) -> Reg {
match r {
Register::RAX => Reg::RAX,
Register::RCX => Reg::RCX,
Register::RDX => Reg::RDX,
Register::RBX => Reg::RBX,
Register::RSP => Reg::RSP,
Register::RBP => Reg::RBP,
Register::RSI => Reg::RSI,
Register::RDI => Reg::RDI,
Register::R8 => Reg::R8,
Register::R9 => Reg::R9,
Register::R10 => Reg::R10,
Register::R11 => Reg::R11,
Register::R12 => Reg::R12,
Register::R13 => Reg::R13,
Register::R14 => Reg::R14,
Register::R15 => Reg::R15,
}
}
impl From<&'_ FunctionEpilogInstruction> for OffsetOrPop {
fn from(value: &'_ FunctionEpilogInstruction) -> Self {
match value {
FunctionEpilogInstruction::AddSP(offset) => {
if let Ok(v) = (offset / 8).try_into() {
OffsetOrPop::OffsetBy8(v)
} else {
OffsetOrPop::None
}
}
FunctionEpilogInstruction::Pop(reg) => OffsetOrPop::Pop(convert_pe_register(*reg)),
_ => OffsetOrPop::None,
}
}
}
impl From<&'_ UnwindOperation> for OffsetOrPop {
fn from(value: &'_ UnwindOperation) -> Self {
match value {
UnwindOperation::UnStackAlloc(offset) => {
if let Ok(v) = (offset / 8).try_into() {
OffsetOrPop::OffsetBy8(v)
} else {
OffsetOrPop::None
}
}
UnwindOperation::PopNonVolatile(reg) => OffsetOrPop::Pop(convert_pe_register(*reg)),
_ => OffsetOrPop::None,
}
}
}
impl PeUnwinding for ArchX86_64 {
fn unwind_frame<F, D>(
sections: PeSections<D>,
address: u32,
regs: &mut Self::UnwindRegs,
is_first_frame: bool,
read_stack: &mut F,
) -> Result<UnwindResult<Self::UnwindRule>, PeUnwinderError>
where
F: FnMut(u64) -> Result<u64, ()>,
D: core::ops::Deref<Target = [u8]>,
{
let entries = FunctionTableEntries::parse(sections.pdata);
let Some(function) = entries.lookup(address) else {
return Ok(UnwindResult::ExecRule(UnwindRuleX86_64::JustReturn));
};
let read_stack_err = |read_stack: &mut F, addr| {
read_stack(addr).map_err(|()| PeUnwinderError::MissingStackData(Some(addr)))
};
let unwind_info_address = function.unwind_info_address.get();
let unwind_info =
UnwindInfo::parse(sections.unwind_info_memory_at_rva(unwind_info_address)?)
.ok_or(PeUnwinderError::UnwindInfoParseError)?;
if is_first_frame {
// Check whether the address is in the function epilog. If so, we need to
// simulate the remaining epilog instructions (unwind codes don't account for
// unwinding from the epilog). We only need to check this for the first unwind info (if
// there are chained infos).
let bytes = (function.end_address.get() - address) as usize;
let instruction = §ions.text_memory_at_rva(address)?[..bytes];
if let Ok(epilog_instructions) =
FunctionEpilogInstruction::parse_sequence(instruction, unwind_info.frame_register())
{
// If the epilog is an optional AddSP followed by Pops, we can return a cache
// rule.
if let Some(rule) =
UnwindRuleX86_64::for_sequence_of_offset_or_pop(epilog_instructions.iter())
{
return Ok(UnwindResult::ExecRule(rule));
}
for instruction in epilog_instructions.iter() {
match instruction {
FunctionEpilogInstruction::AddSP(offset) => {
let rsp = regs.get(Reg::RSP);
regs.set(Reg::RSP, rsp + *offset as u64);
}
FunctionEpilogInstruction::AddSPFromFP(offset) => {
let fp = unwind_info
.frame_register()
.expect("invalid fp register offset");
let fp = convert_pe_register(fp);
let fp = regs.get(fp);
regs.set(Reg::RSP, fp + *offset as u64);
}
FunctionEpilogInstruction::Pop(reg) => {
let rsp = regs.get(Reg::RSP);
let val = read_stack_err(read_stack, rsp)?;
regs.set(convert_pe_register(*reg), val);
regs.set(Reg::RSP, rsp + 8);
}
}
}
let rsp = regs.get(Reg::RSP);
let ra = read_stack_err(read_stack, rsp)?;
regs.set(Reg::RSP, rsp + 8);
return Ok(UnwindResult::Uncacheable(ra));
}
}
// Get all chained UnwindInfo and resolve errors when collecting.
let chained_info = core::iter::successors(Some(Ok(unwind_info)), |info| {
let Ok(info) = info else {
return None;
};
if let Some(UnwindInfoTrailer::ChainedUnwindInfo { chained }) = info.trailer() {
let unwind_info_address = chained.unwind_info_address.get();
Some(
sections
.unwind_info_memory_at_rva(unwind_info_address)
.and_then(|data| {
UnwindInfo::parse(data).ok_or(PeUnwinderError::UnwindInfoParseError)
}),
)
} else {
None
}
})
.collect::<Result<Vec<_>, _>>()?;
// Get all operations across chained UnwindInfo. The first should be filtered to only those
// operations which are before the offset in the function.
let offset = address - function.begin_address.get();
let operations = chained_info.into_iter().enumerate().flat_map(|(i, info)| {
info.unwind_operations()
.skip_while(move |(o, _)| i == 0 && *o as u32 > offset)
.map(|(_, op)| op)
});
// We need to collect operations to first check (without losing ownership) whether an
// unwind rule can be returned.
let operations = operations.collect::<Vec<_>>();
if let Some(rule) = UnwindRuleX86_64::for_sequence_of_offset_or_pop(operations.iter()) {
return Ok(UnwindResult::ExecRule(rule));
}
// Resolve operations to get the return address.
let mut state = State { regs, read_stack };
for op in operations {
if let ControlFlow::Break(ra) = unwind_info
.resolve_operation(&mut state, &op)
.ok_or(PeUnwinderError::MissingStackData(None))?
{
return Ok(UnwindResult::Uncacheable(ra));
}
}
let rsp = regs.get(Reg::RSP);
let ra = read_stack_err(read_stack, rsp)?;
regs.set(Reg::RSP, rsp + 8);
Ok(UnwindResult::Uncacheable(ra))
}
}