From 0cd36f4eb27e68c2337bb13c9f973d16e0f190d3 Mon Sep 17 00:00:00 2001
From: Nathaniel Wesley Filardo <nfilardo@microsoft.com>
Date: Sat, 10 Dec 2022 11:11:07 +0000
Subject: [PATCH] StrictProvenance: plumb Authmaps through backends

No use of them, yet, though.
---
 src/snmalloc/backend/backend.h                |  1 +
 src/snmalloc/backend/fixedglobalconfig.h      | 19 +++++++++++---
 src/snmalloc/backend/globalconfig.h           | 25 ++++++++++++++-----
 .../backend_helpers/defaultpagemapentry.h     |  7 +++++-
 src/test/func/domestication/domestication.cc  | 19 ++++++++++----
 5 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/src/snmalloc/backend/backend.h b/src/snmalloc/backend/backend.h
index 581f67d..0496b9a 100644
--- a/src/snmalloc/backend/backend.h
+++ b/src/snmalloc/backend/backend.h
@@ -12,6 +12,7 @@ namespace snmalloc
     SNMALLOC_CONCEPT(IsPAL) PAL,
     typename PagemapEntry,
     typename Pagemap,
+    typename Authmap,
     typename LocalState>
   class BackendAllocator
   {
diff --git a/src/snmalloc/backend/fixedglobalconfig.h b/src/snmalloc/backend/fixedglobalconfig.h
index d8b0c0b..b45db72 100644
--- a/src/snmalloc/backend/fixedglobalconfig.h
+++ b/src/snmalloc/backend/fixedglobalconfig.h
@@ -20,12 +20,25 @@ namespace snmalloc
 
     using Pagemap = BasicPagemap<PAL, ConcretePagemap, PagemapEntry, true>;
 
+    struct Authmap
+    {
+      static inline capptr::Arena<void> arena;
+
+      template<bool potentially_out_of_range = false>
+      static SNMALLOC_FAST_PATH capptr::Arena<void>
+      amplify(capptr::Alloc<void> c)
+      {
+        return Aal::capptr_rebound(arena, c);
+      }
+    };
+
   public:
     using LocalState = StandardLocalState<PAL, Pagemap>;
 
     using GlobalPoolState = PoolState<CoreAllocator<FixedRangeConfig>>;
 
-    using Backend = BackendAllocator<PAL, PagemapEntry, Pagemap, LocalState>;
+    using Backend =
+      BackendAllocator<PAL, PagemapEntry, Pagemap, Authmap, LocalState>;
     using Pal = PAL;
 
   private:
@@ -74,9 +87,9 @@ namespace snmalloc
       auto [heap_base, heap_length] =
         Pagemap::concretePagemap.init(base, length);
 
-      auto heap_arena = capptr::Arena<void>::unsafe_from(heap_base);
+      Authmap::arena = capptr::Arena<void>::unsafe_from(heap_base);
 
-      Pagemap::register_range(heap_arena, heap_length);
+      Pagemap::register_range(Authmap::arena, heap_length);
 
       // Push memory into the global range.
       range_to_pow_2_blocks<MIN_CHUNK_BITS>(
diff --git a/src/snmalloc/backend/globalconfig.h b/src/snmalloc/backend/globalconfig.h
index 46614f7..0a69c86 100644
--- a/src/snmalloc/backend/globalconfig.h
+++ b/src/snmalloc/backend/globalconfig.h
@@ -55,15 +55,22 @@ namespace snmalloc
 
     using Pagemap = BasicPagemap<Pal, ConcretePagemap, PagemapEntry, false>;
 
+    using ConcreteAuthmap =
+      FlatPagemap<MinBaseSizeBits<Pal>(), capptr::Arena<void>, Pal, false>;
+
+    using Authmap = DefaultAuthmap<ConcreteAuthmap>;
+
     /**
-     * This specifies where this configurations sources memory from.
-     *
-     * Takes account of any platform specific constraints like whether
-     * mmap/virtual alloc calls can be consolidated.
+     * This specifies where this configurations sources memory from and the
+     * pagemap (and authmap) that maintain metadata about underlying OS
+     * allocations.
      * @{
      */
 
-    using Base = Pipe<PalRange<Pal>, PagemapRegisterRange<Pagemap>>;
+    using Base = Pipe<
+      PalRange<Pal>,
+      PagemapRegisterRange<Pagemap>,
+      PagemapRegisterRange<Authmap>>;
 
     /**
      * @}
@@ -81,7 +88,8 @@ namespace snmalloc
     /**
      * Use the default backend.
      */
-    using Backend = BackendAllocator<Pal, PagemapEntry, Pagemap, LocalState>;
+    using Backend =
+      BackendAllocator<Pal, PagemapEntry, Pagemap, Authmap, LocalState>;
 
   private:
     SNMALLOC_REQUIRE_CONSTINIT
@@ -139,6 +147,11 @@ namespace snmalloc
 
       Pagemap::concretePagemap.template init<pagemap_randomize>();
 
+      if constexpr (aal_supports<StrictProvenance>)
+      {
+        Authmap::init();
+      }
+
       initialised = true;
     }
 
diff --git a/src/snmalloc/backend_helpers/defaultpagemapentry.h b/src/snmalloc/backend_helpers/defaultpagemapentry.h
index de5ac30..034a2f0 100644
--- a/src/snmalloc/backend_helpers/defaultpagemapentry.h
+++ b/src/snmalloc/backend_helpers/defaultpagemapentry.h
@@ -24,7 +24,12 @@ namespace snmalloc
     /**
      * The private initialising constructor is usable only by this back end.
      */
-    template<SNMALLOC_CONCEPT(IsPAL) A1, typename A2, typename A3, typename A4>
+    template<
+      SNMALLOC_CONCEPT(IsPAL) A1,
+      typename A2,
+      typename A3,
+      typename A4,
+      typename A5>
     friend class BackendAllocator;
 
     /**
diff --git a/src/test/func/domestication/domestication.cc b/src/test/func/domestication/domestication.cc
index 351643f..f4d2e06 100644
--- a/src/test/func/domestication/domestication.cc
+++ b/src/test/func/domestication/domestication.cc
@@ -32,15 +32,23 @@ namespace snmalloc
   public:
     using Pagemap = BasicPagemap<Pal, ConcretePagemap, PagemapEntry, false>;
 
+    using ConcreteAuthmap =
+      FlatPagemap<MinBaseSizeBits<Pal>(), capptr::Arena<void>, Pal, false>;
+
+    using Authmap = DefaultAuthmap<ConcreteAuthmap>;
+
   public:
-    using LocalState = StandardLocalState<
-      Pal,
-      Pagemap,
-      Pipe<PalRange<Pal>, PagemapRegisterRange<Pagemap>>>;
+    using Base = Pipe<
+      PalRange<Pal>,
+      PagemapRegisterRange<Pagemap>,
+      PagemapRegisterRange<Authmap>>;
+
+    using LocalState = StandardLocalState<Pal, Pagemap, Base>;
 
     using GlobalPoolState = PoolState<CoreAllocator<CustomConfig>>;
 
-    using Backend = BackendAllocator<Pal, PagemapEntry, Pagemap, LocalState>;
+    using Backend =
+      BackendAllocator<Pal, PagemapEntry, Pagemap, Authmap, LocalState>;
 
   private:
     SNMALLOC_REQUIRE_CONSTINIT
@@ -128,6 +136,7 @@ int main()
 #  endif
 
   snmalloc::CustomConfig::Pagemap::concretePagemap.init<pagemap_randomize>();
+  snmalloc::CustomConfig::Authmap::init();
   snmalloc::CustomConfig::domesticate_count = 0;
 
   LocalEntropy entropy;