From 11ae49b7cd4cc5f4948b936e0e32e49435bf12ee Mon Sep 17 00:00:00 2001
From: s1341 <github@shmarya.net>
Date: Sun, 26 Dec 2021 10:44:25 +0200
Subject: [PATCH 1/2] Implement max total allocation size for frida asan

---
 libafl_frida/src/alloc.rs |  9 ++++++++-
 libafl_frida/src/lib.rs   | 12 ++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/libafl_frida/src/alloc.rs b/libafl_frida/src/alloc.rs
index 733a92a70d..22912881fa 100644
--- a/libafl_frida/src/alloc.rs
+++ b/libafl_frida/src/alloc.rs
@@ -28,6 +28,7 @@ pub struct Allocator {
     shadow_pages: RangeSet<usize>,
     allocation_queue: HashMap<usize, Vec<AllocationMetadata>>,
     largest_allocation: usize,
+    total_allocation_size: usize,
     base_mapping_addr: usize,
     current_mapping_addr: usize,
 }
@@ -70,7 +71,7 @@ impl Allocator {
         let mut shadow_bit = 0;
 
         #[cfg(all(target_arch = "aarch64", target_os = "android"))]
-        for try_shadow_bit in &[46usize, 36usize] {
+        for try_shadow_bit in &[44usize, 36usize] {
             let addr: usize = 1 << try_shadow_bit;
             if unsafe {
                 mmap(
@@ -147,6 +148,7 @@ impl Allocator {
             shadow_pages: RangeSet::new(),
             allocation_queue: HashMap::new(),
             largest_allocation: 0,
+            total_allocation_size: 0,
             base_mapping_addr: addr + addr + addr,
             current_mapping_addr: addr + addr + addr,
         }
@@ -205,6 +207,11 @@ impl Allocator {
         }
         let rounded_up_size = self.round_up_to_page(size) + 2 * self.page_size;
 
+        if self.total_allocation_size + rounded_up_size > self.options.asan_max_total_allocation() {
+            return std::ptr::null_mut();
+        }
+        self.total_allocation_size += rounded_up_size;
+
         let metadata = if let Some(mut metadata) = self.find_smallest_fit(rounded_up_size) {
             //println!("reusing allocation at {:x}, (actual mapping starts at {:x}) size {:x}", metadata.address, metadata.address - self.page_size, size);
             metadata.is_malloc_zero = is_malloc_zero;
diff --git a/libafl_frida/src/lib.rs b/libafl_frida/src/lib.rs
index 4c1d8f805c..50ca01c210 100644
--- a/libafl_frida/src/lib.rs
+++ b/libafl_frida/src/lib.rs
@@ -37,6 +37,7 @@ pub struct FridaOptions {
     enable_asan_continue_after_error: bool,
     enable_asan_allocation_backtraces: bool,
     asan_max_allocation: usize,
+    asan_max_total_allocation: usize,
     asan_max_allocation_panics: bool,
     enable_coverage: bool,
     enable_drcov: bool,
@@ -79,6 +80,9 @@ impl FridaOptions {
                     "asan-max-allocation" => {
                         options.asan_max_allocation = value.parse().unwrap();
                     }
+                    "asan-max-total-allocation" => {
+                        options.asan_max_total_allocation = value.parse().unwrap();
+                    }
                     "asan-max-allocation-panics" => {
                         options.asan_max_allocation_panics = value.parse().unwrap();
                     }
@@ -208,6 +212,13 @@ impl FridaOptions {
         self.asan_max_allocation
     }
 
+    /// The maximum total allocation size that the ASAN allocator should allocate
+    #[must_use]
+    #[inline]
+    pub fn asan_max_total_allocation(&self) -> usize {
+        self.asan_max_total_allocation
+    }
+
     /// Should we panic if the max ASAN allocation size is exceeded
     #[must_use]
     #[inline]
@@ -252,6 +263,7 @@ impl Default for FridaOptions {
             enable_asan_continue_after_error: false,
             enable_asan_allocation_backtraces: true,
             asan_max_allocation: 1 << 30,
+            asan_max_total_allocation: 1 << 32,
             asan_max_allocation_panics: false,
             enable_coverage: true,
             enable_drcov: false,

From 2e92a34494295f0ffe0d5b2eaf5b991448423579 Mon Sep 17 00:00:00 2001
From: s1341 <github@shmarya.net>
Date: Sun, 26 Dec 2021 11:17:27 +0200
Subject: [PATCH 2/2] Reset total allocations on reset

---
 libafl_frida/src/alloc.rs | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/libafl_frida/src/alloc.rs b/libafl_frida/src/alloc.rs
index 22912881fa..f9c4ae2bb9 100644
--- a/libafl_frida/src/alloc.rs
+++ b/libafl_frida/src/alloc.rs
@@ -352,6 +352,8 @@ impl Allocator {
         for allocation in tmp_allocations {
             self.allocations.insert(allocation.address, allocation);
         }
+
+        self.total_allocation_size = 0;
     }
 
     pub fn get_usable_size(&self, ptr: *mut c_void) -> usize {