[dfsan] Add thread registration

[stephan.yichao.zhao]

Summary

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
 * https://github.com/llvm/llvm-project/commit/98f5ea0dbae664a2e5f9381a64f2913fe1add208
 * https://github.com/llvm/llvm-project/commit/f653cda2695ac7390fe5663f2c0895213938334d
 * https://github.com/llvm/llvm-project/commit/5a7c3643437c262137bd3dac7f6a0f5b9e8501be

 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.

[stephan.yichao.zhao]

fixed broken builts

[stephan.yichao.zhao]

fixed clang-tidy errors

[morehouse]

compiler-rt/lib/dfsan/dfsan.cpp

434

Isn't the else block dead code? Getting past the first if block implies WillUseFastUnwind is true.

compiler-rt/lib/dfsan/dfsan_custom.cpp

486–488

Since it isn't used by anything else, let's inline dfsan_pthread_create here.

compiler-rt/lib/dfsan/dfsan_thread.cpp

27

nit

70

stack_switching_ is always false.

77

next_stack_ is never set anywhere.

109

Is the cast needed?

compiler-rt/lib/dfsan/dfsan_thread.h

55

If we use the proper function ptr type here and in the constructor, we can remove some unnecessary casting to void * and back.

compiler-rt/test/dfsan/pthread.c

19

Is the cast needed?

21

Nit: let's use a const variable instead of magic number 24.

[stephan.yichao.zhao]

addressed comments

compiler-rt/lib/dfsan/dfsan.cpp

434

Yes. This branch cannot happen. Removed.

compiler-rt/lib/dfsan/dfsan_custom.cpp

486–488

We will use it after origin tracking is added.

compiler-rt/lib/dfsan/dfsan_thread.cpp

70

removed

77

removed it with stack_switching_. In MSan/ASan threads, they expose APIs to allow users tell them when a fiber thread switch contexts. DFSan does not have the case yet.

109

removed.

compiler-rt/lib/dfsan/dfsan_thread.h

55

This start_routine_trampoline_ will have a different type with and without origin tracking like this

typedef void *(*thread_callback_trampoline_t)(void *, void *, dfsan_label,
                                              dfsan_label *);
typedef void *(*thread_callback_origin_trampoline_t)(
    void *, void *, dfsan_label, dfsan_label *, dfsan_origin, dfsan_origin *);

If we use different types here at the C code of dfsan_custom.c, we have to define two functions with similar code.
See the pthread_create wrappers in https://reviews.llvm.org/D95835#change-GvgDqkOr0aQn
Or maybe we can delay how to change this in the late origin tracking related change?

compiler-rt/test/dfsan/pthread.c

19

removed.

21

added kNumThreads=4.

[morehouse]

compiler-rt/lib/dfsan/dfsan_thread.cpp

62

Please address this lint.

[stephan.yichao.zhao]

compiler-rt/lib/dfsan/dfsan_thread.cpp

62

For some reason the local lint conflicts with the one used here, they propose two different styles... maybe they use different versions..?
For the submission, I will try make the local one happy, so it does not block check in. We will change the function in the next diffs. Then we will see if the conflict is gone.

compiler-rt/lib/dfsan/CMakeLists.txt

 include_directories(..)
 
 # Runtime library sources and build flags.
 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})
 append_rtti_flag(OFF DFSAN_COMMON_CFLAGS)
 # Prevent clang from generating libc calls.
 append_list_if(COMPILER_RT_HAS_FFREESTANDING_FLAG -ffreestanding DFSAN_COMMON_CFLAGS)
 
 # Static runtime library.

