[llvm-dev] [PATCH compiler-rt 4/5] asan: better support for swapcontext() and setcontext()

[Paweł Dziepak via llvm-dev]

This patch makes swapcontext() and setcontext() update stack information
in AsanThread each time the stack is changed. It is still far from full
support of these functions but at least makes __asan_handle_no_return()
work correctly on a custom user stack.
---
 lib/asan/asan_interceptors.cc | 30 ++++++++++++++++++++++++++++--
 lib/asan/asan_thread.cc       | 12 ++++++++----
 lib/asan/asan_thread.h        | 32 ++++++++++++++++++++++++++++----
 3 files changed, 64 insertions(+), 10 deletions(-)

diff --git a/lib/asan/asan_interceptors.cc b/lib/asan/asan_interceptors.cc
index 6b1aa00..62ecc7c 100644
--- a/lib/asan/asan_interceptors.cc
+++ b/lib/asan/asan_interceptors.cc
@@ -347,11 +347,17 @@ INTERCEPTOR(int, swapcontext, struct ucontext_t *oucp,
   ReadContextStack(ucp, &stack, &ssize);
   ClearShadowMemoryForContextStack(stack, ssize);
   AsanThread *curr_thread = GetCurrentThread();
+  uptr old_stack, old_ssize;
   if (curr_thread) {
-    WriteContextStack(oucp, curr_thread->stack_bottom(),
-                      curr_thread->stack_size());
+    old_stack = curr_thread->stack_bottom();
+    old_ssize = curr_thread->stack_size();
+    WriteContextStack(oucp, old_stack, old_ssize);
+    curr_thread->SetUserStack(stack, ssize);
   }
   int res = REAL(swapcontext)(oucp, ucp);
+  if (curr_thread) {
+    curr_thread->SetUserStack(old_stack, old_ssize);
+  }
   // swapcontext technically does not return, but program may swap context to
   // "oucp" later, that would look as if swapcontext() returned 0.
   // We need to clear shadow for ucp once again, as it may be in arbitrary
@@ -360,6 +366,25 @@ INTERCEPTOR(int, swapcontext, struct ucontext_t *oucp,
   return res;
 }
 
+INTERCEPTOR(int, setcontext, struct ucontext_t *ucp) {
+  uptr stack, ssize;
+  ReadContextStack(ucp, &stack, &ssize);
+  ClearShadowMemoryForContextStack(stack, ssize);
+
+  AsanThread *curr_thread = GetCurrentThread();
+  uptr old_stack, old_ssize;
+  if (curr_thread) {
+      old_stack = curr_thread->stack_bottom();
+      old_ssize = curr_thread->stack_size();
+      curr_thread->SetUserStack(stack, ssize);
+  }
+  int res = REAL(setcontext)(ucp);
+  if (curr_thread) {
+    curr_thread->SetUserStack(old_stack, old_ssize);
+  }
+  return res;
+}
+
 INTERCEPTOR(int, getcontext, struct ucontext_t *ucp) {
   int res = REAL(getcontext)(ucp);
   AsanThread *curr_thread = GetCurrentThread();
@@ -811,6 +836,7 @@ void InitializeAsanInterceptors() {
 #endif
 #if ASAN_INTERCEPT_SWAPCONTEXT
   ASAN_INTERCEPT_FUNC(swapcontext);
+  ASAN_INTERCEPT_FUNC(setcontext);
   ASAN_INTERCEPT_FUNC(getcontext);
 #endif
 #if ASAN_INTERCEPT__LONGJMP
diff --git a/lib/asan/asan_thread.cc b/lib/asan/asan_thread.cc
index 33fffec..3fa3143 100644
--- a/lib/asan/asan_thread.cc
+++ b/lib/asan/asan_thread.cc
@@ -150,6 +150,9 @@ FakeStack *AsanThread::AsyncSignalSafeLazyInitFakeStack() {
 }
 
 void AsanThread::Init() {
+  temp_stack_ = &stacks_[0];
+  next_stack_ = &stacks_[1];
+  previous_stack_ = &stacks_[2];
   fake_stack_ = nullptr;  // Will be initialized lazily if needed.
   CHECK_EQ(this->stack_size(), 0U);
   SetThreadStackAndTls();
@@ -195,10 +198,11 @@ thread_return_t AsanThread::ThreadStart(
 
 void AsanThread::SetThreadStackAndTls() {
   uptr tls_size = 0;
-  GetThreadStackAndTls(tid() == 0, &current_stack_.stack_bottom,
-                       &current_stack_.stack_size, &tls_begin_, &tls_size);
-  current_stack_.stack_top = current_stack_.stack_bottom
-                             + current_stack_.stack_size;
+  GetThreadStackAndTls(tid() == 0, &next_stack_->stack_bottom,
+                       &next_stack_->stack_size, &tls_begin_, &tls_size);
+  next_stack_->stack_top = next_stack_->stack_bottom + next_stack_->stack_size;
+  previous_stack_->stack_top = previous_stack_->stack_bottom = 0;
+  previous_stack_->stack_size = 0;
   tls_end_ = tls_begin_ + tls_size;
   dtls_ = DTLS_Get();
 
diff --git a/lib/asan/asan_thread.h b/lib/asan/asan_thread.h
index c4670e1..850b04c 100644
--- a/lib/asan/asan_thread.h
+++ b/lib/asan/asan_thread.h
@@ -82,6 +82,17 @@ class AsanThread {
   AsanThreadContext *context() { return context_; }
   void set_context(AsanThreadContext *context) { context_ = context; }
 
+  void SetUserStack(uptr base, uptr size) {
+    temp_stack_->stack_bottom = base;
+    temp_stack_->stack_top = base + size;
+    temp_stack_->stack_size = size;
+
+    StackDescriptor* oprev = previous_stack_;
+    previous_stack_ = next_stack_;
+    next_stack_ = temp_stack_;
+    temp_stack_ = oprev;
+  }
+
   struct StackFrameAccess {
     uptr offset;
     uptr frame_pc;
@@ -94,7 +105,12 @@ class AsanThread {
   }
 
   StackDescriptor *CurrentStack() {
-    return &current_stack_;
+    int local;
+    if (AddrIsInStack(previous_stack_, (uptr)&local)) {
+        return previous_stack_;
+    } else {
+        return next_stack_;
+    }
   }
 
   bool AddrIsInStack(uptr addr) {
@@ -144,9 +160,17 @@ class AsanThread {
   AsanThreadContext *context_;
   thread_callback_t start_routine_;
   void *arg_;
-  // stack_size_ == stack_top_ - stack_bottom_;
-  // It needs to be set in a async-signal-safe manner.
-  StackDescriptor current_stack_;
+
+  // We have three stack descriptors for async-safe stack change. New stack
+  // information is written to temp_stack_. Then previous_stack_ is made to
+  // point to the same descriptor that next_stack_ does. Finally, temp_stack_
+  // is assigned to next_stack_. The result is that at any time either
+  // previous_stack_ or next_stack_ contain the correct stack information.
+  StackDescriptor stacks_[3];
+  StackDescriptor* temp_stack_;
+  StackDescriptor* next_stack_;
+  StackDescriptor* previous_stack_;
+
   uptr tls_begin_;
   uptr tls_end_;
   DTLS *dtls_;
-- 
2.5.0