[lsan] Handle fork() correctly.

Update the main thread's os_id on every pthread_create, and before
initiating leak checking. This ensures that we have the correct os_id even if we
have forked after Init().

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@185815 91177308-0d34-0410-b5e6-96231b3b80d8

10 files changed, 99 insertions, 2 deletions

diff --git a/lib/asan/asan_interceptors.cc b/lib/asan/asan_interceptors.cc
index 32f1ebfa8..374344f27 100644
--- a/lib/asan/asan_interceptors.cc
+++ b/lib/asan/asan_interceptors.cc
@@ -138,6 +138,7 @@ extern "C" int pthread_attr_getdetachstate(void *attr, int *v);
 
 INTERCEPTOR(int, pthread_create, void *thread,
     void *attr, void *(*start_routine)(void*), void *arg) {
+  EnsureMainThreadIDIsCorrect();
   // Strict init-order checking in thread-hostile.
   if (flags()->strict_init_order)
     StopInitOrderChecking();
diff --git a/lib/asan/asan_thread.cc b/lib/asan/asan_thread.cc
index 39153bf87..1253ff196 100644
--- a/lib/asan/asan_thread.cc
+++ b/lib/asan/asan_thread.cc
@@ -211,7 +211,8 @@ static bool ThreadStackContainsAddress(ThreadContextBase *tctx_base,
 }
 
 AsanThread *GetCurrentThread() {
-  AsanThreadContext *context = (AsanThreadContext*)AsanTSDGet();
+  AsanThreadContext *context =
+      reinterpret_cast<AsanThreadContext *>(AsanTSDGet());
   if (!context) {
     if (SANITIZER_ANDROID) {
       // On Android, libc constructor is called _after_ asan_init, and cleans up
@@ -254,6 +255,13 @@ AsanThread *FindThreadByStackAddress(uptr addr) {
                                                    (void *)addr));
   return tctx ? tctx->thread : 0;
 }
+
+void EnsureMainThreadIDIsCorrect() {
+  AsanThreadContext *context =
+      reinterpret_cast<AsanThreadContext *>(AsanTSDGet());
+  if (context && (context->tid == 0))
+    context->os_id = GetTid();
+}
 }  // namespace __asan
 
 // --- Implementation of LSan-specific functions --- {{{1
@@ -283,4 +291,8 @@ void LockThreadRegistry() {
 void UnlockThreadRegistry() {
   __asan::asanThreadRegistry().Unlock();
 }
+
+void EnsureMainThreadIDIsCorrect() {
+  __asan::EnsureMainThreadIDIsCorrect();
+}
 }  // namespace __lsan
diff --git a/lib/asan/asan_thread.h b/lib/asan/asan_thread.h
index 8aac20279..5fe81541e 100644
--- a/lib/asan/asan_thread.h
+++ b/lib/asan/asan_thread.h
@@ -124,6 +124,8 @@ void SetCurrentThread(AsanThread *t);
 u32 GetCurrentTidOrInvalid();
 AsanThread *FindThreadByStackAddress(uptr addr);
 
+// Used to handle fork().
+void EnsureMainThreadIDIsCorrect();
 }  // namespace __asan
 
 #endif  // ASAN_THREAD_H
