From 8a8669f9572329494ae9ab98cf69ebaf7f522b88 Mon Sep 17 00:00:00 2001
From: Matthew Parkinson <mattpark@microsoft.com>
Date: Mon, 19 Jul 2021 11:25:42 +0100
Subject: [PATCH] Add randomised start to pagemap layout

The pagemap contains a lot of important data.  This commit makes the
checked mode overallocate, and then start the pagemap at a random
offset within this range.
---
 src/backend/pagemap.h | 29 +++++++++++++++++++++++++----
 src/mem/entropy.h     |  2 ++
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/backend/pagemap.h b/src/backend/pagemap.h
index f1cda4a..22126c8 100644
--- a/src/backend/pagemap.h
+++ b/src/backend/pagemap.h
@@ -2,6 +2,7 @@
 
 #include "../ds/bits.h"
 #include "../ds/helpers.h"
+#include "../mem/entropy.h"
 #include "../pal/pal.h"
 
 #include <atomic>
@@ -95,6 +96,8 @@ namespace snmalloc
       body = reinterpret_cast<T*>(b);
 
       // Advance by size of pagemap.
+      // Note that base needs to be aligned to GRANULARITY for the rest of the
+      // code to work
       // TODO CHERI capability bound here!
       heap_base = pointer_align_up(
         pointer_offset(b, (size >> SHIFT) * sizeof(T)),
@@ -102,6 +105,7 @@ namespace snmalloc
       base = address_cast(heap_base);
       SNMALLOC_ASSERT(
         base == bits::align_up(base, bits::one_at_bit(GRANULARITY_BITS)));
+
       return {heap_base, pointer_diff(heap_base, end)};
     }
 
@@ -117,15 +121,32 @@ namespace snmalloc
         bits::ADDRESS_BITS - GRANULARITY_BITS;
       static constexpr size_t ENTRIES = bits::one_at_bit(COVERED_BITS);
 
-      // TODO request additional space, and move to random offset.
+      static constexpr size_t REQUIRED_SIZE = ENTRIES * sizeof(T);
 
-      auto new_body_untyped = Pal::reserve(ENTRIES * sizeof(T));
+#ifdef SNMALLOC_CHECK_CLIENT
+      // Allocate a power of two extra to allow the placement of the
+      // pagemap be difficult to guess.
+      size_t additional_size = bits::next_pow2(REQUIRED_SIZE) * 2;
+      size_t request_size = REQUIRED_SIZE + additional_size;
+#else
+      size_t request_size = REQUIRED_SIZE;
+#endif
 
+      auto new_body_untyped = PAL::reserve(request_size);
+
+#ifdef SNMALLOC_CHECK_CLIENT
+      // Begin pagemap at random offset within the additionally allocated space.
+      static_assert(bits::is_pow2(sizeof(T)), "Next line assumes this.");
+      size_t offset = get_entropy64<PAL>() & (additional_size - sizeof(T));
+      auto new_body =
+        reinterpret_cast<T*>(pointer_offset(new_body_untyped, offset));
+#else
       auto new_body = reinterpret_cast<T*>(new_body_untyped);
-
+#endif
       // Ensure bottom page is committed
       // ASSUME: new memory is zeroed.
-      Pal::notify_using<NoZero>(new_body, OS_PAGE_SIZE);
+      PAL::template notify_using<NoZero>(
+        pointer_align_down<OS_PAGE_SIZE>(new_body), OS_PAGE_SIZE);
 
       // Set up zero page
       new_body[0] = body[0];
diff --git a/src/mem/entropy.h b/src/mem/entropy.h
index 80bd6f4..d9b40ae 100644
--- a/src/mem/entropy.h
+++ b/src/mem/entropy.h
@@ -1,3 +1,5 @@
+#pragma once
+
 #include "../ds/address.h"
 #include "../pal/pal.h"