1 files changed, 106 insertions, 5 deletions

diff --git a/lib/asan/asan_thread.cc b/lib/asan/asan_thread.cc
index 69813546f551..d7e2cca65660 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,10 +262,12 @@ 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();
 
   int local;
   CHECK(AddrIsInStack((uptr)&local));
@@ -249,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);
@@ -322,8 +396,8 @@ __asan::AsanThread *GetAsanThreadByOsIDLocked(uptr os_id) {
 // --- Implementation of LSan-specific functions --- {{{1
 namespace __lsan {
 bool GetThreadRangesLocked(uptr os_id, uptr *stack_begin, uptr *stack_end,
-                           uptr *tls_begin, uptr *tls_end,
-                           uptr *cache_begin, uptr *cache_end) {
+                           uptr *tls_begin, uptr *tls_end, uptr *cache_begin,
+                           uptr *cache_end, DTLS **dtls) {
   __asan::AsanThread *t = __asan::GetAsanThreadByOsIDLocked(os_id);
   if (!t) return false;
   *stack_begin = t->stack_bottom();
@@ -333,6 +407,7 @@ bool GetThreadRangesLocked(uptr os_id, uptr *stack_begin, uptr *stack_end,
   // ASan doesn't keep allocator caches in TLS, so these are unused.
   *cache_begin = 0;
   *cache_end = 0;
+  *dtls = t->dtls();
   return true;
 }
 
@@ -355,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);
+}
+}