From 7ee9d213da48c43daf587f0c8c6a5ae4fc2581ce Mon Sep 17 00:00:00 2001 From: Matthew Parkinson Date: Mon, 23 Sep 2024 12:18:19 +0100 Subject: [PATCH] Add a replay allocator for ABA Asan delays reallocation to increase the likelyhood of catching a UAF. This adds a simple free list allocator that tries to reuse allocations quickly to increase the chance of an ABA problem being detected. --- .github/workflows/cmake.yml | 7 ++ src/rt/CMakeLists.txt | 4 + src/rt/debug/harness.h | 1 + src/rt/ds/heap.h | 172 +++++++++++++++++++++++++++ test/func/replayalloc/replayalloc.cc | 64 ++++++++++ 5 files changed, 248 insertions(+) create mode 100644 test/func/replayalloc/replayalloc.cc diff --git a/.github/workflows/cmake.yml b/.github/workflows/cmake.yml index 52387413..4a5b68d8 100644 --- a/.github/workflows/cmake.yml +++ b/.github/workflows/cmake.yml @@ -34,6 +34,13 @@ jobs: cmake-flags: "-G Ninja -DSANITIZER=address,undefined -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_C_COMPILER=clang" dependencies: "sudo apt install ninja-build" + - variant: "Sanitizer with replay" + platform: "ubuntu-latest" + build-type: "Release" + cmake-flags: "-G Ninja -DUSE_REPLAY_ALLOCATOR=On -DSANITIZER=address,undefined -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_C_COMPILER=clang" + dependencies: "sudo apt install ninja-build" + + - platform: "windows-2019" build-type: "Debug" # Exclude perf tests in Debug on Windows as they take too long. diff --git a/src/rt/CMakeLists.txt b/src/rt/CMakeLists.txt index a7845d25..013f198d 100644 --- a/src/rt/CMakeLists.txt +++ b/src/rt/CMakeLists.txt @@ -37,6 +37,10 @@ if(USE_SYSTEMATIC_TESTING) target_compile_definitions(verona_rt INTERFACE USE_SYSTEMATIC_TESTING) endif() +if(USE_REPLAY_ALLOCATOR) + target_compile_definitions(verona_rt INTERFACE USE_REPLAY_ALLOCATOR) +endif() + if(USE_CRASH_LOGGING) target_compile_definitions(verona_rt INTERFACE USE_FLIGHT_RECORDER) endif() diff --git a/src/rt/debug/harness.h b/src/rt/debug/harness.h index 39d9f0ea..3d177776 100644 --- a/src/rt/debug/harness.h +++ b/src/rt/debug/harness.h @@ -146,6 +146,7 @@ class SystematicTestHarness Scheduler& sched = Scheduler::get(); #ifdef USE_SYSTEMATIC_TESTING Systematic::set_seed(seed); + heap::set_seed(seed); if (seed % 2 == 1) { sched.set_fair(true); diff --git a/src/rt/ds/heap.h b/src/rt/ds/heap.h index f7de99ca..a0b87fc8 100644 --- a/src/rt/ds/heap.h +++ b/src/rt/ds/heap.h @@ -3,10 +3,176 @@ #pragma once +#include #include +#if defined(__SANITIZE_ADDRESS__) +# define HAS_ASAN +#elif defined(__has_feature) +# if __has_feature(address_sanitizer) +# define HAS_ASAN +# endif +#endif + +#if defined(HAS_ASAN) +# include +// Asan does not appear to support a __asan_update_deallocation_context +# define VERONA_TRACK_FREE(ptr, size) __asan_poison_memory_region(ptr, size); +# define VERONA_TRACK_ALLOC(ptr, size) __asan_unpoison_memory_region(ptr, size); __asan_update_allocation_context(ptr); +# define VERONA_NO_SANITIZE_ADDRESS __attribute__((no_sanitize_address)) +#else +# define VERONA_TRACK_FREE(ptr, size) +# define VERONA_TRACK_ALLOC(ptr, size) +# define VERONA_NO_SANITIZE_ADDRESS +#endif + namespace verona::rt::heap { +#ifdef USE_REPLAY_ALLOCATOR + class ReplayAllocator + { + struct Node + { + Node* next; + }; + + static inline std::array allocs; + static inline std::array lengths; + static inline snmalloc::FlagWord lock; + static inline PRNG rng; + + public: + static void set_seed(uint64_t seed) + { + rng.set_seed(seed); + } + + VERONA_NO_SANITIZE_ADDRESS + static void* alloc(size_t size) + { + auto sc = snmalloc::size_to_sizeclass_full(size); + auto idx = sc.index(); + { + snmalloc::FlagLock l(lock); + if (lengths[idx] > 0) + { + // Reuse an element if there are at least 16, or + // the randomisation says to. + auto r = rng.next(); +// std::cout << std::hex << r << std::endl; + + auto reuse = + lengths[idx] > 16 || ((r & 0xf) == 0); + if (reuse) + { + auto r = rng.next(); + Node** prev = &allocs[idx]; + for (size_t i = 0; i < r % lengths[idx]; i++) + { + prev = &(*prev)->next; + } + + auto curr = *prev; + auto next = curr->next; + *prev = next; + lengths[idx]--; + VERONA_TRACK_ALLOC(curr, size); + return curr; + } + } + } + + return snmalloc::ThreadAlloc::get().alloc(size); + } + + static void dealloc(void* ptr, size_t size) + { + auto sc = snmalloc::size_to_sizeclass_full(size); + auto idx = sc.index(); + snmalloc::FlagLock l(lock); + auto hd = reinterpret_cast(ptr); + hd->next = allocs[idx]; + allocs[idx] = hd; + lengths[idx]++; + VERONA_TRACK_FREE(ptr, size); + } + + VERONA_NO_SANITIZE_ADDRESS + static void flush() + { + for (size_t i = 0; i < allocs.size(); i++) + { + auto hd = allocs[i]; + size_t count = 0; + while (hd != nullptr) + { + auto next = hd->next; + snmalloc::ThreadAlloc::get().dealloc( + hd, snmalloc::sizeclass_to_size(i)); + hd = next; + count++; + } + assert(count == lengths[i]); + lengths[i] = 0; + allocs[i] = nullptr; + } + } + }; + + inline void* alloc(size_t size) + { + return ReplayAllocator::alloc(size); + } + + template + inline void* alloc() + { + return ReplayAllocator::alloc(size); + } + + inline void* calloc(size_t size) + { + auto obj = ReplayAllocator::alloc(size); + memset(obj, 0, size); + return obj; + } + + template + inline void* calloc() + { + auto obj = ReplayAllocator::alloc(size); + memset(obj, 0, size); + return obj; + } + + inline void dealloc(void* ptr, size_t size) + { + ReplayAllocator::dealloc(ptr, size); + } + + inline void dealloc(void* ptr) + { + auto size = snmalloc::ThreadAlloc::get().alloc_size(ptr); + dealloc(ptr, size); + } + + template + inline void dealloc(void* ptr) + { + ReplayAllocator::dealloc(ptr, size); + } + + inline void debug_check_empty() + { + ReplayAllocator::flush(); + snmalloc::debug_check_empty(); + } + + inline void set_seed(uint64_t seed) + { + ReplayAllocator::set_seed(seed); + } +#else inline void* alloc(size_t size) { return snmalloc::ThreadAlloc::get().alloc(size); @@ -49,4 +215,10 @@ namespace verona::rt::heap { snmalloc::debug_check_empty(); } + + inline void set_seed(uint64_t seed) + { + // Do nothing + } +#endif } // namespace verona::rt::heap \ No newline at end of file diff --git a/test/func/replayalloc/replayalloc.cc b/test/func/replayalloc/replayalloc.cc new file mode 100644 index 00000000..c7483f69 --- /dev/null +++ b/test/func/replayalloc/replayalloc.cc @@ -0,0 +1,64 @@ +// Copyright Microsoft and Project Verona Contributors. +// SPDX-License-Identifier: MIT + +#define USE_REPLAY_ALLOCATOR + +#include +#include + +/** + * Tests heap replay functionality with systematic testing. + * + * We count each time alloc returns the value just freed, this should + * be possible, but not guaranteed. + **/ + +size_t count = 0; + +void test_replay() +{ + size_t size = 16; + void* p = heap::alloc(size); + uintptr_t p_addr = (uintptr_t)p; + heap::dealloc(p, size); + + void* p2 = heap::alloc(size); + uintptr_t p2_addr = (uintptr_t)p2; + heap::dealloc(p2, size); + + if (p_addr == p2_addr) + { + std::cout << "*" << std::flush; + count ++; + } + else + { + std::cout << "." << std::flush; + } +} + + +int main(int argc, char** argv) +{ + auto t = Aal::tick(); + size_t repeats = 1000; + for (size_t i = 0; i < repeats; i++) + { + heap::set_seed(i + t); + test_replay(); + heap::debug_check_empty(); + if (i % 64 == 0) + { + std::cout << std::endl; + } + } + + std::cout << std::endl << "count: " << count << std::endl; + + // We should have at least one replay + check(count != 0); + // We should not have all replays + check(count != repeats); + + return 0; +}