TCMalloc: support userland ASLR on Linux

third_party/tcmalloc/chromium/src/system-alloc.cc

Index: third_party/tcmalloc/chromium/src/system-alloc.cc
diff --git a/third_party/tcmalloc/chromium/src/system-alloc.cc b/third_party/tcmalloc/chromium/src/system-alloc.cc
index b755b3fdc9bb3e0d34e644b03b7e16582a3aa61a..b41c573c337001928dc0e46d22af092e02ffa051 100644
--- a/third_party/tcmalloc/chromium/src/system-alloc.cc
+++ b/third_party/tcmalloc/chromium/src/system-alloc.cc
@@ -100,6 +100,91 @@ template <> bool CheckAddressBits<8 * sizeof(void*)>(uintptr_t ptr) {
   return true;
 }
 
+// From libdieharder, public domain library by Bob Jenkins (rngav.c).
+// Described at http://burtleburtle.net/bob/rand/smallprng.html.
+// Not cryptographically secure, but good enough for what we need.
+typedef uint32_t u4;
+typedef struct ranctx { u4 a; u4 b; u4 c; u4 d; } ranctx;
+
+#define rot(x,k) ((x<<(k))|(x>>(32-(k))))
+
+static u4 ranval(ranctx* x) {
+  /* xxx: the generator being tested */
+  u4 e = x->a - rot(x->b, 27);
+  x->a = x->b ^ rot(x->c, 17);
+  x->b = x->c + x->d;
+  x->c = x->d + e;
+  x->d = e + x->a;
+  return x->d;
+}
+
+static u4 raninit(ranctx* x, u4 seed) {
+  u4 i, e;
+  x->a = x->b = x->c = 0xf1ea5eed;
+  x->d = seed - x->a;
+  for (i=0; i<20; ++i) {
+    e = ranval(x);
+  }
+  return e;
+}
+
+// End PRNG code.
+
+#define ASLR_IS_SUPPORTED \\
+    (defined(OS_LINUX) || defined(OS_CHROMEOS)) && defined(__x86_64__)

[Chris Evans, 2013/01/29 06:08:08]

Put in ARM64 to be future proof? :)

[jln (very slow on Chromium), 2013/01/29 06:23:46]

ARM will be very different in terms of the canonical address space and how the
kernel lays everything out.
I'll be happy to work on ARM64 version when Chrome supports that architecture :)

+
+// Give a random "hint" that is suitable for use with mmap(). This cannot make
+// mmap fail , as the kernel will simply not follow the hint if it can't.
+// However, this will create address space fragmentation.  Currently, we only
+// implement it on x86_64, where we have a 47 bits userland address space and
+// fragmentation is not an issue.
+void* GetRandomAddrHint() {
+#if defined(ASLR_IS_SUPPORTED)
+  // Note: we are protected by the general TCMalloc_SystemAlloc spinlock. Given
+  // the nature of what we're doing, it wouldn't be critical if we weren't.
+  // It's nice to share the state between threads, because scheduling will add
+  // some randomness to the succession of ranval() calls.
+  static ranctx ctx;
+  static bool initialized = false;
+  if (!initialized) {
+    volatile int c;
+    // Pre-initialize our seed, but /dev/urandom should always be available.
+    uint32_t seed = reinterpret_cast<unsigned long>(&c);

[Chris Evans, 2013/01/29 06:08:08]

I think I see what you're doing here -- using the entropy in a stack address
just in case the /dev/urandom thing fails.

On 64-bit Linux, the best 8 bits of entropy are often in the upper 4 bytes of
the pointer (with some more in the lower 4 bytes). But the unsigned long ->
uint32_t truncation loses it.
How about xor'ing the upper and lower 32-bits of the pointer address for form
"seed"?

Also add comment saying "grab a little entropy from a stack pointer in case
/dev/urandom fails"

[jln (very slow on Chromium), 2013/01/29 06:23:46]

It was a quick stopgap measure, but it doesn't hurt to make it better indeed.

Done :)

+    int urandom_fd = open("/dev/urandom", O_RDONLY);
+#if !defined(NDEBUG)

[Chris Evans, 2013/01/29 06:08:08]

Do you need this ifdef? Doesn't ASSERT() compile out in NDEBUG builds
automatically?

[jln (very slow on Chromium), 2013/01/29 06:23:46]

Yeah I didn't really know where ASSERT came from in tcmalloc. But I checked and
you're right.

Done.

+    ASSERT(urandom_fd >= 0);
+#endif
+    if (urandom_fd >= 0) {
+      ssize_t len;
+      len = read(urandom_fd, &seed, sizeof(seed));
+#if !defined(NDEBUG)
+    ASSERT(len == sizeof(seed));

[Chris Evans, 2013/01/29 06:08:08]

Same NDEBUG vs. ASSERT question.
And indentation issue -- should be two to the right?

[jln (very slow on Chromium), 2013/01/29 06:23:46]

Done.

+#endif
+      close(seed);

[Chris Evans, 2013/01/29 06:08:08]

I think you mean close(urandom_fd)? :)

[jln (very slow on Chromium), 2013/01/29 06:23:46]

Ooch. And I *always* check return values normally. Thanks :)

