[lsan] Support ASan's stack-use-after-return mode in LSan.

Treat the fake stack as live memory.

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@192593 91177308-0d34-0410-b5e6-96231b3b80d8

8 files changed, 69 insertions, 4 deletions

diff --git a/lib/asan/asan_fake_stack.cc b/lib/asan/asan_fake_stack.cc
index 81091f18..4f617cac 100644
--- a/lib/asan/asan_fake_stack.cc
+++ b/lib/asan/asan_fake_stack.cc
@@ -134,6 +134,20 @@ NOINLINE void FakeStack::GC(uptr real_stack) {
   needs_gc_ = false;
 }
 
+void FakeStack::ForEachFakeFrame(RangeIteratorCallback callback, void *arg) {
+  for (uptr class_id = 0; class_id < kNumberOfSizeClasses; class_id++) {
+    u8 *flags = GetFlags(stack_size_log(), class_id);
+    for (uptr i = 0, n = NumberOfFrames(stack_size_log(), class_id); i < n;
+         i++) {
+      if (flags[i] == 0) continue;  // not allocated.
+      FakeFrame *ff = reinterpret_cast<FakeFrame *>(
+          GetFrame(stack_size_log(), class_id, i));
+      uptr begin = reinterpret_cast<uptr>(ff);
+      callback(begin, begin + FakeStack::BytesInSizeClass(class_id), arg);
+    }
+  }
+}
+
 #if SANITIZER_LINUX && !SANITIZER_ANDROID
 static THREADLOCAL FakeStack *fake_stack_tls;
 
diff --git a/lib/asan/asan_fake_stack.h b/lib/asan/asan_fake_stack.h
index b4f74d90..f17ee026 100644
--- a/lib/asan/asan_fake_stack.h
+++ b/lib/asan/asan_fake_stack.h
@@ -148,6 +148,8 @@ class FakeStack {
   void HandleNoReturn();
   void GC(uptr real_stack);
 
+  void ForEachFakeFrame(RangeIteratorCallback callback, void *arg);
+
  private:
   FakeStack() { }
   static const uptr kFlagsOffset = 4096;  // This is were the flags begin.
diff --git a/lib/asan/asan_thread.cc b/lib/asan/asan_thread.cc
index 1949cd83..7ffac991 100644
--- a/lib/asan/asan_thread.cc
+++ b/lib/asan/asan_thread.cc
@@ -294,6 +294,13 @@ void EnsureMainThreadIDIsCorrect() {
   if (context && (context->tid == 0))
     context->os_id = GetTid();
 }
+
+__asan::AsanThread *GetAsanThreadByOsIDLocked(uptr os_id) {
+  __asan::AsanThreadContext *context = static_cast<__asan::AsanThreadContext *>(
+      __asan::asanThreadRegistry().FindThreadContextByOsIDLocked(os_id));
+  if (!context) return 0;
+  return context->thread;
+}
 }  // namespace __asan
 
 // --- Implementation of LSan-specific functions --- {{{1
@@ -301,10 +308,7 @@ 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) {
-  __asan::AsanThreadContext *context = static_cast<__asan::AsanThreadContext *>(
-      __asan::asanThreadRegistry().FindThreadContextByOsIDLocked(os_id));
-  if (!context) return false;
-  __asan::AsanThread *t = context->thread;
+  __asan::AsanThread *t = __asan::GetAsanThreadByOsIDLocked(os_id);
   if (!t) return false;
   *stack_begin = t->stack_bottom();
   *stack_end = t->stack_top();
@@ -316,6 +320,13 @@ bool GetThreadRangesLocked(uptr os_id, uptr *stack_begin, uptr *stack_end,
   return true;
 }
 
+void ForEachExtraStackRange(uptr os_id, RangeIteratorCallback callback,
+                            void *arg) {
+  __asan::AsanThread *t = __asan::GetAsanThreadByOsIDLocked(os_id);
+  if (t && t->has_fake_stack())
+    t->fake_stack()->ForEachFakeFrame(callback, arg);
+}
+
 void LockThreadRegistry() {
   __asan::asanThreadRegistry().Lock();
 }
diff --git a/lib/lsan/lit_tests/TestCases/use_after_return.cc b/lib/lsan/lit_tests/TestCases/use_after_return.cc
new file mode 100644
index 00000000..7ccbe435
--- /dev/null
+++ b/lib/lsan/lit_tests/TestCases/use_after_return.cc
@@ -0,0 +1,23 @@
+// Test that fake stack (introduced by ASan's use-after-return mode) is included
+// in the root set.
+// RUN: LSAN_BASE="report_objects=1:use_registers=0"
+// RUN: %clangxx_lsan %s -O2 -o %t
+// RUN: ASAN_OPTIONS=$ASAN_OPTIONS:detect_stack_use_after_return=1 LSAN_OPTIONS=$LSAN_BASE:"use_stacks=0" not %t 2>&1 | FileCheck %s
+// RUN: ASAN_OPTIONS=$ASAN_OPTIONS:detect_stack_use_after_return=1 LSAN_OPTIONS=$LSAN_BASE:"use_stacks=1" %t 2>&1
+// RUN: ASAN_OPTIONS=$ASAN_OPTIONS:detect_stack_use_after_return=1 LSAN_OPTIONS="" %t 2>&1
+
+#include <stdio.h>
+#include <stdlib.h>
+
+int main() {
+  void *stack_var = malloc(1337);
+  fprintf(stderr, "Test alloc: %p.\n", stack_var);
+  // Take pointer to variable, to ensure it's not optimized into a register.
+  fprintf(stderr, "Stack var at: %p.\n", &stack_var);
+  // Do not return from main to prevent the pointer from going out of scope.
+  exit(0);
+}
+// CHECK: Test alloc: [[ADDR:.*]].
+// CHECK: Directly leaked 1337 byte object at [[ADDR]]
+// CHECK: LeakSanitizer: detected memory leaks
+// CHECK: SUMMARY: LeakSanitizer:
diff --git a/lib/lsan/lsan_common.cc b/lib/lsan/lsan_common.cc
index 416828ab..a58e604f 100644
--- a/lib/lsan/lsan_common.cc
+++ b/lib/lsan/lsan_common.cc
@@ -151,6 +151,11 @@ void ScanRangeForPointers(uptr begin, uptr end,
   }
 }
 
