[sanitizers] Fixes strndup API behaviour when intercepted by sanitizers

Sanitizers ignore flag allocator_may_return_null=1 in strndup() calls.
When OOM is emulated, this causes to the unexpected crash.

Committed by pgousseau on behalf of "Kostyantyn Melnik, kmnls.kmnls@gmail.com"

Reviewed by: pgousseau

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

2 files changed, 24 insertions, 7 deletions

diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_common_interceptors.inc b/compiler-rt/lib/sanitizer_common/sanitizer_common_interceptors.inc
index 99bd59a..3b316f5 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_common_interceptors.inc
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_common_interceptors.inc
@@ -385,9 +385,11 @@ extern const short *_tolower_tab_;
   if (common_flags()->intercept_strndup) {                                    \
     COMMON_INTERCEPTOR_READ_STRING(ctx, s, Min(size, copy_length + 1));       \
   }                                                                           \
-  COMMON_INTERCEPTOR_COPY_STRING(ctx, new_mem, s, copy_length);               \
-  internal_memcpy(new_mem, s, copy_length);                                   \
-  new_mem[copy_length] = '\0';                                                \
+  if (new_mem) {                                                              \
+    COMMON_INTERCEPTOR_COPY_STRING(ctx, new_mem, s, copy_length);             \
+    internal_memcpy(new_mem, s, copy_length);                                 \
+    new_mem[copy_length] = '\0';                                              \
+  }                                                                           \
   return new_mem;
 #endif
 
diff --git a/compiler-rt/test/sanitizer_common/TestCases/max_allocation_size.cpp b/compiler-rt/test/sanitizer_common/TestCases/max_allocation_size.cpp
index 9cc799a..fc9bc3a 100644
--- a/compiler-rt/test/sanitizer_common/TestCases/max_allocation_size.cpp
+++ b/compiler-rt/test/sanitizer_common/TestCases/max_allocation_size.cpp
@@ -35,6 +35,10 @@
 // RUN:   | FileCheck %s --check-prefix=CHECK-nnCRASH
 // RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=1 \
 // RUN:   %run %t new-nothrow 2>&1 | FileCheck %s --check-prefix=CHECK-NULL
+// RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=0 \
+// RUN:   not %run %t strndup 2>&1 | FileCheck %s --check-prefix=CHECK-sCRASH
+// RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=1 \
+// RUN:   %run %t strndup 2>&1 | FileCheck %s --check-prefix=CHECK-NULL
 
 // win32 is disabled due to failing errno tests.
 // UNSUPPORTED: ubsan, windows-msvc
@@ -47,6 +51,8 @@
 #include <stdlib.h>
 #include <string.h>
 
+constexpr size_t MaxAllocationSize = size_t{2} << 20;
+
 static void *allocate(const char *Action, size_t Size) {
   if (!strcmp(Action, "malloc"))
     return malloc(Size);
@@ -65,12 +71,21 @@ static void *allocate(const char *Action, size_t Size) {
     return ::operator new(Size);
   if (!strcmp(Action, "new-nothrow"))
     return ::operator new(Size, std::nothrow);
+  if (!strcmp(Action, "strndup")) {
+    static char pstr[MaxAllocationSize + 1] = {'a'};
+    for (size_t i = 0; i < MaxAllocationSize + 1; i++)
+      pstr[i] = 'a';
+    if (Size == MaxAllocationSize)
+      pstr[MaxAllocationSize - 1] = '\0';
+    return strndup(pstr, Size);
+  }
   assert(0);
 }
 
 static void deallocate(const char *Action, void *Ptr) {
   if (!strcmp(Action, "malloc") || !strcmp(Action, "calloc") ||
-      !strcmp(Action, "realloc") || !strcmp(Action, "realloc-after-malloc"))
+      !strcmp(Action, "realloc") || !strcmp(Action, "realloc-after-malloc") ||
+      !strcmp(Action, "strndup"))
     return free(Ptr);
   if (!strcmp(Action, "new"))
     return ::operator delete(Ptr);
@@ -84,8 +99,6 @@ int main(int Argc, char **Argv) {
   const char *Action = Argv[1];
   fprintf(stderr, "%s:\n", Action);
 
-  constexpr size_t MaxAllocationSize = size_t{2} << 20;
-
   // Should succeed when max_allocation_size_mb is set.
   void *volatile P = allocate(Action, MaxAllocationSize);
   assert(P);
@@ -120,8 +133,10 @@ int main(int Argc, char **Argv) {
 // CHECK-nCRASH-OOM: {{SUMMARY: .*Sanitizer: out-of-memory}}
 // CHECK-nnCRASH: new-nothrow:
 // CHECK-nnCRASH: {{SUMMARY: .*Sanitizer: allocation-size-too-big}}
+// CHECK-sCRASH: strndup:
+// CHECK-sCRASH: {{SUMMARY: .*Sanitizer: allocation-size-too-big}}
 
-// CHECK-NULL: {{malloc|calloc|calloc-overflow|realloc|realloc-after-malloc|new-nothrow}}
+// CHECK-NULL: {{malloc|calloc|calloc-overflow|realloc|realloc-after-malloc|new-nothrow|strndup}}
 // CHECK-NULL: errno: 12, P: 0
 //
 // CHECK-NOTNULL-NOT: P: 0