From 33db263260830b70ac8b3763070fdeade0358e21 Mon Sep 17 00:00:00 2001
From: "Dongjia \"toka\" Zhang" <tokazerkje@outlook.com>
Date: Thu, 13 Feb 2025 11:33:24 +0100
Subject: [PATCH] Fix brk() handling for snapshot module (#2970)

* drop grown address

* this clippy lint literally makes 0 sense!

* i hate you rust

* mm

* don't use drop! add comments for why alignment is not necessary
---
 Cargo.toml                                   |  1 +
 libafl_qemu/src/modules/usermode/snapshot.rs | 45 +++++++++++++-------
 2 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index c31440a0bb..475edf6db4 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -162,6 +162,7 @@ module_name_repetitions = "allow"
 unsafe_derive_deserialize = "allow"
 similar_names = "allow"
 too_many_lines = "allow"
+comparison_chain = "allow"          # This lint makes **ZERO** sense
 
 
 [workspace.lints.rustdoc]
diff --git a/libafl_qemu/src/modules/usermode/snapshot.rs b/libafl_qemu/src/modules/usermode/snapshot.rs
index 56b68f708e..cf78f2b25d 100644
--- a/libafl_qemu/src/modules/usermode/snapshot.rs
+++ b/libafl_qemu/src/modules/usermode/snapshot.rs
@@ -407,11 +407,20 @@ impl SnapshotModule {
                 let aligned_new_brk = (new_brk + ((SNAPSHOT_PAGE_SIZE - 1) as GuestAddr))
                     & (!(SNAPSHOT_PAGE_SIZE - 1) as GuestAddr);
                 log::debug!("New brk ({:#x?}) < snapshotted brk ({:#x?})! Mapping back in the target {:#x?} - {:#x?}", new_brk, self.brk, aligned_new_brk, aligned_new_brk + (self.brk - aligned_new_brk));
-                drop(qemu.map_fixed(
+                qemu.map_fixed(
                     aligned_new_brk,
                     (self.brk - aligned_new_brk) as usize,
                     MmapPerms::ReadWrite,
-                ));
+                )
+                .unwrap();
+            } else if new_brk > self.brk {
+                // The heap has grown. so we want to drop those
+                let drop_sz = (new_brk - self.brk) as usize;
+
+                // if self.brk is not aligned this call will return an error
+                // and it will page align this drop_sz too
+                // look at target_munmap in qemu-libafl-bridge
+                qemu.unmap(self.brk, drop_sz).unwrap();
             }
 
             for acc in &mut self.accesses {
@@ -426,11 +435,12 @@ impl SnapshotModule {
                                 .query_mut(*page..(page + SNAPSHOT_PAGE_SIZE as GuestAddr))
                             {
                                 if !entry.value.perms.unwrap_or(MmapPerms::None).writable() {
-                                    drop(qemu.mprotect(
+                                    qemu.mprotect(
                                         entry.interval.start,
                                         (entry.interval.end - entry.interval.start) as usize,
                                         MmapPerms::ReadWrite,
-                                    ));
+                                    )
+                                    .unwrap();
                                     entry.value.changed = true;
                                     entry.value.perms = Some(MmapPerms::ReadWrite);
                                 }
@@ -464,11 +474,12 @@ impl SnapshotModule {
                     if !entry.value.perms.unwrap_or(MmapPerms::None).writable()
                         && !entry.value.changed
                     {
-                        drop(qemu.mprotect(
+                        qemu.mprotect(
                             entry.interval.start,
                             (entry.interval.end - entry.interval.start) as usize,
                             MmapPerms::ReadWrite,
-                        ));
+                        )
+                        .unwrap();
                         entry.value.changed = true;
                     }
                 }
@@ -487,11 +498,12 @@ impl SnapshotModule {
 
         for entry in self.maps.tree.query_mut(0..GuestAddr::MAX) {
             if entry.value.changed {
-                drop(qemu.mprotect(
+                qemu.mprotect(
                     entry.interval.start,
                     (entry.interval.end - entry.interval.start) as usize,
                     entry.value.perms.unwrap(),
-                ));
+                )
+                .unwrap();
                 entry.value.changed = false;
             }
         }
@@ -653,26 +665,29 @@ impl SnapshotModule {
 
             if found.is_empty() {
                 //panic!("A pre-snapshot memory region was unmapped");
-                drop(qemu.map_fixed(
+                qemu.map_fixed(
                     entry.interval.start,
                     (entry.interval.end - entry.interval.start) as usize,
                     entry.value.perms.unwrap(),
-                ));
+                )
+                .unwrap();
             } else if found.len() == 1 && found[0].0 == *entry.interval {
                 if found[0].1 && found[0].2 != entry.value.perms {
-                    drop(qemu.mprotect(
+                    qemu.mprotect(
                         entry.interval.start,
                         (entry.interval.end - entry.interval.start) as usize,
                         entry.value.perms.unwrap(),
-                    ));
+                    )
+                    .unwrap();
                 }
             } else {
                 //  TODO check for holes
-                drop(qemu.mprotect(
+                qemu.mprotect(
                     entry.interval.start,
                     (entry.interval.end - entry.interval.start) as usize,
                     entry.value.perms.unwrap(),
-                ));
+                )
+                .unwrap();
             }
 
             for (interval, ..) in found {
@@ -685,7 +700,7 @@ impl SnapshotModule {
             to_unmap.push((*entry.interval, entry.value.changed, entry.value.perms));
         }
         for (i, ..) in to_unmap {
-            drop(qemu.unmap(i.start, (i.end - i.start) as usize));
+            qemu.unmap(i.start, (i.end - i.start) as usize).unwrap();
             new_maps.tree.delete(i);
         }