[MSAN] Add fiber switching APIs

[justincady]

Summary

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

[vitalybuka]

compiler-rt/lib/msan/msan_thread.cpp

99

Please drop {} for single line if/for for consistency

114

if we don't expect multiple threads access this field, then we don't need atomic
if we do expect that, this atomic does not solve data race.

From what I see stack_top() stack_bottom() only used from BufferedStackTrace::UnwindImpl for current thread.

So please remove remove stack_switching_ and simplify GetStackBounds() (next_stack_top_ == next_stack_bottom_ should be enought? )

compiler-rt/lib/msan/msan_thread.h

64–65

[justincady]

compiler-rt/lib/msan/msan_thread.cpp

114

Yes, you are correct. We do not expect multiple threads here currently. But I don't think this is attempting to prevent a data race (I used this pattern after seeing it in the ASAN implementation). My intent, and I believe the original intent in ASAN, is to prevent incorrect program logic and make sure error reports are always correct.

Regarding incorrect program logic, consider this sequence from a single OS thread:

1. __msan_start_switch_fiber() is invoked
2. __msan_start_switch_fiber() is invoked again

(Or similarly any sequence of events in which start and finish are not paired properly.)

In those cases, the atomic is the breadcrumb that lets MSAN terminate the program immediately and inform the user there was a bad sequence of calls (at the top of StartSwitchFiber and FinishSwitchFiber).

Regarding correct error reports, consider this sequence from a single OS thread:

1. __msan_start_switch_fiber() is invoked
2. Something bad happens before __msan_finish_switch_fiber() is invoked
3. An error report with stack trace is generated

Note that Step 2 is basically never expected to happen. But, if it does, Steps 2 and 3 could occur either before of after the actual fiber switch. The atomic is again used as a breadcrumb to make sure that the error report shows the correct stack, be it pre- or post- fiber switch (in GetStackBounds). If stack_switching_ is 1, we need to create a local stack variable and check if it is in the "old" stack or "new" stack, returning the appropriate stack addresses based on that check.

I think the confusing part is the comment in GetStackBounds that I also brought over from asan_thread.cpp, which incorrectly indicates a data race is being protected against.

But your comment also points out that even if all of the above is true, we don't necessarily need the breadcrumb to be an atomic. It could just be a normal bool, since we don't expect access by multiple threads.

Of course, if I am misunderstanding something please let me know. How would you like to proceed?

[vitalybuka]

compiler-rt/lib/msan/msan_thread.cpp

45–47

This should go into ::Create()
However Create uses MmapOrDie and implicitly reset everything to 0.

89–91

> is not possible
== only for 0,0

94

114

Lsan in Asan can inspect stacks from other threads, so I see why there is atomic.

114–117

maybe

124–127

maybe

128–133

same for the rest

compiler-rt/lib/msan/msan_thread.h

62

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

41–44

I don't understand why the second one is poisoned

[justincady]

compiler-rt/lib/msan/msan_thread.cpp

114

Ah, I see. Thanks for clarifying!

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

41–44

Because main is unsanitized, calls from main do not update the TLS. But foo() is sanitized, so when foo() calls bar() it does update the TLS.

The TLS is left in that state when we make the second foo() call from main. Line 44 exists to validate that understanding before the actual test below, which should result in the TLS being cleared.

I got the half-sanitized trick from here: compiler-rt/test/msan/unpoison_param.cpp

[justincady]

Address outstanding comments. I have re-run all of the tests with the new changes.

[vitalybuka]

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

4

I guess this requirement is copy/pasted, if so just remove it
msan has a different set of supported platforms, so it's likely not needed here

same in another test

[justincady]

Remove required test architectures from unit tests.

[justincady]

@vitalybuka Thanks for the review. I updated the diff with your recommended changes for the unit tests.

