[ASan][Darwin][GCD] Add interceptor for dispatch_apply (#149238)

ASan had a gap in coverage for wqthreads blocks submitted by
dispatch_apply

This adds interceptor for dispatch_apply and dispatch_apply_f and adds a
test that a failure in a dispatch apply block contains thread and stack
info.

rdar://139660648

5 files changed, 107 insertions, 2 deletions

diff --git a/compiler-rt/lib/asan/asan_mac.cpp b/compiler-rt/lib/asan/asan_mac.cpp
index 4bc2c88..1f3c79e 100644
--- a/compiler-rt/lib/asan/asan_mac.cpp
+++ b/compiler-rt/lib/asan/asan_mac.cpp
@@ -103,6 +103,8 @@ void FlushUnneededASanShadowMemory(uptr p, uptr size) {
 //   dispatch_after()
 //   dispatch_group_async_f()
 //   dispatch_group_async()
+//   dispatch_apply()
+//   dispatch_apply_f()
 // TODO(glider): libdispatch API contains other functions that we don't support
 // yet.
 //
@@ -243,13 +245,31 @@ INTERCEPTOR(void, dispatch_group_async_f, dispatch_group_t group,
                                asan_dispatch_call_block_and_release);
 }
 
-#if !defined(MISSING_BLOCKS_SUPPORT)
+extern "C" void asan_dispatch_apply_f_work(void *context, size_t iteration) {
+  GET_STACK_TRACE_THREAD;
+  asan_block_context_t *asan_ctxt = (asan_block_context_t *)context;
+  asan_register_worker_thread(asan_ctxt->parent_tid, &stack);
+  ((void (*)(void *, size_t))asan_ctxt->func)(asan_ctxt->block, iteration);
+}
+
+INTERCEPTOR(void, dispatch_apply_f, size_t iterations, dispatch_queue_t queue,
+            void *ctxt, void (*work)(void *, size_t)) {
+  GET_STACK_TRACE_THREAD;
+  asan_block_context_t *asan_ctxt =
+      alloc_asan_context(ctxt, (dispatch_function_t)work, &stack);
+  REAL(dispatch_apply_f)(iterations, queue, (void *)asan_ctxt,
+                         asan_dispatch_apply_f_work);
+}
+
+#  if !defined(MISSING_BLOCKS_SUPPORT)
 extern "C" {
 void dispatch_async(dispatch_queue_t dq, void(^work)(void));
 void dispatch_group_async(dispatch_group_t dg, dispatch_queue_t dq,
                           void(^work)(void));
 void dispatch_after(dispatch_time_t when, dispatch_queue_t queue,
                     void(^work)(void));
+void dispatch_apply(size_t iterations, dispatch_queue_t queue,
+                    void (^block)(size_t iteration));
 void dispatch_source_set_cancel_handler(dispatch_source_t ds,
                                         void(^work)(void));
 void dispatch_source_set_event_handler(dispatch_source_t ds, void(^work)(void));
@@ -332,6 +352,20 @@ INTERCEPTOR(void *, dispatch_mach_create_f, const char *label,
       });
 }
 
-#endif
+INTERCEPTOR(void, dispatch_apply, size_t iterations, dispatch_queue_t queue,
+            void (^block)(size_t iteration)) {
+  ENABLE_FRAME_POINTER;
+  int parent_tid = GetCurrentTidOrInvalid();
+
+  void (^asan_block)(size_t) = ^(size_t iteration) {
+    GET_STACK_TRACE_THREAD;
+    asan_register_worker_thread(parent_tid, &stack);
+    block(iteration);
+  };
+
+  REAL(dispatch_apply)(iterations, queue, asan_block);
+}
+
+#  endif
 
 #endif  // SANITIZER_APPLE
diff --git a/compiler-rt/lib/asan/tests/asan_mac_test.cpp b/compiler-rt/lib/asan/tests/asan_mac_test.cpp
index bd36089..4b21f12 100644
--- a/compiler-rt/lib/asan/tests/asan_mac_test.cpp
+++ b/compiler-rt/lib/asan/tests/asan_mac_test.cpp
@@ -116,6 +116,12 @@ TEST(AddressSanitizerMac, GCDDispatchAfter) {
   EXPECT_DEATH(TestGCDDispatchAfter(), "Shadow byte legend");
 }
 
