Fix hook_func and implement deinit (#2180)

* Hook fix

* Implement deinit to pass tests

* Fix aarch64 errors as well as signficantly speed up startup on macos

* fmt

* Fix cmplog_rt

* Fix windows ci

* ci fix 2

* Fix clippy

* Revert hook implementation to new impl

* Add nolinkage option for hooks and fix some comments

* fmt

* Move hooking mechanism to entirely static mut

* fmt

* Fix clippy

* fix windows

* fmt

* Overall the function hook macro to ensure that the original function is correct and not the same function in a different library. Also change static muts to OnceLock

libafl_frida/src/alloc.rs

@@ -556,12 +556,14 @@ impl Allocator {
             &mut |range: &RangeDetails| -> bool {
                 let start = range.memory_range().base_address().0 as usize;
                 let end = start + range.memory_range().size();
-
-                if !self.is_managed(start as *mut c_void) {
+                //the shadow region should be the highest valid userspace region, so don't continue after
+                if self.is_managed(start as *mut c_void) {
+                    false
+                } else {
+                    log::trace!("Unpoisoning: {:#x}-{:#x}", start, end);
                     self.map_shadow_for_region(start, end, true);
-                }
-
                     true
+                }
             },
         );
     }
@@ -587,6 +589,14 @@ impl Allocator {
             &mut |range: &RangeDetails| -> bool {
                 let start = range.memory_range().base_address().0 as usize;
                 let end = start + range.memory_range().size();
+                log::trace!("New range: {:#x}-{:#x}", start, end);
+
+                #[cfg(target_vendor = "apple")]
+                if start >= 0x600000000000 {
+                    //this is the MALLOC_NANO region. There is no point in spending time tracking this region as we hook malloc
+                    return false;
+                }
+
                 occupied_ranges.push((start, end));
                 // On x64, if end > 2**48, then that's in vsyscall or something.
                 #[cfg(all(unix, target_arch = "x86_64"))]
@@ -614,15 +624,17 @@ impl Allocator {
         #[cfg(windows)]
         let maxbit = 63;
         #[cfg(unix)]
-        for power in 1..64 {
+        for power in 44..64 {
             if 2_usize.pow(power) > userspace_max {
                 maxbit = power;
                 break;
             }
         }
 
+        log::trace!("max bit: {}", maxbit);
+
         {
-            for try_shadow_bit in 44..maxbit {
+            for try_shadow_bit in 44..=maxbit {
                 let addr: usize = 1 << try_shadow_bit;
                 let shadow_start = addr;
                 let shadow_end = addr + addr + addr;

libafl_frida/src/cmplog_rt.rs

@@ -128,6 +128,8 @@ impl FridaRuntime for CmpLogRuntime {
         self.generate_instrumentation_blobs();
     }
 
+    fn deinit(&mut self, _gum: &frida_gum::Gum) {}
+
     fn pre_exec<I: Input + HasTargetBytes>(&mut self, _input: &I) -> Result<(), Error> {
         Ok(())
     }

libafl_frida/src/coverage_rt.rs

@@ -42,6 +42,8 @@ impl FridaRuntime for CoverageRuntime {
     ) {
     }
 
+    fn deinit(&mut self, _gum: &frida_gum::Gum) {}
+
     fn pre_exec<I: libafl::inputs::Input + libafl::inputs::HasTargetBytes>(
         &mut self,
         _input: &I,

libafl_frida/src/drcov_rt.rs

@@ -40,6 +40,8 @@ impl FridaRuntime for DrCovRuntime {
             .expect("failed to create directory for coverage files");
     }
 
+    fn deinit(&mut self, _gum: &frida_gum::Gum) {}
+
     /// Called before execution, does nothing
     fn pre_exec<I: Input + HasTargetBytes>(&mut self, _input: &I) -> Result<(), Error> {
         Ok(())

libafl_frida/src/helper.rs

@@ -45,6 +45,8 @@ pub trait FridaRuntime: 'static + Debug {
         ranges: &RangeMap<usize, (u16, String)>,
         module_map: &Rc<ModuleMap>,
     );
+    /// Deinitialization
+    fn deinit(&mut self, gum: &Gum);
 
     /// Method called before execution
     fn pre_exec<I: Input + HasTargetBytes>(&mut self, input: &I) -> Result<(), Error>;
@@ -63,6 +65,9 @@ pub trait FridaRuntimeTuple: MatchFirstType + Debug {
         module_map: &Rc<ModuleMap>,
     );
 
+    /// Deinitialization
+    fn deinit_all(&mut self, gum: &Gum);
+
     /// Method called before execution
     fn pre_exec_all<I: Input + HasTargetBytes>(&mut self, input: &I) -> Result<(), Error>;
 
@@ -78,6 +83,8 @@ impl FridaRuntimeTuple for () {
         _module_map: &Rc<ModuleMap>,
     ) {
     }
+    fn deinit_all(&mut self, _gum: &Gum) {}
+
     fn pre_exec_all<I: Input + HasTargetBytes>(&mut self, _input: &I) -> Result<(), Error> {
         Ok(())
     }
@@ -101,6 +108,11 @@ where
         self.1.init_all(gum, ranges, module_map);
     }
 
+    fn deinit_all(&mut self, gum: &Gum) {
+        self.0.deinit(gum);
+        self.1.deinit_all(gum);
+    }
+
     fn pre_exec_all<I: Input + HasTargetBytes>(&mut self, input: &I) -> Result<(), Error> {
         self.0.pre_exec(input)?;
         self.1.pre_exec_all(input)
@@ -575,6 +587,11 @@ where
         }
     }
 
+    /// Clean up all runtimes
+    pub fn deinit(&mut self, gum: &Gum) {
+        (*self.runtimes).borrow_mut().deinit_all(gum);
+    }
+
     /*
     /// Return the runtime
     pub fn runtime<R>(&self) -> Option<&R>

libafl_frida/src/lib.rs

@@ -529,6 +529,8 @@ mod tests {
                 }
             }
         }
+
+        frida_helper.deinit(GUM.get().expect("Gum uninitialized"));
     }
 
     #[test]