[lsan][Darwin] Scan libdispatch and Foundation memory regions

libdispatch uses its own heap (_dispatch_main_heap) for some allocations, including the dispatch_continuation_t that holds a dispatch source's event handler.
Objective-C block trampolines (creating methods at runtime with a block as the implementations) use the VM_MEMORY_FOUNDATION region (see https://github.com/apple-oss-distributions/objc4/blob/8701d5672d3fd3cd817aeb84db1077aafe1a1604/runtime/objc-block-trampolines.mm#L371).

This change scans both regions to fix false positives. See tests for details; unfortunately I was unable to reduce the trampoline example with imp_implementationWithBlock on a new class, so I'm resorting to something close to the bug as seen in the wild.

Differential Revision: https://reviews.llvm.org/D129385

3 files changed, 96 insertions, 20 deletions

diff --git a/compiler-rt/lib/lsan/lsan_common_mac.cpp b/compiler-rt/lib/lsan/lsan_common_mac.cpp
index 26b623f..b6b1509 100644
--- a/compiler-rt/lib/lsan/lsan_common_mac.cpp
+++ b/compiler-rt/lib/lsan/lsan_common_mac.cpp
@@ -17,21 +17,36 @@
 
 #if CAN_SANITIZE_LEAKS && SANITIZER_APPLE
 
-#include "sanitizer_common/sanitizer_allocator_internal.h"
-#include "lsan_allocator.h"
+#  include <mach/mach.h>
+#  include <mach/vm_statistics.h>
+#  include <pthread.h>
 
-#include <pthread.h>
+#  include "lsan_allocator.h"
+#  include "sanitizer_common/sanitizer_allocator_internal.h"
+namespace __lsan {
 
-#include <mach/mach.h>
+enum class SeenRegion {
+  None = 0,
+  AllocOnce = 1 << 0,
+  LibDispatch = 1 << 1,
+  Foundation = 1 << 2,
+  All = AllocOnce | LibDispatch | Foundation
+};
+
+inline SeenRegion operator|(SeenRegion left, SeenRegion right) {
+  return static_cast<SeenRegion>(static_cast<int>(left) |
+                                 static_cast<int>(right));
+}
 
-// Only introduced in Mac OS X 10.9.
-#ifdef VM_MEMORY_OS_ALLOC_ONCE
-static const int kSanitizerVmMemoryOsAllocOnce = VM_MEMORY_OS_ALLOC_ONCE;
-#else
-static const int kSanitizerVmMemoryOsAllocOnce = 73;
-#endif
+inline SeenRegion &operator|=(SeenRegion &left, const SeenRegion &right) {
+  left = left | right;
+  return left;
+}
 
-namespace __lsan {
+struct RegionScanState {
+  SeenRegion seen_regions = SeenRegion::None;
+  bool in_libdispatch = false;
+};
 
 typedef struct {
   int disable_counter;
@@ -148,6 +163,7 @@ void ProcessPlatformSpecificAllocations(Frontier *frontier) {
 
   InternalMmapVectorNoCtor<RootRegion> const *root_regions = GetRootRegions();
 
+  RegionScanState scan_state;
   while (err == KERN_SUCCESS) {
     vm_size_t size = 0;
     unsigned depth = 1;
@@ -157,17 +173,35 @@ void ProcessPlatformSpecificAllocations(Frontier *frontier) {
                                (vm_region_info_t)&info, &count);
 
     uptr end_address = address + size;
-
-    // libxpc stashes some pointers in the Kernel Alloc Once page,
-    // make sure not to report those as leaks.
-    if (info.user_tag == kSanitizerVmMemoryOsAllocOnce) {
+    if (info.user_tag == VM_MEMORY_OS_ALLOC_ONCE) {
+      // libxpc stashes some pointers in the Kernel Alloc Once page,
+      // make sure not to report those as leaks.
+      scan_state.seen_regions |= SeenRegion::AllocOnce;
       ScanRangeForPointers(address, end_address, frontier, "GLOBAL",
                            kReachable);
+    } else if (info.user_tag == VM_MEMORY_FOUNDATION) {
+      // Objective-C block trampolines use the Foundation region.
+      scan_state.seen_regions |= SeenRegion::Foundation;
+      ScanRangeForPointers(address, end_address, frontier, "GLOBAL",
+                           kReachable);
+    } else if (info.user_tag == VM_MEMORY_LIBDISPATCH) {
+      // Dispatch continuations use the libdispatch region. Empirically, there
+      // can be more than one region with this tag, so we'll optimistically
+      // assume that they're continguous. Otherwise, we would need to scan every
+      // region to ensure we find them all.
+      scan_state.in_libdispatch = true;
+      ScanRangeForPointers(address, end_address, frontier, "GLOBAL",
+                           kReachable);
+    } else if (scan_state.in_libdispatch) {
+      scan_state.seen_regions |= SeenRegion::LibDispatch;
+      scan_state.in_libdispatch = false;
+    }
 
-      // Recursing over the full memory map is very slow, break out
-      // early if we don't need the full iteration.
-      if (!flags()->use_root_regions || !root_regions->size())
-        break;
+    // Recursing over the full memory map is very slow, break out
+    // early if we don't need the full iteration.
+    if (scan_state.seen_regions == SeenRegion::All &&
+        !(flags()->use_root_regions && root_regions->size() > 0)) {
+      break;
     }
 
     // This additional root region scan is required on Darwin in order to
@@ -199,6 +233,6 @@ void LockStuffAndStopTheWorld(StopTheWorldCallback callback,
   StopTheWorld(callback, argument);
 }
 
-} // namespace __lsan
+}  // namespace __lsan
 
 #endif // CAN_SANITIZE_LEAKS && SANITIZER_APPLE
diff --git a/compiler-rt/test/lsan/TestCases/Darwin/dispatch_continuations.mm b/compiler-rt/test/lsan/TestCases/Darwin/dispatch_continuations.mm
new file mode 100644
index 0000000..d0420d8
--- /dev/null
+++ b/compiler-rt/test/lsan/TestCases/Darwin/dispatch_continuations.mm
@@ -0,0 +1,24 @@
+// Test that dispatch continuation memory region is scanned.
+// RUN: %clangxx_lsan %s  -o %t -framework Foundation
+// RUN: %env_lsan_opts="report_objects=1" %run %t 2>&1 && echo "" | FileCheck %s
+
+#include <dispatch/dispatch.h>
+#include <sanitizer/lsan_interface.h>
+
+int main() {
+  // Reduced from `CFRunLoopCreate`
+  dispatch_queue_t fake_rl_queue = dispatch_get_global_queue(2, 0);
+  dispatch_source_t timer =
+      dispatch_source_create(DISPATCH_SOURCE_TYPE_TIMER, 0, 0, fake_rl_queue);
+  dispatch_source_set_event_handler(timer, ^{
+                                    });
+  dispatch_source_set_timer(timer, DISPATCH_TIME_FOREVER, DISPATCH_TIME_FOREVER,
+                            321);
+  dispatch_resume(timer);
+  __lsan_do_leak_check();
+  dispatch_source_cancel(timer);
+  dispatch_release(timer);
+  return 0;
+}
+
+// CHECK-NOT: LeakSanitizer: detected memory leaks
diff --git a/compiler-rt/test/lsan/TestCases/Darwin/trampoline.mm b/compiler-rt/test/lsan/TestCases/Darwin/trampoline.mm
new file mode 100644
index 0000000..ce3182f
--- /dev/null
+++ b/compiler-rt/test/lsan/TestCases/Darwin/trampoline.mm
@@ -0,0 +1,18 @@
+// Test that the memory region  that contains Objective-C block trampolines
+// is scanned.
+// FIXME: Find a way to reduce this without AppKit to remove Mac requirement.
+// UNSUPPORTED: ios
+// RUN: %clangxx_lsan %s  -o %t -framework Cocoa -fno-objc-arc
+// RUN: %env_lsan_opts="report_objects=1" %run %t 2>&1  && echo "" | FileCheck %s
+
+#import <Cocoa/Cocoa.h>
+
+#include <sanitizer/lsan_interface.h>
+
+int main() {
+  NSView *view =
+      [[[NSView alloc] initWithFrame:CGRectMake(0, 0, 20, 20)] autorelease];
+  __lsan_do_leak_check();
+  return 0;
+}
+// CHECK-NOT: LeakSanitizer: detected memory leaks