[MSAN] Add fiber switching APIs

Add functions exposed via the MSAN interface to enable MSAN within
binaries that perform manual stack switching (e.g. through using fibers
or coroutines).

This functionality is analogous to the fiber APIs available for ASAN and TSAN.

Fixes google/sanitizers#1232

Reviewed By: vitalybuka

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

compiler-rt/include/sanitizer/msan_interface.h

@@ -114,6 +114,9 @@ extern "C" {
      call to __msan_scoped_disable_interceptor_checks. */
   void __msan_scoped_enable_interceptor_checks(void);
 
+  void __msan_start_switch_fiber(const void *bottom, size_t size);
+  void __msan_finish_switch_fiber(const void **bottom_old, size_t *size_old);
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif

compiler-rt/lib/msan/msan.cpp

@@ -695,6 +695,37 @@ void __msan_set_death_callback(void (*callback)(void)) {
   SetUserDieCallback(callback);
 }
 
+void __msan_start_switch_fiber(const void *bottom, uptr size) {
+  MsanThread *t = GetCurrentThread();
+  if (!t) {
+    VReport(1, "__msan_start_switch_fiber called from unknown thread\n");
+    return;
+  }
+  t->StartSwitchFiber((uptr)bottom, size);
+}
+
+void __msan_finish_switch_fiber(const void **bottom_old, uptr *size_old) {
+  MsanThread *t = GetCurrentThread();
+  if (!t) {
+    VReport(1, "__msan_finish_switch_fiber called from unknown thread\n");
+    return;
+  }
+  t->FinishSwitchFiber((uptr *)bottom_old, (uptr *)size_old);
+
+  internal_memset(__msan_param_tls, 0, sizeof(__msan_param_tls));
+  internal_memset(__msan_retval_tls, 0, sizeof(__msan_retval_tls));
+  internal_memset(__msan_va_arg_tls, 0, sizeof(__msan_va_arg_tls));
+
+  if (__msan_get_track_origins()) {
+    internal_memset(__msan_param_origin_tls, 0,
+                    sizeof(__msan_param_origin_tls));
+    internal_memset(&__msan_retval_origin_tls, 0,
+                    sizeof(__msan_retval_origin_tls));
+    internal_memset(__msan_va_arg_origin_tls, 0,
+                    sizeof(__msan_va_arg_origin_tls));
+  }
+}
+
 #if !SANITIZER_SUPPORTS_WEAK_HOOKS
 extern "C" {
 SANITIZER_INTERFACE_ATTRIBUTE SANITIZER_WEAK_ATTRIBUTE

compiler-rt/lib/msan/msan_interface_internal.h

@@ -187,6 +187,12 @@ void __msan_scoped_disable_interceptor_checks();
 
 SANITIZER_INTERFACE_ATTRIBUTE
 void __msan_scoped_enable_interceptor_checks();
+
+SANITIZER_INTERFACE_ATTRIBUTE
+void __msan_start_switch_fiber(const void *bottom, uptr size);
+
+SANITIZER_INTERFACE_ATTRIBUTE
+void __msan_finish_switch_fiber(const void **bottom_old, uptr *size_old);
 }  // extern "C"
 
 #endif  // MSAN_INTERFACE_INTERNAL_H

compiler-rt/lib/msan/msan_thread.cpp

@@ -22,9 +22,9 @@ MsanThread *MsanThread::Create(thread_callback_t start_routine,
 void MsanThread::SetThreadStackAndTls() {
   uptr tls_size = 0;
   uptr stack_size = 0;
-  GetThreadStackAndTls(IsMainThread(), &stack_bottom_, &stack_size,
-                       &tls_begin_, &tls_size);
-  stack_top_ = stack_bottom_ + stack_size;
+  GetThreadStackAndTls(IsMainThread(), &stack_.bottom, &stack_size, &tls_begin_,
+                       &tls_size);
+  stack_.top = stack_.bottom + stack_size;
   tls_end_ = tls_begin_ + tls_size;
 
   int local;
@@ -32,7 +32,7 @@ void MsanThread::SetThreadStackAndTls() {
 }
 
 void MsanThread::ClearShadowForThreadStackAndTLS() {
-  __msan_unpoison((void *)stack_bottom_, stack_top_ - stack_bottom_);
+  __msan_unpoison((void *)stack_.bottom, stack_.top - stack_.bottom);
   if (tls_begin_ != tls_end_)
     __msan_unpoison((void *)tls_begin_, tls_end_ - tls_begin_);
   DTLS *dtls = DTLS_Get();
@@ -43,8 +43,8 @@ void MsanThread::ClearShadowForThreadStackAndTLS() {
 
 void MsanThread::Init() {
   SetThreadStackAndTls();
-  CHECK(MEM_IS_APP(stack_bottom_));
-  CHECK(MEM_IS_APP(stack_top_ - 1));
+  CHECK(MEM_IS_APP(stack_.bottom));
+  CHECK(MEM_IS_APP(stack_.top - 1));
   ClearShadowForThreadStackAndTLS();
 }
 
@@ -79,4 +79,45 @@ thread_return_t MsanThread::ThreadStart() {
   return res;
 }
 
+MsanThread::StackBounds MsanThread::GetStackBounds() const {
+  if (!stack_switching_)
+    return {stack_.bottom, stack_.top};
+  const uptr cur_stack = GET_CURRENT_FRAME();
+  // 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 MsanThread::stack_top() { return GetStackBounds().top; }
+
+uptr MsanThread::stack_bottom() { return GetStackBounds().bottom; }
+
+bool MsanThread::AddrIsInStack(uptr addr) {
+  const auto bounds = GetStackBounds();
+  return addr >= bounds.bottom && addr < bounds.top;
+}
+
+void MsanThread::StartSwitchFiber(uptr bottom, uptr size) {
+  CHECK(!stack_switching_);
+  next_stack_.bottom = bottom;
+  next_stack_.top = bottom + size;
+  stack_switching_ = true;
+}
+
+void MsanThread::FinishSwitchFiber(uptr *bottom_old, uptr *size_old) {
+  CHECK(stack_switching_);
+  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;
+  stack_switching_ = false;
+  next_stack_.top = 0;
+  next_stack_.bottom = 0;
+}
+
 } // namespace __msan

compiler-rt/lib/msan/msan_thread.h

@@ -27,20 +27,21 @@ class MsanThread {
   void Init();  // Should be called from the thread itself.
   thread_return_t ThreadStart();
 
-  uptr stack_top() { return stack_top_; }
-  uptr stack_bottom() { return stack_bottom_; }
+  uptr stack_top();
+  uptr stack_bottom();
   uptr tls_begin() { return tls_begin_; }
   uptr tls_end() { return tls_end_; }
   bool IsMainThread() { return start_routine_ == nullptr; }
 
-  bool AddrIsInStack(uptr addr) {
-    return addr >= stack_bottom_ && addr < stack_top_;
-  }
+  bool AddrIsInStack(uptr addr);
 
   bool InSignalHandler() { return in_signal_handler_; }
   void EnterSignalHandler() { in_signal_handler_++; }
   void LeaveSignalHandler() { in_signal_handler_--; }
 
+  void StartSwitchFiber(uptr bottom, uptr size);
+  void FinishSwitchFiber(uptr *bottom_old, uptr *size_old);
+
   MsanThreadLocalMallocStorage &malloc_storage() { return malloc_storage_; }
 
   int destructor_iterations_;
@@ -50,10 +51,19 @@ class MsanThread {
   // via mmap() and *must* be valid in zero-initialized state.
   void SetThreadStackAndTls();
   void ClearShadowForThreadStackAndTLS();
+  struct StackBounds {
+    uptr bottom;
+    uptr top;
+  };
+  StackBounds GetStackBounds() const;
   thread_callback_t start_routine_;
   void *arg_;
-  uptr stack_top_;
-  uptr stack_bottom_;
+
+  bool stack_switching_;
+
+  StackBounds stack_;
+  StackBounds next_stack_;
+
   uptr tls_begin_;
   uptr tls_end_;
 

compiler-rt/test/msan/Linux/swapcontext_annotation.cpp

@@ -0,0 +1,68 @@
+// RUN: %clangxx_msan -O0 %s -o %t && %run %t
+
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <ucontext.h>
+#include <unistd.h>
+
+#include <sanitizer/msan_interface.h>
+
+namespace {
+
+const int kStackSize = 1 << 20;
+char fiber_stack[kStackSize] = {};
+
+ucontext_t main_ctx;
+ucontext_t fiber_ctx;
+
+void fiber() {
+  printf("%s: entering fiber\n", __FUNCTION__);
+
+  // This fiber was switched into from main. Verify the details of main's stack
+  // have been populated by MSAN.
+  const void *previous_stack_bottom = nullptr;
+  size_t previous_stack_size = 0;
+  __msan_finish_switch_fiber(&previous_stack_bottom, &previous_stack_size);
+  assert(previous_stack_bottom != nullptr);
+  assert(previous_stack_size != 0);
+
+  printf("%s: implicitly swapcontext to main\n", __FUNCTION__);
+  __msan_start_switch_fiber(previous_stack_bottom, previous_stack_size);
+}
+
+} // namespace
+
+// Set up a fiber, switch to it, and switch back, invoking __msan_*_switch_fiber
+// functions along the way. At each step, validate the correct stack addresses and
+// sizes are returned from those functions.
+int main(int argc, char **argv) {
+  if (getcontext(&fiber_ctx) == -1) {
+    perror("getcontext");
+    _exit(1);
+  }
+  fiber_ctx.uc_stack.ss_sp = fiber_stack;
+  fiber_ctx.uc_stack.ss_size = sizeof(fiber_stack);
+  fiber_ctx.uc_link = &main_ctx;
+  makecontext(&fiber_ctx, fiber, 0);
+
+  // Tell MSAN a fiber switch is about to occur, then perform the switch
+  printf("%s: swapcontext to fiber\n", __FUNCTION__);
+  __msan_start_switch_fiber(fiber_stack, kStackSize);
+  if (swapcontext(&main_ctx, &fiber_ctx) == -1) {
+    perror("swapcontext");
+    _exit(1);
+  }
+
+  // The fiber switched to above now switched back here. Tell MSAN that switch
+  // is complete and verify the fiber details return by MSAN are correct.
+  const void *previous_stack_bottom = nullptr;
+  size_t previous_stack_size = 0;
+  __msan_finish_switch_fiber(&previous_stack_bottom, &previous_stack_size);
+  assert(previous_stack_bottom == fiber_stack);
+  assert(previous_stack_size == kStackSize);
+
+  printf("%s: exiting\n", __FUNCTION__);
+
+  return 0;
+}

compiler-rt/test/msan/Linux/swapcontext_annotation_reset.cpp

@@ -0,0 +1,65 @@
+// RUN: %clangxx_msan -fno-sanitize=memory -c %s -o %t-main.o
+// RUN: %clangxx_msan %t-main.o %s -o %t
+// RUN: %run %t
+
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <ucontext.h>
+#include <unistd.h>
+
+#include <sanitizer/msan_interface.h>
+
+#if __has_feature(memory_sanitizer)
+
+__attribute__((noinline)) int bar(int a, int b) {
+  volatile int zero = 0;
+  return zero;
+}
+
+void foo(int x, int y, int expected) {
+  assert(__msan_test_shadow(&x, sizeof(x)) == expected);
+  assert(__msan_test_shadow(&y, sizeof(y)) == expected);
+
+  // Poisons parameter shadow in TLS so that the next call (to foo) from
+  // uninstrumented main has params 1 and 2 poisoned no matter what.
+  int a, b;
+  (void)bar(a, b);
+}
+
+#else
+
+// This code is not instrumented by MemorySanitizer to prevent it from modifying
+// MSAN TLS data for this test.
+
+int foo(int, int, int);
+
+int main(int argc, char **argv) {
+  int x, y;
+  // The parameters should _not_ be poisoned; this is the first call to foo.
+  foo(x, y, -1);
+  // The parameters should be poisoned; the prior call to foo left them so.
+  foo(x, y, 0);
+
+  ucontext_t ctx;
+  if (getcontext(&ctx) == -1) {
+    perror("getcontext");
+    _exit(1);
+  }
+
+  // Simulate a fiber switch occurring from MSAN's perspective (though no switch
+  // actually occurs).
+  const void *previous_stack_bottom = nullptr;
+  size_t previous_stack_size = 0;
+  __msan_start_switch_fiber(ctx.uc_stack.ss_sp, ctx.uc_stack.ss_size);
+  __msan_finish_switch_fiber(&previous_stack_bottom, &previous_stack_size);
+
+  // The simulated fiber switch will reset the TLS parameter shadow. So even
+  // though the most recent call to foo left the parameter shadow poisoned, the
+  // parameters are _not_ expected to be poisoned now.
+  foo(x, y, -1);
+
+  return 0;
+}
+
+#endif