[ASAN] Use ThreadArgRetval in ASAN

Fixes false leaks on thread retval.

Reviewed By: thurston

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

6 files changed, 65 insertions, 44 deletions

diff --git a/compiler-rt/lib/asan/asan_interceptors.cpp b/compiler-rt/lib/asan/asan_interceptors.cpp
index 2df596ed7fe1..7d658e527e17 100644
--- a/compiler-rt/lib/asan/asan_interceptors.cpp
+++ b/compiler-rt/lib/asan/asan_interceptors.cpp
@@ -198,11 +198,16 @@ DECLARE_REAL_AND_INTERCEPTOR(void, free, void *)
 static thread_return_t THREAD_CALLING_CONV asan_thread_start(void *arg) {
   AsanThread *t = (AsanThread *)arg;
   SetCurrentThread(t);
-  return t->ThreadStart(GetTid());
+  auto self = GetThreadSelf();
+  auto args = asanThreadArgRetval().GetArgs(self);
+  thread_return_t retval = t->ThreadStart(GetTid());
+  asanThreadArgRetval().Finish(self, retval);
+  CHECK_EQ(args.arg_retval, t->get_arg());
+  return retval;
 }
 
-INTERCEPTOR(int, pthread_create, void *thread,
-    void *attr, void *(*start_routine)(void*), void *arg) {
+INTERCEPTOR(int, pthread_create, void *thread, void *attr,
+            void *(*start_routine)(void *), void *arg) {
   EnsureMainThreadIDIsCorrect();
   // Strict init-order checking is thread-hostile.
   if (flags()->strict_init_order)
@@ -222,10 +227,13 @@ INTERCEPTOR(int, pthread_create, void *thread,
     // stored by pthread for future reuse even after thread destruction, and
     // the linked list it's stored in doesn't even hold valid pointers to the
     // objects, the latter are calculated by obscure pointer arithmetic.
-#if CAN_SANITIZE_LEAKS
+#    if CAN_SANITIZE_LEAKS
     __lsan::ScopedInterceptorDisabler disabler;
-#endif
-    result = REAL(pthread_create)(thread, attr, asan_thread_start, t);
+#    endif
+    asanThreadArgRetval().Create(detached, {start_routine, arg}, [&]() -> uptr {
+      result = REAL(pthread_create)(thread, attr, asan_thread_start, t);
+      return result ? 0 : *(uptr *)(thread);
+    });
   }
   if (result != 0) {
     // If the thread didn't start delete the AsanThread to avoid leaking it.
@@ -237,27 +245,48 @@ INTERCEPTOR(int, pthread_create, void *thread,
 }
 
 INTERCEPTOR(int, pthread_join, void *thread, void **retval) {
-  return REAL(pthread_join)(thread, retval);
+  int result;
+  asanThreadArgRetval().Join((uptr)thread, [&]() {
+    result = REAL(pthread_join)(thread, retval);
+    return !result;
+  });
+  return result;
 }
 
 INTERCEPTOR(int, pthread_detach, void *thread) {
-  return REAL(pthread_detach)(thread);
+  int result;
+  asanThreadArgRetval().Detach((uptr)thread, [&](){
+    result = REAL(pthread_detach)(thread);
+    return !result;
+  });
+  return result;
 }
 
 INTERCEPTOR(int, pthread_exit, void *retval) {
+  asanThreadArgRetval().Finish(GetThreadSelf(), retval);
   return REAL(pthread_exit)(retval);
 }
 
 #    if ASAN_INTERCEPT_TRYJOIN
 INTERCEPTOR(int, pthread_tryjoin_np, void *thread, void **ret) {
-  return REAL(pthread_tryjoin_np)(thread, ret);
+  int result;
+  asanThreadArgRetval().Join((uptr)thread, [&]() {
+    result = REAL(pthread_tryjoin_np)(thread, ret);
+    return !result;
+  });
+  return result;
 }
 #    endif
 
 #    if ASAN_INTERCEPT_TIMEDJOIN
 INTERCEPTOR(int, pthread_timedjoin_np, void *thread, void **ret,
             const struct timespec *abstime) {
-  return REAL(pthread_timedjoin_np)(thread, ret, abstime);
+  int result;
+  asanThreadArgRetval().Join((uptr)thread, [&]() {
+    result = REAL(pthread_timedjoin_np)(thread, ret, abstime);
+    return !result;
+  });
+  return result;
 }
 #    endif
 
diff --git a/compiler-rt/lib/asan/asan_thread.cpp b/compiler-rt/lib/asan/asan_thread.cpp
index a1734a6f896d..343fd079fa02 100644
--- a/compiler-rt/lib/asan/asan_thread.cpp
+++ b/compiler-rt/lib/asan/asan_thread.cpp
@@ -41,6 +41,8 @@ void AsanThreadContext::OnFinished() {
 }
 
 static ThreadRegistry *asan_thread_registry;
+static ThreadArgRetval *thread_data;
+
 static Mutex mu_for_thread_context;
 static LowLevelAllocator allocator_for_thread_context;
 
@@ -63,9 +65,12 @@ static void InitThreads() {
   // MIPS requires aligned address
   static ALIGNED(alignof(
       ThreadRegistry)) char thread_registry_placeholder[sizeof(ThreadRegistry)];
+  static ALIGNED(alignof(
+      ThreadArgRetval)) char thread_data_placeholder[sizeof(ThreadArgRetval)];
 
   asan_thread_registry =
       new (thread_registry_placeholder) ThreadRegistry(GetAsanThreadContext);
+  thread_data = new (thread_data_placeholder) ThreadArgRetval();
   initialized = true;
 }
 
@@ -74,6 +79,11 @@ ThreadRegistry &asanThreadRegistry() {
   return *asan_thread_registry;
 }
 
+ThreadArgRetval &asanThreadArgRetval() {
+  InitThreads();
+  return *thread_data;
+}
+
 AsanThreadContext *GetThreadContextByTidLocked(u32 tid) {
   return static_cast<AsanThreadContext *>(
       asanThreadRegistry().GetThreadLocked(tid));
@@ -484,9 +494,15 @@ __asan::AsanThread *GetAsanThreadByOsIDLocked(tid_t os_id) {
 
 // --- Implementation of LSan-specific functions --- {{{1
 namespace __lsan {
-void LockThreadRegistry() { __asan::asanThreadRegistry().Lock(); }
+void LockThreadRegistry() {
+  __asan::asanThreadRegistry().Lock();
+  __asan::asanThreadArgRetval().Lock();
+}
 
-void UnlockThreadRegistry() { __asan::asanThreadRegistry().Unlock(); }
+void UnlockThreadRegistry() {
+  __asan::asanThreadArgRetval().Unlock();
+  __asan::asanThreadRegistry().Unlock();
+}
 
 static ThreadRegistry *GetAsanThreadRegistryLocked() {
   __asan::asanThreadRegistry().CheckLocked();
@@ -541,33 +557,7 @@ void GetThreadExtraStackRangesLocked(InternalMmapVector<Range> *ranges) {
 }
 
 void GetAdditionalThreadContextPtrsLocked(InternalMmapVector<uptr> *ptrs) {
-  GetAsanThreadRegistryLocked()->RunCallbackForEachThreadLocked(
-      [](ThreadContextBase *tctx, void *ptrs) {
-        // Look for the arg pointer of threads that have been created or are
-        // running. This is necessary to prevent false positive leaks due to the
-        // AsanThread holding the only live reference to a heap object.  This
-        // can happen because the `pthread_create()` interceptor doesn't wait
-        // for the child thread to start before returning and thus loosing the
-        // the only live reference to the heap object on the stack.
-
-        __asan::AsanThreadContext *atctx =
-            static_cast<__asan::AsanThreadContext *>(tctx);
-
-        // Note ThreadStatusRunning is required because there is a small window
-        // where the thread status switches to `ThreadStatusRunning` but the
-        // `arg` pointer still isn't on the stack yet.
-        if (atctx->status != ThreadStatusCreated &&
-            atctx->status != ThreadStatusRunning)
-          return;
-
-        uptr thread_arg = reinterpret_cast<uptr>(atctx->thread->get_arg());
-        if (!thread_arg)
-          return;
-
-        auto ptrsVec = reinterpret_cast<InternalMmapVector<uptr> *>(ptrs);
-        ptrsVec->push_back(thread_arg);
-      },
-      ptrs);
+  __asan::asanThreadArgRetval().GetAllPtrsLocked(ptrs);
 }
 
 void GetRunningThreadsLocked(InternalMmapVector<tid_t> *threads) {
diff --git a/compiler-rt/lib/asan/asan_thread.h b/compiler-rt/lib/asan/asan_thread.h
index dff1a0fd85d5..c131dd40d864 100644
--- a/compiler-rt/lib/asan/asan_thread.h
+++ b/compiler-rt/lib/asan/asan_thread.h
@@ -20,6 +20,7 @@
 #include "asan_stats.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "sanitizer_common/sanitizer_libc.h"
+#include "sanitizer_common/sanitizer_thread_arg_retval.h"
 #include "sanitizer_common/sanitizer_thread_registry.h"
 
 namespace __sanitizer {
@@ -171,6 +172,7 @@ class AsanThread {
 
 // Returns a single instance of registry.
 ThreadRegistry &asanThreadRegistry();
+ThreadArgRetval &asanThreadArgRetval();
 
 // Must be called under ThreadRegistryLock.
 AsanThreadContext *GetThreadContextByTidLocked(u32 tid);
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.cpp b/compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.cpp
index fd53208bd27a..85362e8a9721 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.cpp
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.cpp
@@ -74,7 +74,7 @@ void ThreadArgRetval::DetachLocked(uptr thread) {
   CHECK(t);
   CHECK(!t->second.detached);
   if (t->second.done) {
-    // Detached thread has no use after it started and returned args.
+    // We can't retrive retval after detached thread finished.
     data_.erase(t);
     return;
   }
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.h b/compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.h
index f2e5af208474..c77021beb67d 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.h
@@ -69,8 +69,8 @@ class SANITIZER_MUTEX ThreadArgRetval {
   template <typename JoinFn /* returns true on success */>
   void Join(uptr thread, const JoinFn& fn) {
     // Remember internal id of the thread to prevent re-use of the thread
-    // between fn() and DetachLocked() calls. We can't just lock JoinFn
-    // like in Detach() implementation.
+    // between fn() and AfterJoin() calls. Locking JoinFn, like in
+    // Detach(), implementation can cause deadlock.
     auto gen = BeforeJoin(thread);
     if (fn())
       AfterJoin(thread, gen);
diff --git a/compiler-rt/test/lsan/TestCases/create_thread_leak.cpp b/compiler-rt/test/lsan/TestCases/create_thread_leak.cpp
index 10c377658e61..4074cd4e540e 100644
--- a/compiler-rt/test/lsan/TestCases/create_thread_leak.cpp
+++ b/compiler-rt/test/lsan/TestCases/create_thread_leak.cpp
@@ -6,8 +6,8 @@
 // RUN: %run not %t 10 0 0 1 2>&1 | FileCheck %s --check-prefixes=LEAK,LEAK234
 
 // FIXME: Remove "not". There is no leak.
-// False LEAK234 is broken for ASAN, LSAN.
-// RUN: %run %if asan %{ not %} %if lsan-standalone %{ not %} %t 10 0 0 0
+// False LEAK234 is broken for LSAN.
+// RUN: %run %if lsan-standalone %{ not %} %t 10 0 0 0
 
 #include <pthread.h>
 #include <stdlib.h>