Add PE support for x86_64

Cargo.toml

@@ -17,6 +17,7 @@ gimli = "0.27.0"
 object = { version = ">=0.30", optional = true }
 thiserror = "1.0.30"
 macho-unwind-info = "0.3.0"
+pe-unwind-info = "0.1.0"
 fallible-iterator = "0.2.0"
 
 [dev-dependencies]

src/aarch64/mod.rs

@@ -3,6 +3,7 @@ mod cache;
 mod dwarf;
 mod instruction_analysis;
 mod macho;
+mod pe;
 mod unwind_rule;
 mod unwinder;
 mod unwindregs;
@@ -12,6 +13,7 @@ pub use cache::*;
 pub use dwarf::*;
 pub use instruction_analysis::*;
 pub use macho::*;
+pub use pe::*;
 pub use unwind_rule::*;
 pub use unwinder::*;
 pub use unwindregs::*;

src/aarch64/pe.rs

@@ -0,0 +1,18 @@
+use super::arch::ArchAarch64;
+use crate::pe::{PeSections, PeUnwinderError, PeUnwinding};
+use crate::unwind_result::UnwindResult;
+
+impl PeUnwinding for ArchAarch64 {
+    fn unwind_frame<F, D>(
+        _sections: PeSections<D>,
+        _address: u32,
+        _regs: &mut Self::UnwindRegs,
+        _read_stack: &mut F,
+    ) -> Result<UnwindResult<Self::UnwindRule>, PeUnwinderError>
+    where
+        F: FnMut(u64) -> Result<u64, ()>,
+        D: std::ops::Deref<Target = [u8]>,
+    {
+        Err(PeUnwinderError::Aarch64Unsupported)
+    }
+}

src/error.rs

@@ -1,5 +1,6 @@
 use crate::dwarf::DwarfUnwinderError;
 use crate::macho::CompactUnwindInfoUnwinderError;
+use crate::pe::PeUnwinderError;
 
 /// The error type used in this crate.
 #[derive(thiserror::Error, Debug, Clone, Copy, PartialEq, Eq)]
@@ -28,6 +29,9 @@ pub enum UnwinderError {
     #[error("DWARF unwinding failed: {0}")]
     Dwarf(#[from] DwarfUnwinderError),
 
+    #[error("PE unwinding failed: {0}")]
+    Pe(#[from] PeUnwinderError),
+
     #[error("__unwind_info referred to DWARF FDE but we do not have __eh_frame data")]
     NoDwarfData,
 

src/lib.rs

@@ -120,6 +120,7 @@ mod dwarf;
 mod error;
 mod instruction_analysis;
 mod macho;
+mod pe;
 mod rule_cache;
 mod unwind_result;
 mod unwind_rule;

src/pe.rs

@@ -0,0 +1,65 @@
+use crate::{arch::Arch, unwind_result::UnwindResult};
+use std::ops::Range;
+
+#[derive(thiserror::Error, Clone, Copy, Debug, PartialEq, Eq)]
+pub enum PeUnwinderError {
+    #[error("failed to read unwind info memory at RVA {0:x}")]
+    MissingUnwindInfoData(u32),
+    #[error("failed to read instruction memory at RVA {0:x}")]
+    MissingInstructionData(u32),
+    #[error("failed to read stack{}", .0.map(|a| format!(" at address {a:x}")).unwrap_or_default())]
+    MissingStackData(Option<u64>),
+    #[error("failed to parse UnwindInfo")]
+    UnwindInfoParseError,
+    #[error("AArch64 is not yet supported")]
+    Aarch64Unsupported,
+}
+
+pub struct PeSections<'a, D> {
+    pub pdata: &'a D,
+    pub rdata: Option<&'a (Range<u32>, D)>,
+    pub xdata: Option<&'a (Range<u32>, D)>,
+    pub text: Option<&'a (Range<u32>, D)>,
+}
+
+impl<'a, D> PeSections<'a, D>
+where
+    D: std::ops::Deref<Target = [u8]>,
+{
+    pub fn unwind_info_memory_at_rva(&self, rva: u32) -> Result<&'a [u8], PeUnwinderError> {
+        [&self.rdata, &self.xdata]
+            .into_iter()
+            .find_map(|o| o.and_then(|m| memory_at_rva(m, rva)))
+            .ok_or(PeUnwinderError::MissingUnwindInfoData(rva))
+    }
+
+    pub fn text_memory_at_rva(&self, rva: u32) -> Result<&'a [u8], PeUnwinderError> {
+        self.text
+            .and_then(|m| memory_at_rva(m, rva))
+            .ok_or(PeUnwinderError::MissingInstructionData(rva))
+    }
+}
+
+fn memory_at_rva<D: std::ops::Deref<Target = [u8]>>(
+    (range, data): &(Range<u32>, D),
+    address: u32,
+) -> Option<&[u8]> {
+    if range.contains(&address) {
+        let offset = address - range.start;
+        Some(&data[(offset as usize)..])
+    } else {
+        None
+    }
+}
+
+pub trait PeUnwinding: Arch {
+    fn unwind_frame<F, D>(
+        sections: PeSections<D>,
+        address: u32,
+        regs: &mut Self::UnwindRegs,
+        read_stack: &mut F,
+    ) -> Result<UnwindResult<Self::UnwindRule>, PeUnwinderError>
+    where
+        F: FnMut(u64) -> Result<u64, ()>,
+        D: std::ops::Deref<Target = [u8]>;
+}

