From 208ab9a8e892305c7ec37cc648a60af26da5dee6 Mon Sep 17 00:00:00 2001
From: Matthew Parkinson <mattpark@microsoft.com>
Date: Fri, 9 Apr 2021 08:35:00 +0100
Subject: [PATCH] Rederive allocator id for remotes.

Storing a pointer to an allocator id in an unused object could be a
gadget for escallating priviledge of an attacker, by enabling
use-after-free to corrupt the allocator structure, and then create more
damage.

This commit adds an alternative implementation that does not cache the
allocator id.
---
 src/mem/alloc.h           |  9 ++-----
 src/mem/remoteallocator.h | 52 ++++++++++++++++++++++++++++++++++++---
 2 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/src/mem/alloc.h b/src/mem/alloc.h
index 66540c3..b628aa1 100644
--- a/src/mem/alloc.h
+++ b/src/mem/alloc.h
@@ -734,16 +734,11 @@ namespace snmalloc
 
     SNMALLOC_FAST_PATH void handle_dealloc_remote(CapPtr<Remote, CBAlloc> p)
     {
-      if (likely(p->trunc_target_id() == get_trunc_id()))
+      if (likely(Remote::trunc_target_id(p, &large_allocator) == get_trunc_id()))
       {
         // Destined for my slabs
         auto p_auth = large_allocator.template capptr_amplify<Remote>(p);
         auto super = Superslab::get(p_auth);
-
-        check_client(
-          p->trunc_target_id() == super->get_allocator()->trunc_id(),
-          "Detected memory corruption.  Potential use-after-free");
-
         dealloc_not_large_local(super, p, p->sizeclass());
       }
       else
@@ -751,7 +746,7 @@ namespace snmalloc
         // Merely routing; despite the cast here, p is going to be cast right
         // back to a Remote.
         remote.dealloc<Allocator>(
-          p->trunc_target_id(),
+          Remote::trunc_target_id(p, &large_allocator),
           p.template as_reinterpret<FreeObject>(),
           p->sizeclass());
       }
diff --git a/src/mem/remoteallocator.h b/src/mem/remoteallocator.h
index bc3e743..26d6867 100644
--- a/src/mem/remoteallocator.h
+++ b/src/mem/remoteallocator.h
@@ -8,6 +8,10 @@
 
 #include <atomic>
 
+#ifdef CHECK_CLIENT
+#  define SNMALLOC_DONT_CACHE_ALLOCATOR_PTR
+#endif
+
 namespace snmalloc
 {
   /*
@@ -24,7 +28,21 @@ namespace snmalloc
       AtomicCapPtr<Remote, CBAlloc> next{nullptr};
     };
 
-    /*
+#ifdef SNMALLOC_DONT_CACHE_ALLOCATOR_PTR
+    /**
+     * Cache the size class of the object to improve performance.
+     *
+     * This implementation does not cache the allocator id due to security
+     * concerns. Alternative implementations may store the allocator
+     * id, so that amplification costs can be mitigated on CHERI with MTE.
+     */
+    sizeclass_t sizeclasscache;
+#else
+    /* This implementation assumes that storing the allocator ID in a freed
+     * object is not a security concern.  Either we trust the code running on
+     * top of the allocator, or additional security measure are in place such
+     * as MTE + CHERI.
+     *
      * We embed the size class in the bottom 8 bits of an allocator ID (i.e.,
      * the address of an Alloc's remote_alloc's message_queue; in practice we
      * only need 7 bits, but using 8 is conjectured to be faster).  The hashing
@@ -38,20 +56,45 @@ namespace snmalloc
      * and use SNMALLOC_ASSERT to verify that they do not exist in debug builds.
      */
     alloc_id_t alloc_id_and_sizeclass;
+#endif
 
+    /**
+     * Set up a remote object.  Potentially cache sizeclass and allocator id.
+     */
     void set_info(alloc_id_t id, sizeclass_t sc)
     {
+#ifdef SNMALLOC_DONT_CACHE_ALLOCATOR_PTR
+      UNUSED(id);
+      sizeclasscache = sc;
+#else
       alloc_id_and_sizeclass = (id & ~SIZECLASS_MASK) | sc;
+#endif
     }
 
-    alloc_id_t trunc_target_id()
+    /**
+     * Return allocator for this object.  This may perform amplification.
+     */
+    template<typename LargeAlloc>
+    static alloc_id_t trunc_target_id(CapPtr<Remote, CBAlloc> r, LargeAlloc* large_allocator)
     {
-      return alloc_id_and_sizeclass & ~SIZECLASS_MASK;
+#ifdef SNMALLOC_DONT_CACHE_ALLOCATOR_PTR
+      // Rederive allocator id.
+      auto r_auth = large_allocator->template capptr_amplify<Remote>(r);
+      auto super = Superslab::get(r_auth);
+      return super->get_allocator()->trunc_id();
+#else
+      UNUSED(large_allocator);
+      return r->alloc_id_and_sizeclass & ~SIZECLASS_MASK;
+#endif
     }
 
     sizeclass_t sizeclass()
     {
+#ifdef SNMALLOC_DONT_CACHE_ALLOCATOR_PTR
+      return sizeclasscache;
+#else
       return alloc_id_and_sizeclass & SIZECLASS_MASK;
+#endif
     }
 
     /** Zero out a Remote tracking structure, return pointer to object base */
@@ -204,7 +247,8 @@ namespace snmalloc
         {
           // Use the next N bits to spread out remote deallocs in our own
           // slot.
-          size_t slot = get_slot<Alloc>(r->trunc_target_id(), post_round);
+          size_t slot = get_slot<Alloc>(
+            Remote::trunc_target_id(r, &allocator->large_allocator), post_round);
           RemoteList* l = &list[slot];
           l->last->non_atomic_next = r;
           l->last = r;