From e3eaa7bb35b2981d848e9806fdfc98585c456cb6 Mon Sep 17 00:00:00 2001
From: Evgeniy Stepanov <eugeni.stepanov@gmail.com>
Date: Mon, 2 Sep 2013 09:24:53 +0000
Subject: [PATCH] [msan] Intercept memalign, valloc, pvalloc.

PR17039

llvm-svn: 189750
---
 compiler-rt/lib/msan/msan_interceptors.cc | 28 +++++++++++++++++++++++
 compiler-rt/lib/msan/tests/msan_test.cc   | 28 +++++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/compiler-rt/lib/msan/msan_interceptors.cc b/compiler-rt/lib/msan/msan_interceptors.cc
index 52c1913c6ece..88c9f86eb370 100644
--- a/compiler-rt/lib/msan/msan_interceptors.cc
+++ b/compiler-rt/lib/msan/msan_interceptors.cc
@@ -139,6 +139,31 @@ INTERCEPTOR(int, posix_memalign, void **memptr, SIZE_T alignment, SIZE_T size) {
   return 0;
 }
 
+INTERCEPTOR(void *, memalign, SIZE_T boundary, SIZE_T size) {
+  GET_MALLOC_STACK_TRACE;
+  CHECK_EQ(boundary & (boundary - 1), 0);
+  void *ptr = MsanReallocate(&stack, 0, size, boundary, false);
+  return ptr;
+}
+
+INTERCEPTOR(void *, valloc, SIZE_T size) {
+  GET_MALLOC_STACK_TRACE;
+  void *ptr = MsanReallocate(&stack, 0, size, GetPageSizeCached(), false);
+  return ptr;
+}
+
+INTERCEPTOR(void *, pvalloc, SIZE_T size) {
+  GET_MALLOC_STACK_TRACE;
+  uptr PageSize = GetPageSizeCached();
+  size = RoundUpTo(size, PageSize);
+  if (size == 0) {
+    // pvalloc(0) should allocate one page.
+    size = PageSize;
+  }
+  void *ptr = MsanReallocate(&stack, 0, size, PageSize, false);
+  return ptr;
+}
+
 INTERCEPTOR(void, free, void *ptr) {
   ENSURE_MSAN_INITED();
   if (ptr == 0) return;
@@ -1188,6 +1213,9 @@ void InitializeInterceptors() {
   INTERCEPT_FUNCTION(mmap);
   INTERCEPT_FUNCTION(mmap64);
   INTERCEPT_FUNCTION(posix_memalign);
+  INTERCEPT_FUNCTION(memalign);
+  INTERCEPT_FUNCTION(valloc);
+  INTERCEPT_FUNCTION(pvalloc);
   INTERCEPT_FUNCTION(malloc);
   INTERCEPT_FUNCTION(calloc);
   INTERCEPT_FUNCTION(realloc);
diff --git a/compiler-rt/lib/msan/tests/msan_test.cc b/compiler-rt/lib/msan/tests/msan_test.cc
index 717c8930cf5a..244e3e9a8d12 100644
--- a/compiler-rt/lib/msan/tests/msan_test.cc
+++ b/compiler-rt/lib/msan/tests/msan_test.cc
@@ -26,6 +26,7 @@
 #include <assert.h>
 #include <wchar.h>
 #include <math.h>
+#include <malloc.h>
 
 #include <arpa/inet.h>
 #include <dlfcn.h>
@@ -56,6 +57,8 @@
 # define MSAN_HAS_M128 0
 #endif
 
+static const int kPageSize = 4096;
+
 typedef unsigned char      U1;
 typedef unsigned short     U2;  // NOLINT
 typedef unsigned int       U4;
@@ -2253,6 +2256,31 @@ TEST(MemorySanitizer, posix_memalign) {
   int res = posix_memalign(&p, 4096, 13);
   ASSERT_EQ(0, res);
   EXPECT_NOT_POISONED(p);
+  EXPECT_EQ(0U, (uintptr_t)p % 4096);
+  free(p);
+}
+
+TEST(MemorySanitizer, memalign) {
+  void *p = memalign(4096, 13);
+  EXPECT_EQ(0U, (uintptr_t)p % kPageSize);
+  free(p);
+}
+
+TEST(MemorySanitizer, valloc) {
+  void *a = valloc(100);
+  EXPECT_EQ(0U, (uintptr_t)a % kPageSize);
+  free(a);
+}
+
+TEST(MemorySanitizer, pvalloc) {
+  void *p = pvalloc(kPageSize + 100);
+  EXPECT_EQ(0U, (uintptr_t)p % kPageSize);
+  EXPECT_EQ(2 * kPageSize, __msan_get_allocated_size(p));
+  free(p);
+
+  p = pvalloc(0);  // pvalloc(0) should allocate at least one page.
+  EXPECT_EQ(0U, (uintptr_t)p % kPageSize);
+  EXPECT_EQ(kPageSize, __msan_get_allocated_size(p));
   free(p);
 }