[asan] implement __asan_set_death_callback

llvm-svn: 150414

6 files changed, 49 insertions, 24 deletions

diff --git a/compiler-rt/lib/asan/asan_interface.h b/compiler-rt/lib/asan/asan_interface.h
index eb7d87d09b7..bfbbd09e9aa 100644
--- a/compiler-rt/lib/asan/asan_interface.h
+++ b/compiler-rt/lib/asan/asan_interface.h
@@ -81,7 +81,7 @@ extern "C" {
 
   // Performs cleanup before a NoReturn function. Must be called before things
   // like _exit and execl to avoid false positives on stack.
-  void __asan_handle_no_return();
+  void __asan_handle_no_return() ASAN_INTERFACE_FUNCTION_ATTRIBUTE;
 
 // User code should use macro instead of functions.
 #if defined(__has_feature) && __has_feature(address_sanitizer)
@@ -109,7 +109,13 @@ extern "C" {
 
   // Sets the exit code to use when reporting an error.
   // Returns the old value.
-  int __asan_set_error_exit_code(int exit_code);
+  int __asan_set_error_exit_code(int exit_code)
+      ASAN_INTERFACE_FUNCTION_ATTRIBUTE;
+
+  // Sets the callback to be called right before death on error.
+  // Passing NULL will unset the callback.
+  void __asan_set_death_callback(void (*callback)(void))
+      ASAN_INTERFACE_FUNCTION_ATTRIBUTE;
 
   // Returns the estimated number of bytes that will be reserved by allocator
   // for request of "size" bytes. If ASan allocator can't allocate that much
diff --git a/compiler-rt/lib/asan/asan_internal.h b/compiler-rt/lib/asan/asan_internal.h
index 67ed869d633..6440e8f36f2 100644
--- a/compiler-rt/lib/asan/asan_internal.h
+++ b/compiler-rt/lib/asan/asan_internal.h
@@ -210,6 +210,8 @@ extern bool asan_init_is_running;
 enum LinkerInitialized { LINKER_INITIALIZED = 0 };
 
 void AsanDie();
+void SleepForSeconds(int seconds);
+void Exit(int exitcode);
 
 #define CHECK(cond) do { if (!(cond)) { \
   CheckFailed(#cond, __FILE__, __LINE__); \
diff --git a/compiler-rt/lib/asan/asan_posix.cc b/compiler-rt/lib/asan/asan_posix.cc
index 78ec98d671b..1b8b7c2dda1 100644
--- a/compiler-rt/lib/asan/asan_posix.cc
+++ b/compiler-rt/lib/asan/asan_posix.cc
@@ -69,14 +69,6 @@ void AsanDisableCoreDumper() {
   setrlimit(RLIMIT_CORE, &nocore);
 }
 
-void AsanDie() {
-  if (FLAG_sleep_before_dying) {
-    Report("Sleeping for %d second(s)\n", FLAG_sleep_before_dying);
-    sleep(FLAG_sleep_before_dying);
-  }
-  _exit(FLAG_exitcode);
-}
-
 int GetPid() {
   return getpid();
 }
@@ -85,6 +77,14 @@ uintptr_t GetThreadSelf() {
   return (uintptr_t)pthread_self();
 }
 
+void SleepForSeconds(int seconds) {
+  sleep(seconds);
+}
+
+void Exit(int exitcode) {
+  return _exit(exitcode);
+}
+
 int AtomicInc(int *a) {
 #ifdef ANDROID
   return __atomic_inc(a) + 1;
diff --git a/compiler-rt/lib/asan/asan_rtl.cc b/compiler-rt/lib/asan/asan_rtl.cc
index 715edb086c9..4d3eb6e3770 100644
--- a/compiler-rt/lib/asan/asan_rtl.cc
+++ b/compiler-rt/lib/asan/asan_rtl.cc
@@ -52,6 +52,7 @@ int    FLAG_sleep_before_dying;
 // -------------------------- Globals --------------------- {{{1
 int asan_inited;
 bool asan_init_is_running;
+static void (*death_callback)(void);
 
 // -------------------------- Misc ---------------- {{{1
 void ShowStatsAndAbort() {
@@ -100,6 +101,16 @@ size_t ReadFileToBuffer(const char *file_name, char **buff,
   return read_len;
 }
 
+void AsanDie() {
+  if (FLAG_sleep_before_dying) {
+    Report("Sleeping for %d second(s)\n", FLAG_sleep_before_dying);
+    SleepForSeconds(FLAG_sleep_before_dying);
+  }
+  if (death_callback)
+    death_callback();
+  Exit(FLAG_exitcode);
+}
+
 // ---------------------- mmap -------------------- {{{1
 void OutOfMemoryMessageAndDie(const char *mem_type, size_t size) {
   Report("ERROR: AddressSanitizer failed to allocate "
@@ -283,7 +294,6 @@ int __asan_set_error_exit_code(int exit_code) {
   return old;
 }
 
-NOINLINE ASAN_INTERFACE_ATTRIBUTE
 void __asan_handle_no_return() {
   int local_stack;
   AsanThread *curr_thread = asanThreadRegistry().GetCurrent();
@@ -293,6 +303,10 @@ void __asan_handle_no_return() {
   PoisonShadow(bottom, top - bottom, 0);
 }
 
+void __asan_set_death_callback(void (*callback)(void)) {
+  death_callback = callback;
+}
+
 void __asan_report_error(uintptr_t pc, uintptr_t bp, uintptr_t sp,
                          uintptr_t addr, bool is_write, size_t access_size) {
   // Do not print more than one report, otherwise they will mix up.
diff --git a/compiler-rt/lib/asan/asan_win.cc b/compiler-rt/lib/asan/asan_win.cc
index 7de5d589799..b35f40efe60 100644
--- a/compiler-rt/lib/asan/asan_win.cc
+++ b/compiler-rt/lib/asan/asan_win.cc
@@ -48,7 +48,6 @@ void *AsanMprotect(uintptr_t fixed_addr, size_t size) {
 void AsanUnmapOrDie(void *addr, size_t size) {
   UNIMPLEMENTED();
 }
-// }}}
 
 // ---------------------- IO ---------------- {{{1
 size_t AsanWrite(int fd, const void *buf, size_t count) {
@@ -75,7 +74,6 @@ int AsanClose(int fd) {
   UNIMPLEMENTED();
   return -1;
 }
-// }}}
 
 // ---------------------- Stacktraces, symbols, etc. ---------------- {{{1
 static AsanLock dbghelp_lock(LINKER_INITIALIZED);
@@ -156,7 +154,6 @@ bool WinSymbolize(const void *addr, char *out_buffer, int buffer_size) {
   }
   return true;
 }
-// }}}
 
 // ---------------------- AsanLock ---------------- {{{1
 enum LockState {
@@ -193,7 +190,6 @@ void AsanLock::Unlock() {
   owner_ = LOCK_READY;
   LeaveCriticalSection((LPCRITICAL_SECTION)opaque_storage_);
 }
-// }}}
 
 // ---------------------- TSD ---------------- {{{1
 static bool tsd_key_inited = false;
@@ -215,7 +211,6 @@ void AsanTSDSet(void *tsd) {
   CHECK(tsd_key_inited);
   fake_tsd = tsd;
 }
-// }}}
 
 // ---------------------- Various stuff ---------------- {{{1
 void *AsanDoesNotSupportStaticLinkage() {
@@ -253,15 +248,13 @@ void AsanDisableCoreDumper() {
   UNIMPLEMENTED();
 }
 
-void AsanDie() {
-  // FIXME: AsanDie() should be the same on all platforms.
-  if (FLAG_sleep_before_dying) {
-    Report("Sleeping for %d second(s)\n", FLAG_sleep_before_dying);
-    Sleep(FLAG_sleep_before_dying * 1000);
-  }
-  _exit(FLAG_exitcode);
+void SleepForSeconds(int seconds) {
+  Sleep(seconds * 1000);
+}
+
+void Exit(int exitcode) {
+  _exit(exitcode);
 }
-// }}}
 
 }  // namespace __asan
 
diff --git a/compiler-rt/lib/asan/tests/asan_interface_test.cc b/compiler-rt/lib/asan/tests/asan_interface_test.cc
index d17353d0211..782466dd37f 100644
--- a/compiler-rt/lib/asan/tests/asan_interface_test.cc
+++ b/compiler-rt/lib/asan/tests/asan_interface_test.cc
@@ -199,6 +199,16 @@ TEST(AddressSanitizerInterface, ExitCode) {
               ::testing::ExitedWithCode(original_exit_code), "");
 }
 
+static void MyDeathCallback() {
+  fprintf(stderr, "MyDeathCallback\n");
+}
+
+TEST(AddressSanitizerInterface, DeathCallbackTest) {
+  __asan_set_death_callback(MyDeathCallback);
+  EXPECT_DEATH(DoDoubleFree(), "MyDeathCallback");
+  __asan_set_death_callback(NULL);
+}
+
 static const char* kUseAfterPoisonErrorMessage = "use-after-poison";
 
 #define ACCESS(ptr, offset) Ident(*(ptr + offset))