[dfsan] Add thread registration

This is a part of https://reviews.llvm.org/D95835.

This change is to address two problems
1) When recording stacks in origin tracking, libunwind is not async signal safe. Inside signal callbacks, we need
to use fast unwind. Fast unwind needs threads
2) StackDepot used by origin tracking is not async signal safe, we set a flag per thread inside
a signal callback to prevent from using it.

The thread registration is similar to ASan and MSan.

Related MSan changes are
* 98f5ea0dba
* f653cda269
* 5a7c364343

Some changes in the diff are used in the next diffs
1) The test case pthread.c is not very interesting for now. It will be
  extended to test origin tracking later.
2) DFsanThread::InSignalHandler will be used by origin tracking later.

Reviewed-by: morehouse

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

compiler-rt/lib/dfsan/CMakeLists.txt

@@ -5,12 +5,14 @@ set(DFSAN_RTL_SOURCES
   dfsan.cpp
   dfsan_custom.cpp
   dfsan_interceptors.cpp
+  dfsan_thread.cpp
   )
 
 set(DFSAN_RTL_HEADERS
   dfsan.h
   dfsan_flags.inc
   dfsan_platform.h
+  dfsan_thread.h
   )
 
 set(DFSAN_COMMON_CFLAGS ${SANITIZER_COMMON_CFLAGS})

compiler-rt/lib/dfsan/dfsan.cpp

@@ -20,6 +20,7 @@
 
 #include "dfsan/dfsan.h"
 
+#include "dfsan/dfsan_thread.h"
 #include "sanitizer_common/sanitizer_atomic.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "sanitizer_common/sanitizer_file.h"
