From 888d182baccb3288e41172fec3e0501158418613 Mon Sep 17 00:00:00 2001
From: Matthew Parkinson <mattpark@microsoft.com>
Date: Tue, 10 May 2022 10:45:59 +0100
Subject: [PATCH] Check size is correct if specified

Some secure allocators check that the C++ supplied size is correct
relative to the meta-data. This adds a check to the secure version of
snmalloc to do that.
---
 src/snmalloc/mem/localalloc.h     | 20 ++++++++++++++++++--
 src/snmalloc/mem/sizeclasstable.h |  5 +++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/src/snmalloc/mem/localalloc.h b/src/snmalloc/mem/localalloc.h
index 086fd8a..cd1b6e5 100644
--- a/src/snmalloc/mem/localalloc.h
+++ b/src/snmalloc/mem/localalloc.h
@@ -675,16 +675,32 @@ namespace snmalloc
 #endif
     }
 
+    void check_size(void* p, size_t size)
+    {
+#ifdef SNMALLOC_CHECK_CLIENT
+      size = size == 0 ? 1 : size;
+      auto sc = size_to_sizeclass_full(size);
+      auto pm_sc =
+        Backend::Pagemap::get_metaentry(address_cast(p)).get_sizeclass();
+      auto rsize = sizeclass_full_to_size(sc);
+      auto pm_size = sizeclass_full_to_size(pm_sc);
+      snmalloc_check_client(
+        sc == pm_sc, "Dealloc rounded size mismatch: {} != {}", rsize, pm_size);
+#else
+      UNUSED(p, size);
+#endif
+    }
+
     SNMALLOC_FAST_PATH void dealloc(void* p, size_t s)
     {
-      UNUSED(s);
+      check_size(p, s);
       dealloc(p);
     }
 
     template<size_t size>
     SNMALLOC_FAST_PATH void dealloc(void* p)
     {
-      UNUSED(size);
+      check_size(p, size);
       dealloc(p);
     }
 
diff --git a/src/snmalloc/mem/sizeclasstable.h b/src/snmalloc/mem/sizeclasstable.h
index b257d07..fa85d94 100644
--- a/src/snmalloc/mem/sizeclasstable.h
+++ b/src/snmalloc/mem/sizeclasstable.h
@@ -120,6 +120,11 @@ namespace snmalloc
     {
       return value == 0;
     }
+
+    constexpr bool operator==(sizeclass_t other)
+    {
+      return value == other.value;
+    }
   };
 
   using sizeclass_compress_t = uint8_t;