wmemcheck support for additional allocation functions and granular memory

crates/cranelift/src/func_environ.rs

@@ -585,6 +585,112 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
         builder.ins().call(check_malloc, &[vmctx, retval, len]);
     }
 
+    #[cfg(feature = "wmemcheck")]
+    fn hook_calloc_exit(&mut self, builder: &mut FunctionBuilder, retvals: &[ir::Value]) {
+        let check_calloc = self.builtin_functions.check_calloc(builder.func);
+        let vmctx = self.vmctx_val(&mut builder.cursor());
+        let func_args = builder
+            .func
+            .dfg
+            .block_params(builder.func.layout.entry_block().unwrap());
+        let (count, size) = if func_args.len() < 4 {
+            return;
+        } else {
+            (func_args[2], func_args[3])
+        };
+        let retval = if retvals.len() < 1 {
+            return;
+        } else {
+            retvals[0]
+        };
+        builder.ins().call(check_calloc, &[vmctx, retval, count, size]);
+    }
+
+    #[cfg(feature = "wmemcheck")]
+    fn hook_realloc_exit(&mut self, builder: &mut FunctionBuilder, retvals: &[ir::Value]) {
+        let check_realloc = self.builtin_functions.check_realloc(builder.func);
+        let vmctx = self.vmctx_val(&mut builder.cursor());
+        let func_args = builder
+            .func
+            .dfg
+            .block_params(builder.func.layout.entry_block().unwrap());
+        let (ptr, len) = if func_args.len() < 4 {
+            return;
+        } else {
+            // If a function named `realloc` has at least two arguments, we assume the
+            // first arguments are the pointer and requested allocation size.
+            (func_args[2], func_args[3])
+        };
+        let retval = if retvals.len() < 1 {
+            return;
+        } else {
+            retvals[0]
+        };
+        builder.ins().call(check_realloc, &[vmctx, retval, ptr, len]);
+    }
+
+    #[cfg(feature = "wmemcheck")]
+    fn hook_malloc_usable_size_exit(&mut self, builder: &mut FunctionBuilder, retvals: &[ir::Value]) {
+        let check_malloc_usable_size = self.builtin_functions.check_malloc_usable_size(builder.func);
+        let vmctx = self.vmctx_val(&mut builder.cursor());
+        let func_args = builder
+            .func
+            .dfg
+            .block_params(builder.func.layout.entry_block().unwrap());
+        let ptr = if func_args.len() < 3 {
+            return;
+        } else {
+            func_args[2]
+        };
+        let retval = if retvals.len() < 1 {
+            return;
+        } else {
+            retvals[0]
+        };
+        builder.ins().call(check_malloc_usable_size, &[vmctx, retval, ptr]);
+    }
+
+    #[cfg(feature = "wmemcheck")]
+    fn hook_posix_memalign_exit(&mut self, builder: &mut FunctionBuilder) {
+        let check_posix_memalign = self.builtin_functions.check_posix_memalign(builder.func);
+        let vmctx = self.vmctx_val(&mut builder.cursor());
+        let func_args = builder
+            .func
+            .dfg
+            .block_params(builder.func.layout.entry_block().unwrap());
+        let (outptr, _alignment, size) = if func_args.len() < 5 {
+            return;
+        } else {
+            // If a function named `malloc` has at least one argument, we assume the
+            // first argument is the requested allocation size.
+            (func_args[2], func_args[3], func_args[4])
+        };
+        builder.ins().call(check_posix_memalign, &[vmctx, outptr, size]);
+    }
+
+    #[cfg(feature = "wmemcheck")]
+    fn hook_aligned_alloc(&mut self, builder: &mut FunctionBuilder, retvals: &[ir::Value]) {
+        let check_malloc = self.builtin_functions.check_malloc(builder.func);
+        let vmctx = self.vmctx_val(&mut builder.cursor());
+        let func_args = builder
+            .func
+            .dfg
+            .block_params(builder.func.layout.entry_block().unwrap());
+        let (_alignment, size) = if func_args.len() < 4 {
+            return;
+        } else {
+            // If a function named `malloc` has at least one argument, we assume the
+            // first argument is the requested allocation size.
+            (func_args[2], func_args[3])
+        };
+        let retval = if retvals.len() < 1 {
+            return;
+        } else {
+            retvals[0]
+        };
+        builder.ins().call(check_malloc, &[vmctx, retval, size]);
+    }
+
     #[cfg(feature = "wmemcheck")]
     fn hook_free_exit(&mut self, builder: &mut FunctionBuilder) {
         let check_free = self.builtin_functions.check_free(builder.func);
@@ -927,17 +1033,10 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
     }
 
     #[cfg(feature = "wmemcheck")]
