Performance optimizations for address tracking

libresolve/src/lib.rs

@@ -2,6 +2,9 @@
 // LGPL-3; See LICENSE.txt in the repo root for details.
 
 #![feature(btree_cursors)]
+#![feature(macro_metavar_expr_concat)]
+#![feature(test)]
+extern crate test;
 
 mod remediate;
 mod shadowobjs;

libresolve/src/remediate.rs

@@ -4,7 +4,9 @@ use libc::{
     c_char, c_void, calloc, free, malloc, realloc, strdup, strlen, strndup, strnlen,
 };
 
-use crate::shadowobjs::{ALIVE_OBJ_LIST, AllocType, FREED_OBJ_LIST, ShadowObject, Vaddr};
+use crate::shadowobjs::{
+    ALIVE_OBJ_LIST, AllocType, FREED_OBJ_LIST, STACK_OBJ_LIST, ShadowObject, Vaddr,
+};
 
 use log::{error, info, trace, warn};
 
@@ -24,26 +26,55 @@ use log::{error, info, trace, warn};
 #[unsafe(no_mangle)]
 pub extern "C" fn resolve_stack_obj(ptr: *mut c_void, size: usize) -> () {
     let base = ptr as Vaddr;
-    {
-        let mut obj_list = ALIVE_OBJ_LIST.lock();
-        obj_list.add_shadow_object(AllocType::Stack, base, size);
-    }
+
+    STACK_OBJ_LIST.with_borrow_mut(|l| {
+        l.add_shadow_object(AllocType::Stack, base, size);
+    });
 
     info!("[STACK] Object allocated with size: {size}, address: 0x{base:x}");
 }
 
+/*
 #[unsafe(no_mangle)]
+pub extern "C" fn resolve_invalidate_stack(base: *mut c_void) {
 pub extern "C" fn resolve_invalidate_stack(base: *mut c_void) {
     let base = base as Vaddr;
 
-    {
-        let mut obj_list = ALIVE_OBJ_LIST.lock();
-        obj_list.invalidate_at(base);
-    }
+    STACK_OBJ_LIST.with_borrow_mut(|l| {
+        l.invalidate_at(base);
+    });
+
+    info!("[STACK] Free range 0x{base:x}");
+}
+*/
+
+macro_rules! invalidate_stacks {
+    ($name:ident; $($ns:literal),+) => {
+
+#[unsafe(no_mangle)]
+pub extern "C" fn $name($(${concat(base_, $ns)}: *mut c_void, )+) {
+
+    STACK_OBJ_LIST.with_borrow_mut(|l| {
+        $(
+        l.invalidate_at(${concat(base_, $ns)} as Vaddr);
+        info!("[STACK] Free address 0x{:x}", ${concat(base_, $ns)} as Vaddr);
+        )+
+    });
+
+}
 
-    info!("[STACK] Free addr 0x{base:x}");
+    }
 }
 
+// the x64 ABI allows up to 6 arguments to be passed via register,
+// so provide that many functions as we cannot define a variadic
+invalidate_stacks!(resolve_invalidate_stack; 1);
+invalidate_stacks!(resolve_invalidate_stack_2; 1, 2);
+invalidate_stacks!(resolve_invalidate_stack_3; 1, 2, 3);
+invalidate_stacks!(resolve_invalidate_stack_4; 1, 2, 3, 4);
+invalidate_stacks!(resolve_invalidate_stack_5; 1, 2, 3, 4, 5);
+invalidate_stacks!(resolve_invalidate_stack_6; 1, 2, 3, 4, 5, 6);
+
 /**
  * @brief - Allocator logging interface for malloc
  * @input - size of the allocation in bytes
@@ -81,10 +112,37 @@ pub extern "C" fn resolve_malloc(size: usize) -> *mut c_void {
  */
 
 #[unsafe(no_mangle)]
