From c68b30ae2a15c770135fbd918093ff4521334299 Mon Sep 17 00:00:00 2001
From: WorksButNotTested <62701594+WorksButNotTested@users.noreply.github.com>
Date: Fri, 28 Mar 2025 19:52:57 +0000
Subject: [PATCH] Add iniitialize feature to librasan (#3113)

* Optimize memset

* Make the initialization of memory buffers an optional feature

---------

Co-authored-by: Your Name <you@example.com>
---
 libafl_qemu/librasan/asan/Cargo.toml          |  2 ++
 .../asan/src/allocator/frontend/default.rs    |  9 ++++--
 libafl_qemu/librasan/asan/src/mem.rs          | 30 ++++++++++++++++---
 3 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/libafl_qemu/librasan/asan/Cargo.toml b/libafl_qemu/librasan/asan/Cargo.toml
index d48a51733d..1a2fec4697 100644
--- a/libafl_qemu/librasan/asan/Cargo.toml
+++ b/libafl_qemu/librasan/asan/Cargo.toml
@@ -34,6 +34,8 @@ guest = []
 hooks = []
 ## Enable support for shadow memory and tracking in the host
 host = ["dep:syscalls"]
+## Intialize all allocations with 0xff
+initialize = []
 ## Enable use of the `libc` library to support creation of mappings, read/write, logging etc (more OS agnostic)
 libc = ["dep:libc"]
 ## Enable the use of direct syscalls (supported by `rustix`) to interact with the operating system (Linux specific).
diff --git a/libafl_qemu/librasan/asan/src/allocator/frontend/default.rs b/libafl_qemu/librasan/asan/src/allocator/frontend/default.rs
index 493a99c04e..58e8ab8b41 100644
--- a/libafl_qemu/librasan/asan/src/allocator/frontend/default.rs
+++ b/libafl_qemu/librasan/asan/src/allocator/frontend/default.rs
@@ -13,7 +13,8 @@ use alloc::{
     collections::{BTreeMap, VecDeque},
     fmt::Debug,
 };
-use core::slice::from_raw_parts_mut;
+#[cfg(feature = "initialize")]
+use core::ptr::write_bytes;
 
 use log::debug;
 use thiserror::Error;
@@ -107,8 +108,10 @@ impl<B: GlobalAlloc + Send, S: Shadow, T: Tracking> AllocatorFrontend for Defaul
             .poison(data + len, poison_len, PoisonType::AsanStackRightRz)
             .map_err(|e| DefaultFrontendError::ShadowError(e))?;
 
-        let buffer = unsafe { from_raw_parts_mut(data as *mut u8, len) };
-        buffer.iter_mut().for_each(|b| *b = 0xff);
+        #[cfg(feature = "initialize")]
+        unsafe {
+            write_bytes(data as *mut u8, 0xff, len)
+        };
         Ok(data)
     }
 
diff --git a/libafl_qemu/librasan/asan/src/mem.rs b/libafl_qemu/librasan/asan/src/mem.rs
index 21b5b3bbf7..0a33793b37 100644
--- a/libafl_qemu/librasan/asan/src/mem.rs
+++ b/libafl_qemu/librasan/asan/src/mem.rs
@@ -65,10 +65,32 @@ pub unsafe extern "C" fn memcpy(dest: *mut u8, src: *const u8, count: usize) {
 
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn memset(dest: *mut u8, value: u8, count: usize) {
-    let dest_slice = unsafe { from_raw_parts_mut(dest, count) };
-    #[allow(clippy::needless_range_loop)]
-    for i in 0..count {
-        dest_slice[i] = value;
+    unsafe {
+        let mut cursor = dest;
+        let word_value = match value {
+            u8::MIN => Some(usize::MIN),
+            u8::MAX => Some(usize::MAX),
+            _ => None,
+        };
+
+        if let Some(word_value) = word_value {
+            let num_words = count / size_of::<usize>();
+            for _ in 0..num_words {
+                *(cursor as *mut usize) = word_value;
+                cursor = cursor.wrapping_add(size_of::<usize>());
+            }
+
+            let num_bytes = count % size_of::<usize>();
+            for _ in 0..num_bytes {
+                *cursor = value;
+                cursor = cursor.wrapping_add(1);
+            }
+        } else {
+            for _ in 0..count {
+                *cursor = value;
+                cursor = cursor.wrapping_add(1);
+            }
+        }
     }
 }