[ASan/Win] Fix a CHECK failure when an exception is thrown from a callback passed to QueueUserWorkItem

llvm-svn: 231947

compiler-rt/lib/asan/asan_interceptors.cc

@@ -832,9 +832,55 @@ INTERCEPTOR_WINAPI(DWORD, CreateThread,
                             asan_thread_start, t, thr_flags, tid);
 }
 
+struct UserWorkItemInfo {
+  DWORD (__stdcall *function)(void *arg);
+  void *arg;
+  u32 parent_tid;
+};
+
+static BlockingMutex mu_for_thread_tracking(LINKER_INITIALIZED);
+
+// QueueUserWorkItem may silently create a thread we should keep track of.
+// We achieve this by wrapping the user-supplied work items with our function.
+static DWORD __stdcall QueueUserWorkItemWrapper(void *arg) {
+  UserWorkItemInfo *item = (UserWorkItemInfo *)arg;
+
+  {
+    // FIXME: GetCurrentThread relies on TSD, which might not play well with
+    // system thread pools.  We might want to use something like reference
+    // counting to zero out GetCurrentThread() underlying storage when the last
+    // work item finishes?  Or can we disable reclaiming of threads in the pool?
+    BlockingMutexLock l(&mu_for_thread_tracking);
+    AsanThread *t = GetCurrentThread();
+    if (!t) {
+      GET_STACK_TRACE_THREAD;
+      t = AsanThread::Create(/* start_routine */ nullptr, /* arg */ nullptr,
+                             item->parent_tid, &stack, /* detached */ true);
+      t->Init();
+      asanThreadRegistry().StartThread(t->tid(), 0, 0);
+      SetCurrentThread(t);
+    }
+  }
+
+  DWORD ret = item->function(item->arg);
+  delete item;
+  return ret;
+}
+
+INTERCEPTOR_WINAPI(DWORD, QueueUserWorkItem, DWORD(__stdcall *function)(void *),
+                   void *arg, DWORD flags) {
+  UserWorkItemInfo *work_item_info = new UserWorkItemInfo;
+  work_item_info->function = function;
+  work_item_info->arg = arg;
+  work_item_info->parent_tid = GetCurrentTidOrInvalid();
+  return REAL(QueueUserWorkItem)(QueueUserWorkItemWrapper, work_item_info,
+                                 flags);
+}
+
 namespace __asan {
 void InitializeWindowsInterceptors() {
   ASAN_INTERCEPT_FUNC(CreateThread);
+  ASAN_INTERCEPT_FUNC(QueueUserWorkItem);
   ASAN_INTERCEPT_FUNC(RaiseException);
   ASAN_INTERCEPT_FUNC(_except_handler3);
   ASAN_INTERCEPT_FUNC(_except_handler4);

compiler-rt/test/asan/TestCases/Windows/queue_user_work_item.cc

@@ -0,0 +1,55 @@
+// Make sure we can throw exceptions from work items executed via
+// QueueUserWorkItem.
+//
+// Clang doesn't support exceptions on Windows yet, so for the time being we
+// build this program in two parts: the code with exceptions is built with CL,
+// the rest is built with Clang.  This represents the typical scenario when we
+// build a large project using "clang-cl -fallback -fsanitize=address".
+//
+// RUN: cl -c %s -Fo%t.obj
+// RUN: %clangxx_asan -o %t.exe %s %t.obj
+// RUN: %run %t.exe 2>&1 | FileCheck %s
+
+#include <windows.h>
+#include <stdio.h>
+
+void ThrowAndCatch();
+
+#if !defined(__clang__)
+__declspec(noinline)
+void Throw() {
+  fprintf(stderr, "Throw\n");
+// CHECK: Throw
+  throw 1;
+}
+
+void ThrowAndCatch() {
+  int local;
+  try {
+    Throw();
+  } catch(...) {
+    fprintf(stderr, "Catch\n");
+// CHECK: Catch
+  }
+}
+#else
+
+HANDLE done;
+
+DWORD CALLBACK work_item(LPVOID) {
+  ThrowAndCatch();
+  SetEvent(done);
+  return 0;
+}
+
+int main(int argc, char **argv) {
+  done = CreateEvent(0, false, false, "job is done");
+  if (!done)
+    return 1;
+  QueueUserWorkItem(&work_item, nullptr, 0);
+  if (WAIT_OBJECT_0 != WaitForSingleObject(done, INFINITE))
+    return 2;
+  fprintf(stderr, "Done!\n");
+// CHECK: Done!
+}
+#endif