Implementing start_switch_fiber and finish_switch_fiber for uthread t…

…o enable address sanitizer

async_simple/CMakeLists.txt

@@ -11,7 +11,7 @@ if(${UTHREAD})
     # for the user thread.
     # FIXME2: the new added sanitize action: address-use-after-return
     # can't handle stackful coroutine well too. So disable it as a workaround now.
-    set_property(SOURCE ${uthread_src} PROPERTY COMPILE_OPTIONS "-fno-sanitize=address;-fsanitize-address-use-after-return=never")
+    # set_property(SOURCE ${uthread_src} PROPERTY COMPILE_OPTIONS "-fno-sanitize=address;-fsanitize-address-use-after-return=never")
     #endif()
   endif()
 file(GLOB uthread_asm_src "uthread/internal/${CMAKE_SYSTEM_NAME}/${CMAKE_SYSTEM_PROCESSOR}/*.S")

async_simple/Common.h

@@ -32,6 +32,18 @@
 #define AS_INLINE __attribute__((__always_inline__)) inline
 #endif
 
+#ifdef __clang__
+#if __has_feature(address_sanitizer)
+#define AS_INTERNAL_USE_ASAN 1
+#endif  // __has_feature(address_sanitizer)
+#endif  // __clang__
+
+#ifdef __GNUC__
+#ifdef __SANITIZE_ADDRESS__  // GCC
+#define AS_INTERNAL_USE_ASAN 1
+#endif  // __SANITIZE_ADDRESS__
+#endif  // __GNUC__
+
 namespace async_simple {
 // Different from assert, logicAssert is meaningful in
 // release mode. logicAssert should be used in case that

async_simple/uthread/internal/thread.cc

@@ -30,6 +30,36 @@ namespace async_simple {
 namespace uthread {
 namespace internal {
 
+#ifdef AS_INTERNAL_USE_ASAN
+
+extern "C" {
+void __sanitizer_start_switch_fiber(void** fake_stack_save,
+                                    const void* stack_bottom,
+                                    size_t stack_size);
+void __sanitizer_finish_switch_fiber(void* fake_stack_save,
+                                     const void** stack_bottom_old,
+                                     size_t* stack_size_old);
+}
+
+inline void start_switch_fiber(jmp_buf_link* context) {
+    __sanitizer_start_switch_fiber(&context->asan_fack_stack,
+                                   context->asan_stack_bottom,
+                                   context->asan_stack_size);
+}
+
+inline void finish_switch_fiber(jmp_buf_link* context) {
+    __sanitizer_finish_switch_fiber(context->asan_fack_stack,
+                                    &context->asan_stack_bottom,
+                                    &context->asan_stack_size);
+}
+
+#else
+
+inline void start_switch_fiber(jmp_buf_link* context) {}
+inline void finish_switch_fiber(jmp_buf_link* context) {}
+
+#endif  // AS_INTERNAL_USE_ASAN
+
 thread_local jmp_buf_link g_unthreaded_context;
 thread_local jmp_buf_link* g_current_context = nullptr;
 
@@ -55,16 +85,41 @@ inline void jmp_buf_link::switch_in() {
     link = std::exchange(g_current_context, this);
     if (!link)
         AS_UNLIKELY { link = &g_unthreaded_context; }
+    start_switch_fiber(this);
     // `thread` is currently only used in `s_main`
     fcontext = _fl_jump_fcontext(fcontext, thread).fctx;
+    finish_switch_fiber(this);
 }
 
 inline void jmp_buf_link::switch_out() {
     g_current_context = link;
+    start_switch_fiber(link);
     link->fcontext = _fl_jump_fcontext(link->fcontext, thread).fctx;
+    finish_switch_fiber(link);
 }
 
-inline void jmp_buf_link::initial_switch_in_completed() {}
+inline void jmp_buf_link::initial_switch_in_completed() {
+#ifdef AS_INTERNAL_USE_ASAN
+    // This is a new thread and it doesn't have the fake stack yet. ASan will
+    // create it lazily, for now just pass nullptr.
+    __sanitizer_finish_switch_fiber(nullptr, &link->asan_stack_bottom,
+                                    &link->asan_stack_size);
+#endif
+}
+
+inline void jmp_buf_link::final_switch_out() {
+    g_current_context = link;
+#ifdef AS_INTERNAL_USE_ASAN
+    // Since the thread is about to die we pass nullptr as fake_stack_save
+    // argument so that ASan knows it can destroy the fake stack if it exists.
+    __sanitizer_start_switch_fiber(nullptr, link->asan_stack_bottom,
+                                   link->asan_stack_size);
+#endif
+    _fl_jump_fcontext(link->fcontext, thread);
+
+    // never reach here
+    assert(false);
+}
 
 thread_context::thread_context(std::function<void()> func, size_t stack_size)
     : stack_size_(stack_size ? stack_size : get_base_stack_size()),
@@ -87,6 +142,10 @@ void thread_context::setup() {
     context_.fcontext = _fl_make_fcontext(stack_.get() + stack_size_,
                                           stack_size_, thread_context::s_main);
     context_.thread = this;
+#ifdef AS_INTERNAL_USE_ASAN
+    context_.asan_stack_bottom = stack_.get();
+    context_.asan_stack_size = stack_size_;
+#endif
     context_.switch_in();
 }
 
@@ -122,10 +181,7 @@ void thread_context::main() {
         done_.setException(std::current_exception());
     }
 
-    context_.switch_out();
-
-    // never reach here
-    assert(false);
+    context_.final_switch_out();
 }
 
 namespace thread_impl {

async_simple/uthread/internal/thread_impl.h

@@ -23,6 +23,8 @@
 #ifndef ASYNC_SIMPLE_UTHREAD_INTERNAL_UTHREAD_IMPL_H
 #define ASYNC_SIMPLE_UTHREAD_INTERNAL_UTHREAD_IMPL_H
 
+#include "async_simple/Common.h"
+
 namespace async_simple {
 namespace uthread {
 namespace internal {
@@ -44,10 +46,17 @@ struct jmp_buf_link {
     jmp_buf_link* link = nullptr;
     thread_context* thread = nullptr;
 
+#ifdef AS_INTERNAL_USE_ASAN
+    const void* asan_stack_bottom = nullptr;
+    std::size_t asan_stack_size = 0;
+    void* asan_fack_stack = nullptr;
+#endif
+
 public:
     void switch_in();
     void switch_out();
     void initial_switch_in_completed();
+    void final_switch_out();
 };
 
 extern thread_local jmp_buf_link* g_current_context;

async_simple/uthread/test/CMakeLists.txt

@@ -8,7 +8,7 @@ if (CMAKE_BUILD_TYPE STREQUAL "Debug" AND CMAKE_CXX_COMPILER_ID STREQUAL "Clang"
   # for the user thread.
   # FIXME2: the new added sanitize action: address-use-after-return
   # can't handle stackful coroutine well too. So disable it as a workaround now.
-  set_property(SOURCE ${uthread_test_src} PROPERTY COMPILE_OPTIONS "-fno-sanitize=address;-fsanitize-address-use-after-return=never")
+  # set_property(SOURCE ${uthread_test_src} PROPERTY COMPILE_OPTIONS "-fno-sanitize=address;-fsanitize-address-use-after-return=never")
 endif()
 
 add_executable(async_simple_uthread_test ${uthread_test_src} ${PROJECT_SOURCE_DIR}/async_simple/test/dotest.cpp)