Fix fibers asan integration for new asan APIs

Summary:
Updating fibers code to use https://github.com/llvm-mirror/compiler-rt/commit/b0477747dfa8a9706f2c902e877e616aca51e06f

Patch by andriigrynenko!

Reviewed By: andrewcox

Differential Revision: D3500482

fbshipit-source-id: f51f4bb4ebf0d95a898eb1d4098459aa691acd61

folly/fibers/Fiber.cpp

@@ -126,6 +126,11 @@ void Fiber::fiberFuncHelper(intptr_t fiber) {
 }
 
 void Fiber::fiberFunc() {
+#ifdef FOLLY_SANITIZE_ADDRESS
+  fiberManager_.registerFinishSwitchStackWithAsan(
+      nullptr, &asanMainStackBase_, &asanMainStackSize_);
+#endif
+
   while (true) {
     DCHECK_EQ(state_, NOT_STARTED);
 

folly/fibers/Fiber.h

@@ -184,6 +184,12 @@ class Fiber {
                                            queues */
   folly::IntrusiveListHook globalListHook_; /**< list hook for global list */
   std::thread::id threadId_{};
+
+#ifdef FOLLY_SANITIZE_ADDRESS
+  void* asanFakeStack_{nullptr};
+  const void* asanMainStackBase_{nullptr};
+  size_t asanMainStackSize_{0};
+#endif
 };
 }
 }

folly/fibers/FiberManager.cpp

@@ -33,18 +33,22 @@
 
 #include <dlfcn.h>
 
