[gwp-asan] Implement malloc_iterate.

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

Author: eugenis

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),

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".

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

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);
+}

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() {