-pub extern "C" fn resolve_gep(ptr: *mut c_void, derived: *mut c_void) -> *mut c_void {
+pub extern "C" fn resolve_gep(ptr: *mut c_void, derived: *mut c_void, max_access: usize) -> *mut c_void {
     let base = ptr as Vaddr;
     let derived = derived as Vaddr;
 
+
+    let contains_or_err = |obj: &ShadowObject| {
+        // If shadow object exists then check if the access is within bounds
+        if obj.contains(ShadowObject::limit(derived, max_access)) {
+            trace!(
+                "[GEP] ptr {max_access}x{derived:x} valid for base 0x{base:x}, obj: {}@0x{:x}",
+                obj.size(),
+                obj.base
+            );
+            return derived as *mut c_void;
+        }
+
+        error!(
+            "[GEP] ptr {max_access}x{derived:x} not valid for base 0x{base:x}, obj: {}@0x{:x}",
+            obj.size(),
+            obj.base
+        );
+
+        // Pointer known-invalid, return sentinel to indicate failure
+        0 as *mut c_void
+    };
+
+    // Check the local stack list first since it is cheaper.
+    if let Some(sobj) = STACK_OBJ_LIST.with_borrow(|l| l.search_intersection(base).cloned()) {
+        return contains_or_err(&sobj);
+    };
+
     let sobj_table = ALIVE_OBJ_LIST.lock();
 
     // Look up the shadow object corresponding to this access.
@@ -101,24 +159,7 @@ pub extern "C" fn resolve_gep(ptr: *mut c_void, derived: *mut c_void) -> *mut c_
         return derived as *mut c_void;
     };
 
-    // If shadow object exists then check if the access is within bounds
-    if sobj.contains(derived as Vaddr) {
-        info!(
-            "[GEP] ptr 0x{derived:x} valid for base 0x{base:x}, obj: {}@0x{:x}",
-            sobj.size(),
-            sobj.base
-        );
-        return derived as *mut c_void;
-    }
-
-    error!(
-        "[GEP] ptr 0x{derived:x} not valid for base 0x{base:x}, obj: {}@0x{:x}",
-        sobj.size(),
-        sobj.base
-    );
-
-    // Return 1-past limit of allocation @ ptr
-    sobj.past_limit() as *mut c_void
+    contains_or_err(sobj)
 }
 
 
