[gwp-asan] Implement malloc_iterate.

eugenis · eugenis · commit 46044a6900e8 · 2020-01-27T14:37:33.000-08:00
Summary: This is an Android-specific interface for iterating over all live allocations in a memory range. Reviewers: hctim, cferris Subscribers: mgorny, mgrang, #sanitizers, llvm-commits Tags: #sanitizers, #llvm Differential Revision: https://reviews.llvm.org/D73305
diff --git a/compiler-rt/lib/gwp_asan/guarded_pool_allocator.cpp b/compiler-rt/lib/gwp_asan/guarded_pool_allocator.cpp
@@ -182,6 +182,17 @@ void GuardedPoolAllocator::disable() { PoolMutex.lock(); }
 
 void GuardedPoolAllocator::enable() { PoolMutex.unlock(); }
 
+void GuardedPoolAllocator::iterate(void *Base, size_t Size, iterate_callback Cb,
+                                   void *Arg) {
+  uintptr_t Start = reinterpret_cast<uintptr_t>(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<void *>(GuardedPagePool),
diff --git a/compiler-rt/lib/gwp_asan/guarded_pool_allocator.h b/compiler-rt/lib/gwp_asan/guarded_pool_allocator.h
@@ -106,6 +106,12 @@ class GuardedPoolAllocator {
   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/tests/CMakeLists.txt b/compiler-rt/lib/gwp_asan/tests/CMakeLists.txt
@@ -14,6 +14,7 @@ set(GWP_ASAN_UNITTESTS
   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
@@ -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<std::pair<void *, size_t>> 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<size_t>(Allocated.back().first) -
+                reinterpret_cast<size_t>(Base) + 1;
+  std::vector<std::pair<void *, size_t>> Found;
+  GPA.iterate(
+      Base, Size,
+      [](uintptr_t Addr, size_t Size, void *Arg) {
+        reinterpret_cast<std::vector<std::pair<void *, size_t>> *>(Arg)
+            ->push_back({(void *)Addr, Size});
+      },
+      reinterpret_cast<void *>(&Found));
+  GPA.enable();
+
+  std::sort(Found.begin(), Found.end());
+  EXPECT_EQ(Allocated, Found);
+
+  // Now without the last allocation.
+  GPA.disable();
+  Size = reinterpret_cast<size_t>(Allocated.back().first) -
+         reinterpret_cast<size_t>(Base); // Allocated.back() is out of range.
+  Found.clear();
+  GPA.iterate(
+      Base, Size,
+      [](uintptr_t Addr, size_t Size, void *Arg) {
+        reinterpret_cast<std::vector<std::pair<void *, size_t>> *>(Arg)
+            ->push_back({(void *)Addr, Size});
+      },
+      reinterpret_cast<void *>(&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
@@ -592,6 +592,9 @@ class Allocator {
     };
     Primary.iterateOverBlocks(Lambda);
     Secondary.iterateOverBlocks(Lambda);
+#ifdef GWP_ASAN_HOOKS
+    GuardedAlloc.iterate(reinterpret_cast<void *>(Base), Size, Callback, Arg);
+#endif
   }
 
   bool canReturnNull() {

Original file line number	Diff line number	Diff line change
`@@ -592,6 +592,9 @@ class Allocator {`
`592`	`592`	`};`
`593`	`593`	`Primary.iterateOverBlocks(Lambda);`
`594`	`594`	`Secondary.iterateOverBlocks(Lambda);`
	`595`	`+#ifdef GWP_ASAN_HOOKS`
	`596`	`+ GuardedAlloc.iterate(reinterpret_cast<void *>(Base), Size, Callback, Arg);`
	`597`	`+#endif`
`595`	`598`	`}`
`596`	`599`
`597`	`600`	`bool canReturnNull() {`