-    fn check_malloc_start(&mut self, builder: &mut FunctionBuilder) {
-        let malloc_start = self.builtin_functions.malloc_start(builder.func);
-        let vmctx = self.vmctx_val(&mut builder.cursor());
-        builder.ins().call(malloc_start, &[vmctx]);
-    }
-
-    #[cfg(feature = "wmemcheck")]
-    fn check_free_start(&mut self, builder: &mut FunctionBuilder) {
-        let free_start = self.builtin_functions.free_start(builder.func);
+    fn hook_memcheck_off(&mut self, builder: &mut FunctionBuilder) {
+        let memcheck_off = self.builtin_functions.memcheck_off(builder.func);
         let vmctx = self.vmctx_val(&mut builder.cursor());
-        builder.ins().call(free_start, &[vmctx]);
+        builder.ins().call(memcheck_off, &[vmctx]);
     }
 
     #[cfg(feature = "wmemcheck")]
@@ -3088,11 +3187,22 @@ impl<'module_environment> crate::translate::FuncEnvironment
 
         #[cfg(feature = "wmemcheck")]
         if self.wmemcheck {
-            let func_name = self.current_func_name(builder);
-            if func_name == Some("malloc") {
-                self.check_malloc_start(builder);
-            } else if func_name == Some("free") {
-                self.check_free_start(builder);
+            match self.current_func_name(builder) {
+                Some("__wrap_malloc") | Some("malloc") =>
+                    self.hook_memcheck_off(builder),
+                Some("__wrap_calloc") | Some("calloc") =>
+                    self.hook_memcheck_off(builder),
+                Some("__wrap_realloc") | Some("realloc") =>
+                    self.hook_memcheck_off(builder),
+                Some("__wrap_malloc_usable_size") | Some("malloc_usable_size") =>
+                    self.hook_memcheck_off(builder),
+                Some("__wrap_posix_memalign") | Some("posix_memalign") =>
+                    self.hook_memcheck_off(builder),
+                Some("__wrap_aligned_alloc") | Some("aligned_alloc") =>
+                    self.hook_memcheck_off(builder),
+                Some("__wrap_free") | Some("free") =>
+                    self.hook_memcheck_off(builder),
+                _ => ()
             }
         }
 
@@ -3141,11 +3251,23 @@ impl<'module_environment> crate::translate::FuncEnvironment
     #[cfg(feature = "wmemcheck")]
     fn handle_before_return(&mut self, retvals: &[ir::Value], builder: &mut FunctionBuilder) {
         if self.wmemcheck {
-            let func_name = self.current_func_name(builder);
-            if func_name == Some("malloc") {
-                self.hook_malloc_exit(builder, retvals);
-            } else if func_name == Some("free") {
-                self.hook_free_exit(builder);
+            let name = self.current_func_name(builder);
+            match name {
+                Some("__wrap_malloc") | Some("malloc") =>
+                    self.hook_malloc_exit(builder, retvals),
+                Some("__wrap_calloc") | Some("calloc") =>
+                    self.hook_calloc_exit(builder, retvals),
+                Some("__wrap_realloc") | Some("realloc") =>
+                    self.hook_realloc_exit(builder, retvals),
+                Some("__wrap_malloc_usable_size") | Some("malloc_usable_size") =>
+                    self.hook_malloc_usable_size_exit(builder, retvals),
+                Some("__wrap_posix_memalign") | Some("posix_memalign") =>
+                    self.hook_posix_memalign_exit(builder),
+                Some("__wrap_aligned_alloc") | Some("aligned_alloc") =>
+                    self.hook_aligned_alloc(builder, retvals),
+                Some("__wrap_free") | Some("free") =>
+                    self.hook_free_exit(builder),
+                _ => ()
             }
         }
     }

crates/environ/src/builtin.rs

@@ -44,6 +44,18 @@ macro_rules! foreach_builtin_function {
             // Invoked before malloc returns.
             #[cfg(feature = "wmemcheck")]
             check_malloc(vmctx: vmctx, addr: i32, len: i32) -> i32;
+            // Invoked before calloc returns.
+            #[cfg(feature = "wmemcheck")]
+            check_calloc(vmctx: vmctx, addr: i32, count: i32, size: i32) -> i32;
+            // Invoked before realloc returns.
+            #[cfg(feature = "wmemcheck")]
+            check_realloc(vmctx: vmctx, end_addr: i32, start_addr: i32, len: i32) -> i32;
+            // Invoked before malloc_usable_size returns.
+            #[cfg(feature = "wmemcheck")]
+            check_malloc_usable_size(vmctx: vmctx, len: i32, addr: i32) -> i32;
+            // Invoked before posix_memalign returns.
+            #[cfg(feature = "wmemcheck")]
+            check_posix_memalign(vmctx: vmctx, outptr: i32, size: i32) -> i32;
             // Invoked before the free returns.
             #[cfg(feature = "wmemcheck")]
             check_free(vmctx: vmctx, addr: i32) -> i32;
@@ -53,19 +65,17 @@ macro_rules! foreach_builtin_function {
             // Invoked before a store is executed.
             #[cfg(feature = "wmemcheck")]
             check_store(vmctx: vmctx, num_bytes: i32, addr: i32, offset: i32) -> i32;
-            // Invoked after malloc is called.
-            #[cfg(feature = "wmemcheck")]
-            malloc_start(vmctx: vmctx);
-            // Invoked after free is called.
-            #[cfg(feature = "wmemcheck")]
-            free_start(vmctx: vmctx);
             // Invoked when wasm stack pointer is updated.
             #[cfg(feature = "wmemcheck")]
             update_stack_pointer(vmctx: vmctx, value: i32);
             // Invoked before memory.grow is called.
             #[cfg(feature = "wmemcheck")]
             update_mem_size(vmctx: vmctx, num_bytes: i32);
 
+            // Invoked before stuff is called.
+            #[cfg(feature = "wmemcheck")]
+            memcheck_off(vmctx: vmctx);
+
             // Drop a non-stack GC reference (eg an overwritten table entry)
             // once it will no longer be used again. (Note: `val` is not of type
             // `reference` because it needn't appear in any stack maps, as it

crates/wasmtime/src/runtime/vm/libcalls.rs

@@ -1137,6 +1137,40 @@ unsafe fn check_malloc(
     Ok(0)
 }
 
+// Hook for validating calloc using wmemcheck_state.
+#[cfg(feature = "wmemcheck")]
+unsafe fn check_calloc(store: &mut dyn VMStore, instance: &mut Instance, addr: u32, count: u32, size: u32) -> Result<u32> {
+    check_malloc(store, instance, addr, count * size)
+}
+
+// Hook for validating realloc using wmemcheck_state.
+#[cfg(feature = "wmemcheck")]
+unsafe fn check_realloc(store: &mut dyn VMStore, instance: &mut Instance, end_addr: u32, start_addr: u32, len: u32) -> Result<u32> {
+    check_free(store, instance, start_addr)?;
+    check_malloc(store, instance, end_addr, len)
+}
+
+// Hook for validating malloc_usable_size using wmemcheck_state.
+#[cfg(feature = "wmemcheck")]
+unsafe fn check_malloc_usable_size(store: &mut dyn VMStore, instance: &mut Instance, len: u32, addr: u32) -> Result<u32> {
+    check_free(store, instance, addr)?;
+    check_malloc(store, instance, addr, len)
+}
+
+
+// Hook for validating posix_memalign using wmemcheck_state.
+#[cfg(feature = "wmemcheck")]
+unsafe fn check_posix_memalign(store: &mut dyn VMStore, instance: &mut Instance, outptr: u32, size: u32) -> Result<u32> {
+    for (_, entry) in instance.exports() {
+        if let wasmtime_environ::EntityIndex::Memory(mem_index) = entry {
+            let mem = instance.get_memory(*mem_index);
+            let out = *(mem.base.offset(outptr as isize) as *mut u32);
+            return check_malloc(store, instance, out, size)
+        }
+    }
+    todo!("Why is there no memory?")
+}
+
 // Hook for validating free using wmemcheck_state.
 #[cfg(feature = "wmemcheck")]
 unsafe fn check_free(_store: &mut dyn VMStore, instance: &mut Instance, addr: u32) -> Result<u32> {
@@ -1158,6 +1192,30 @@ unsafe fn check_free(_store: &mut dyn VMStore, instance: &mut Instance, addr: u3
     Ok(0)
 }
 
+#[cfg(feature = "wmemcheck")]
+fn log_allocation_previous_to(instance: &mut Instance, addr: usize) {
+    if let Some(wmemcheck_state) = &mut instance.wmemcheck_state {
+        if let Some((prev_malloc, prev_len)) = wmemcheck_state.malloc_previous_to(addr) {
+            println!("previous malloc'd range was {:#x}..{:#x}", prev_malloc, prev_malloc + prev_len);
+            for (_, entry) in instance.exports() {
+                if let wasmtime_environ::EntityIndex::Memory(mem_index) = entry {
+                    let mem = instance.get_memory(*mem_index);
+                    for i in 0..prev_len {
+                        if i > 0 && i % 40 == 0 {
+                            println!();
+                        }
+                        unsafe {
+                            print!("{:02x} ", *mem.base.offset((prev_malloc + i) as isize));
+                        }
+                    }
+                    println!();
+                    break
+                }
+            }
+        }
+    }
+}
+
 // Hook for validating load using wmemcheck_state.
 #[cfg(feature = "wmemcheck")]
 fn check_load(
@@ -1174,6 +1232,7 @@ fn check_load(
                 return Ok(0);
             }
             Err(InvalidRead { addr, len }) => {
+                log_allocation_previous_to(instance, addr);
                 bail!("Invalid load at addr {:#x} of size {}", addr, len);
             }
             Err(OutOfBounds { addr, len }) => {
@@ -1203,6 +1262,7 @@ fn check_store(
                 return Ok(0);
             }
             Err(InvalidWrite { addr, len }) => {
+                log_allocation_previous_to(instance, addr);
                 bail!("Invalid store at addr {:#x} of size {}", addr, len)
             }
             Err(OutOfBounds { addr, len }) => {
@@ -1216,17 +1276,8 @@ fn check_store(
     Ok(0)
 }
 
-// Hook for turning wmemcheck load/store validation off when entering a malloc function.
-#[cfg(feature = "wmemcheck")]
-fn malloc_start(_store: &mut dyn VMStore, instance: &mut Instance) {
-    if let Some(wmemcheck_state) = &mut instance.wmemcheck_state {
-        wmemcheck_state.memcheck_off();
-    }
-}
-
-// Hook for turning wmemcheck load/store validation off when entering a free function.
 #[cfg(feature = "wmemcheck")]
-fn free_start(_store: &mut dyn VMStore, instance: &mut Instance) {
+fn memcheck_off(_store: &mut dyn VMStore, instance: &mut Instance) {
     if let Some(wmemcheck_state) = &mut instance.wmemcheck_state {
         wmemcheck_state.memcheck_off();
     }

crates/wmemcheck/src/lib.rs

@@ -10,6 +10,21 @@ pub struct Wmemcheck {
     pub flag: bool,
 }
 
+impl Wmemcheck {
+    pub fn malloc_previous_to(&self, addr: usize) -> Option<(usize, usize)> {
+        let mut best: Option<(usize, usize)> = None;
+        for (base, len) in self.mallocs.iter() {
+            if let Some((prev_base, _)) = best {
+                if prev_base < *base && *base <= addr {
+                    best = Some((*base, *len));
+                }
+            } else {
+                best = Some((*base, *len));
+            }
+        }
+        best
+    }
+}
 /// Error types for memory checker.
 #[derive(Debug, PartialEq)]
 pub enum AccessError {
@@ -51,7 +66,10 @@ impl Wmemcheck {
     }
 
     /// Updates memory checker memory state metadata when malloc is called.
-    pub fn malloc(&mut self, addr: usize, len: usize) -> Result<(), AccessError> {
+    pub fn malloc(&mut self, addr: usize, start_len: usize) -> Result<(), AccessError> {
+        // round up to multiple of 4
+        let len = (start_len + 3) & !3;
+
         if !self.is_in_bounds_heap(addr, len) {
             return Err(AccessError::OutOfBounds {
                 addr: addr,
@@ -101,12 +119,12 @@ impl Wmemcheck {
                         len: len,
                     });
                 }
-                MemState::ValidToWrite => {
+                /* MemState::ValidToWrite => {
                     return Err(AccessError::InvalidRead {
                         addr: addr,
                         len: len,
                     });
-                }
+                } */
                 _ => {}
             }
         }
@@ -140,6 +158,9 @@ impl Wmemcheck {
 
     /// Updates memory checker memory state metadata when free is called.
     pub fn free(&mut self, addr: usize) -> Result<(), AccessError> {
+        if addr == 0 {
+            return Ok(());
+        }
         if !self.mallocs.contains_key(&addr) {
             return Err(AccessError::InvalidFree { addr: addr });
         }