@@ -317,27 +358,39 @@ pub extern "C" fn resolve_strndup(ptr: *mut c_char, size: usize) -> *mut c_char
 pub extern "C" fn resolve_check_bounds(base_ptr: *mut c_void, size: usize) -> bool {
     let base = base_ptr as Vaddr;
 
-    let sobj_table = ALIVE_OBJ_LIST.lock();
+    // Nullptr clearly are not valid, and may be returned by resolve_gep if it is an invalid call.
+    if base == 0 {
+        return false;
+    }
 
-    // Look up the shadow object corresponding to this access
-    if let Some(sobj) = sobj_table.search_intersection(base) {
+    let contains_or_err = |obj: &ShadowObject| {
         // If shadow object exists then check if the access is within bounds
-        if sobj.contains(ShadowObject::limit(base, size)) {
+        if obj.contains(ShadowObject::limit(base, size)) {
             // Access in Bounds
             trace!(
                 "[BOUNDS] Access allowed {size}@0x{base:x} for allocation {}@0x{:x}",
-                sobj.size(),
-                sobj.base
+                obj.size(),
+                obj.base
             );
             return true;
         } else {
             error!(
                 "[BOUNDS] OOB access at {size}@0x{base:x} too big for allocation {}@0x{:x}",
-                sobj.size(),
-                sobj.base
+                obj.size(),
+                obj.base
             );
             return false;
         }
+    };
+
+    if let Some(sobj) = STACK_OBJ_LIST.with_borrow(|l| l.search_intersection(base).cloned()) {
+        return contains_or_err(&sobj);
+    }
+
+    // Look up the shadow object corresponding to this access
+    let sobj_table = ALIVE_OBJ_LIST.lock();
+    if let Some(sobj) = sobj_table.search_intersection(base) {
+        return contains_or_err(sobj);
     }
 
     // Check if this is an invalid pointer for one of the known shadow objects
@@ -350,6 +403,8 @@ pub extern "C" fn resolve_check_bounds(base_ptr: *mut c_void, size: usize) -> bo
         return false;
     }
 
+    warn!("[BOUNDS] unknown pointer 0x:{base:x}");
+
     // Not a tracked pointer, assume good to avoid trapping on otherwise valid pointers
     // TODO: add a strict mode to reject here / add extra tracking.
     true
@@ -359,13 +414,21 @@ pub extern "C" fn resolve_check_bounds(base_ptr: *mut c_void, size: usize) -> bo
 pub extern "C" fn resolve_obj_type(base_ptr: *mut c_void) -> AllocType {
     let base = base_ptr as Vaddr;
 
-    let find_in = |table: &crate::MutexWrap<crate::shadowobjs::ShadowObjectTable>| {
-        let t = table.lock();
-        t.search_intersection(base).map(|o| o.alloc_type)
+    let find_in = |table: &crate::shadowobjs::ShadowObjectTable| {
+        table.search_intersection(base).map(|o| o.alloc_type)
     };
 
     // Why does this search freed before alive?
-    let alloc_type = find_in(&FREED_OBJ_LIST).or_else(|| find_in(&ALIVE_OBJ_LIST));
+    let alloc_type = STACK_OBJ_LIST
+        .with_borrow(|l| find_in(l))
+        .or_else(|| {
+            let l = FREED_OBJ_LIST.lock();
+            find_in(&l)
+        })
+        .or_else(|| {
+            let l = ALIVE_OBJ_LIST.lock();
+            find_in(&l)
+        });
 
     alloc_type.unwrap_or(AllocType::Unknown)
 }
@@ -378,7 +441,7 @@ pub extern "C" fn resolve_obj_type(base_ptr: *mut c_void) -> AllocType {
  */
 #[unsafe(no_mangle)]
 pub extern "C" fn resolve_report_sanitize_mem_inst_triggered(ptr: *mut c_void) {
-    info!(
+    error!(
         "[SANITIZE] Applying sanitizer to address 0x{:x}",
         ptr as Vaddr
     );
@@ -396,6 +459,7 @@ pub extern "C" fn resolve_report_sanitizer_triggered() -> () {
 mod tests {
     use super::*;
     use crate::{resolve_init, shadowobjs::AllocType};
+    use test::Bencher;
 
     #[test]
     fn test_malloc_free() {
@@ -479,4 +543,24 @@ mod tests {
             resolve_free(p);
         }
     }
+
+
+    #[bench]
+    fn bench_resolve_stack(b: &mut Bencher) {
+        resolve_init();
+
+        let addrs: Vec<_> = (0x7FFF_0000_0000_0000..0x7FFF_0000_0001_0000)
+            .map(|a: usize| a as *mut c_void)
+            .collect();
+
+        b.iter(|| {
+            addrs.iter().for_each(|a| resolve_stack_obj(*a, 1));
+
+            addrs.iter().for_each(|&a| {
+                let _ = resolve_gep(a, a, 1);
+            });
+
+            addrs.iter().for_each(|a| resolve_invalidate_stack(*a));
+        });
+    }
 }

libresolve/src/shadowobjs.rs

@@ -2,6 +2,7 @@
 // LGPL-3; See LICENSE.txt in the repo root for details.
 
 use crate::MutexWrap;
+use std::cell::RefCell;
 use std::collections::BTreeMap;
 use std::ops::RangeInclusive;
 use std::ops::Bound::Included;
@@ -131,6 +132,9 @@ impl ShadowObjectTable {
 }
 
 // static object lists to store all objects
+thread_local! {
+    pub static STACK_OBJ_LIST: RefCell<ShadowObjectTable> = RefCell::new(ShadowObjectTable::new());
+}
 pub static ALIVE_OBJ_LIST: MutexWrap<ShadowObjectTable> = MutexWrap::new(ShadowObjectTable::new());
 pub static FREED_OBJ_LIST: MutexWrap<ShadowObjectTable> = MutexWrap::new(ShadowObjectTable::new());
 

llvm-plugin/Makefile

@@ -1,10 +1,11 @@
-all: build test
 
-.PHONY: clangformat test
+.PHONY: all build clangformat test
+
+all: build test
 
 build: src
 	mkdir -p build
-	cd build && cmake ../src && make -j4
+	cmake -S ./src -B ./build  && cd build && make -j4
 
 clangformat: src
 	cd build && make clangformat

llvm-plugin/src/CMakeLists.txt

@@ -13,7 +13,7 @@ if(NOT CMAKE_BUILD_TYPE)
     set(CMAKE_BUILD_TYPE RelWithDebInfo)
 endif()
 
-# set(CMAKE_CXX_STANDARD 17 STRING "")
+set(CMAKE_CXX_STANDARD 20 STRING "")
 
 # LLVM is normally built without RTTI. Be consistent with that.
 if(NOT LLVM_ENABLE_RTTI)