Fix early drop for frida transformer (#992)

* Fix early drop for frida transformer

* clippy

fuzzers/frida_libpng/src/fuzzer.rs

@@ -54,16 +54,16 @@ pub fn main() {
     let options = parse_args();
 
     unsafe {
-        match fuzz(options) {
+        match fuzz(&options) {
             Ok(()) | Err(Error::ShuttingDown) => println!("\nFinished fuzzing. Good bye."),
-            Err(e) => panic!("Error during fuzzing: {:?}", e),
+            Err(e) => panic!("Error during fuzzing: {e:?}"),
         }
     }
 }
 
 /// The actual fuzzer
 #[allow(clippy::too_many_lines, clippy::too_many_arguments)]
-unsafe fn fuzz(options: FuzzerOptions) -> Result<(), Error> {
+unsafe fn fuzz(options: &FuzzerOptions) -> Result<(), Error> {
     // 'While the stats are state, they are usually used in the broker - which is likely never restarted
     let monitor = MultiMonitor::new(|s| println!("{s}"));
 
@@ -96,7 +96,7 @@ unsafe fn fuzz(options: FuzzerOptions) -> Result<(), Error> {
 
                 #[cfg(unix)]
                 let mut frida_helper =
-                    FridaInstrumentationHelper::new(&gum, &options, tuple_list!(coverage, asan));
+                    FridaInstrumentationHelper::new(&gum, options, tuple_list!(coverage, asan));
                 #[cfg(windows)]
                 let mut frida_helper =
                     FridaInstrumentationHelper::new(&gum, &options, tuple_list!(coverage));
@@ -141,7 +141,7 @@ unsafe fn fuzz(options: FuzzerOptions) -> Result<(), Error> {
                         // Corpus in which we store solutions (crashes in this example),
                         // on disk so the user can get them after stopping the fuzzer
                         OnDiskCorpus::new_save_meta(
-                            options.output.to_path_buf(),
+                            options.output.clone(),
                             Some(OnDiskMetadataFormat::JsonPretty),
                         )
                         .unwrap(),
@@ -219,7 +219,7 @@ unsafe fn fuzz(options: FuzzerOptions) -> Result<(), Error> {
                 let cmplog = CmpLogRuntime::new();
 
                 let mut frida_helper =
-                    FridaInstrumentationHelper::new(&gum, &options, tuple_list!(coverage, cmplog));
+                    FridaInstrumentationHelper::new(&gum, options, tuple_list!(coverage, cmplog));
 
                 // Create an observation channel using the coverage map
                 let edges_observer = HitcountsMapObserver::new(StdMapObserver::from_mut_ptr(
@@ -259,7 +259,7 @@ unsafe fn fuzz(options: FuzzerOptions) -> Result<(), Error> {
                         // Corpus in which we store solutions (crashes in this example),
                         // on disk so the user can get them after stopping the fuzzer
                         OnDiskCorpus::new_save_meta(
-                            options.output.to_path_buf(),
+                            options.output.clone(),
                             Some(OnDiskMetadataFormat::JsonPretty),
                         )
                         .unwrap(),
@@ -352,7 +352,7 @@ unsafe fn fuzz(options: FuzzerOptions) -> Result<(), Error> {
                 let coverage = CoverageRuntime::new();
 
                 let mut frida_helper =
-                    FridaInstrumentationHelper::new(&gum, &options, tuple_list!(coverage));
+                    FridaInstrumentationHelper::new(&gum, options, tuple_list!(coverage));
 
                 // Create an observation channel using the coverage map
                 let edges_observer = HitcountsMapObserver::new(StdMapObserver::from_mut_ptr(
@@ -392,7 +392,7 @@ unsafe fn fuzz(options: FuzzerOptions) -> Result<(), Error> {
                         // Corpus in which we store solutions (crashes in this example),
                         // on disk so the user can get them after stopping the fuzzer
                         OnDiskCorpus::new_save_meta(
-                            options.output.to_path_buf(),
+                            options.output.clone(),
                             Some(OnDiskMetadataFormat::JsonPretty),
                         )
                         .unwrap(),

libafl_frida/src/executor.rs

@@ -39,7 +39,6 @@ where
     /// User provided callback for instrumentation
     helper: &'c mut FridaInstrumentationHelper<'b, RT>,
     followed: bool,
-    gum: &'b Gum,
     _phantom: PhantomData<&'b u8>,
 }
 
@@ -85,8 +84,8 @@ where
                 self.stalker.activate(NativePointer(core::ptr::null_mut()));
             } else {
                 self.followed = true;
-                let transformer = self.helper.transformer(self.gum);
+                let transformer = self.helper.transformer();
-                self.stalker.follow_me::<NoneEventSink>(&transformer, None);
+                self.stalker.follow_me::<NoneEventSink>(transformer, None);
             }
         }
         let res = self.base.run_target(fuzzer, state, mgr, input);
@@ -190,7 +189,6 @@ where
             base,
             stalker,
             helper,
-            gum,
             followed: false,
             _phantom: PhantomData,
         }

libafl_frida/src/helper.rs

@@ -120,6 +120,7 @@ pub struct FridaInstrumentationHelper<'a, RT> {
     ranges: RangeMap<usize, (u16, String)>,
     module_map: ModuleMap,
     options: &'a FuzzerOptions,
+    transformer: Option<Transformer<'a>>,
     runtimes: RT,
 }
 
@@ -220,6 +221,7 @@ where
             module_map: ModuleMap::new_from_names(gum, &modules_to_instrument),
             options,
             runtimes,
+            transformer: None,
         };
 
         if options.cmplog || options.asan || !options.disable_coverage {
@@ -252,6 +254,117 @@ where
                 .runtimes
                 .init_all(gum, &helper.ranges, &modules_to_instrument);
         }
+
+        let transformer = 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 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
     }
 
@@ -272,113 +385,12 @@ where
     }
 
     /// Returns ref to the Transformer
-    pub fn transformer(&mut self, gum: &'a Gum) -> Transformer<'a> {
+    pub fn transformer(&mut self) -> &Transformer<'a> {
-        Transformer::from_callback(gum, |basic_block, output| {
+        // the Transformer is always initialized on `new`. We can safely unwrap.
-            let mut first = true;
+        self.transformer.as_ref().unwrap()
-            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
+    /// Initialize all
     pub fn init(
         &mut self,
         gum: &'a Gum,