4 files changed, 58 insertions, 0 deletions
diff --git a/compiler-rt/include/sanitizer/msan_interface.h b/compiler-rt/include/sanitizer/msan_interface.h
index 5c5363f5ca6..d40c556a46d 100644
--- a/compiler-rt/include/sanitizer/msan_interface.h
+++ b/compiler-rt/include/sanitizer/msan_interface.h
@@ -42,6 +42,9 @@ extern "C" {
      contents). */
   void __msan_unpoison_string(const volatile char *a);
 
+  /* Make first n parameters of the next function call fully initialized. */
+  void __msan_unpoison_param(size_t n);
+
   /* Make memory region fully uninitialized (without changing its contents).
      This is a legacy interface that does not update origin information. Use
      __msan_allocated_memory() instead. */
diff --git a/compiler-rt/lib/msan/msan_interceptors.cc b/compiler-rt/lib/msan/msan_interceptors.cc
index af9d1402976..eed9bb5b8fd 100644
--- a/compiler-rt/lib/msan/msan_interceptors.cc
+++ b/compiler-rt/lib/msan/msan_interceptors.cc
@@ -1543,6 +1543,8 @@ void __msan_poison_stack(void *a, uptr size) {
   SetShadow(a, size, __msan::flags()->poison_stack_with_zeroes ? 0 : -1);
 }
 
+void __msan_unpoison_param(uptr n) { UnpoisonParam(n); }
+
 void __msan_clear_and_unpoison(void *a, uptr size) {
   REAL(memset)(a, 0, size);
   SetShadow(a, size, 0);
diff --git a/compiler-rt/lib/msan/msan_interface_internal.h b/compiler-rt/lib/msan/msan_interface_internal.h
index 09680f67a96..1abbba018c9 100644
--- a/compiler-rt/lib/msan/msan_interface_internal.h
+++ b/compiler-rt/lib/msan/msan_interface_internal.h
@@ -69,6 +69,8 @@ void __msan_unpoison(const void *a, uptr size);
 SANITIZER_INTERFACE_ATTRIBUTE
 void __msan_unpoison_string(const char *s);
 SANITIZER_INTERFACE_ATTRIBUTE
+void __msan_unpoison_param(uptr n);
+SANITIZER_INTERFACE_ATTRIBUTE
 void __msan_clear_and_unpoison(void *a, uptr size);
 SANITIZER_INTERFACE_ATTRIBUTE
 void* __msan_memcpy(void *dst, const void *src, uptr size);
diff --git a/compiler-rt/test/msan/unpoison_param.cc b/compiler-rt/test/msan/unpoison_param.cc
new file mode 100644
index 00000000000..f59d2cbe231
--- /dev/null
+++ b/compiler-rt/test/msan/unpoison_param.cc
@@ -0,0 +1,51 @@
+// Tests that __msan_unpoison_param() works as specified.  To prevent MSan
+// instrumentation from modifying parameter shadow before each call to foo(), we
+// compile main() without MSan.
+
+// RUN: %clangxx_msan -fno-sanitize=memory -c %s -o %t-main.o
+// RUN: %clangxx_msan %t-main.o %s -o %t
+// RUN: %run %t
+
+#include <assert.h>
+#include <sanitizer/msan_interface.h>
+
+#if __has_feature(memory_sanitizer)
+
+__attribute__((noinline)) int bar(int a, int b) {
+  volatile int zero = 0;
+  return zero;
+}
+
+int foo(int a, int b, int unpoisoned_params) {
+  if (unpoisoned_params == 0) {
+    assert(__msan_test_shadow(&a, sizeof(a)) == 0);
+    assert(__msan_test_shadow(&b, sizeof(b)) == 0);
+  } else if (unpoisoned_params == 1) {
+    assert(__msan_test_shadow(&a, sizeof(a)) == -1);
+    assert(__msan_test_shadow(&b, sizeof(b)) == 0);
+  } else if (unpoisoned_params == 2) {
+    assert(__msan_test_shadow(&a, sizeof(a)) == -1);
+    assert(__msan_test_shadow(&b, sizeof(b)) == -1);
+  }
+
+  // Poisons parameter shadow in TLS so that the next call from uninstrumented
+  // main has params 1 and 2 poisoned no matter what.
+  int x, y;
+  return bar(x, y);
+}
+
+#else
+
+int foo(int, int, int);
+
+int main() {
+  foo(0, 0, 2); // Poison parameters for next call.
+  foo(0, 0, 0); // Check that both params are poisoned.
+  __msan_unpoison_param(1);
+  foo(0, 0, 1); // Check that only first param is unpoisoned.
+  __msan_unpoison_param(2);
+  foo(0, 0, 2); // Check that first and second params are unpoisoned.
+  return 0;
+}
+
+#endif