+void ForEachExtraStackRangeCb(uptr begin, uptr end, void* arg) {
+  Frontier *frontier = reinterpret_cast<Frontier *>(arg);
+  ScanRangeForPointers(begin, end, frontier, "FAKE STACK", kReachable);
+}
+
 // Scans thread data (stacks and TLS) for heap pointers.
 static void ProcessThreads(SuspendedThreadsList const &suspended_threads,
                            Frontier *frontier) {
@@ -199,6 +204,7 @@ static void ProcessThreads(SuspendedThreadsList const &suspended_threads,
       }
       ScanRangeForPointers(stack_begin, stack_end, frontier, "STACK",
                            kReachable);
+      ForEachExtraStackRange(os_id, ForEachExtraStackRangeCb, frontier);
     }
 
     if (flags()->use_tls) {
diff --git a/lib/lsan/lsan_common.h b/lib/lsan/lsan_common.h
index 971f7486..d490f8ba 100644
--- a/lib/lsan/lsan_common.h
+++ b/lib/lsan/lsan_common.h
@@ -135,6 +135,8 @@ void UnlockThreadRegistry();
 bool GetThreadRangesLocked(uptr os_id, uptr *stack_begin, uptr *stack_end,
                            uptr *tls_begin, uptr *tls_end,
                            uptr *cache_begin, uptr *cache_end);
+void ForEachExtraStackRange(uptr os_id, RangeIteratorCallback callback,
+                            void *arg);
 // If called from the main thread, updates the main thread's TID in the thread
 // registry. We need this to handle processes that fork() without a subsequent
 // exec(), which invalidates the recorded TID. To update it, we must call
diff --git a/lib/lsan/lsan_thread.cc b/lib/lsan/lsan_thread.cc
index 3bfccdf7..0f8efc09 100644
--- a/lib/lsan/lsan_thread.cc
+++ b/lib/lsan/lsan_thread.cc
@@ -145,6 +145,10 @@ bool GetThreadRangesLocked(uptr os_id, uptr *stack_begin, uptr *stack_end,
   return true;
 }
 
+void ForEachExtraStackRange(uptr os_id, RangeIteratorCallback callback,
+                            void *arg) {
+}
+
 void LockThreadRegistry() {
   thread_registry->Lock();
 }
diff --git a/lib/sanitizer_common/sanitizer_common.h b/lib/sanitizer_common/sanitizer_common.h
index f8f9f036..af1e4f10 100644
--- a/lib/sanitizer_common/sanitizer_common.h
+++ b/lib/sanitizer_common/sanitizer_common.h
@@ -446,6 +446,9 @@ const uptr kPthreadDestructorIterations = 4;
 // Unused on Windows.
 const uptr kPthreadDestructorIterations = 0;
 #endif
+
+// Callback type for iterating over a set of memory ranges.
+typedef void (*RangeIteratorCallback)(uptr begin, uptr end, void *arg);
 }  // namespace __sanitizer
 
 #endif  // SANITIZER_COMMON_H