diff --git a/lib/lsan/lit_tests/TestCases/fork.cc b/lib/lsan/lit_tests/TestCases/fork.cc
new file mode 100644
index 000000000..69258d9a0
--- /dev/null
+++ b/lib/lsan/lit_tests/TestCases/fork.cc
@@ -0,0 +1,24 @@
+// Test that thread local data is handled correctly after forking without exec().
+// RUN: %clangxx_lsan %s -o %t
+// RUN: %t 2>&1
+
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/wait.h>
+#include <unistd.h>
+
+__thread void *thread_local_var;
+
+int main() {
+  int status = 0;
+  thread_local_var = malloc(1337);
+  pid_t pid = fork();
+  assert(pid >= 0);
+  if (pid > 0) {
+    waitpid(pid, &status, 0);
+    assert(WIFEXITED(status));
+    return WEXITSTATUS(status);
+  }
+  return 0;
+}
diff --git a/lib/lsan/lit_tests/TestCases/fork_threaded.cc b/lib/lsan/lit_tests/TestCases/fork_threaded.cc
new file mode 100644
index 000000000..b162de458
--- /dev/null
+++ b/lib/lsan/lit_tests/TestCases/fork_threaded.cc
@@ -0,0 +1,44 @@
+// Test that thread local data is handled correctly after forking without
+// exec(). In this test leak checking is initiated from a non-main thread.
+// RUN: %clangxx_lsan %s -o %t
+// RUN: %t 2>&1
+
+#include <assert.h>
+#include <pthread.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/wait.h>
+#include <unistd.h>
+
+__thread void *thread_local_var;
+
+void *exit_thread_func(void *arg) {
+  exit(0);
+}
+
+void ExitFromThread() {
+  pthread_t tid;
+  int res;
+  res = pthread_create(&tid, 0, exit_thread_func, 0);
+  assert(res == 0);
+  res = pthread_join(tid, 0);
+  assert(res == 0);
+}
+
+int main() {
+  int status = 0;
+  thread_local_var = malloc(1337);
+  pid_t pid = fork();
+  assert(pid >= 0);
+  if (pid > 0) {
+    waitpid(pid, &status, 0);
+    assert(WIFEXITED(status));
+    return WEXITSTATUS(status);
+  } else {
+    // Spawn a thread and call exit() from there, to check that we track main
+    // thread's pid correctly even if leak checking is initiated from another
+    // thread.
+    ExitFromThread();
+  }
+  return 0;
+}
diff --git a/lib/lsan/lsan_common.cc b/lib/lsan/lsan_common.cc
index af5df4cdb..534f154e8 100644
--- a/lib/lsan/lsan_common.cc
+++ b/lib/lsan/lsan_common.cc
@@ -350,6 +350,7 @@ static void DoLeakCheckCallback(const SuspendedThreadsList &suspended_threads,
 }
 
 void DoLeakCheck() {
+  EnsureMainThreadIDIsCorrect();
   BlockingMutexLock l(&global_mutex);
   static bool already_done;
   CHECK(!already_done);
diff --git a/lib/lsan/lsan_common.h b/lib/lsan/lsan_common.h
index 74e6a8175..926dea80f 100644
--- a/lib/lsan/lsan_common.h
+++ b/lib/lsan/lsan_common.h
@@ -135,6 +135,13 @@ void UnlockThreadRegistry();
 bool GetThreadRangesLocked(uptr os_id, uptr *stack_begin, uptr *stack_end,
                            uptr *tls_begin, uptr *tls_end,
                            uptr *cache_begin, uptr *cache_end);
+// If called from the main thread, updates the main thread's TID in the thread
+// registry. We need this to handle processes that fork() without a subsequent
+// exec(), which invalidates the recorded TID. To update it, we must call
+// gettid() from the main thread. Our solution is to call this function before
+// leak checking and also before every call to pthread_create() (to handle cases
+// where leak checking is initiated from a non-main thread).
+void EnsureMainThreadIDIsCorrect();
 // If p points into a chunk that has been allocated to the user, returns its
 // user-visible address. Otherwise, returns 0.
 uptr PointsIntoChunk(void *p);
diff --git a/lib/lsan/lsan_interceptors.cc b/lib/lsan/lsan_interceptors.cc
index e9e2dd2da..454b5561c 100644
--- a/lib/lsan/lsan_interceptors.cc
+++ b/lib/lsan/lsan_interceptors.cc
@@ -218,6 +218,7 @@ extern "C" void *__lsan_thread_start_func(void *arg) {
 INTERCEPTOR(int, pthread_create, void *th, void *attr,
             void *(*callback)(void *), void *param) {
   Init();
+  EnsureMainThreadIDIsCorrect();
   __sanitizer_pthread_attr_t myattr;
   if (attr == 0) {
     pthread_attr_init(&myattr);
diff --git a/lib/lsan/lsan_thread.cc b/lib/lsan/lsan_thread.cc
index 3e28dee13..3bfccdf71 100644
--- a/lib/lsan/lsan_thread.cc
+++ b/lib/lsan/lsan_thread.cc
@@ -123,6 +123,11 @@ void ThreadJoin(u32 tid) {
   thread_registry->JoinThread(tid, /* arg */0);
 }
 
+void EnsureMainThreadIDIsCorrect() {
+  if (GetCurrentThread() == 0)
+    CurrentThreadContext()->os_id = GetTid();
+}
+
 ///// Interface to the common LSan module. /////
 
 bool GetThreadRangesLocked(uptr os_id, uptr *stack_begin, uptr *stack_end,
diff --git a/lib/lsan/lsan_thread.h b/lib/lsan/lsan_thread.h
index b62f04b8e..4641b32ed 100644
--- a/lib/lsan/lsan_thread.h
+++ b/lib/lsan/lsan_thread.h
@@ -47,7 +47,7 @@ u32 ThreadTid(uptr uid);
 u32 GetCurrentThread();
 void SetCurrentThread(u32 tid);
 ThreadContext *CurrentThreadContext();
-
+void EnsureMainThreadIDIsCorrect();
 }  // namespace __lsan
 
 #endif  // LSAN_THREAD_H