wmemcheck support for additional allocation functions and granular memory

crates/cranelift/src/func_environ.rs

@@ -563,44 +563,26 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
     }
 
     #[cfg(feature = "wmemcheck")]
-    fn hook_malloc_exit(&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 len = if func_args.len() < 3 {
-            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]
-        };
-        let retval = if retvals.len() < 1 {
-            return;
-        } else {
-            retvals[0]
-        };
-        builder.ins().call(check_malloc, &[vmctx, retval, len]);
-    }
-
-    #[cfg(feature = "wmemcheck")]
-    fn hook_free_exit(&mut self, builder: &mut FunctionBuilder) {
-        let check_free = self.builtin_functions.check_free(builder.func);
+    fn hook_memcheck_exit(
+        &mut self,
+        builder: &mut FunctionBuilder,
+        builtin_memcheck_check: ir::FuncRef,
+        expected_arg_count: usize,
+        retvals: &[ir::Value],
+    ) {
         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 {
+        if func_args.len() < expected_arg_count + 2 {
+            // Assume the first n arguments are the expected ones
             return;
-        } else {
-            // If a function named `free` has at least one argument, we assume the
-            // first argument is a pointer to memory.
-            func_args[2]
-        };
-        builder.ins().call(check_free, &[vmctx, ptr]);
+        }
+        let mut args = vec![vmctx];
+        args.extend_from_slice(&retvals);
+        args.extend_from_slice(&func_args[2..2 + expected_arg_count]);
+        builder.ins().call(builtin_memcheck_check, &args);
     }
 
     fn epoch_ptr(&mut self, builder: &mut FunctionBuilder<'_>) -> ir::Value {
@@ -927,17 +909,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 check_allocator_start(&mut self, builder: &mut FunctionBuilder) {
+        let allocator_start = self.builtin_functions.allocator_start(builder.func);
         let vmctx = self.vmctx_val(&mut builder.cursor());
-        builder.ins().call(free_start, &[vmctx]);
+        builder.ins().call(allocator_start, &[vmctx]);
     }
 
     #[cfg(feature = "wmemcheck")]
@@ -3088,11 +3063,15 @@ 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("malloc")
+                | Some("free")
+                | Some("calloc")
+                | Some("realloc")
+                | Some("posix_memalign")
+                | Some("aligned_alloc")
+                | Some("malloc_usable_size") => self.check_allocator_start(builder),
+                _ => (),
             }
         }
 
@@ -3141,11 +3120,41 @@ 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("malloc") => {
+                    let check_malloc = self.builtin_functions.check_malloc(builder.func);
+                    self.hook_memcheck_exit(builder, check_malloc, 1, retvals)
+                }
+                Some("free") => {
+                    let check_free = self.builtin_functions.check_free(builder.func);
+                    self.hook_memcheck_exit(builder, check_free, 1, &[])
+                }
+                Some("calloc") => {
+                    let check_calloc = self.builtin_functions.check_calloc(builder.func);
+                    self.hook_memcheck_exit(builder, check_calloc, 2, retvals)
+                }
+                Some("realloc") => {
+                    let check_realloc = self.builtin_functions.check_realloc(builder.func);
+                    self.hook_memcheck_exit(builder, check_realloc, 2, retvals)
+                }
+                Some("posix_memalign") => {
+                    let check_posix_memalign =
+                        self.builtin_functions.check_posix_memalign(builder.func);
+                    self.hook_memcheck_exit(builder, check_posix_memalign, 3, retvals)
+                }
+                Some("aligned_alloc") => {
+                    let check_aligned_alloc =
+                        self.builtin_functions.check_aligned_alloc(builder.func);
+                    self.hook_memcheck_exit(builder, check_aligned_alloc, 2, retvals)
+                }
+                Some("malloc_usable_size") => {
+                    let check_malloc_usable_size = self
+                        .builtin_functions
+                        .check_malloc_usable_size(builder.func);
+                    self.hook_memcheck_exit(builder, check_malloc_usable_size, 1, retvals)
+                }
+                _ => (),
             }
         }
     }

