diff --git a/compiler-rt/lib/gwp_asan/guarded_pool_allocator.h b/compiler-rt/lib/gwp_asan/guarded_pool_allocator.h --- a/compiler-rt/lib/gwp_asan/guarded_pool_allocator.h +++ b/compiler-rt/lib/gwp_asan/guarded_pool_allocator.h @@ -103,6 +103,12 @@ void disable(); void enable(); + typedef void (*iterate_callback)(uintptr_t base, size_t size, void *arg); + // Execute the callback Cb for every allocation the lies in [Base, Base + Size). + // Must be called while the allocator is disabled. The callback can not + // allocate. + void iterate(void *Base, size_t Size, iterate_callback Cb, void *Arg); + // Return whether the allocation should be randomly chosen for sampling. GWP_ASAN_ALWAYS_INLINE bool shouldSample() { // NextSampleCounter == 0 means we "should regenerate the counter". diff --git a/compiler-rt/lib/gwp_asan/guarded_pool_allocator.cpp b/compiler-rt/lib/gwp_asan/guarded_pool_allocator.cpp --- a/compiler-rt/lib/gwp_asan/guarded_pool_allocator.cpp +++ b/compiler-rt/lib/gwp_asan/guarded_pool_allocator.cpp @@ -180,6 +180,17 @@ void GuardedPoolAllocator::enable() { PoolMutex.unlock(); } +void GuardedPoolAllocator::iterate(void *Base, size_t Size, iterate_callback Cb, + void *Arg) { + uintptr_t Start = reinterpret_cast(Base); + for (size_t i = 0; i < MaxSimultaneousAllocations; ++i) { + const AllocationMetadata &Meta = Metadata[i]; + if (Meta.Addr && !Meta.IsDeallocated && Meta.Addr >= Start && + Meta.Addr < Start + Size) + Cb(Meta.Addr, Meta.Size, Arg); + } +} + void GuardedPoolAllocator::uninitTestOnly() { if (GuardedPagePool) { unmapMemory(reinterpret_cast(GuardedPagePool), diff --git a/compiler-rt/lib/gwp_asan/tests/CMakeLists.txt b/compiler-rt/lib/gwp_asan/tests/CMakeLists.txt --- a/compiler-rt/lib/gwp_asan/tests/CMakeLists.txt +++ b/compiler-rt/lib/gwp_asan/tests/CMakeLists.txt @@ -14,6 +14,7 @@ backtrace.cpp basic.cpp compression.cpp + iterate.cpp driver.cpp mutex_test.cpp slot_reuse.cpp diff --git a/compiler-rt/lib/gwp_asan/tests/iterate.cpp b/compiler-rt/lib/gwp_asan/tests/iterate.cpp new file mode 100644 --- /dev/null +++ b/compiler-rt/lib/gwp_asan/tests/iterate.cpp @@ -0,0 +1,66 @@ +//===-- iterate.cpp ---------------------------------------------*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "gwp_asan/tests/harness.h" + +TEST_F(CustomGuardedPoolAllocator, Iterate) { + InitNumSlots(7); + std::vector> Allocated; + auto alloc = [&](size_t size) { + Allocated.push_back({GPA.allocate(size), size}); + }; + + void *Ptr = GPA.allocate(5); + alloc(2); + alloc(1); + alloc(100); + GPA.deallocate(Ptr); + alloc(42); + std::sort(Allocated.begin(), Allocated.end()); + + GPA.disable(); + void *Base = Allocated[0].first; + size_t Size = reinterpret_cast(Allocated.back().first) - + reinterpret_cast(Base) + 1; + std::vector> Found; + GPA.iterate( + Base, Size, + [](uintptr_t Addr, size_t Size, void *Arg) { + reinterpret_cast> *>(Arg) + ->push_back({(void *)Addr, Size}); + }, + reinterpret_cast(&Found)); + GPA.enable(); + + std::sort(Found.begin(), Found.end()); + EXPECT_EQ(Allocated, Found); + + // Now without the last allocation. + GPA.disable(); + Size = reinterpret_cast(Allocated.back().first) - + reinterpret_cast(Base); // Allocated.back() is out of range. + Found.clear(); + GPA.iterate( + Base, Size, + [](uintptr_t Addr, size_t Size, void *Arg) { + reinterpret_cast> *>(Arg) + ->push_back({(void *)Addr, Size}); + }, + reinterpret_cast(&Found)); + GPA.enable(); + + // We should have found every allocation but the last. + // Remove it and compare the rest. + std::sort(Found.begin(), Found.end()); + GPA.deallocate(Allocated.back().first); + Allocated.pop_back(); + EXPECT_EQ(Allocated, Found); + + for (auto PS : Allocated) + GPA.deallocate(PS.first); +} diff --git a/compiler-rt/lib/scudo/standalone/combined.h b/compiler-rt/lib/scudo/standalone/combined.h --- a/compiler-rt/lib/scudo/standalone/combined.h +++ b/compiler-rt/lib/scudo/standalone/combined.h @@ -592,6 +592,9 @@ }; Primary.iterateOverBlocks(Lambda); Secondary.iterateOverBlocks(Lambda); +#ifdef GWP_ASAN_HOOKS + GuardedAlloc.iterate(reinterpret_cast(Base), Size, Callback, Arg); +#endif } bool canReturnNull() {