From bb5027b45428ad04849dac422eb3ba69da5092e9 Mon Sep 17 00:00:00 2001
From: Matthew Parkinson <mattpark@microsoft.com>
Date: Mon, 21 Jan 2019 17:44:09 +0000
Subject: [PATCH] Harden sizeclass table

If a sizeclass in the metadata is corrupted, then this can be used to
force an index beyond the end of these tables. This extends the tables
to the next power of two, and uses a mask on the index, so they are
always either a valid piece of data, or zero.
---
 src/ds/helpers.h         | 19 +++++++++++++++++++
 src/mem/sizeclasstable.h | 14 ++++++++------
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/src/ds/helpers.h b/src/ds/helpers.h
index fc4461c..44d61f3 100644
--- a/src/ds/helpers.h
+++ b/src/ds/helpers.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include "bits.h"
 #include "flaglock.h"
 
 namespace snmalloc
@@ -30,4 +31,22 @@ namespace snmalloc
       return obj;
     }
   };
+
+  template <size_t length, typename T>
+  class ModArray
+  {
+    static constexpr size_t rlength = bits::next_pow2_const(length);
+    T array[rlength];
+  
+  public:
+    constexpr const T &operator[] (const size_t i) const
+    {
+      return array[i & (rlength - 1)];
+    } 
+
+    constexpr T &operator[] (const size_t i)
+    {
+      return array[i & (rlength - 1)];
+    }  
+  };
 }
diff --git a/src/mem/sizeclasstable.h b/src/mem/sizeclasstable.h
index 099b465..21aa913 100644
--- a/src/mem/sizeclasstable.h
+++ b/src/mem/sizeclasstable.h
@@ -1,16 +1,17 @@
 #pragma once
 
 #include "superslab.h"
+#include "../ds/helpers.h"
 
 namespace snmalloc
 {
   struct SizeClassTable
   {
-    size_t size[NUM_SIZECLASSES];
-    uint16_t bump_ptr_start[NUM_SMALL_CLASSES];
-    uint16_t short_bump_ptr_start[NUM_SMALL_CLASSES];
-    uint16_t count_per_slab[NUM_SMALL_CLASSES];
-    uint16_t medium_slab_slots[NUM_MEDIUM_CLASSES];
+    ModArray<NUM_SIZECLASSES, size_t> size;
+    ModArray<NUM_SMALL_CLASSES, uint16_t> bump_ptr_start;
+    ModArray<NUM_SMALL_CLASSES, uint16_t> short_bump_ptr_start;
+    ModArray<NUM_SMALL_CLASSES, uint16_t> count_per_slab;
+    ModArray<NUM_MEDIUM_CLASSES, uint16_t> medium_slab_slots;
 
     constexpr SizeClassTable()
     : size(),
@@ -66,6 +67,7 @@ namespace snmalloc
 
   constexpr static inline uint16_t medium_slab_free(uint8_t sizeclass)
   {
-    return sizeclass_metadata.medium_slab_slots[sizeclass - NUM_SMALL_CLASSES];
+    return sizeclass_metadata.medium_slab_slots
+      [(sizeclass - NUM_SMALL_CLASSES)];
   }
 }