[WIP] Mimalloc Integration

resolve-cveassert/libresolve/CMakeLists.txt

@@ -3,6 +3,31 @@
 
 include(GNUInstallDirs)
 include(ExternalProject)
+include(FetchContent)
+
+#### Fetching and building mimalloc project 
+FetchContent_Declare(
+    mimalloc
+    GIT_REPOSITORY https://github.com/microsoft/mimalloc.git
+    GIT_TAG v3.3.2
+)
+
+## Force static build 
+set(MI_BUILD_SHARED OFF CACHE BOOL "" FORCE)
+set(MI_BUILD_STATIC ON CACHE BOOL "" FORCE)
+
+FetchContent_MakeAvailable(mimalloc)
+
+# Add the shim source to the mimalloc proj
+target_sources(mimalloc-static PRIVATE
+    ${CMAKE_CURRENT_SOURCE_DIR}/mimalloc_shadow.c
+)
+
+## mimalloc proj must be PIC to ensure compatibility with libresolve
+set_target_properties(mimalloc-static PROPERTIES
+    POSITION_INDEPENDENT_CODE ON
+    OUTPUT_NAME mimalloc
+)
 
 # Map CMake build type to Cargo flags
 if(CMAKE_BUILD_TYPE STREQUAL "Debug" OR CMAKE_BUILD_TYPE STREQUAL "")