Could you please land this for me? If I should contact someone else for that or we should wait for another review first, please let me know. (I'm assuming arc land isn't going to work, since I don't have commit access.)

Thanks again!

compiler-rt/include/sanitizer/msan_interface.h

#endif
 
   /* Disables uninitialized memory checks in interceptors. */
   void __msan_scoped_disable_interceptor_checks(void);
 
   /* Re-enables uninitialized memory checks in interceptors after a previous
      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
 
 #endif

compiler-rt/lib/msan/msan.cpp

if (uu32 o = __msan_param_origin_tls[2])
       SetOriginIfPoisoned((uptr)p, (uptr)&s, sizeof(s), o);
   *p = x;
 }
 
 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
 const char* __msan_default_options() { return ""; }
 }  // extern "C"
 #endif
 
 extern "C" {
 SANITIZER_INTERFACE_ATTRIBUTE
 void __sanitizer_print_stack_trace() {
   GET_FATAL_STACK_TRACE_PC_BP(StackTrace::GetCurrentPc(), GET_CURRENT_FRAME());
   stack.Print();
 }
 } // extern "C"

compiler-rt/lib/msan/msan_interface_internal.h

 SANITIZER_INTERFACE_ATTRIBUTE
 void __msan_copy_shadow(void *dst, const void *src, uptr size);
 
 SANITIZER_INTERFACE_ATTRIBUTE
 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.h

  public:
   static MsanThread *Create(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() { 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_;
 
  private:
   // NOTE: There is no MsanThread constructor. It is allocated
   // 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_;
+
+  atomic_uint8_t stack_switching_;
+
   uptr stack_top_;
   uptr stack_bottom_;
+
+  uptr next_stack_top_;
+  uptr next_stack_bottom_;
+
   uptr tls_begin_;
   uptr tls_end_;
 
   unsigned in_signal_handler_;
 
   MsanThreadLocalMallocStorage malloc_storage_;
 };
 
 MsanThread *GetCurrentThread();
 void SetCurrentThread(MsanThread *t);
 
 } // namespace __msan
 
 #endif // MSAN_THREAD_H

compiler-rt/lib/msan/msan_thread.cpp

if (tls_begin_ != tls_end_)
     __msan_unpoison((void *)tls_begin_, tls_end_ - tls_begin_);
   DTLS *dtls = DTLS_Get();
   CHECK_NE(dtls, 0);
   for (uptr i = 0; i < dtls->dtv_size; ++i)
     __msan_unpoison((void *)(dtls->dtv[i].beg), dtls->dtv[i].size);
 }
 
 void MsanThread::Init() {
+  atomic_store(&stack_switching_, 0, memory_order_release);
+  next_stack_top_ = 0;
+  next_stack_bottom_ = 0;
+
   SetThreadStackAndTls();
   CHECK(MEM_IS_APP(stack_bottom_));
   CHECK(MEM_IS_APP(stack_top_ - 1));
   ClearShadowForThreadStackAndTLS();
 }
 
 void MsanThread::TSDDtor(void *tsd) {
   MsanThread *t = (MsanThread*)tsd;
if (!start_routine_) {
     return 0;
   }
 
   thread_return_t res = start_routine_(arg_);
 
   return res;
 }
 
+MsanThread::StackBounds MsanThread::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_};
+  }
+  char local;
+  const uptr cur_stack = (uptr)&local;
+  // 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) {
+  if (atomic_load(&stack_switching_, memory_order_relaxed)) {
+    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 MsanThread::FinishSwitchFiber(uptr *bottom_old, uptr *size_old) {
+  if (!atomic_load(&stack_switching_, memory_order_relaxed)) {
+    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;
+}
+
 } // namespace __msan

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

+// RUN: %clangxx_msan -O0 %s -o %t && %run %t
+
+// REQUIRES: x86-target-arch && !android
+
+#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

+// RUN: %clangxx_msan -fno-sanitize=memory -c %s -o %t-main.o
+// RUN: %clangxx_msan %t-main.o %s -o %t
+// RUN: %run %t
+
+// REQUIRES: x86-target-arch && !android
+
+#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