[compiler-rt] 1d3ef5f - [MSAN] Add fiber switching APIs

Vitaly Buka via llvm-commits llvm-commits at lists.llvm.org
Thu Aug 27 19:30:52 PDT 2020


Author: Justin Cady
Date: 2020-08-27T19:30:40-07:00
New Revision: 1d3ef5f122fe0b25283b7d72d5664474eabfb61e

URL: https://github.com/llvm/llvm-project/commit/1d3ef5f122fe0b25283b7d72d5664474eabfb61e
DIFF: https://github.com/llvm/llvm-project/commit/1d3ef5f122fe0b25283b7d72d5664474eabfb61e.diff

LOG: [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

Added: 
    compiler-rt/test/msan/Linux/swapcontext_annotation.cpp
    compiler-rt/test/msan/Linux/swapcontext_annotation_reset.cpp

Modified: 
    compiler-rt/include/sanitizer/msan_interface.h
    compiler-rt/lib/msan/msan.cpp
    compiler-rt/lib/msan/msan_interface_internal.h
    compiler-rt/lib/msan/msan_thread.cpp
    compiler-rt/lib/msan/msan_thread.h

Removed: 
    


################################################################################
diff  --git a/compiler-rt/include/sanitizer/msan_interface.h b/compiler-rt/include/sanitizer/msan_interface.h
index d40c556a46d9..eeb39fbed8b4 100644
--- a/compiler-rt/include/sanitizer/msan_interface.h
+++ b/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

diff  --git a/compiler-rt/lib/msan/msan.cpp b/compiler-rt/lib/msan/msan.cpp
index a1ad5c4f1abc..3028f79f041c 100644
--- a/compiler-rt/lib/msan/msan.cpp
+++ b/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

diff  --git a/compiler-rt/lib/msan/msan_interface_internal.h b/compiler-rt/lib/msan/msan_interface_internal.h
index 9e3db06bd64d..17922a888b9c 100644
--- a/compiler-rt/lib/msan/msan_interface_internal.h
+++ b/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

diff  --git a/compiler-rt/lib/msan/msan_thread.cpp b/compiler-rt/lib/msan/msan_thread.cpp
index 0ba499350064..456901c6789b 100644
--- a/compiler-rt/lib/msan/msan_thread.cpp
+++ b/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

diff  --git a/compiler-rt/lib/msan/msan_thread.h b/compiler-rt/lib/msan/msan_thread.h
index 808780cd57b9..fe795e3a547a 100644
--- a/compiler-rt/lib/msan/msan_thread.h
+++ b/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_;
 

diff  --git a/compiler-rt/test/msan/Linux/swapcontext_annotation.cpp b/compiler-rt/test/msan/Linux/swapcontext_annotation.cpp
new file mode 100644
index 000000000000..16bdd28fdec9
--- /dev/null
+++ b/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;
+}

diff  --git a/compiler-rt/test/msan/Linux/swapcontext_annotation_reset.cpp b/compiler-rt/test/msan/Linux/swapcontext_annotation_reset.cpp
new file mode 100644
index 000000000000..342ef735dacc
--- /dev/null
+++ b/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


        


More information about the llvm-commits mailing list