compiler-rt/lib/dfsan/dfsan.cpp

 //
 // The public interface is defined in include/sanitizer/dfsan_interface.h whose
 // functions are prefixed dfsan_ while the compiler interface functions are
 // prefixed __dfsan_.
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "sanitizer_common/sanitizer_flag_parser.h"
 #include "sanitizer_common/sanitizer_flags.h"
 #include "sanitizer_common/sanitizer_internal_defs.h"
 #include "sanitizer_common/sanitizer_libc.h"
 #include "sanitizer_common/sanitizer_stacktrace.h"

 #define GET_FATAL_STACK_TRACE_PC_BP(pc, bp) \
   BufferedStackTrace stack;                 \
   stack.Unwind(pc, bp, nullptr, common_flags()->fast_unwind_on_fatal);
 
 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() {
   GET_FATAL_STACK_TRACE_PC_BP(StackTrace::GetCurrentPc(), GET_CURRENT_FRAME());
   stack.Print();
 }
 
 void Flags::SetDefaults() {
 #define DFSAN_FLAG(Type, Name, DefaultValue, Description) Name = DefaultValue;
 #include "dfsan_flags.inc"
 #undef DFSAN_FLAG
static void dfsan_init(int argc, char **argv, char **envp) {
 
   InitializeInterceptors();
 
   // Register the fini callback to run when the program terminates successfully
   // or it is killed by the runtime.
   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>";
 }
 
 #if SANITIZER_CAN_USE_PREINIT_ARRAY
 __attribute__((section(".preinit_array"), used))
 static void (*dfsan_init_ptr)(int, char **, char **) = dfsan_init;
 #endif

compiler-rt/lib/dfsan/dfsan_custom.cpp

 #include <sys/socket.h>
 #include <sys/stat.h>
 #include <sys/time.h>
 #include <sys/types.h>
 #include <time.h>
 #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"
 
 using namespace __dfsan;
 
 #define CALL_WEAK_INTERCEPTOR_HOOK(f, ...)                                     \
   do {                                                                         \
__dfsw_dlopen(const char *filename, int flag, dfsan_label filename_label,
   void *handle = dlopen(filename, flag);
   link_map *map = GET_LINK_MAP_BY_DLOPEN_HANDLE(handle);
   if (map)
     ForEachMappedRegion(map, unpoison);
   *ret_label = 0;
   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(
     pthread_t *thread, const pthread_attr_t *attr,
     void *(*start_routine_trampoline)(void *, void *, dfsan_label,
                                       dfsan_label *),
     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,
                                                       void **retval,
                                                       dfsan_label thread_label,
                                                       dfsan_label retval_label,
                                                       dfsan_label *ret_label) {
   int ret = pthread_join(thread, retval);
   if (ret == 0 && retval)

 SANITIZER_INTERFACE_ATTRIBUTE
 int __dfsw_sigemptyset(sigset_t *set, dfsan_label set_label,
                        dfsan_label *ret_label) {
   int ret = sigemptyset(set);
   dfsan_set_label(0, set, sizeof(sigset_t));
   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
 // instances of this class may live on the signal handler stack.
 //
 // DFSan uses TLS to pass metadata of arguments and return values. When an
 // instrumented function accesses the TLS, if a signal callback happens, and the
 // callback calls other instrumented functions with updating the same TLS, the
public:
   ~ScopedClearThreadLocalState() { dfsan_clear_thread_local_state(); }
 };
 
 // SignalSpinLocker::sigactions_mu guarantees atomicity of sigaction() calls.
 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,
   // and uses the trampoline to clear shadows. However sigaction does not use
   // a function pointer directly, so we have to do this manually.
   dfsan_clear_arg_tls(0, sizeof(dfsan_label));
 
   typedef void (*signal_cb)(int x);
   signal_cb cb =
       (signal_cb)atomic_load(&sigactions[signo], memory_order_relaxed);
   cb(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));
   dfsan_set_label(0, si, sizeof(*si));
   dfsan_set_label(0, uc, sizeof(ucontext_t));
 
   typedef void (*sigaction_cb)(int, siginfo_t *, void *);
   sigaction_cb cb =

compiler-rt/lib/dfsan/dfsan_thread.h

+//===-- 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/lib/dfsan/dfsan_thread.cpp

+#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/test/dfsan/pthread.c

+// 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;
+}