1 files changed, 175 insertions, 0 deletions
diff --git a/xpcom/tests/gtest/TestAllocReplacement.cpp b/xpcom/tests/gtest/TestAllocReplacement.cpp
new file mode 100644
index 000000000..7bcf3d4ad
--- /dev/null
+++ b/xpcom/tests/gtest/TestAllocReplacement.cpp
@@ -0,0 +1,175 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8 sts=2 et sw=2 tw=80: */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "mozilla/Attributes.h"
+#include "mozilla/Function.h"
+#include "mozilla/mozalloc.h"
+#include "mozilla/ScopeExit.h"
+#include "nsCOMPtr.h"
+#include "nsIMemoryReporter.h"
+#include "nsServiceManagerUtils.h"
+#include "gtest/gtest.h"
+
+// We want to ensure that various functions are hooked properly and that
+// allocations are getting routed through jemalloc.  The strategy
+// pursued below relies on jemalloc's statistics tracking: we measure
+// the size of the jemalloc heap using nsIMemoryReporterManager,
+// allocate a chunk of memory with whatever function is supposed to be
+// hooked, and then ask for the size of the jemalloc heap again.  If the
+// function has been hooked correctly, then the heap size should be
+// different between the two measurements.  We can also check the
+// hooking of |free| and similar functions: once we free() the returned
+// pointer, we can measure the jemalloc heap size again, expecting it to
+// be identical to the size prior to the allocation.
+//
+// If we're not using jemalloc, then nsIMemoryReporterManager will
+// simply report an error, and we will ignore the entire test.
+//
+// This strategy is not perfect: it relies on GTests being
+// single-threaded, which they are, and no other threads doing
+// allocation during the test, which is uncertain, as XPCOM has started
+// up during gtests, and who knows what might be going on behind the
+// scenes.  This latter assumption, however, does not seem to be a
+// problem in practice.
+#if defined(MOZ_MEMORY)
+#define ALLOCATION_ASSERT(b) ASSERT_TRUE((b))
+#else
+#define ALLOCATION_ASSERT(b) (void)(b)
+#endif
+
+#define ASSERT_ALLOCATION_HAPPENED(lambda)      \
+  ALLOCATION_ASSERT(ValidateHookedAllocation(lambda, free));
+
+// We do run the risk of OOM'ing when we allocate something...all we can
+// do is try to allocate something so small that OOM'ing is unlikely.
+const size_t kAllocAmount = 16;
+
+// We declare this function MOZ_NEVER_INLINE to work around optimizing
+// compilers.  If we permitted inlining here, then the compiler might
+// inline both this function and the calls to the function pointers we
+// pass in, giving something like:
+//
+//   void* p = malloc(...);
+//   ...do nothing with p except check nullptr-ness...
+//   free(p);
+//
+// and the optimizer can delete the calls to malloc and free entirely,
+// which would make checking that the jemalloc heap had never changed
+// difficult.
+static MOZ_NEVER_INLINE bool
+ValidateHookedAllocation(void* (*aAllocator)(void),
+                         void (*aFreeFunction)(void*))
+{
+  nsCOMPtr<nsIMemoryReporterManager> manager =
+    do_GetService("@mozilla.org/memory-reporter-manager;1");
+
+  int64_t before = 0;
+  nsresult rv = manager->GetHeapAllocated(&before);
+  if (NS_FAILED(rv)) {
+    return false;
+  }
+
+  {
+    void* p = aAllocator();
+
+    if (!p) {
+      return false;
+    }
+
+    int64_t after = 0;
+    rv = manager->GetHeapAllocated(&after);
+
+    // Regardless of whether that call succeeded or failed, we are done with
+    // the allocated buffer now.
+    aFreeFunction(p);
+
+    if (NS_FAILED(rv)) {
+      return false;
+    }
+
+    // Verify that our heap stats have changed.
+    if ((before + int64_t(kAllocAmount)) != after) {
+      return false;
+    }
+  }
+
+  // Verify that freeing the allocated pointer resets our heap to what it
+  // was before.
+  int64_t after = 0;
+  rv = manager->GetHeapAllocated(&after);
+  if (NS_FAILED(rv)) {
+    return false;
+  }
+
+  return before == after;
+}
+
+TEST(AllocReplacement, malloc_check)
+{
+  ASSERT_ALLOCATION_HAPPENED([] {
+    return malloc(kAllocAmount);
+  });
+}
+
+TEST(AllocReplacement, calloc_check)
+{
+  ASSERT_ALLOCATION_HAPPENED([] {
+    return calloc(1, kAllocAmount);
+  });
+}
+
+TEST(AllocReplacement, realloc_check)
+{
+  ASSERT_ALLOCATION_HAPPENED([] {
+    return realloc(nullptr, kAllocAmount);
+  });
+}
+
+#if defined(HAVE_POSIX_MEMALIGN)
+TEST(AllocReplacement, posix_memalign_check)
+{
+  ASSERT_ALLOCATION_HAPPENED([] {
+    void* p = nullptr;
+    int result = posix_memalign(&p, sizeof(void*), kAllocAmount);
+    if (result != 0) {
+      return static_cast<void*>(nullptr);
+    }
+    return p;
+  });
+}
+#endif
+
+#if defined(XP_WIN)
+#include <windows.h>
+
+#undef ASSERT_ALLOCATION_HAPPENED
+#define ASSERT_ALLOCATION_HAPPENED(lambda)      \
+  ALLOCATION_ASSERT(ValidateHookedAllocation(lambda, [](void* p) { \
+    HeapFree(GetProcessHeap(), 0, p); \
+  }));
+
+TEST(AllocReplacement, HeapAlloc_check)
+{
+  ASSERT_ALLOCATION_HAPPENED([] {
+    HANDLE h = GetProcessHeap();
+    return HeapAlloc(h, 0, kAllocAmount);
+  });
+}
+
+TEST(AllocReplacement, HeapReAlloc_check)
+{
+  ASSERT_ALLOCATION_HAPPENED([] {
+    HANDLE h = GetProcessHeap();
+    void *p = HeapAlloc(h, 0, kAllocAmount / 2);
+
+    if (!p) {
+      return static_cast<void*>(nullptr);
+    }
+
+    return HeapReAlloc(h, 0, p, kAllocAmount);
+  });
+}
+#endif