2 files changed, 126 insertions, 13 deletions

diff --git a/lib/asan/asan_thread.cc b/lib/asan/asan_thread.cc
index 526ef3dc5..d7e2cca65 100644
--- a/lib/asan/asan_thread.cc
+++ b/lib/asan/asan_thread.cc
@@ -120,6 +120,71 @@ void AsanThread::Destroy() {
   DTLS_Destroy();
 }
 
+void AsanThread::StartSwitchFiber(FakeStack **fake_stack_save, 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);
+
+  FakeStack *current_fake_stack = fake_stack_;
+  if (fake_stack_save)
+    *fake_stack_save = fake_stack_;
+  fake_stack_ = nullptr;
+  SetTLSFakeStack(nullptr);
+  // if fake_stack_save is null, the fiber will die, delete the fakestack
+  if (!fake_stack_save && current_fake_stack)
+    current_fake_stack->Destroy(this->tid());
+}
+
+void AsanThread::FinishSwitchFiber(FakeStack *fake_stack_save) {
+  if (!atomic_load(&stack_switching_, memory_order_relaxed)) {
+    Report("ERROR: finishing a fiber switch that has not started\n");
+    Die();
+  }
+
+  if (fake_stack_save) {
+    SetTLSFakeStack(fake_stack_save);
+    fake_stack_ = fake_stack_save;
+  }
+
+  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 AsanThread::StackBounds AsanThread::GetStackBounds() const {
+  if (!atomic_load(&stack_switching_, memory_order_acquire))
+    return StackBounds{stack_bottom_, stack_top_};  // NOLINT
+  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 StackBounds{next_stack_bottom_, next_stack_top_};  // NOLINT
+  return StackBounds{stack_bottom_, stack_top_};              // NOLINT
+}
+
+uptr AsanThread::stack_top() {
+  return GetStackBounds().top;
+}
+
+uptr AsanThread::stack_bottom() {
+  return GetStackBounds().bottom;
+}
+
+uptr AsanThread::stack_size() {
+  const auto bounds = GetStackBounds();
+  return bounds.top - bounds.bottom;
+}
+
 // We want to create the FakeStack lazyly on the first use, but not eralier
 // than the stack size is known and the procedure has to be async-signal safe.
 FakeStack *AsanThread::AsyncSignalSafeLazyInitFakeStack() {
@@ -150,6 +215,8 @@ FakeStack *AsanThread::AsyncSignalSafeLazyInitFakeStack() {
 }
 
 void AsanThread::Init() {
+  next_stack_top_ = next_stack_bottom_ = 0;
+  atomic_store(&stack_switching_, false, memory_order_release);
   fake_stack_ = nullptr;  // Will be initialized lazily if needed.
   CHECK_EQ(this->stack_size(), 0U);
   SetThreadStackAndTls();
@@ -195,9 +262,10 @@ thread_return_t AsanThread::ThreadStart(
 
 void AsanThread::SetThreadStackAndTls() {
   uptr tls_size = 0;
-  GetThreadStackAndTls(tid() == 0, &stack_bottom_, &stack_size_, &tls_begin_,
-                       &tls_size);
-  stack_top_ = stack_bottom_ + stack_size_;
+  uptr stack_size = 0;
+  GetThreadStackAndTls(tid() == 0, const_cast<uptr *>(&stack_bottom_),
+                       const_cast<uptr *>(&stack_size), &tls_begin_, &tls_size);
+  stack_top_ = stack_bottom_ + stack_size;
   tls_end_ = tls_begin_ + tls_size;
   dtls_ = DTLS_Get();
 
@@ -250,6 +318,11 @@ bool AsanThread::GetStackFrameAccessByAddr(uptr addr,
   return true;
 }
 
+bool AsanThread::AddrIsInStack(uptr addr) {
+  const auto bounds = GetStackBounds();
+  return addr >= bounds.bottom && addr < bounds.top;
+}
+
 static bool ThreadStackContainsAddress(ThreadContextBase *tctx_base,
                                        void *addr) {
   AsanThreadContext *tctx = static_cast<AsanThreadContext*>(tctx_base);
@@ -357,3 +430,29 @@ void EnsureMainThreadIDIsCorrect() {
   __asan::EnsureMainThreadIDIsCorrect();
 }
 } // namespace __lsan
+
+// ---------------------- Interface ---------------- {{{1
+using namespace __asan;  // NOLINT
+
+extern "C" {
+SANITIZER_INTERFACE_ATTRIBUTE
+void __sanitizer_start_switch_fiber(void **fakestacksave, const void *bottom,
+                                    uptr size) {
+  AsanThread *t = GetCurrentThread();
+  if (!t) {
+    VReport(1, "__asan_start_switch_fiber called from unknown thread\n");
+    return;
+  }
+  t->StartSwitchFiber((FakeStack**)fakestacksave, (uptr)bottom, size);
+}
+
+SANITIZER_INTERFACE_ATTRIBUTE
+void __sanitizer_finish_switch_fiber(void* fakestack) {
+  AsanThread *t = GetCurrentThread();
+  if (!t) {
+    VReport(1, "__asan_finish_switch_fiber called from unknown thread\n");
+    return;
+  }
+  t->FinishSwitchFiber((FakeStack*)fakestack);
+}
+}
diff --git a/lib/asan/asan_thread.h b/lib/asan/asan_thread.h
index b05d720d6..92a92a2e8 100644
--- a/lib/asan/asan_thread.h
+++ b/lib/asan/asan_thread.h
@@ -66,9 +66,9 @@ class AsanThread {
   thread_return_t ThreadStart(uptr os_id,
                               atomic_uintptr_t *signal_thread_is_registered);
 
-  uptr stack_top() { return stack_top_; }
-  uptr stack_bottom() { return stack_bottom_; }
-  uptr stack_size() { return stack_size_; }
+  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_; }
@@ -83,9 +83,7 @@ class AsanThread {
   };
   bool GetStackFrameAccessByAddr(uptr addr, StackFrameAccess *access);
 
-  bool AddrIsInStack(uptr addr) {
-    return addr >= stack_bottom_ && addr < stack_top_;
-  }
+  bool AddrIsInStack(uptr addr);
 
   void DeleteFakeStack(int tid) {
     if (!fake_stack_) return;
@@ -95,13 +93,19 @@ class AsanThread {
     t->Destroy(tid);
   }
 
+  void StartSwitchFiber(FakeStack **fake_stack_save, uptr bottom, uptr size);
+  void FinishSwitchFiber(FakeStack *fake_stack_save);
+
   bool has_fake_stack() {
-    return (reinterpret_cast<uptr>(fake_stack_) > 1);
+    return !atomic_load(&stack_switching_, memory_order_relaxed) &&
+           (reinterpret_cast<uptr>(fake_stack_) > 1);
   }
 
   FakeStack *fake_stack() {
     if (!__asan_option_detect_stack_use_after_return)
       return nullptr;
+    if (atomic_load(&stack_switching_, memory_order_relaxed))
+      return nullptr;
     if (!has_fake_stack())
       return AsyncSignalSafeLazyInitFakeStack();
     return fake_stack_;
@@ -127,14 +131,24 @@ class AsanThread {
   void ClearShadowForThreadStackAndTLS();
   FakeStack *AsyncSignalSafeLazyInitFakeStack();
 
+  struct StackBounds {
+    uptr bottom;
+    uptr top;
+  };
+  StackBounds GetStackBounds() const;
+
   AsanThreadContext *context_;
   thread_callback_t start_routine_;
   void *arg_;
+
   uptr stack_top_;
   uptr stack_bottom_;
-  // stack_size_ == stack_top_ - stack_bottom_;
-  // It needs to be set in a async-signal-safe manner.
-  uptr stack_size_;
+  // 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_;