Fix FridaInstrumentationHelper bugs caused by moving it after creation (#931)

* move Transformer out of FridaInstrumentationHelper's fields and create it dynamically instead; wrap CoverageRuntime in Pin<Arc<RefCell>>

* Update helper.rs

* run cargo fmt

* switch Arc for Rc

libafl_frida/src/coverage_rt.rs

@@ -1,5 +1,12 @@
 //! Functionality regarding binary-only coverage collection.
 use core::ptr::addr_of_mut;
+use std::{
+    cell::{Ref, RefCell},
+    marker::PhantomPinned,
+    ops::Deref,
+    pin::Pin,
+    rc::Rc,
+};
 
 use dynasmrt::{dynasm, DynasmApi, DynasmLabelApi};
 #[cfg(target_arch = "x86_64")]
@@ -15,15 +22,19 @@ use crate::helper::FridaRuntime;
 /// (Default) map size for frida coverage reporting
 pub const MAP_SIZE: usize = 64 * 1024;
 
-/// Frida binary-only coverage
 #[derive(Debug)]
-pub struct CoverageRuntime {
+struct CoverageRuntimeInner {
     map: [u8; MAP_SIZE],
     previous_pc: u64,
     current_log_impl: u64,
     blob_maybe_log: Option<Box<[u8]>>,
+    _pinned: PhantomPinned,
 }
 
+/// Frida binary-only coverage
+#[derive(Debug)]
+pub struct CoverageRuntime(Pin<Rc<RefCell<CoverageRuntimeInner>>>);
+
 impl Default for CoverageRuntime {
     fn default() -> Self {
         Self::new()
@@ -32,6 +43,7 @@ impl Default for CoverageRuntime {
 
 impl FridaRuntime for CoverageRuntime {
     /// Initialize the coverage runtime
+    /// The struct MUST NOT be moved after this function is called, as the generated assembly references it
     fn init(
         &mut self,
         _gum: &frida_gum::Gum,
@@ -60,23 +72,24 @@ impl CoverageRuntime {
     /// Create a new coverage runtime
     #[must_use]
     pub fn new() -> Self {
-        Self {
+        Self(Rc::pin(RefCell::new(CoverageRuntimeInner {
             map: [0_u8; MAP_SIZE],
             previous_pc: 0,
             current_log_impl: 0,
             blob_maybe_log: None,
-        }
+            _pinned: PhantomPinned,
+        })))
     }
 
     /// Retrieve the coverage map pointer
     pub fn map_ptr_mut(&mut self) -> *mut u8 {
-        self.map.as_mut_ptr()
+        self.0.borrow_mut().map.as_mut_ptr()
     }
 
     /// Retrieve the `maybe_log` code blob, that will write coverage into the map
     #[must_use]
-    pub fn blob_maybe_log(&self) -> &[u8] {
+    pub fn blob_maybe_log(&self) -> impl Deref<Target = Box<[u8]>> + '_ {
-        self.blob_maybe_log.as_ref().unwrap()
+        Ref::map(self.0.borrow(), |s| s.blob_maybe_log.as_ref().unwrap())
     }
 
     /// A minimal `maybe_log` implementation. We insert this into the transformed instruction stream
@@ -104,12 +117,13 @@ impl CoverageRuntime {
             ;   ldp x1, x2, [sp], #0x10
             ;   ret
             ;map_addr:
-            ;.qword addr_of_mut!(self.map) as i64
+            ;.qword addr_of_mut!(self.0.borrow_mut().map) as i64
             ;previous_loc:
             ;.qword 0
         );
         let ops_vec = ops.finalize().unwrap();
-        self.blob_maybe_log = Some(ops_vec[..ops_vec.len() - 8].to_vec().into_boxed_slice());
+        self.0.borrow_mut().blob_maybe_log =
+            Some(ops_vec[..ops_vec.len() - 8].to_vec().into_boxed_slice());
     }
 
     /// A minimal `maybe_log` implementation. We insert this into the transformed instruction stream
@@ -140,12 +154,13 @@ impl CoverageRuntime {
             ;   popfq
             ;   ret
             ;map_addr:
-            ;.qword addr_of_mut!(self.map) as i64
+            ;.qword addr_of_mut!(self.0.borrow_mut().map) as i64
             ;previous_loc:
             ;.qword 0
         );
         let ops_vec = ops.finalize().unwrap();
-        self.blob_maybe_log = Some(ops_vec[..ops_vec.len() - 8].to_vec().into_boxed_slice());
+        self.0.borrow_mut().blob_maybe_log =
+            Some(ops_vec[..ops_vec.len() - 8].to_vec().into_boxed_slice());
     }
 
     /// Emits coverage mapping into the current basic block.
@@ -156,9 +171,10 @@ impl CoverageRuntime {
         let writer = output.writer();
         #[allow(clippy::cast_possible_wrap)] // gum redzone size is u32, we need an offset as i32.
         let redzone_size = i64::from(frida_gum_sys::GUM_RED_ZONE_SIZE);
-        if self.current_log_impl == 0
+        if self.0.borrow().current_log_impl == 0
-            || !writer.can_branch_directly_to(self.current_log_impl)
+            || !writer.can_branch_directly_to(self.0.borrow().current_log_impl)
-            || !writer.can_branch_directly_between(writer.pc() + 128, self.current_log_impl)
+            || !writer
+                .can_branch_directly_between(writer.pc() + 128, self.0.borrow().current_log_impl)
         {
             let after_log_impl = writer.code_offset() + 1;
 
@@ -167,9 +183,9 @@ impl CoverageRuntime {
             #[cfg(target_arch = "aarch64")]
             writer.put_b_label(after_log_impl);
 
-            self.current_log_impl = writer.pc();
+            self.0.borrow_mut().current_log_impl = writer.pc();
-            writer.put_bytes(self.blob_maybe_log());
+            writer.put_bytes(&self.blob_maybe_log());
-            let prev_loc_pointer = addr_of_mut!(self.previous_pc) as u64; // Get the pointer to self.previous_pc
+            let prev_loc_pointer = addr_of_mut!(self.0.borrow_mut().previous_pc) as u64; // Get the pointer to self.previous_pc
 
             writer.put_bytes(&prev_loc_pointer.to_ne_bytes());
 
@@ -180,7 +196,7 @@ impl CoverageRuntime {
             writer.put_lea_reg_reg_offset(X86Register::Rsp, X86Register::Rsp, -(redzone_size));
             writer.put_push_reg(X86Register::Rdi);
             writer.put_mov_reg_address(X86Register::Rdi, h64 & (MAP_SIZE as u64 - 1));
-            writer.put_call_address(self.current_log_impl);
+            writer.put_call_address(self.0.borrow().current_log_impl);
             writer.put_pop_reg(X86Register::Rdi);
             writer.put_lea_reg_reg_offset(X86Register::Rsp, X86Register::Rsp, redzone_size);
         }
@@ -195,7 +211,7 @@ impl CoverageRuntime {
             );
             writer.put_ldr_reg_u64(Aarch64Register::X0, h64 & (MAP_SIZE as u64 - 1));
 
-            writer.put_bl_imm(self.current_log_impl);
+            writer.put_bl_imm(self.0.borrow().current_log_impl);
             writer.put_ldp_reg_reg_reg_offset(
                 Aarch64Register::Lr,
                 Aarch64Register::X0,

libafl_frida/src/executor.rs

@@ -31,6 +31,7 @@ where
     S::Input: HasTargetBytes,
     S: UsesInput,
     OT: ObserversTuple<S>,
+    'a: 'b,
 {
     base: InProcessExecutor<'a, H, OT, S>,
     /// Frida's dynamic rewriting engine
@@ -38,6 +39,7 @@ where
     /// User provided callback for instrumentation
     helper: &'c mut FridaInstrumentationHelper<'b, RT>,
     followed: bool,
+    gum: &'b Gum,
     _phantom: PhantomData<&'b u8>,
 }
 
@@ -83,8 +85,8 @@ where
                 self.stalker.activate(NativePointer(core::ptr::null_mut()));
             } else {
                 self.followed = true;
-                self.stalker
+                let transformer = self.helper.transformer(self.gum);
-                    .follow_me::<NoneEventSink>(self.helper.transformer(), None);
+                self.stalker.follow_me::<NoneEventSink>(&transformer, None);
             }
         }
         let res = self.base.run_target(fuzzer, state, mgr, input);
@@ -185,6 +187,7 @@ where
             base,
             stalker,
             helper,
+            gum,
             followed: false,
             _phantom: PhantomData,
         }

libafl_frida/src/helper.rs

@@ -115,8 +115,6 @@ where
 
 /// An helper that feeds `FridaInProcessExecutor` with edge-coverage instrumentation
 pub struct FridaInstrumentationHelper<'a, RT> {
-    /// Transformer that has to be passed to FridaInProcessExecutor
-    transformer: Option<Transformer<'a>>,
     #[cfg(unix)]
     capstone: Capstone,
     ranges: RangeMap<usize, (u16, String)>,
@@ -204,7 +202,6 @@ where
             modules_to_instrument.iter().map(AsRef::as_ref).collect();
 
         let mut helper = Self {
-            transformer: None,
             #[cfg(target_arch = "aarch64")]
             capstone: Capstone::new()
                 .arm64()
@@ -245,122 +242,12 @@ where
                 }
             }
 
-            let transformer = Transformer::from_callback(gum, |basic_block, output| {
+            // make sure we aren't in the instrumented list, as it would cause recursions
-                let mut first = true;
+            assert!(
-                for instruction in basic_block {
+                !helper.ranges.contains_key(&(Self::new as usize)),
-                    let instr = instruction.instr();
+                "instrumented libraries must not include the fuzzer"
-                    #[cfg(unix)]
+            );
-                    let instr_size = instr.bytes().len();
-                    let address = instr.address();
-                    //println!("block @ {:x} transformed to {:x}", address, output.writer().pc());
 
-                    //println!(
-                    //"address: {:x} contains: {:?}",
-                    //address,
-                    //helper.ranges.contains_key(&(address as usize))
-                    //);
-
-                    // println!("Ranges: {:#?}", helper.ranges);
-                    if helper.ranges.contains_key(&(address as usize)) {
-                        if first {
-                            first = false;
-                            //println!("block @ {:x} transformed to {:x}", address, output.writer().pc());
-                            if let Some(rt) = helper.runtime_mut::<CoverageRuntime>() {
-                                rt.emit_coverage_mapping(address, &output);
-                            }
-
-                            #[cfg(unix)]
-                            if let Some(rt) = helper.runtime_mut::<DrCovRuntime>() {
-                                instruction.put_callout(|context| {
-                                    let real_address = rt.real_address_for_stalked(pc(&context));
-                                    //let (range, (id, name)) = helper.ranges.get_key_value(&real_address).unwrap();
-                                    //println!("{}:0x{:016x}", name, real_address - range.start);
-                                    rt.drcov_basic_blocks.push(DrCovBasicBlock::new(
-                                        real_address,
-                                        real_address + instr_size,
-                                    ));
-                                });
-                            }
-                        }
-
-                        #[cfg(unix)]
-                        let res = if let Some(_rt) = helper.runtime::<AsanRuntime>() {
-                            AsanRuntime::asan_is_interesting_instruction(
-                                &helper.capstone,
-                                address,
-                                instr,
-                            )
-                        } else {
-                            None
-                        };
-
-                        #[cfg(all(target_arch = "x86_64", unix))]
-                        if let Some((segment, width, basereg, indexreg, scale, disp)) = res {
-                            if let Some(rt) = helper.runtime_mut::<AsanRuntime>() {
-                                rt.emit_shadow_check(
-                                    address, &output, segment, width, basereg, indexreg, scale,
-                                    disp,
-                                );
-                            }
-                        }
-
-                        #[cfg(target_arch = "aarch64")]
-                        if let Some((basereg, indexreg, displacement, width, shift, extender)) = res
-                        {
-                            if let Some(rt) = helper.runtime_mut::<AsanRuntime>() {
-                                rt.emit_shadow_check(
-                                    address,
-                                    &output,
-                                    basereg,
-                                    indexreg,
-                                    displacement,
-                                    width,
-                                    shift,
-                                    extender,
-                                );
-                            }
-                        }
-
-                        #[cfg(all(feature = "cmplog", target_arch = "aarch64"))]
-                        if let Some(rt) = helper.runtime::<CmpLogRuntime>() {
-                            if let Some((op1, op2, special_case)) =
-                                CmpLogRuntime::cmplog_is_interesting_instruction(
-                                    &helper.capstone,
-                                    address,
-                                    instr,
-                                )
-                            {
-                                //emit code that saves the relevant data in runtime(passes it to x0, x1)
-                                rt.emit_comparison_handling(
-                                    address,
-                                    &output,
-                                    &op1,
-                                    &op2,
-                                    special_case,
-                                );
-                            }
-                        }
-
-                        #[cfg(unix)]
-                        if let Some(rt) = helper.runtime_mut::<AsanRuntime>() {
-                            rt.add_stalked_address(
-                                output.writer().pc() as usize - instr_size,
-                                address as usize,
-                            );
-                        }
-
-                        #[cfg(unix)]
-                        if let Some(rt) = helper.runtime_mut::<DrCovRuntime>() {
-                            rt.add_stalked_address(
-                                output.writer().pc() as usize - instr_size,
-                                address as usize,
-                            );
-                        }
-                    }
-                    instruction.keep();
-                }
-            });
-            helper.transformer = Some(transformer);
             helper
                 .runtimes
                 .init_all(gum, &helper.ranges, &modules_to_instrument);
@@ -385,8 +272,110 @@ where
     }
 
     /// Returns ref to the Transformer
-    pub fn transformer(&self) -> &Transformer<'a> {
+    pub fn transformer(&mut self, gum: &'a Gum) -> Transformer<'a> {
-        self.transformer.as_ref().unwrap()
+        Transformer::from_callback(gum, |basic_block, output| {
+            let mut first = true;
+            for instruction in basic_block {
+                let instr = instruction.instr();
+                #[cfg(unix)]
+                let instr_size = instr.bytes().len();
+                let address = instr.address();
+                //println!("block @ {:x} transformed to {:x}", address, output.writer().pc());
+
+                //println!(
+                //"address: {:x} contains: {:?}",
+                //address,
+                //self.ranges().contains_key(&(address as usize))
+                //);
+
+                // println!("Ranges: {:#?}", self.ranges());
+                if self.ranges().contains_key(&(address as usize)) {
+                    if first {
+                        first = false;
+                        //println!("block @ {:x} transformed to {:x}", address, output.writer().pc());
+                        if let Some(rt) = self.runtime_mut::<CoverageRuntime>() {
+                            rt.emit_coverage_mapping(address, &output);
+                        }
+
+                        #[cfg(unix)]
+                        if let Some(rt) = self.runtime_mut::<DrCovRuntime>() {
+                            instruction.put_callout(|context| {
+                                let real_address = rt.real_address_for_stalked(pc(&context));
+                                //let (range, (id, name)) = helper.ranges.get_key_value(&real_address).unwrap();
+                                //println!("{}:0x{:016x}", name, real_address - range.start);
+                                rt.drcov_basic_blocks.push(DrCovBasicBlock::new(
+                                    real_address,
+                                    real_address + instr_size,
+                                ));
+                            });
+                        }
+                    }
+
+                    #[cfg(unix)]
+                    let res = if let Some(_rt) = self.runtime::<AsanRuntime>() {
+                        AsanRuntime::asan_is_interesting_instruction(&self.capstone, address, instr)
+                    } else {
+                        None
+                    };
+
+                    #[cfg(all(target_arch = "x86_64", unix))]
+                    if let Some((segment, width, basereg, indexreg, scale, disp)) = res {
+                        if let Some(rt) = self.runtime_mut::<AsanRuntime>() {
+                            rt.emit_shadow_check(
+                                address, &output, segment, width, basereg, indexreg, scale, disp,
+                            );
+                        }
+                    }
+
+                    #[cfg(target_arch = "aarch64")]
+                    if let Some((basereg, indexreg, displacement, width, shift, extender)) = res {
+                        if let Some(rt) = self.runtime_mut::<AsanRuntime>() {
+                            rt.emit_shadow_check(
+                                address,
+                                &output,
+                                basereg,
+                                indexreg,
+                                displacement,
+                                width,
+                                shift,
+                                extender,
+                            );
+                        }
+                    }
+
+                    #[cfg(all(feature = "cmplog", target_arch = "aarch64"))]
+                    if let Some(rt) = self.runtime::<CmpLogRuntime>() {
+                        if let Some((op1, op2, special_case)) =
+                            CmpLogRuntime::cmplog_is_interesting_instruction(
+                                &self.capstone,
+                                address,
+                                instr,
+                            )
+                        {
+                            //emit code that saves the relevant data in runtime(passes it to x0, x1)
+                            rt.emit_comparison_handling(address, &output, &op1, &op2, special_case);
+                        }
+                    }
+
+                    #[cfg(unix)]
+                    if let Some(rt) = self.runtime_mut::<AsanRuntime>() {
+                        rt.add_stalked_address(
+                            output.writer().pc() as usize - instr_size,
+                            address as usize,
+                        );
+                    }
+
+                    #[cfg(unix)]
+                    if let Some(rt) = self.runtime_mut::<DrCovRuntime>() {
+                        rt.add_stalked_address(
+                            output.writer().pc() as usize - instr_size,
+                            address as usize,
+                        );
+                    }
+                }
+                instruction.keep();
+            }
+        })
     }
 
     /// Initializa all