[MSAN] add interceptor for stpncpy

Reviewed By: vitalybuka

Differential Revision: https://reviews.llvm.org/D138386

GitOrigin-RevId: e49a6c3d7f50c7e8426e9b5166e8567fced1aa57
diff --git a/lib/msan/msan_interceptors.cpp b/lib/msan/msan_interceptors.cpp
index b3d5484..058c10a 100644
--- a/lib/msan/msan_interceptors.cpp
+++ b/lib/msan/msan_interceptors.cpp
@@ -310,9 +310,21 @@
   CopyShadowAndOrigin(dest, src, n + 1, &stack);
   return res;
 }
-#define MSAN_MAYBE_INTERCEPT_STPCPY INTERCEPT_FUNCTION(stpcpy)
+
+INTERCEPTOR(char *, stpncpy, char *dest, const char *src, SIZE_T n) {
+  ENSURE_MSAN_INITED();
+  GET_STORE_STACK_TRACE;
+  SIZE_T copy_size = Min(n, internal_strnlen(src, n) + 1);
+  char *res = REAL(stpncpy)(dest, src, n);
+  CopyShadowAndOrigin(dest, src, copy_size, &stack);
+  __msan_unpoison(dest + copy_size, n - copy_size);
+  return res;
+}
+#  define MSAN_MAYBE_INTERCEPT_STPCPY INTERCEPT_FUNCTION(stpcpy)
+#  define MSAN_MAYBE_INTERCEPT_STPNCPY INTERCEPT_FUNCTION(stpncpy)
 #else
 #define MSAN_MAYBE_INTERCEPT_STPCPY
+#  define MSAN_MAYBE_INTERCEPT_STPNCPY
 #endif
 
 INTERCEPTOR(char *, strdup, char *src) {
@@ -1686,6 +1698,7 @@
   INTERCEPT_FUNCTION(wmemmove);
   INTERCEPT_FUNCTION(strcpy);
   MSAN_MAYBE_INTERCEPT_STPCPY;
+  MSAN_MAYBE_INTERCEPT_STPNCPY;
   INTERCEPT_FUNCTION(strdup);
   MSAN_MAYBE_INTERCEPT___STRDUP;
   INTERCEPT_FUNCTION(strncpy);
diff --git a/lib/msan/tests/msan_test.cpp b/lib/msan/tests/msan_test.cpp
index 8babe87..e3ad9bf 100644
--- a/lib/msan/tests/msan_test.cpp
+++ b/lib/msan/tests/msan_test.cpp
@@ -1687,6 +1687,21 @@
   EXPECT_NOT_POISONED(y[2]);
 }
 
+TEST(MemorySanitizer, stpncpy) {
+  char *x = new char[3];
+  char *y = new char[5];
+  x[0] = 'a';
+  x[1] = *GetPoisoned<char>(1, 1);
+  x[2] = '\0';
+  char *res = stpncpy(y, x, 4);
+  ASSERT_EQ(res, y + 2);
+  EXPECT_NOT_POISONED(y[0]);
+  EXPECT_POISONED(y[1]);
+  EXPECT_NOT_POISONED(y[2]);
+  EXPECT_NOT_POISONED(y[3]);
+  EXPECT_POISONED(y[4]);
+}
+
 TEST(MemorySanitizer, strcat) {
   char a[10];
   char b[] = "def";