@@ -24,17 +49,24 @@ file(GLOB_RECURSE RUST_SOURCES CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/sr
 add_custom_command(
     OUTPUT ${RUST_LIB}
     COMMAND 
-        ${CMAKE_COMMAND} -E env CARGO_TARGET_DIR=${RUST_OUT_DIR} cargo build ${CARGO_FLAGS} 
+        ${CMAKE_COMMAND} -E env
+        CARGO_TARGET_DIR=${RUST_OUT_DIR}
+        MIMALLOC_LIB_DIR=$<TARGET_FILE_DIR:mimalloc-static>
+        cargo build ${CARGO_FLAGS} 
     WORKING_DIRECTORY ${RUST_CRATE_DIR}
     COMMENT "Building libresolve.so" VERBATIM
-    DEPENDS ${RUST_SOURCES} ${CMAKE_CURRENT_SOURCE_DIR}/Cargo.lock ${CMAKE_CURRENT_SOURCE_DIR}/Cargo.toml ${CMAKE_CURRENT_SOURCE_DIR}/rust-toolchain.toml
+    DEPENDS mimalloc-static ${RUST_SOURCES} ${CMAKE_CURRENT_SOURCE_DIR}/Cargo.lock ${CMAKE_CURRENT_SOURCE_DIR}/Cargo.toml ${CMAKE_CURRENT_SOURCE_DIR}/rust-toolchain.toml
 )
 
 add_custom_target(test-libresolve
-    COMMAND cargo test
+    COMMAND 
+        ${CMAKE_COMMAND} -E env
+        MIMALLOC_LIB_DIR=${MIMALLOC_LIB_DIR} # Note: Pass the static mimalloc loc so that build.rs does not panic
+        cargo test
     WORKING_DIRECTORY ${RUST_CRATE_DIR}
     COMMENT "Running regression tests for libresolve"
 )
 
 add_custom_target(libresolve ALL DEPENDS ${RUST_LIB})
+add_dependencies(libresolve mimalloc-static)
 install(FILES ${RUST_LIB} DESTINATION ${CMAKE_INSTALL_LIBDIR})

resolve-cveassert/libresolve/Cargo.toml

@@ -1,12 +1,14 @@
 [package]
 name = "resolve"
+build = "src/build.rs"
 version = "0.1.0"
 edition = "2024"
 
 [lib]
 crate-type = ["cdylib"]
 
 [dependencies]
+cc = "1.2.63"
 env_logger = "0.11.8"
 libc = "0.2.174"
 log = "0.4.29"
@@ -18,3 +20,5 @@ debug = true
 [profile.dev]
 panic = "abort"
 
+[toolchain]
+channel = "nightly"

resolve-cveassert/libresolve/mimalloc_shadow.c

@@ -0,0 +1,82 @@
+#include "mimalloc.h"
+#include "mimalloc/internal.h"
+#include <stdarg.h>
+#include <stdint.h>
+#include <stdio.h>
+
+extern void* __resolve_malloc(size_t);
+extern void __resolve_free(void*);
+
+
+typedef struct {
+  void *base;
+  void *limit;
+  size_t block_size;
+  size_t block_index;
+} bounds_info_t;
+
+
+bounds_info_t mi_resolve_ptr(void* p) {
+  // Can return null if ptr is not owned by mimalloc
+  mi_page_t *page = _mi_ptr_page(p);
+
+  const size_t block_size = page->block_size;
+
+  uintptr_t page_start = (uintptr_t)page->page_start;
+  uintptr_t ptr = (uintptr_t)p;
+
+  size_t block_index = (ptr - page_start) / block_size; 
+  uintptr_t base_addr = page_start + block_index * block_size;
+
+  bounds_info_t bounds;
+  bounds.base = (void*)base_addr;
+  bounds.limit = (void*)(base_addr + block_size);
+  bounds.block_size = block_size;
+  bounds.block_index = block_index;
+  return bounds;
+}
+
+bool mi_is_heap_owned(const void* p) {
+  return _mi_ptr_page(p) != NULL;
+}
+
+int __vasprintf(char **strp, const char *fmt, va_list ap)
+{
+  va_list ap_copy;
+  va_copy(ap_copy, ap);
+
+  int len = vsnprintf(NULL, 0, fmt, ap_copy);
+  va_end(ap_copy);
+
+  if (len < 0) { 
+    // to match glibc behavior
+    *strp = NULL;
+    return -1; 
+  }
+
+  char *buf = __resolve_malloc((size_t)len + 1);
+  if (!buf) { return -1; }
+
+  va_copy(ap_copy, ap);
+
+  int written = vsnprintf(buf, (size_t)len + 1, fmt, ap_copy);
+
+  va_end(ap_copy);
+
+  if (written < 0) {
+    __resolve_free(buf);
+    return -1;
+  }
+
+  *strp = buf;
+  return written;
+}
+
+void *resolve_return_address(unsigned level) {
+  switch(level) {
+    case 0: return __builtin_return_address(0);
+    case 1: return __builtin_return_address(1);
+    case 2: return __builtin_return_address(2);
+    default: return NULL;
+  }
+}

resolve-cveassert/libresolve/src/build.rs

@@ -0,0 +1,6 @@
+fn main() {
+    let dir = std::env::var("MIMALLOC_LIB_DIR").unwrap();
+    println!("cargo::warning=LINKING_MIMALLOC_ARCHIVE");
+    println!("cargo::rustc-link-search=native={}", dir);
+    println!("cargo::rustc-link-lib=static=mimalloc");
+}

resolve-cveassert/libresolve/src/lib.rs

@@ -2,7 +2,7 @@
 // LGPL-3; See LICENSE.txt in the repo root for details.
 
 #![feature(btree_cursors)]
-
+#![feature(c_variadic)]
 mod remediate;
 mod shadowobjs;
 mod trace;
@@ -11,8 +11,8 @@ use libc::{Dl_info, atexit, c_void, dladdr, dlsym};
 use std::ffi::{CStr, OsString};
 use std::fmt::Display;
 use std::fs::{self, File};
-use std::path::PathBuf;
 use std::io::{self, Seek, Write};
+use std::path::PathBuf;
 use std::sync::{LazyLock, Mutex};
 use std::{env, process};
 
@@ -34,23 +34,23 @@ impl<T> MutexWrap<T> {
 }
 
 fn idify_file_path(path: &mut PathBuf, id: impl Display) {
-    let file_name = path.file_name()
+    let file_name = path
+        .file_name()
         .expect("Path could not be found in file system.")
         .to_owned();
 
     let mut updated_file_name = OsString::new();
 
     updated_file_name.push(file_name);
     updated_file_name.push("-");
-    updated_file_name.push(id.to_string()); 
+    updated_file_name.push(id.to_string());
 
     path.set_file_name(updated_file_name);
 }
 
 /// File for "resolve_dlsym.json"
 pub static DLSYM_LOG_FILE: LazyLock<MutexWrap<File>> = LazyLock::new(|| {
-    let log_dir = env::var("RESOLVE_DLSYM_LOG_DIR")
-        .unwrap_or_else(|_| ".".to_string());
+    let log_dir = env::var("RESOLVE_DLSYM_LOG_DIR").unwrap_or_else(|_| ".".to_string());
 
     let mut path = PathBuf::from(log_dir);
 
@@ -84,27 +84,26 @@ pub extern "C" fn resolve_init() {
     if cfg!(test) {
         builder.is_test(true);
     } else {
-        let file = open_resolve_log_file().unwrap_or_else(|err| { 
+        let file = open_resolve_log_file().unwrap_or_else(|err| {
             eprintln!("Libresolve log file could not be created.");
             eprintln!("Error: {err:?}");
             process::exit(12);
         });
-        
+
         builder.target(env_logger::Target::Pipe(Box::new(file)));
     }
 
     let _ = builder.try_init();
 }
 
 fn open_resolve_log_file() -> Result<File, io::Error> {
-    let log_dir = env::var("RESOLVE_RUNTIME_LOG_DIR")
-        .unwrap_or_else(|_| ".".to_string());
+    let log_dir = env::var("RESOLVE_RUNTIME_LOG_DIR").unwrap_or_else(|_| ".".to_string());
 
     let mut path = PathBuf::from(log_dir);
 
     // Ensure the parent directories exist
     fs::create_dir_all(&path)?;
-    
+
     // Append the file name
     path.push("resolve_log.out");