+    }
+    raninit(&ctx, seed);
+    initialized = true;
+  }
+  uint64_t random_address = static_cast<uint64_t>(ranval(&ctx)) << 32 |
+                            ranval(&ctx);
+  // If the kernel cannot honor the hint in arch_get_unmapped_area_topdown, it
+  // will simply ignore it. So we give a hint that has a good chance of
+  // working.
+  // The mmap top-down allocator will normally allocate below TASK_SIZE - gap,
+  // with a gap that depends on the max stack size. See x86/mm/mmap.c. We
+  // should make allocations that are below this area, which would be
+  // 0x7ffbf8000000.
+  // We use 0x3ffffffff000 as the mask so that we only "pollute" half of the
+  // address space. In the unlikely case where fragmentation would become an
+  // issue, the kernel will still have another half to use.
+  // A a bit-wise "and" won't bias our random distribution.
+  random_address &= 0x3ffffffff000;

[Chris Evans, 2013/01/29 06:08:08]

Postfix constant with ULL? Some compilers might otherwise throw a fit.

[jln (very slow on Chromium), 2013/01/29 06:23:46]

Given that it's a &=, I think it's guaranteed ok, but it doesn't hurt. Done.

+  return reinterpret_cast<void*>(random_address);
+#else
+  return NULL;
+#endif  // ASLR_IS_SUPPORTED
+}
+
 }  // Anonymous namespace to avoid name conflicts on "CheckAddressBits".
 
 COMPILE_ASSERT(kAddressBits <= 8 * sizeof(void*),
@@ -139,6 +224,15 @@ DEFINE_bool(malloc_skip_mmap,
             EnvToBool("TCMALLOC_SKIP_MMAP", false),
             "Whether mmap can be used to obtain memory.");
 
+DEFINE_bool(malloc_random_allocator,
+            EnvToBool("TCMALLOC_ASLR",
+#if defined(ASLR_IS_SUPPORTED)
+                      true),
+#else
+                      false),
+#endif
+            "Whether to randomize the address space via mmap().");
+
 // static allocators
 class SbrkSysAllocator : public SysAllocator {
 public:
@@ -304,7 +398,11 @@ void* MmapSysAllocator::Alloc(size_t size, size_t *actual_size,
   //            size + alignment < (1<<NBITS).
   // and        extra <= alignment
   // therefore  size + extra < (1<<NBITS)
-  void* result = mmap(NULL, size + extra,
+  void* address_hint = NULL;
+  if (FLAGS_malloc_random_allocator) {
+    address_hint = GetRandomAddrHint();
+  }
+  void* result = mmap(address_hint, size + extra,
                       PROT_READ|PROT_WRITE,
                       MAP_PRIVATE|MAP_ANONYMOUS,
                       -1, 0);
@@ -445,6 +543,14 @@ void InitSystemAllocators(void) {
   MmapSysAllocator *mmap = new (mmap_space) MmapSysAllocator();
   SbrkSysAllocator *sbrk = new (sbrk_space) SbrkSysAllocator();
 
+  DefaultSysAllocator *sdef = new (default_space) DefaultSysAllocator();
+
+  // Unfortunately, this code runs before flags are initialized. So
+  // we can't use FLAGS_malloc_random_allocator.
+#if defined(ASLR_IS_SUPPORTED)
+  // Our only random allocator is mmap.
+  sdef->SetChildAllocator(mmap, 0, mmap_name);
+#else
   // In 64-bit debug mode, place the mmap allocator first since it
   // allocates pointers that do not fit in 32 bits and therefore gives
   // us better testing of code's 64-bit correctness.  It also leads to
@@ -452,7 +558,6 @@ void InitSystemAllocators(void) {
   // likely to look like pointers and therefore the conservative gc in
   // the heap-checker is less likely to misinterpret a number as a
   // pointer).
-  DefaultSysAllocator *sdef = new (default_space) DefaultSysAllocator();
   if (kDebugMode && sizeof(void*) > 4) {
     sdef->SetChildAllocator(mmap, 0, mmap_name);
     sdef->SetChildAllocator(sbrk, 1, sbrk_name);
@@ -460,6 +565,7 @@ void InitSystemAllocators(void) {
     sdef->SetChildAllocator(sbrk, 0, sbrk_name);
     sdef->SetChildAllocator(mmap, 1, mmap_name);
   }
+#endif  // ASLR_IS_SUPPORTED
   sys_alloc = sdef;
 }