[msan] Intercept wcsncpy, wcsnlen.

llvm-svn: 297793

2 files changed, 59 insertions, 0 deletions

diff --git a/compiler-rt/lib/msan/msan_interceptors.cc b/compiler-rt/lib/msan/msan_interceptors.cc
index 6447bb1b270..f3d5f4a0911 100644
--- a/compiler-rt/lib/msan/msan_interceptors.cc
+++ b/compiler-rt/lib/msan/msan_interceptors.cc
@@ -580,6 +580,13 @@ INTERCEPTOR(SIZE_T, wcslen, const wchar_t *s) {
   return res;
 }
 
+INTERCEPTOR(SIZE_T, wcsnlen, const wchar_t *s, SIZE_T n) {
+  ENSURE_MSAN_INITED();
+  SIZE_T res = REAL(wcsnlen)(s, n);
+  CHECK_UNPOISONED(s, sizeof(wchar_t) * Min(res + 1, n));
+  return res;
+}
+
 // wchar_t *wcschr(const wchar_t *wcs, wchar_t wc);
 INTERCEPTOR(wchar_t *, wcschr, void *s, wchar_t wc, void *ps) {
   ENSURE_MSAN_INITED();
@@ -597,6 +604,18 @@ INTERCEPTOR(wchar_t *, wcscpy, wchar_t *dest, const wchar_t *src) {
   return res;
 }
 
+INTERCEPTOR(wchar_t *, wcsncpy, wchar_t *dest, const wchar_t *src,
+            SIZE_T n) {  // NOLINT
+  ENSURE_MSAN_INITED();
+  GET_STORE_STACK_TRACE;
+  SIZE_T copy_size = REAL(wcsnlen)(src, n);
+  if (copy_size < n) copy_size++;           // trailing \0
+  wchar_t *res = REAL(wcsncpy)(dest, src, n);  // NOLINT
+  CopyShadowAndOrigin(dest, src, copy_size * sizeof(wchar_t), &stack);
+  __msan_unpoison(dest + copy_size, (n - copy_size) * sizeof(wchar_t));
+  return res;
+}
+
 // wchar_t *wmemcpy(wchar_t *dest, const wchar_t *src, SIZE_T n);
 INTERCEPTOR(wchar_t *, wmemcpy, wchar_t *dest, const wchar_t *src, SIZE_T n) {
   ENSURE_MSAN_INITED();
@@ -1565,8 +1584,10 @@ void InitializeInterceptors() {
   INTERCEPT_FUNCTION(mbtowc);
   INTERCEPT_FUNCTION(mbrtowc);
   INTERCEPT_FUNCTION(wcslen);
+  INTERCEPT_FUNCTION(wcsnlen);
   INTERCEPT_FUNCTION(wcschr);
   INTERCEPT_FUNCTION(wcscpy);
+  INTERCEPT_FUNCTION(wcsncpy);
   INTERCEPT_FUNCTION(wcscmp);
   INTERCEPT_FUNCTION(getenv);
   INTERCEPT_FUNCTION(setenv);
diff --git a/compiler-rt/test/msan/wcsncpy.cc b/compiler-rt/test/msan/wcsncpy.cc
new file mode 100644
index 00000000000..8f809697ecc
--- /dev/null
+++ b/compiler-rt/test/msan/wcsncpy.cc
@@ -0,0 +1,38 @@
+// RUN: %clangxx_msan -fsanitize-memory-track-origins -O0 %s -o %t && not %run %t >%t.out 2>&1
+// RUN: FileCheck %s < %t.out && FileCheck %s < %t.out
+
+#include <assert.h>
+#include <wchar.h>
+
+#include <sanitizer/msan_interface.h>
+
+int main() {
+  const wchar_t *s = L"abc";
+  assert(wcslen(s) == 3);
+
+  wchar_t s2[5];
+  assert(wcsncpy(s2, s, 3) == s2);
+  assert(__msan_test_shadow(&s2, 5 * sizeof(wchar_t)) == 3 * sizeof(wchar_t));
+  assert(wcsncpy(s2, s, 5) == s2);
+  assert(__msan_test_shadow(&s2, 5 * sizeof(wchar_t)) == -1);
+
+  wchar_t s3[5];
+  assert(wcsncpy(s3, s, 2) == s3);
+  assert(__msan_test_shadow(&s3, 5 * sizeof(wchar_t)) == 2 * sizeof(wchar_t));
+
+  __msan_allocated_memory(&s2[1], sizeof(wchar_t));
+  wchar_t s4[5];
+  assert(wcsncpy(s4, s2, 3) == s4);
+  __msan_check_mem_is_initialized(&s4, sizeof(s4));
+}
+// CHECK:  Uninitialized bytes in __msan_check_mem_is_initialized
+// CHECK:  WARNING: MemorySanitizer: use-of-uninitialized-value
+// CHECK:    in main {{.*}}wcsncpy.cc:26
+
+// CHECK:  Uninitialized value was stored to memory at
+// CHECK:    in wcsncpy
+// CHECK:    in main {{.*}}wcsncpy.cc:25
+
+// CHECK:  Memory was marked as uninitialized
+// CHECK:    in __msan_allocated_memory
+// CHECK:    in main {{.*}}wcsncpy.cc:23