src/unwinder.rs

@@ -10,6 +10,7 @@ use crate::instruction_analysis::InstructionAnalysis;
 use crate::macho::{
     CompactUnwindInfoUnwinder, CompactUnwindInfoUnwinding, CuiUnwindResult, TextBytes,
 };
+use crate::pe::PeUnwinding;
 use crate::rule_cache::CacheResult;
 use crate::unwind_result::UnwindResult;
 use crate::unwind_rule::UnwindRule;
@@ -201,7 +202,7 @@ fn next_global_modules_generation() -> u16 {
 
 pub struct UnwinderInternal<
     D: Deref<Target = [u8]>,
-    A: Arch + DwarfUnwinding + CompactUnwindInfoUnwinding + InstructionAnalysis,
+    A: Arch + DwarfUnwinding + CompactUnwindInfoUnwinding + PeUnwinding + InstructionAnalysis,
     P: AllocationPolicy<D>,
 > {
     /// sorted by avma_range.start
@@ -214,7 +215,7 @@ pub struct UnwinderInternal<
 
 impl<
         D: Deref<Target = [u8]>,
-        A: Arch + DwarfUnwinding + CompactUnwindInfoUnwinding + InstructionAnalysis,
+        A: Arch + DwarfUnwinding + CompactUnwindInfoUnwinding + PeUnwinding + InstructionAnalysis,
         P: AllocationPolicy<D>,
     > Default for UnwinderInternal<D, A, P>
 {
@@ -225,7 +226,7 @@ impl<
 
 impl<
         D: Deref<Target = [u8]>,
-        A: Arch + DwarfUnwinding + CompactUnwindInfoUnwinding + InstructionAnalysis,
+        A: Arch + DwarfUnwinding + CompactUnwindInfoUnwinding + PeUnwinding + InstructionAnalysis,
         P: AllocationPolicy<D>,
     > UnwinderInternal<D, A, P>
 {
@@ -522,6 +523,22 @@ impl<
                     read_stack,
                 )?
             }
+            ModuleUnwindDataInternal::PeUnwindInfo {
+                pdata,
+                rdata,
+                xdata,
+                text,
+            } => <A as PeUnwinding>::unwind_frame(
+                crate::pe::PeSections {
+                    pdata,
+                    rdata: rdata.as_ref(),
+                    xdata: xdata.as_ref(),
+                    text: text.as_ref(),
+                },
+                rel_lookup_address,
+                regs,
+                read_stack,
+            )?,
             ModuleUnwindDataInternal::None => return Err(UnwinderError::NoModuleUnwindData),
         };
         Ok(unwind_result)
@@ -575,6 +592,13 @@ enum ModuleUnwindDataInternal<D: Deref<Target = [u8]>> {
         debug_frame: Arc<D>,
         base_addresses: crate::dwarf::BaseAddresses,
     },
+    /// Used with PE binaries (Windows).
+    PeUnwindInfo {
+        pdata: D,
+        rdata: Option<(Range<u32>, D)>,
+        xdata: Option<(Range<u32>, D)>,
+        text: Option<(Range<u32>, D)>,
+    },
     /// No unwind information is used. Unwinding in this module will use a fallback rule
     /// (usually frame pointer unwinding).
     None,
@@ -611,6 +635,24 @@ impl<D: Deref<Target = [u8]>> ModuleUnwindDataInternal<D> {
                 base_addresses: base_addresses_for_sections(section_info),
                 text_data,
             }
+        } else if let Some(pdata) = section_info.section_data(b".pdata") {
+            let mut range_and_data = |name| {
+                section_info
+                    .section_svma_range(name)
+                    .and_then(|range| {
+                        Some(Range {
+                            start: (range.start - section_info.base_svma()).try_into().ok()?,
+                            end: (range.end - section_info.base_svma()).try_into().ok()?,
+                        })
+                    })
+                    .zip(section_info.section_data(name))
+            };
+            ModuleUnwindDataInternal::PeUnwindInfo {
+                pdata,
+                rdata: range_and_data(b".rdata"),
+                xdata: range_and_data(b".xdata"),
+                text: range_and_data(b".text"),
+            }
         } else if let Some(eh_frame) = section_info
             .section_data(b".eh_frame")
             .or_else(|| section_info.section_data(b"__eh_frame"))

