wip: Implement is_first_frame for PE

src/aarch64/dwarf.rs

@@ -11,6 +11,7 @@ use crate::dwarf::{
     eval_cfa_rule, eval_register_rule, ConversionError, DwarfUnwindRegs, DwarfUnwinderError,
     DwarfUnwinding,
 };
+use crate::FrameAddress;
 
 impl DwarfUnwindRegs for UnwindRegsAarch64 {
     fn get(&self, register: Register) -> Option<u64> {
@@ -79,7 +80,9 @@ impl DwarfUnwinding for ArchAarch64 {
         regs.set_sp(cfa);
         regs.set_lr(lr);
 
-        Ok(UnwindResult::Uncacheable(lr))
+        Ok(UnwindResult::Uncacheable(
+            FrameAddress::from_return_address(lr),
+        ))
     }
 
     fn rule_if_uncovered_by_fde() -> Self::UnwindRule {

src/aarch64/pe.rs

@@ -4,10 +4,11 @@ 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,
+        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, ()>,

src/aarch64/unwinder.rs

@@ -54,7 +54,7 @@ impl<D: Deref<Target = [u8]>, P: AllocationPolicy<D>> Unwinder for UnwinderAarch
         regs: &mut UnwindRegsAarch64,
         cache: &mut CacheAarch64<D, P>,
         read_stack: &mut F,
-    ) -> Result<Option<u64>, Error>
+    ) -> Result<Option<FrameAddress>, Error>
     where
         F: FnMut(u64) -> Result<u64, ()>,
     {

src/pe.rs

@@ -71,6 +71,7 @@ pub trait PeUnwinding: Arch {
         sections: PeSections<D>,
         address: u32,
         regs: &mut Self::UnwindRegs,
+        is_first_frame: bool,
         read_stack: &mut F,
     ) -> Result<UnwindResult<Self::UnwindRule>, PeUnwinderError>
     where

src/unwind_result.rs

@@ -1,5 +1,7 @@
+use crate::FrameAddress;
+
 #[derive(Debug, Clone)]
 pub enum UnwindResult<R> {
     ExecRule(R),
-    Uncacheable(u64),
+    Uncacheable(Option<FrameAddress>),
 }

src/unwinder.rs

@@ -66,7 +66,7 @@ pub trait Unwinder {
         regs: &mut Self::UnwindRegs,
         cache: &mut Self::Cache,
         read_stack: &mut F,
-    ) -> Result<Option<u64>, Error>
+    ) -> Result<Option<FrameAddress>, Error>
     where
         F: FnMut(u64) -> Result<u64, ()>;
 
@@ -164,8 +164,6 @@ impl<'u, 'c, 'r, U: Unwinder + ?Sized, F: FnMut(u64) -> Result<u64, ()>>
         };
         match next {
             Some(return_address) => {
-                let return_address = FrameAddress::from_return_address(return_address)
-                    .ok_or(Error::ReturnAddressIsNull)?;
                 self.state = UnwindIteratorState::Unwinding(return_address);
                 Ok(Some(return_address))
             }
@@ -308,7 +306,7 @@ impl<
         cache: &mut Cache<D, A::UnwindRule, P>,
         read_stack: &mut F,
         callback: G,
-    ) -> Result<Option<u64>, Error>
+    ) -> Result<Option<FrameAddress>, Error>
     where
         F: FnMut(u64) -> Result<u64, ()>,
         G: FnOnce(
@@ -327,7 +325,9 @@ impl<
             .lookup(lookup_address, self.modules_generation)
         {
             CacheResult::Hit(unwind_rule) => {
-                return unwind_rule.exec(is_first_frame, regs, read_stack);
+                return Ok(unwind_rule
+                    .exec(is_first_frame, regs, read_stack)?
+                    .and_then(FrameAddress::from_return_address));
             }
             CacheResult::Miss(handle) => handle,
         };
@@ -345,9 +345,7 @@ impl<
                     read_stack,
                 ) {
                     Ok(UnwindResult::ExecRule(rule)) => rule,
-                    Ok(UnwindResult::Uncacheable(return_address)) => {
-                        return Ok(Some(return_address))
-                    }
+                    Ok(UnwindResult::Uncacheable(return_address)) => return Ok(return_address),
                     Err(_err) => {
                         // eprintln!("Unwinder error: {}", err);
                         A::UnwindRule::fallback_rule()
@@ -356,7 +354,9 @@ impl<
             }
         };
         cache.rule_cache.insert(cache_handle, unwind_rule);
