[compiler-rt][hwasan] Add fiber switch for HwASan

compiler-rt/lib/hwasan/hwasan_interface_internal.h

@@ -247,6 +247,13 @@ void *__hwasan_memmove_match_all(void *dest, const void *src, uptr n, u8);
 
 SANITIZER_INTERFACE_ATTRIBUTE
 void __hwasan_set_error_report_callback(void (*callback)(const char *));
+
+// hwasan does not need fake stack, so we leave it empty here.
+SANITIZER_INTERFACE_ATTRIBUTE
+void __sanitizer_start_switch_fiber(void **, const void *bottom, uptr size);
+SANITIZER_INTERFACE_ATTRIBUTE
+void __sanitizer_finish_switch_fiber(void *, const void **bottom_old,
+                                     uptr *size_old);
 }  // extern "C"
 
 #endif  // HWASAN_INTERFACE_INTERNAL_H

compiler-rt/lib/hwasan/hwasan_thread.cpp

@@ -119,6 +119,59 @@ void Thread::Destroy() {
   *GetCurrentThreadLongPtr() = 0;
 }
 
+void Thread::StartSwitchFiber(uptr bottom, uptr size) {
+  if (atomic_load(&stack_switching_, memory_order_acquire)) {
+    Report("ERROR: starting fiber switch while in fiber switch\n");
+    Die();
+  }
+
+  next_stack_bottom_ = bottom;
+  next_stack_top_ = bottom + size;
+  atomic_store(&stack_switching_, 1, memory_order_release);
+}
+
+void Thread::FinishSwitchFiber(uptr *bottom_old, uptr *size_old) {
+  if (!atomic_load(&stack_switching_, memory_order_acquire)) {
+    Report("ERROR: finishing a fiber switch that has not started\n");
+    Die();
+  }
+
+  if (bottom_old)
+    *bottom_old = stack_bottom_;
+  if (size_old)
+    *size_old = stack_top_ - stack_bottom_;
+  stack_bottom_ = next_stack_bottom_;
+  stack_top_ = next_stack_top_;
+  atomic_store(&stack_switching_, 0, memory_order_release);
+  next_stack_top_ = 0;
+  next_stack_bottom_ = 0;
+}
+
+inline Thread::StackBounds Thread::GetStackBounds() const {
+  if (!atomic_load(&stack_switching_, memory_order_acquire)) {
+    // Make sure the stack bounds are fully initialized.
+    if (stack_bottom_ >= stack_top_)
+      return {0, 0};
+    return {stack_bottom_, stack_top_};
+  }
+  const uptr cur_stack = (uptr)__builtin_frame_address(0);
+  // Note: need to check next stack first, because FinishSwitchFiber
+  // may be in process of overwriting stack_top_/bottom_. But in such case
+  // we are already on the next stack.
+  if (cur_stack >= next_stack_bottom_ && cur_stack < next_stack_top_)
+    return {next_stack_bottom_, next_stack_top_};
+  return {stack_bottom_, stack_top_};
+}
+
+uptr Thread::stack_top() { return GetStackBounds().top; }
+
+uptr Thread::stack_bottom() { return GetStackBounds().bottom; }
+
+uptr Thread::stack_size() {
+  const auto bounds = GetStackBounds();
+  return bounds.top - bounds.bottom;
+}
+
 void Thread::Print(const char *Prefix) {
   Printf("%sT%zd %p stack: [%p,%p) sz: %zd tls: [%p,%p)\n", Prefix, unique_id_,
          (void *)this, stack_bottom(), stack_top(),
@@ -226,3 +279,25 @@ void PrintThreads() {
 }
 
 }  // namespace __lsan
+
+// ---------------------- Interface ---------------- {{{1
+using namespace __hwasan;
+
+extern "C" {
+SANITIZER_INTERFACE_ATTRIBUTE
+void __sanitizer_start_switch_fiber(void **, const void *bottom, uptr size) {
+  if (auto *t = GetCurrentThread())
+    t->StartSwitchFiber((uptr)bottom, size);
+  else
+    VReport(1, "__hwasan_start_switch_fiber called from unknown thread\n");
+}
+
+SANITIZER_INTERFACE_ATTRIBUTE
+void __sanitizer_finish_switch_fiber(void *, const void **bottom_old,
+                                     uptr *size_old) {
+  if (auto *t = GetCurrentThread())
+    t->FinishSwitchFiber((uptr *)bottom_old, size_old);
+  else
+    VReport(1, "__hwasan_finish_switch_fiber called from unknown thread\n");
+}
+}

compiler-rt/lib/hwasan/hwasan_thread.h

@@ -41,9 +41,9 @@ class Thread {
 
   void Destroy();
 
-  uptr stack_top() { return stack_top_; }
-  uptr stack_bottom() { return stack_bottom_; }
-  uptr stack_size() { return stack_top() - stack_bottom(); }
+  uptr stack_top();
+  uptr stack_bottom();
+  uptr stack_size();
   uptr tls_begin() { return tls_begin_; }
   uptr tls_end() { return tls_end_; }
   DTLS *dtls() { return dtls_; }
@@ -53,6 +53,9 @@ class Thread {
     return addr >= stack_bottom_ && addr < stack_top_;
   }
 
+  void StartSwitchFiber(uptr bottom, uptr size);
+  void FinishSwitchFiber(uptr *bottom_old, uptr *size_old);
+
   AllocatorCache *allocator_cache() { return &allocator_cache_; }
   HeapAllocationsRingBuffer *heap_allocations() { return heap_allocations_; }
   StackAllocationsRingBuffer *stack_allocations() { return stack_allocations_; }
@@ -80,9 +83,22 @@ class Thread {
   void ClearShadowForThreadStackAndTLS();
   void Print(const char *prefix);
   void InitRandomState();
+
+  struct StackBounds {
+    uptr bottom;
+    uptr top;
+  };
+  StackBounds GetStackBounds() const;
+
   uptr vfork_spill_;
   uptr stack_top_;
   uptr stack_bottom_;
+  // these variables are used when the thread is about to switch stack
+  uptr next_stack_top_;
+  uptr next_stack_bottom_;
+  // true if switching is in progress
+  atomic_uint8_t stack_switching_;
+
   uptr tls_begin_;
   uptr tls_end_;
   DTLS *dtls_;

compiler-rt/test/hwasan/TestCases/Linux/swapcontext_annotation.cpp

@@ -0,0 +1,197 @@
+// Test hwasan __sanitizer_start_switch_fiber and __sanitizer_finish_switch_fiber interface.
+
+// RUN: %clangxx_hwasan -std=c++11 -lpthread -O0 %s -o %t && %run %t 2>&1 | FileCheck %s
+// RUN: %clangxx_hwasan -std=c++11 -lpthread -O1 %s -o %t && %run %t 2>&1 | FileCheck %s
+// RUN: %clangxx_hwasan -std=c++11 -lpthread -O2 %s -o %t && %run %t 2>&1 | FileCheck %s
+// RUN: %clangxx_hwasan -std=c++11 -lpthread -O3 %s -o %t && %run %t 2>&1 | FileCheck %s
+// RUN: seq 30 | xargs -i -- grep LOOPCHECK %s > %t.checks
+// RUN: %clangxx_hwasan -std=c++11 -lpthread -O0 %s -o %t && %run %t 2>&1 | FileCheck %t.checks --check-prefix LOOPCHECK
+// RUN: %clangxx_hwasan -std=c++11 -lpthread -O1 %s -o %t && %run %t 2>&1 | FileCheck %t.checks --check-prefix LOOPCHECK
+// RUN: %clangxx_hwasan -std=c++11 -lpthread -O2 %s -o %t && %run %t 2>&1 | FileCheck %t.checks --check-prefix LOOPCHECK
+// RUN: %clangxx_hwasan -std=c++11 -lpthread -O3 %s -o %t && %run %t 2>&1 | FileCheck %t.checks --check-prefix LOOPCHECK
+
+//
+// Android and musl do not support swapcontext.
+// REQUIRES: glibc-2.27
+
+#include <pthread.h>
+#include <setjmp.h>
+#include <signal.h>
+#include <stdio.h>
+#include <sys/time.h>
+#include <ucontext.h>
+#include <unistd.h>
+
+#include <sanitizer/common_interface_defs.h>
+
+ucontext_t orig_context;
+ucontext_t child_context;
+ucontext_t next_child_context;
+
+char *next_child_stack;
+
+const int kStackSize = 1 << 20;
+
+const void *main_thread_stack;
+size_t main_thread_stacksize;
+
+const void *from_stack;
+size_t from_stacksize;
+
+// hwasan does not support longjmp with tagged stack pointer, make sure it is not tagged.
+char __attribute__((no_sanitize("hwaddress"))) allocated_stack[kStackSize + 1];
+char __attribute__((no_sanitize("hwaddress"))) allocated_child_stack[kStackSize + 1];
+
+__attribute__((noinline, noreturn)) void LongJump(jmp_buf env) {
+  longjmp(env, 1);
+  _exit(1);
+}
+
+// Simulate __asan_handle_no_return().
+__attribute__((noinline)) void CallNoReturn() {
+  jmp_buf env;
+  if (setjmp(env) != 0)
+    return;
+
+  LongJump(env);
+  _exit(1);
+}
+
+void NextChild() {
+  CallNoReturn();
+  __sanitizer_finish_switch_fiber(nullptr, &from_stack, &from_stacksize);
+
+  printf("NextChild from: %p %zu\n", from_stack, from_stacksize);
+
+  char x[32] = {0}; // Stack gets poisoned.
+  printf("NextChild: %p\n", x);
+
+  CallNoReturn();
+
+  __sanitizer_start_switch_fiber(nullptr, main_thread_stack,
+                                 main_thread_stacksize);
+  CallNoReturn();
+  if (swapcontext(&next_child_context, &orig_context) < 0) {
+    perror("swapcontext");
+    _exit(1);
+  }
+}
+
+void Child(int mode) {
+  CallNoReturn();
+  __sanitizer_finish_switch_fiber(nullptr, &main_thread_stack,
+                                  &main_thread_stacksize);
+  char x[32] = {0}; // Stack gets poisoned.
+  printf("Child: %p\n", x);
+  CallNoReturn();
+  // (a) Do nothing, just return to parent function.
+  // (b) Jump into the original function. Stack remains poisoned unless we do
+  //     something.
+  // (c) Jump to another function which will then jump back to the main function
+  if (mode == 0) {
+    __sanitizer_start_switch_fiber(nullptr, main_thread_stack,
+                                   main_thread_stacksize);
+    CallNoReturn();
+  } else if (mode == 1) {
+    __sanitizer_start_switch_fiber(nullptr, main_thread_stack,
+                                   main_thread_stacksize);
+    CallNoReturn();
+    if (swapcontext(&child_context, &orig_context) < 0) {
+      perror("swapcontext");
+      _exit(1);
+    }
+  } else if (mode == 2) {
+    printf("NextChild stack: %p\n", next_child_stack);
+
+    getcontext(&next_child_context);
+    next_child_context.uc_stack.ss_sp = next_child_stack;
+    next_child_context.uc_stack.ss_size = kStackSize / 2;
+    makecontext(&next_child_context, (void (*)())NextChild, 0);
+    __sanitizer_start_switch_fiber(nullptr, next_child_context.uc_stack.ss_sp,
+                                   next_child_context.uc_stack.ss_size);
+    CallNoReturn();
+    if (swapcontext(&child_context, &next_child_context) < 0) {
+      perror("swapcontext");
+      _exit(1);
+    }
+  }
+}
+
+int Run(int arg, int mode, char *child_stack) {
+  printf("Child stack: %p\n", child_stack);
+  // Setup child context.
+  getcontext(&child_context);
+  child_context.uc_stack.ss_sp = child_stack;
+  child_context.uc_stack.ss_size = kStackSize / 2;
+  if (mode == 0) {
+    child_context.uc_link = &orig_context;
+  }
+  makecontext(&child_context, (void (*)())Child, 1, mode);
+  CallNoReturn();
+  __sanitizer_start_switch_fiber(nullptr, child_context.uc_stack.ss_sp,
+                                 child_context.uc_stack.ss_size);
+  CallNoReturn();
+  if (swapcontext(&orig_context, &child_context) < 0) {
+    perror("swapcontext");
+    _exit(1);
+  }
+  CallNoReturn();
+  __sanitizer_finish_switch_fiber(nullptr, &from_stack, &from_stacksize);
+  CallNoReturn();
+  printf("Main context from: %p %zu\n", from_stack, from_stacksize);
+
+  return child_stack[arg];
+}
+
+void handler(int sig) { CallNoReturn(); }
+
+int main(int argc, char **argv) {
+  // This testcase is copied from ASan's swapcontext_annotation.cpp testcase
+  // and adapted to HWASan:
+  // 1. removed huge stack test since hwasan has no huge stack limitations
+  // 2. stack allocations are now done with original malloc/free instead of
+  //   hwasan interceptor, since HWASan does not support tagged stack pointer
+  //   in longjmp (see __hwasan_handle_longjmp)
+
+  // set up a signal that will spam and trigger __hwasan_handle_vfork at
+  // tricky moments
+  struct sigaction act = {};
+  act.sa_handler = &handler;
+  if (sigaction(SIGPROF, &act, 0)) {
+    perror("sigaction");
+    _exit(1);
+  }
+
+  itimerval t;
+  t.it_interval.tv_sec = 0;
+  t.it_interval.tv_usec = 10;
+  t.it_value = t.it_interval;
+  if (setitimer(ITIMER_PROF, &t, 0)) {
+    perror("setitimer");
+    _exit(1);
+  }
+
+  char *heap = allocated_stack;
+  next_child_stack = allocated_child_stack;
+  int ret = 0;
+  // CHECK-NOT: WARNING: HWASan is ignoring requested __hwasan_handle_vfork
+  for (unsigned int i = 0; i < 30; ++i) {
+    // LOOPCHECK: Child stack: [[CHILD_STACK:0x[0-9a-f]*]]
+    // LOOPCHECK: Main context from: [[CHILD_STACK]] 524288
+    ret += Run(argc - 1, 0, heap);
+    // LOOPCHECK: Child stack: [[CHILD_STACK:0x[0-9a-f]*]]
+    // LOOPCHECK: Main context from: [[CHILD_STACK]] 524288
+    ret += Run(argc - 1, 1, heap);
+    // LOOPCHECK: Child stack: [[CHILD_STACK:0x[0-9a-f]*]]
+    // LOOPCHECK: NextChild stack: [[NEXT_CHILD_STACK:0x[0-9a-f]*]]
+    // LOOPCHECK: NextChild from: [[CHILD_STACK]] 524288
+    // LOOPCHECK: Main context from: [[NEXT_CHILD_STACK]] 524288
+    ret += Run(argc - 1, 2, heap);
+    printf("Iteration %d passed\n", i);
+  }
+
+  // CHECK: Test passed
+  printf("Test passed\n");
+
+  return ret;
+}