crates/environ/src/builtin.rs

@@ -41,24 +41,36 @@ macro_rules! foreach_builtin_function {
             out_of_gas(vmctx: vmctx);
             // Invoked when we reach a new epoch.
             new_epoch(vmctx: vmctx) -> i64;
+            // Invoked before memory allocation functions are called.
+            #[cfg(feature = "wmemcheck")]
+            allocator_start(vmctx: vmctx);
             // Invoked before malloc returns.
             #[cfg(feature = "wmemcheck")]
             check_malloc(vmctx: vmctx, addr: i32, len: i32) -> i32;
             // Invoked before the free returns.
             #[cfg(feature = "wmemcheck")]
             check_free(vmctx: vmctx, addr: 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 posix_memalign returns.
+            #[cfg(feature = "wmemcheck")]
+            check_posix_memalign(vmctx: vmctx, result: i32, outptr: i32, alignment: i32, size: i32) -> i32;
+            // Invoked before aligned_alloc returns.
+            #[cfg(feature = "wmemcheck")]
+            check_aligned_alloc(vmctx: vmctx, outptr: i32, alignment: i32, size: i32) -> i32;
+            // Invoked before malloc_usable_size returns.
+            #[cfg(feature = "wmemcheck")]
+            check_malloc_usable_size(vmctx: vmctx, len: i32, addr: i32) -> i32;
             // Invoked before a load is executed.
             #[cfg(feature = "wmemcheck")]
             check_load(vmctx: vmctx, num_bytes: i32, addr: i32, offset: i32) -> i32;
             // 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);

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

@@ -336,7 +336,7 @@ impl Instance {
                             .unwrap_or(0)
                             * 64
                             * 1024;
-                        Some(Wmemcheck::new(size as usize))
+                        Some(Wmemcheck::new(size as usize, 4, false))
                     } else {
                         None
                     }

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

@@ -63,8 +63,11 @@ use core::time::Duration;
 use wasmtime_environ::Unsigned;
 use wasmtime_environ::{DataIndex, ElemIndex, FuncIndex, MemoryIndex, TableIndex, Trap};
 #[cfg(feature = "wmemcheck")]
-use wasmtime_wmemcheck::AccessError::{
-    DoubleMalloc, InvalidFree, InvalidRead, InvalidWrite, OutOfBounds,
+use wasmtime_wmemcheck::{
+    AccessError,
+    AccessError::{
+        DoubleMalloc, InvalidFree, InvalidRead, InvalidRealloc, InvalidWrite, OutOfBounds,
+    },
 };
 
 /// Raw functions which are actually called from compiled code.
@@ -1108,6 +1111,28 @@ fn new_epoch(store: &mut dyn VMStore, _instance: &mut Instance) -> Result<u64> {
     store.new_epoch()
 }
 
+#[cfg(feature = "wmemcheck")]
+fn check_memcheck_result(result: Result<(), AccessError>) -> Result<u32> {
+    match result {
+        Ok(()) => Ok(0),
+        Err(DoubleMalloc { addr, len }) => {
+            bail!("Double malloc at addr {:#x} of size {}", addr, len)
+        }
+        Err(OutOfBounds { addr, len }) => {
+            bail!("Malloc out of bounds at addr {:#x} of size {}", addr, len)
+        }
+        Err(InvalidFree { addr }) => {
+            bail!("Invalid free at addr {:#x}", addr)
+        }
+        Err(InvalidRealloc { addr }) => {
+            bail!("Invalid realloc at addr {:#x}", addr)
+        }
+        _ => {
+            panic!("unreachable")
+        }
+    }
+}
+
 // Hook for validating malloc using wmemcheck_state.
 #[cfg(feature = "wmemcheck")]
 unsafe fn check_malloc(
@@ -1117,22 +1142,9 @@ unsafe fn check_malloc(
     len: u32,
 ) -> Result<u32> {
     if let Some(wmemcheck_state) = &mut instance.wmemcheck_state {
-        let result = wmemcheck_state.malloc(addr as usize, len as usize);
+        let result = wmemcheck_state.allocate(addr as usize, len as usize, false);
         wmemcheck_state.memcheck_on();
-        match result {
-            Ok(()) => {
-                return Ok(0);
-            }
-            Err(DoubleMalloc { addr, len }) => {
-                bail!("Double malloc at addr {:#x} of size {}", addr, len)
-            }
-            Err(OutOfBounds { addr, len }) => {
-                bail!("Malloc out of bounds at addr {:#x} of size {}", addr, len);
-            }
-            _ => {
-                panic!("unreachable")
-            }
-        }
+        return check_memcheck_result(result);
     }
     Ok(0)
 }
@@ -1143,21 +1155,95 @@ unsafe fn check_free(_store: &mut dyn VMStore, instance: &mut Instance, addr: u3
     if let Some(wmemcheck_state) = &mut instance.wmemcheck_state {
         let result = wmemcheck_state.free(addr as usize);
         wmemcheck_state.memcheck_on();
-        match result {
-            Ok(()) => {
-                return Ok(0);
-            }
-            Err(InvalidFree { addr }) => {
-                bail!("Invalid free at addr {:#x}", addr)
-            }
-            _ => {
-                panic!("unreachable")
+        return check_memcheck_result(result);
+    }
+    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> {
+    if let Some(wmemcheck_state) = &mut instance.wmemcheck_state {
+        let result = wmemcheck_state.allocate(addr as usize, (count * size) as usize, true);
+        wmemcheck_state.memcheck_on();
+        return check_memcheck_result(result);
+    }
+    Ok(0)
+}
+
+// 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> {
+    if let Some(wmemcheck_state) = &mut instance.wmemcheck_state {
+        let result = wmemcheck_state.realloc(end_addr as usize, start_addr as usize, len as usize);
+        wmemcheck_state.memcheck_on();
+        return check_memcheck_result(result);
+    }
+    Ok(0)
+}
+
+// Hook for validating posix_memalign using wmemcheck_state.
+#[cfg(feature = "wmemcheck")]
+unsafe fn check_posix_memalign(
+    store: &mut dyn VMStore,
+    instance: &mut Instance,
+    error: u32,
+    outptr: u32,
+    _alignment: u32,
+    size: u32,
+) -> Result<u32> {
+    if let Some(_) = &mut instance.wmemcheck_state {
+        if error != 0 {
+            return Ok(0);
+        }
+        for (_, entry) in instance.exports() {
+            if let wasmtime_environ::EntityIndex::Memory(mem_index) = entry {
+                let mem = instance.get_memory(*mem_index);
+                let out_ptr = *(mem.base.offset(outptr as isize) as *mut u32);
+                return check_malloc(store, instance, out_ptr, size);
             }
         }
+        todo!("Why is there no memory?")
     }
     Ok(0)
 }
 
+// Hook for validating aligned_alloc using wmemcheck_state.
+#[cfg(feature = "wmemcheck")]
+unsafe fn check_aligned_alloc(
+    store: &mut dyn VMStore,
+    instance: &mut Instance,
+    addr: u32,
+    _alignment: u32,
+    size: u32,
+) -> Result<u32> {
+    check_malloc(store, instance, addr, size)
+}
+
+// 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> {
+    // Since the wasm program has checked that the entire allocation is usable, mark it as allocated, similar to realloc
+    check_realloc(store, instance, addr, addr, len)
+}
+
 // Hook for validating load using wmemcheck_state.
 #[cfg(feature = "wmemcheck")]
 fn check_load(
@@ -1216,17 +1302,9 @@ 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.
+// Hook for turning wmemcheck load/store validation off when entering an allocator function.
 #[cfg(feature = "wmemcheck")]
-fn free_start(_store: &mut dyn VMStore, instance: &mut Instance) {
+fn allocator_start(_store: &mut dyn VMStore, instance: &mut Instance) {
     if let Some(wmemcheck_state) = &mut instance.wmemcheck_state {
         wmemcheck_state.memcheck_off();
     }