-        unwind_rule.exec(is_first_frame, regs, read_stack)
+        Ok(unwind_rule
+            .exec(is_first_frame, regs, read_stack)?
+            .and_then(FrameAddress::from_return_address))
     }
 
     pub fn unwind_frame<F>(
@@ -365,7 +365,7 @@ impl<
         regs: &mut A::UnwindRegs,
         cache: &mut Cache<D, A::UnwindRule, P>,
         read_stack: &mut F,
-    ) -> Result<Option<u64>, Error>
+    ) -> Result<Option<FrameAddress>, Error>
     where
         F: FnMut(u64) -> Result<u64, ()>,
     {
@@ -538,6 +538,7 @@ impl<
                 },
                 rel_lookup_address,
                 regs,
+                is_first_frame,
                 read_stack,
             )?,
             ModuleUnwindDataInternal::None => return Err(UnwinderError::NoModuleUnwindData),

src/x86_64/dwarf.rs

@@ -9,6 +9,7 @@ use crate::dwarf::{
     DwarfUnwinding,
 };
 use crate::unwind_result::UnwindResult;
+use crate::FrameAddress;
 
 impl DwarfUnwindRegs for UnwindRegsX86_64 {
     fn get(&self, register: Register) -> Option<u64> {
@@ -74,7 +75,9 @@ impl DwarfUnwinding for ArchX86_64 {
         regs.set_bp(new_bp);
         regs.set_sp(cfa);
 
-        Ok(UnwindResult::Uncacheable(return_address))
+        Ok(UnwindResult::Uncacheable(
+            FrameAddress::from_return_address(return_address),
+        ))
     }
 
     fn rule_if_uncovered_by_fde() -> Self::UnwindRule {

src/x86_64/pe.rs

@@ -8,6 +8,7 @@ use crate::pe::{PeSections, PeUnwinderError, PeUnwinding};
 use crate::unwind_result::UnwindResult;
 use std::ops::ControlFlow;
 
+use crate::FrameAddress;
 use pe_unwind_info::x86_64::{
     FunctionEpilogInstruction, FunctionTableEntries, Register, UnwindInfo, UnwindInfoTrailer,
     UnwindOperation, UnwindState,
@@ -97,6 +98,7 @@ impl PeUnwinding for ArchX86_64 {
         sections: PeSections<D>,
         address: u32,
         regs: &mut Self::UnwindRegs,
+        is_first_frame: bool,
         read_stack: &mut F,
     ) -> Result<UnwindResult<Self::UnwindRule>, PeUnwinderError>
     where
@@ -117,53 +119,58 @@ impl PeUnwinding for ArchX86_64 {
             UnwindInfo::parse(sections.unwind_info_memory_at_rva(unwind_info_address)?)
                 .ok_or(PeUnwinderError::UnwindInfoParseError)?;
 
-        // 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 = &sections.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())
+        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 = &sections.text_memory_at_rva(address)?[..bytes];
+            if let Ok(epilog_instructions) =
+                FunctionEpilogInstruction::parse_sequence(instruction, unwind_info.frame_register())
             {
-                return Ok(UnwindResult::ExecRule(rule));
-            }
+                // 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);
+                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);
+                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));
+                return Ok(UnwindResult::Uncacheable(
+                    FrameAddress::from_return_address(ra),
+                ));
+            }
         }
 
+
         // Get all chained UnwindInfo and resolve errors when collecting.
         let chained_info = std::iter::successors(Some(Ok(unwind_info)), |info| {
             let Ok(info) = info else {
@@ -207,14 +214,18 @@ impl PeUnwinding for ArchX86_64 {
                 .resolve_operation(&mut state, &op)
                 .ok_or(PeUnwinderError::MissingStackData(None))?
             {
-                return Ok(UnwindResult::Uncacheable(ra));
+                return Ok(UnwindResult::Uncacheable(Some(
+                    FrameAddress::from_instruction_pointer(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))
+        Ok(UnwindResult::Uncacheable(
+            FrameAddress::from_return_address(ra),
+        ))
     }
 }

src/x86_64/unwinder.rs

@@ -55,7 +55,7 @@ impl<D: Deref<Target = [u8]>, P: AllocationPolicy<D>> Unwinder for UnwinderX86_6
         regs: &mut UnwindRegsX86_64,
         cache: &mut CacheX86_64<D, P>,
         read_stack: &mut F,
-    ) -> Result<Option<u64>, Error>
+    ) -> Result<Option<FrameAddress>, Error>
     where
         F: FnMut(u64) -> Result<u64, ()>,
     {