-static void __asan_enter_fiber_weak(
+static void __sanitizer_start_switch_fiber_weak(
+    void** fake_stack_save,
     void const* fiber_stack_base,
     size_t fiber_stack_extent)
-    __attribute__((__weakref__("__asan_enter_fiber")));
-static void __asan_exit_fiber_weak()
-    __attribute__((__weakref__("__asan_exit_fiber")));
+    __attribute__((__weakref__("__sanitizer_start_switch_fiber")));
+static void __sanitizer_finish_switch_fiber_weak(
+    void* fake_stack_save,
+    void const** old_stack_base,
+    size_t* old_stack_extent)
+    __attribute__((__weakref__("__sanitizer_finish_switch_fiber")));
 static void __asan_unpoison_memory_region_weak(
     void const /* nolint */ volatile* addr,
     size_t size) __attribute__((__weakref__("__asan_unpoison_memory_region")));
 
-typedef void (*AsanEnterFiberFuncPtr)(void const*, size_t);
-typedef void (*AsanExitFiberFuncPtr)();
+typedef void (*AsanStartSwitchStackFuncPtr)(void**, void const*, size_t);
+typedef void (*AsanFinishSwitchStackFuncPtr)(void*, void const**, size_t*);
 typedef void (*AsanUnpoisonMemoryRegionFuncPtr)(
     void const /* nolint */ volatile*,
     size_t);
@@ -52,8 +56,8 @@ typedef void (*AsanUnpoisonMemoryRegionFuncPtr)(
 namespace folly {
 namespace fibers {
 
-static AsanEnterFiberFuncPtr getEnterFiberFunc();
-static AsanExitFiberFuncPtr getExitFiberFunc();
+static AsanStartSwitchStackFuncPtr getStartSwitchStackFunc();
+static AsanFinishSwitchStackFuncPtr getFinishSwitchStackFunc();
 static AsanUnpoisonMemoryRegionFuncPtr getUnpoisonMemoryRegionFunc();
 }
 }
@@ -187,30 +191,29 @@ void FiberManager::FibersPoolResizer::operator()() {
 
 #ifdef FOLLY_SANITIZE_ADDRESS
 
-void FiberManager::registerFiberActivationWithAsan(Fiber* fiber) {
-  auto context = &fiber->fcontext_;
-  void* top = context->stackBase();
-  void* bottom = context->stackLimit();
-  size_t extent = static_cast<char*>(top) - static_cast<char*>(bottom);
-
+void FiberManager::registerStartSwitchStackWithAsan(
+    void** saveFakeStack,
+    const void* stackBottom,
+    size_t stackSize) {
   // Check if we can find a fiber enter function and call it if we find one
-  static AsanEnterFiberFuncPtr fn = getEnterFiberFunc();
+  static AsanStartSwitchStackFuncPtr fn = getStartSwitchStackFunc();
   if (fn == nullptr) {
     LOG(FATAL) << "The version of ASAN in use doesn't support fibers";
   } else {
-    fn(bottom, extent);
+    fn(saveFakeStack, stackBottom, stackSize);
   }
 }
 
-void FiberManager::registerFiberDeactivationWithAsan(Fiber* fiber) {
-  (void)fiber; // currently unused
-
+void FiberManager::registerFinishSwitchStackWithAsan(
+    void* saveFakeStack,
+    const void** saveStackBottom,
+    size_t* saveStackSize) {
   // Check if we can find a fiber exit function and call it if we find one
-  static AsanExitFiberFuncPtr fn = getExitFiberFunc();
+  static AsanFinishSwitchStackFuncPtr fn = getFinishSwitchStackFunc();
   if (fn == nullptr) {
     LOG(FATAL) << "The version of ASAN in use doesn't support fibers";
   } else {
-    fn();
+    fn(saveFakeStack, saveStackBottom, saveStackSize);
   }
 }
 
@@ -226,17 +229,17 @@ void FiberManager::unpoisonFiberStack(const Fiber* fiber) {
   }
 }
 
-static AsanEnterFiberFuncPtr getEnterFiberFunc() {
-  AsanEnterFiberFuncPtr fn{nullptr};
+static AsanStartSwitchStackFuncPtr getStartSwitchStackFunc() {
+  AsanStartSwitchStackFuncPtr fn{nullptr};
 
   // Check whether weak reference points to statically linked enter function
-  if (nullptr != (fn = &::__asan_enter_fiber_weak)) {
+  if (nullptr != (fn = &::__sanitizer_start_switch_fiber_weak)) {
     return fn;
   }
 
   // Check whether we can find a dynamically linked enter function
-  if (nullptr !=
-      (fn = (AsanEnterFiberFuncPtr)dlsym(RTLD_DEFAULT, "__asan_enter_fiber"))) {
+  if (nullptr != (fn = (AsanStartSwitchStackFuncPtr)dlsym(
+                      RTLD_DEFAULT, "__sanitizer_start_switch_fiber"))) {
     return fn;
   }
 
@@ -244,17 +247,17 @@ static AsanEnterFiberFuncPtr getEnterFiberFunc() {
   return nullptr;
 }
 
-static AsanExitFiberFuncPtr getExitFiberFunc() {
-  AsanExitFiberFuncPtr fn{nullptr};
+static AsanFinishSwitchStackFuncPtr getFinishSwitchStackFunc() {
+  AsanFinishSwitchStackFuncPtr fn{nullptr};
 
   // Check whether weak reference points to statically linked exit function
-  if (nullptr != (fn = &::__asan_exit_fiber_weak)) {
+  if (nullptr != (fn = &::__sanitizer_finish_switch_fiber_weak)) {
     return fn;
   }
 
   // Check whether we can find a dynamically linked exit function
-  if (nullptr !=
-      (fn = (AsanExitFiberFuncPtr)dlsym(RTLD_DEFAULT, "__asan_exit_fiber"))) {
+  if (nullptr != (fn = (AsanFinishSwitchStackFuncPtr)dlsym(
+                      RTLD_DEFAULT, "__sanitizer_finish_switch_fiber"))) {
     return fn;
   }
 

folly/fibers/FiberManagerInternal-inl.h

@@ -62,7 +62,16 @@ inline intptr_t FiberManager::activateFiber(Fiber* fiber) {
   DCHECK_EQ(activeFiber_, (Fiber*)nullptr);
 
 #ifdef FOLLY_SANITIZE_ADDRESS
-  registerFiberActivationWithAsan(fiber);
+  DCHECK(!fiber->asanMainStackBase_);
+  DCHECK(!fiber->asanMainStackSize_);
+  auto stack = fiber->getStack();
+  void* asanFakeStack;
+  registerStartSwitchStackWithAsan(&asanFakeStack, stack.first, stack.second);
+  SCOPE_EXIT {
+    registerFinishSwitchStackWithAsan(asanFakeStack, nullptr, nullptr);
+    fiber->asanMainStackBase_ = nullptr;
+    fiber->asanMainStackSize_ = 0;
+  };
 #endif
 
   activeFiber_ = fiber;
@@ -73,7 +82,21 @@ inline intptr_t FiberManager::deactivateFiber(Fiber* fiber) {
   DCHECK_EQ(activeFiber_, fiber);
 
 #ifdef FOLLY_SANITIZE_ADDRESS
-  registerFiberDeactivationWithAsan(fiber);
+  DCHECK(fiber->asanMainStackBase_);
+  DCHECK(fiber->asanMainStackSize_);
+
+  // Release fake stack if fiber is completed
+  auto saveFakeStackPtr =
+      fiber->state_ == Fiber::INVALID ? nullptr : &fiber->asanFakeStack_;
+  registerStartSwitchStackWithAsan(
+      saveFakeStackPtr, fiber->asanMainStackBase_, fiber->asanMainStackSize_);
+  SCOPE_EXIT {
+    registerFinishSwitchStackWithAsan(
+        fiber->asanFakeStack_,
+        &fiber->asanMainStackBase_,
+        &fiber->asanMainStackSize_);
+    fiber->asanFakeStack_ = nullptr;
+  };
 #endif
 
   activeFiber_ = nullptr;

folly/fibers/FiberManagerInternal.h

@@ -468,8 +468,14 @@ class FiberManager : public ::folly::Executor {
   // These methods notify ASAN when a fiber is entered/exited so that ASAN can
   // find the right stack extents when it needs to poison/unpoison the stack.
 
-  void registerFiberActivationWithAsan(Fiber* fiber);
-  void registerFiberDeactivationWithAsan(Fiber* fiber);
+  void registerStartSwitchStackWithAsan(
+      void** saveFakeStack,
+      const void* stackBase,
+      size_t stackSize);
+  void registerFinishSwitchStackWithAsan(
+      void* fakeStack,
+      const void** saveStackBase,
+      size_t* saveStackSize);
   void unpoisonFiberStack(const Fiber* fiber);
 
 #endif // FOLLY_SANITIZE_ADDRESS