src/x86_64/mod.rs

@@ -3,6 +3,7 @@ mod cache;
 mod dwarf;
 mod instruction_analysis;
 mod macho;
+mod pe;
 mod unwind_rule;
 mod unwinder;
 mod unwindregs;
@@ -12,6 +13,7 @@ pub use cache::*;
 pub use dwarf::*;
 pub use instruction_analysis::*;
 pub use macho::*;
+pub use pe::*;
 pub use unwind_rule::*;
 pub use unwinder::*;
 pub use unwindregs::*;

src/x86_64/pe.rs

@@ -0,0 +1,149 @@
+use super::{arch::ArchX86_64, unwind_rule::UnwindRuleX86_64, unwindregs::Reg};
+use crate::arch::Arch;
+use crate::pe::{PeSections, PeUnwinderError, PeUnwinding};
+use crate::unwind_result::UnwindResult;
+use std::ops::ControlFlow;
+
+use pe_unwind_info::x86_64::{
+    FunctionEpilogInstruction, FunctionEpilogParser, FunctionTableEntries, Register, UnwindInfo,
+    UnwindInfoTrailer, 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 PeUnwinding for ArchX86_64 {
+    fn unwind_frame<F, D>(
+        sections: PeSections<D>,
+        address: u32,
+        regs: &mut Self::UnwindRegs,
+        read_stack: &mut F,
+    ) -> Result<UnwindResult<Self::UnwindRule>, PeUnwinderError>
+    where
+        F: FnMut(u64) -> Result<u64, ()>,
+        D: std::ops::Deref<Target = [u8]>,
+    {
+        let entries = FunctionTableEntries::parse(sections.pdata);
+        let Some(mut 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 offset = address - function.begin_address.get();
+        let mut is_chained = false;
+        loop {
+            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_chained {
+                // 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).
+                let bytes = (function.end_address.get() - address) as usize;
+                let instruction = &sections.text_memory_at_rva(address)?[..bytes];
+                let mut epilog_parser = FunctionEpilogParser::new();
+                if let Some(epilog_instructions) =
+                    epilog_parser.is_function_epilog(instruction, unwind_info.frame_register())
+                {
+                    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);
+                            }
+                        }
+                    }
+                    break;
+                }
+            }
+
+            let mut state = State { regs, read_stack };
+            for (_, op) in unwind_info
+                .unwind_operations()
+                .skip_while(|(o, _)| !is_chained && *o as u32 > offset)
+            {
+                if let ControlFlow::Break(ra) = unwind_info
+                    .resolve_operation(&mut state, &op)
+                    .ok_or(PeUnwinderError::MissingStackData(None))?
+                {
+                    return Ok(UnwindResult::Uncacheable(ra));
+                }
+            }
+            if let Some(UnwindInfoTrailer::ChainedUnwindInfo { chained }) = unwind_info.trailer() {
+                is_chained = true;
+                function = chained;
+            } else {
+                break;
+            }
+        }
+
+        let rsp = regs.get(Reg::RSP);
+        let ra = read_stack_err(read_stack, rsp)?;
+        regs.set(Reg::RSP, rsp + 8);
+
+        Ok(UnwindResult::Uncacheable(ra))
+    }
+}