[compiler-rt][hwasan] Add fiber switch for HwASan (#153822)

Currently HwASan has no fiber switch interface for coroutines. This PR
adds fiber switch interfaces similar to ASan which helps to pass sp
check correctly on unwinding.

The only difference is HwASan does not need a fake stack since tags can
do the same thing (e.g., detect UAR). Interfaces are made identical with
ASan's.

Also adds unit test which is similar to ASan with minor adjustments:

1. change `__asan_handle_no_return` to `__hwasan_handle_vfork`
2. remove huge stack test since `__hwasan_handle_vfork` has no stack
size limitation.
3. use uninstrumented globals to simulate allocation since hwasan do not
support tagged pointer while using `longjmp`

The testcase is tested on both x86 with alias mode enabled and aarch64.

Author: Tomahawkd

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