4 files changed, 58 insertions, 29 deletions

diff --git a/lib/msan/lit_tests/signal_stress_test.cc b/lib/msan/lit_tests/signal_stress_test.cc
index e632cde3..ea75eae1 100644
--- a/lib/msan/lit_tests/signal_stress_test.cc
+++ b/lib/msan/lit_tests/signal_stress_test.cc
@@ -1,4 +1,4 @@
-// RUN: %clangxx_msan -O0 %s -o %t && %t
+// RUN: %clangxx_msan -std=c++11 -O0 %s -o %t && %t
 
 // Test that va_arg shadow from a signal handler does not leak outside.
 
@@ -9,19 +9,14 @@
 #include <sys/time.h>
 #include <stdio.h>
 
-const int kArgCnt = 20;
-const int kSigCnt = 100;
+const int kSigCnt = 200;
 
-volatile int z;
-
-void f(bool poisoned, ...) {
+void f(bool poisoned, int n, ...) {
   va_list vl;
-  va_start(vl, poisoned);
-  for (int i = 0; i < kArgCnt; ++i) {
+  va_start(vl, n);
+  for (int i = 0; i < n; ++i) {
     void *p = va_arg(vl, void *);
-    if (poisoned)
-      assert(__msan_test_shadow(&p, sizeof(p)) == 0);
-    else 
+    if (!poisoned)
       assert(__msan_test_shadow(&p, sizeof(p)) == -1);
   }
   va_end(vl);
@@ -32,13 +27,10 @@ int sigcnt;
 void SignalHandler(int signo) {
   assert(signo == SIGPROF);
   void *p;
-  void ** volatile q = &p;
-  f(true,
-      *q, *q, *q, *q, *q,
-      *q, *q, *q, *q, *q,
-      *q, *q, *q, *q, *q,
-      *q, *q, *q, *q, *q,
-      *q, *q, *q, *q, *q);
+  void **volatile q = &p;
+  f(true, 10,
+    *q, *q, *q, *q, *q,
+    *q, *q, *q, *q, *q);
   ++sigcnt;
 }
 
@@ -52,12 +44,20 @@ int main() {
   itv.it_value.tv_usec = 100;
   setitimer(ITIMER_PROF, &itv, NULL);
 
+  void *p;
+  void **volatile q = &p;
+
   do {
-    f(false,
-        0, 0, 0, 0, 0,
-        0, 0, 0, 0, 0,
-        0, 0, 0, 0, 0,
-        0, 0, 0, 0, 0);
+    f(false, 20,
+      nullptr, nullptr, nullptr, nullptr, nullptr,
+      nullptr, nullptr, nullptr, nullptr, nullptr,
+      nullptr, nullptr, nullptr, nullptr, nullptr,
+      nullptr, nullptr, nullptr, nullptr, nullptr);
+    f(true, 20,
+      *q, *q, *q, *q, *q,
+      *q, *q, *q, *q, *q,
+      *q, *q, *q, *q, *q,
+      *q, *q, *q, *q, *q);
   } while (sigcnt < kSigCnt);
 
   itv.it_interval.tv_sec = 0;
diff --git a/lib/msan/msan.cc b/lib/msan/msan.cc
index d7912fa8..c441a29d 100644
--- a/lib/msan/msan.cc
+++ b/lib/msan/msan.cc
@@ -230,11 +230,29 @@ void UnpoisonParam(uptr n) {
   internal_memset(__msan_param_tls, 0, n * sizeof(*__msan_param_tls));
 }
 
-void UnpoisonThreadLocalState() {
+// Backup MSan runtime TLS state.
+// Implementation must be async-signal-safe.
+// Instances of this class may live on the signal handler stack, and data size
+// may be an issue.
+void ScopedThreadLocalStateBackup::Backup() {
+  va_arg_overflow_size_tls = __msan_va_arg_overflow_size_tls;
+}
+
+void ScopedThreadLocalStateBackup::Restore() {
+  // A lame implementation that only keeps essential state and resets the rest.
+  __msan_va_arg_overflow_size_tls = va_arg_overflow_size_tls;
+
   internal_memset(__msan_param_tls, 0, sizeof(__msan_param_tls));
   internal_memset(__msan_retval_tls, 0, sizeof(__msan_retval_tls));
   internal_memset(__msan_va_arg_tls, 0, sizeof(__msan_va_arg_tls));
-  __msan_va_arg_overflow_size_tls = 0;
+
+  if (__msan_get_track_origins()) {
+    internal_memset(&__msan_retval_origin_tls, 0, sizeof(__msan_retval_tls));
+    internal_memset(__msan_param_origin_tls, 0, sizeof(__msan_param_origin_tls));
+  }
+}
+
+void UnpoisonThreadLocalState() {
 }
 
 }  // namespace __msan
diff --git a/lib/msan/msan.h b/lib/msan/msan.h
index e95ac9e6..5d5ca570 100644
--- a/lib/msan/msan.h
+++ b/lib/msan/msan.h
@@ -86,6 +86,15 @@ void UnpoisonThreadLocalState();
         StackTrace::GetCurrentPc(), GET_CURRENT_FRAME(),           \
         common_flags()->fast_unwind_on_malloc)
 
+class ScopedThreadLocalStateBackup {
+public:
+  ScopedThreadLocalStateBackup() { Backup(); }
+  ~ScopedThreadLocalStateBackup() { Restore(); }
+  void Backup();
+  void Restore();
+private:
+  u64 va_arg_overflow_size_tls;
+};
 }  // namespace __msan
 
 #define MSAN_MALLOC_HOOK(ptr, size) \
diff --git a/lib/msan/msan_interceptors.cc b/lib/msan/msan_interceptors.cc
index eb6888a3..8d39e541 100644
--- a/lib/msan/msan_interceptors.cc
+++ b/lib/msan/msan_interceptors.cc
@@ -911,17 +911,20 @@ static atomic_uintptr_t sigactions[kMaxSignals];
 static StaticSpinMutex sigactions_mu;
 
 static void SignalHandler(int signo) {
+  ScopedThreadLocalStateBackup stlsb;
+  stlsb.Backup();
   UnpoisonParam(1);
 
   typedef void (*signal_cb)(int x);
   signal_cb cb =
       (signal_cb)atomic_load(&sigactions[signo], memory_order_relaxed);
   cb(signo);
-
-  UnpoisonThreadLocalState();
+  stlsb.Restore();
 }
 
 static void SignalAction(int signo, void *si, void *uc) {
+  ScopedThreadLocalStateBackup stlsb;
+  stlsb.Backup();
   UnpoisonParam(3);
   __msan_unpoison(si, sizeof(__sanitizer_sigaction));
   __msan_unpoison(uc, __sanitizer::ucontext_t_sz);
@@ -930,8 +933,7 @@ static void SignalAction(int signo, void *si, void *uc) {
   sigaction_cb cb =
       (sigaction_cb)atomic_load(&sigactions[signo], memory_order_relaxed);
   cb(signo, si, uc);
-
-  UnpoisonThreadLocalState();
+  stlsb.Restore();
 }
 
 INTERCEPTOR(int, sigaction, int signo, const __sanitizer_sigaction *act,