+TEST(AddressSanitizerMac, GCDDispatchApply) {
+  // Make sure the whole ASan report is printed, i.e. that we don't die
+  // on a CHECK.
+  EXPECT_DEATH(TestGCDDispatchApply(), "Shadow byte legend");
+}
+
 TEST(AddressSanitizerMac, GCDSourceEvent) {
   // Make sure the whole ASan report is printed, i.e. that we don't die
   // on a CHECK.
diff --git a/compiler-rt/lib/asan/tests/asan_mac_test.h b/compiler-rt/lib/asan/tests/asan_mac_test.h
index 441547a..ec71546 100644
--- a/compiler-rt/lib/asan/tests/asan_mac_test.h
+++ b/compiler-rt/lib/asan/tests/asan_mac_test.h
@@ -9,6 +9,7 @@ extern "C" {
   void TestGCDReuseWqthreadsAsync();
   void TestGCDReuseWqthreadsSync();
   void TestGCDDispatchAfter();
+  void TestGCDDispatchApply();
   void TestGCDInTSDDestructor();
   void TestGCDSourceEvent();
   void TestGCDSourceCancel();
diff --git a/compiler-rt/lib/asan/tests/asan_mac_test_helpers.mm b/compiler-rt/lib/asan/tests/asan_mac_test_helpers.mm
index 3f8fa26d..ddb50f8 100644
--- a/compiler-rt/lib/asan/tests/asan_mac_test_helpers.mm
+++ b/compiler-rt/lib/asan/tests/asan_mac_test_helpers.mm
@@ -148,6 +148,16 @@ void TestGCDDispatchAfter() {
   wait_forever();
 }
 
+void TestGCDDispatchApply() {
+  dispatch_queue_t queue = dispatch_get_global_queue(0, 0);
+  __block char *buffer = (char *)malloc(4);
+  dispatch_apply(8, queue, ^(size_t i) {
+    access_memory(&buffer[i]);
+  });
+
+  free(buffer);  // not reached
+}
+
 void worker_do_deallocate(void *ptr) {
   free(ptr);
 }
diff --git a/compiler-rt/test/asan/TestCases/Darwin/dispatch_apply_threadno.c b/compiler-rt/test/asan/TestCases/Darwin/dispatch_apply_threadno.c
new file mode 100644
index 0000000..5e06615
--- /dev/null
+++ b/compiler-rt/test/asan/TestCases/Darwin/dispatch_apply_threadno.c
@@ -0,0 +1,54 @@
+// Bugs caught within missing GCD dispatch blocks result in thread being reported as T-1
+// with an empty stack.
+// This tests that dispatch_apply blocks can capture valid thread number and stack.
+
+// RUN: %clang_asan %s -o %t
+// RUN: not %run %t func 2>&1 | FileCheck %s --check-prefixes=CHECK-FUNC,CHECK
+// RUN: not %run %t block 2>&1 | FileCheck %s --check-prefixes=CHECK-BLOCK,CHECK
+
+#include <dispatch/dispatch.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+__attribute__((noinline)) void access_memory_frame(char *x) { *x = 0; }
+
+__attribute__((noinline)) void test_dispatch_apply() {
+  char *x = (char *)malloc(4);
+  dispatch_apply(8, dispatch_get_global_queue(0, 0), ^(size_t i) {
+    access_memory_frame(&x[i]);
+  });
+}
+
+typedef struct {
+  char *data;
+} Context;
+
+void da_func(void *ctx, size_t i) {
+  Context *c = (Context *)ctx;
+  access_memory_frame(&c->data[i]);
+}
+
+__attribute__((noinline)) void test_dispatch_apply_f() {
+  Context *ctx = (Context *)malloc(sizeof(Context));
+  ctx->data = (char *)malloc(4);
+  dispatch_apply_f(8, dispatch_get_global_queue(0, 0), ctx, da_func);
+}
+
+int main(int argc, const char *argv[]) {
+  if (strcmp(argv[1], "func") == 0) {
+    fprintf(stderr, "Test dispatch_apply with function\n");
+    // CHECK-FUNC: dispatch_apply with function
+    test_dispatch_apply_f();
+  } else if (strcmp(argv[1], "block") == 0) {
+    fprintf(stderr, "Test dispatch_apply with block\n");
+    // CHECK-BLOCK: dispatch_apply with block
+    test_dispatch_apply();
+  } else {
+    abort();
+  }
+  return 0;
+}
+
+// CHECK: ERROR: AddressSanitizer: heap-buffer-overflow
+// CHECK: #0 0x{{.*}} in {{.*}}access_memory_frame
+// CHECK-NOT: T-1