From af8ab2daf6455fa1f9d9daedeb65f6abc2d2d50e Mon Sep 17 00:00:00 2001
From: Robert Norton <robert.norton@microsoft.com>
Date: Wed, 9 Feb 2022 16:56:53 +0000
Subject: [PATCH] Clear freelist pointers on allocation for CHERI or
 CHECK_CLIENT builds.

This is especially important on CHERI to avoid leaking capabilities to
the freelist. In the CHERI case we also zero in clear_slab (see comment).
Also add a check in the malloc functional test that there are no valid
capabilities in the returned allocation.
---
 src/mem/corealloc.h            | 11 +++++++++++
 src/mem/freelist.h             | 17 ++++++++++++++++-
 src/test/func/malloc/malloc.cc | 18 ++++++++++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/src/mem/corealloc.h b/src/mem/corealloc.h
index 760ca21..f621e0c 100644
--- a/src/mem/corealloc.h
+++ b/src/mem/corealloc.h
@@ -339,6 +339,17 @@ namespace snmalloc
         address_cast(start_of_slab) ==
         address_cast(chunk_record->meta_common.chunk));
 
+#if defined(__CHERI_PURE_CAPABILITY__) && !defined(SNMALLOC_CHECK_CLIENT)
+      // Zero the whole slab. For CHERI we at least need to clear the freelist
+      // pointers to avoid leaking capabilities but we do not need to do it in
+      // the freelist order as for SNMALLOC_CHECK_CLIENT. Zeroing the whole slab
+      // may be more friendly to hw because it does not involve pointer chasing
+      // and is amenable to prefetching.
+      SharedStateHandle::Pal::zero(
+        chunk_record->meta_common.chunk.unsafe_ptr(),
+        snmalloc::sizeclass_to_slab_size(sizeclass));
+#endif
+
 #ifdef SNMALLOC_TRACING
       std::cout << "Slab " << start_of_slab.unsafe_ptr()
                 << " is unused, Object sizeclass " << sizeclass << std::endl;
diff --git a/src/mem/freelist.h b/src/mem/freelist.h
index fbd8509..bac781c 100644
--- a/src/mem/freelist.h
+++ b/src/mem/freelist.h
@@ -185,6 +185,21 @@ namespace snmalloc
             signed_prev == this->prev_encoded,
             "Heap corruption - free list corrupted!");
         }
+
+        /**
+         * Clean up this object when removing it from the list. This is
+         * important on CHERI to avoid leaking capabilities. On CHECK_CLIENT
+         * builds it might increase the difficulty to bypass the checks.
+         */
+        void cleanup()
+        {
+#if defined(__CHERI_PURE_CAPABILITY__) || defined(SNMALLOC_CHECK_CLIENT)
+          this->next_object = nullptr;
+#  ifdef SNMALLOC_CHECK_CLIENT
+          this->prev_encoded = 0;
+#  endif
+#endif
+        }
       };
 
       // Note the inverted template argument order, since BView is inferable.
@@ -467,7 +482,7 @@ namespace snmalloc
 #else
         UNUSED(key);
 #endif
-
+        c->cleanup();
         return c;
       }
     };
diff --git a/src/test/func/malloc/malloc.cc b/src/test/func/malloc/malloc.cc
index 43fda4a..367be7c 100644
--- a/src/test/func/malloc/malloc.cc
+++ b/src/test/func/malloc/malloc.cc
@@ -60,6 +60,24 @@ void check_result(size_t size, size_t align, void* p, int err, bool null)
       "Cheri size is %zu, but required to be %zu.\n", cheri_size, alloc_size);
     failed = true;
   }
+  if (p != nullptr)
+  {
+    /*
+     * Scan the allocation for any tagged capabilities. Since this test doesn't
+     * use the allocated memory if there is a valid cap it must have leaked from
+     * the allocator, which is bad.
+     */
+    void** vp = static_cast<void**>(p);
+    for (size_t n = 0; n < alloc_size / sizeof(*vp); vp++, n++)
+    {
+      void* c = *vp;
+      if (__builtin_cheri_tag_get(c))
+      {
+        printf("Found cap tag set in alloc: %#p at %#p\n", c, vp);
+        failed = true;
+      }
+    }
+  }
 #endif
   if (exact_size && (alloc_size != expected_size) && (size != 0))
   {