@@ -422,7 +423,12 @@ void __sanitizer::BufferedStackTrace::UnwindImpl(uptr pc, uptr bp,
                                                  void *context,
                                                  bool request_fast,
                                                  u32 max_depth) {
-  Unwind(max_depth, pc, bp, context, 0, 0, false);
+  using namespace __dfsan;
+  DFsanThread *t = GetCurrentThread();
+  if (!t || !StackTrace::WillUseFastUnwind(request_fast)) {
+    return Unwind(max_depth, pc, bp, context, 0, 0, false);
+  }
+  Unwind(max_depth, pc, bp, nullptr, t->stack_top(), t->stack_bottom(), true);
 }
 
 extern "C" SANITIZER_INTERFACE_ATTRIBUTE void __sanitizer_print_stack_trace() {
@@ -530,6 +536,12 @@ static void dfsan_init(int argc, char **argv, char **envp) {
   Atexit(dfsan_fini);
   AddDieCallback(dfsan_fini);
 
+  // Set up threads
+  DFsanTSDInit(DFsanTSDDtor);
+  DFsanThread *main_thread = DFsanThread::Create(nullptr, nullptr, nullptr);
+  SetCurrentThread(main_thread);
+  main_thread->ThreadStart();
+
   __dfsan_label_info[kInitializingLabel].desc = "<init label>";
 }
 

compiler-rt/lib/dfsan/dfsan_custom.cpp

@@ -37,6 +37,7 @@
 #include <unistd.h>
 
 #include "dfsan/dfsan.h"
+#include "dfsan/dfsan_thread.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "sanitizer_common/sanitizer_internal_defs.h"
 #include "sanitizer_common/sanitizer_linux.h"
@@ -446,18 +447,33 @@ __dfsw_dlopen(const char *filename, int flag, dfsan_label filename_label,
   return handle;
 }
 
-struct pthread_create_info {
-  void *(*start_routine_trampoline)(void *, void *, dfsan_label, dfsan_label *);
-  void *start_routine;
-  void *arg;
-};
+static void *DFsanThreadStartFunc(void *arg) {
+  DFsanThread *t = (DFsanThread *)arg;
+  SetCurrentThread(t);
+  return t->ThreadStart();
+}
 
-static void *pthread_create_cb(void *p) {
-  pthread_create_info pci(*(pthread_create_info *)p);
-  free(p);
-  dfsan_label ret_label;
-  return pci.start_routine_trampoline(pci.start_routine, pci.arg, 0,
-                                      &ret_label);
+static int dfsan_pthread_create(pthread_t *thread, const pthread_attr_t *attr,
+                                void *start_routine_trampoline,
+                                void *start_routine, void *arg,
+                                dfsan_label *ret_label) {
+  pthread_attr_t myattr;
+  if (!attr) {
+    pthread_attr_init(&myattr);
+    attr = &myattr;
+  }
+
+  // Ensure that the thread stack is large enough to hold all TLS data.
+  AdjustStackSize((void *)(const_cast<pthread_attr_t *>(attr)));
+
+  DFsanThread *t = DFsanThread::Create(start_routine_trampoline,
+                                       (thread_callback_t)start_routine, arg);
+  int res = pthread_create(thread, attr, DFsanThreadStartFunc, t);
+
+  if (attr == &myattr)
+    pthread_attr_destroy(&myattr);
+  *ret_label = 0;
+  return res;
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE int __dfsw_pthread_create(
@@ -467,16 +483,8 @@ SANITIZER_INTERFACE_ATTRIBUTE int __dfsw_pthread_create(
     void *start_routine, void *arg, dfsan_label thread_label,
     dfsan_label attr_label, dfsan_label start_routine_label,
     dfsan_label arg_label, dfsan_label *ret_label) {
-  pthread_create_info *pci =
-      (pthread_create_info *)malloc(sizeof(pthread_create_info));
-  pci->start_routine_trampoline = start_routine_trampoline;
-  pci->start_routine = start_routine;
-  pci->arg = arg;
-  int rv = pthread_create(thread, attr, pthread_create_cb, (void *)pci);
-  if (rv != 0)
-    free(pci);
-  *ret_label = 0;
-  return rv;
+  return dfsan_pthread_create(thread, attr, (void *)start_routine_trampoline,
+                              start_routine, arg, ret_label);
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE int __dfsw_pthread_join(pthread_t thread,
@@ -864,6 +872,18 @@ int __dfsw_sigemptyset(sigset_t *set, dfsan_label set_label,
   return ret;
 }
 
+class SignalHandlerScope {
+ public:
+  SignalHandlerScope() {
+    if (DFsanThread *t = GetCurrentThread())
+      t->EnterSignalHandler();
+  }
+  ~SignalHandlerScope() {
+    if (DFsanThread *t = GetCurrentThread())
+      t->LeaveSignalHandler();
+  }
+};
+
 // Clear DFSan runtime TLS state at the end of a scope.
 //
 // Implementation must be async-signal-safe and use small data size, because
@@ -891,7 +911,8 @@ const int kMaxSignals = 1024;
 static atomic_uintptr_t sigactions[kMaxSignals];
 
 static void SignalHandler(int signo) {
-  ScopedClearThreadLocalState stlsb;
+  SignalHandlerScope signal_handler_scope;
+  ScopedClearThreadLocalState scoped_clear_tls;
 
   // Clear shadows for all inputs provided by system. This is why DFSan
   // instrumentation generates a trampoline function to each function pointer,
@@ -906,7 +927,8 @@ static void SignalHandler(int signo) {
 }
 
 static void SignalAction(int signo, siginfo_t *si, void *uc) {
-  ScopedClearThreadLocalState stlsb;
+  SignalHandlerScope signal_handler_scope;
+  ScopedClearThreadLocalState scoped_clear_tls;
 
   // Clear shadows for all inputs provided by system. Similar to SignalHandler.
   dfsan_clear_arg_tls(0, 3 * sizeof(dfsan_label));

compiler-rt/lib/dfsan/dfsan_thread.cpp

@@ -0,0 +1,115 @@
+#include "dfsan_thread.h"
+
+#include <pthread.h>
+
+#include "dfsan.h"
+
+namespace __dfsan {
+
+DFsanThread *DFsanThread::Create(void *start_routine_trampoline,
+                                 thread_callback_t start_routine, void *arg) {
+  uptr PageSize = GetPageSizeCached();
+  uptr size = RoundUpTo(sizeof(DFsanThread), PageSize);
+  DFsanThread *thread = (DFsanThread *)MmapOrDie(size, __func__);
+  thread->start_routine_trampoline_ = start_routine_trampoline;
+  thread->start_routine_ = start_routine;
+  thread->arg_ = arg;
+  thread->destructor_iterations_ = GetPthreadDestructorIterations();
+
+  return thread;
+}
+
+void DFsanThread::SetThreadStackAndTls() {
+  uptr tls_size = 0;
+  uptr stack_size = 0;
+  uptr tls_begin;
+  GetThreadStackAndTls(IsMainThread(), &stack_.bottom, &stack_size, &tls_begin,
+                       &tls_size);
+  stack_.top = stack_.bottom + stack_size;
+
+  int local;
+  CHECK(AddrIsInStack((uptr)&local));
+}
+
+void DFsanThread::Init() { SetThreadStackAndTls(); }
+
+void DFsanThread::TSDDtor(void *tsd) {
+  DFsanThread *t = (DFsanThread *)tsd;
+  t->Destroy();
+}
+
+void DFsanThread::Destroy() {
+  uptr size = RoundUpTo(sizeof(DFsanThread), GetPageSizeCached());
+  UnmapOrDie(this, size);
+}
+
+thread_return_t DFsanThread::ThreadStart() {
+  Init();
+
+  if (!start_routine_) {
+    // start_routine_ == 0 if we're on the main thread or on one of the
+    // OS X libdispatch worker threads. But nobody is supposed to call
+    // ThreadStart() for the worker threads.
+    return 0;
+  }
+
+  CHECK(start_routine_trampoline_);
+
+  typedef void *(*thread_callback_trampoline_t)(void *, void *, dfsan_label,
+                                                dfsan_label *);
+
+  dfsan_label ret_label;
+  return ((thread_callback_trampoline_t)
+              start_routine_trampoline_)((void *)start_routine_, arg_, 0,
+                                         &ret_label);
+}
+
+DFsanThread::StackBounds DFsanThread::GetStackBounds() const {
+  return {stack_.bottom, stack_.top};
+}
+
+uptr DFsanThread::stack_top() { return GetStackBounds().top; }
+
+uptr DFsanThread::stack_bottom() { return GetStackBounds().bottom; }
+
+bool DFsanThread::AddrIsInStack(uptr addr) {
+  const auto bounds = GetStackBounds();
+  return addr >= bounds.bottom && addr < bounds.top;
+}
+
+static pthread_key_t tsd_key;
+static bool tsd_key_inited = false;
+
+void DFsanTSDInit(void (*destructor)(void *tsd)) {
+  CHECK(!tsd_key_inited);
+  tsd_key_inited = true;
+  CHECK_EQ(0, pthread_key_create(&tsd_key, destructor));
+}
+
+static THREADLOCAL DFsanThread *dfsan_current_thread;
+
+DFsanThread *GetCurrentThread() { return dfsan_current_thread; }
+
+void SetCurrentThread(DFsanThread *t) {
+  // Make sure we do not reset the current DFsanThread.
+  CHECK_EQ(0, dfsan_current_thread);
+  dfsan_current_thread = t;
+  // Make sure that DFsanTSDDtor gets called at the end.
+  CHECK(tsd_key_inited);
+  pthread_setspecific(tsd_key, t);
+}
+
+void DFsanTSDDtor(void *tsd) {
+  DFsanThread *t = (DFsanThread *)tsd;
+  if (t->destructor_iterations_ > 1) {
+    t->destructor_iterations_--;
+    CHECK_EQ(0, pthread_setspecific(tsd_key, tsd));
+    return;
+  }
+  dfsan_current_thread = nullptr;
+  // Make sure that signal handler can not see a stale current thread pointer.
+  atomic_signal_fence(memory_order_seq_cst);
+  DFsanThread::TSDDtor(tsd);
+}
+
+}  // namespace __dfsan

compiler-rt/lib/dfsan/dfsan_thread.h

@@ -0,0 +1,70 @@
+//===-- dfsan_thread.h -------------------------------------------*- C++
+//-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is a part of DataFlowSanitizer.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef DFSAN_THREAD_H
+#define DFSAN_THREAD_H
+
+#include "sanitizer_common/sanitizer_common.h"
+
+namespace __dfsan {
+
+class DFsanThread {
+ public:
+  // NOTE: There is no DFsanThread constructor. It is allocated
+  // via mmap() and *must* be valid in zero-initialized state.
+
+  static DFsanThread *Create(void *start_routine_trampoline,
+                             thread_callback_t start_routine, void *arg);
+  static void TSDDtor(void *tsd);
+  void Destroy();
+
+  void Init();  // Should be called from the thread itself.
+  thread_return_t ThreadStart();
+
+  uptr stack_top();
+  uptr stack_bottom();
+  bool IsMainThread() { return start_routine_ == nullptr; }
+
+  bool InSignalHandler() { return in_signal_handler_; }
+  void EnterSignalHandler() { in_signal_handler_++; }
+  void LeaveSignalHandler() { in_signal_handler_--; }
+
+  int destructor_iterations_;
+
+ private:
+  void SetThreadStackAndTls();
+  struct StackBounds {
+    uptr bottom;
+    uptr top;
+  };
+  StackBounds GetStackBounds() const;
+
+  bool AddrIsInStack(uptr addr);
+
+  void *start_routine_trampoline_;
+  thread_callback_t start_routine_;
+  void *arg_;
+
+  StackBounds stack_;
+
+  unsigned in_signal_handler_;
+};
+
+DFsanThread *GetCurrentThread();
+void SetCurrentThread(DFsanThread *t);
+void DFsanTSDInit(void (*destructor)(void *tsd));
+void DFsanTSDDtor(void *tsd);
+
+}  // namespace __dfsan
+
+#endif  // DFSAN_THREAD_H

compiler-rt/test/dfsan/pthread.c

@@ -0,0 +1,29 @@
+// RUN: %clang_dfsan -mllvm -dfsan-fast-16-labels=true %s -o %t && %run %t
+
+#include <sanitizer/dfsan_interface.h>
+
+#include <assert.h>
+#include <pthread.h>
+
+int volatile x;
+int __thread y;
+
+static void *ThreadFn(void *a) {
+  y = x;
+  assert(dfsan_get_label(y) == 8);
+  return 0;
+}
+
+int main(void) {
+  dfsan_set_label(8, &x, sizeof(x));
+
+  const int kNumThreads = 24;
+  pthread_t t[kNumThreads];
+  for (size_t i = 0; i < kNumThreads; ++i) {
+    pthread_create(&t[i], 0, ThreadFn, (void *)i);
+  }
+  for (size_t i = 0; i < kNumThreads; ++i) {
+    pthread_join(t[i], 0);
+  }
+  return 0;
+}