From ca69fe0dd375ba1091c64d28a48145e2e03a852d Mon Sep 17 00:00:00 2001
From: Nathaniel Wesley Filardo <nfilardo@microsoft.com>
Date: Sat, 10 Dec 2022 11:02:48 +0000
Subject: [PATCH] StrictProvenance: Add "authority map" implementations

---
 src/snmalloc/backend_helpers/authmap.h        | 78 +++++++++++++++++++
 .../backend_helpers/backend_helpers.h         |  1 +
 2 files changed, 79 insertions(+)
 create mode 100644 src/snmalloc/backend_helpers/authmap.h

diff --git a/src/snmalloc/backend_helpers/authmap.h b/src/snmalloc/backend_helpers/authmap.h
new file mode 100644
index 0000000..9ad4f06
--- /dev/null
+++ b/src/snmalloc/backend_helpers/authmap.h
@@ -0,0 +1,78 @@
+#pragma once
+
+#include "commonconfig.h"
+#include "pagemapregisterrange.h"
+
+namespace snmalloc
+{
+  /**
+   * A dummy authmap that simply passes pointers through. For use on
+   * non-StrictProvenance architectures.
+   */
+  struct DummyAuthmap
+  {
+    static SNMALLOC_FAST_PATH void init() {}
+
+    static SNMALLOC_FAST_PATH void register_range(capptr::Arena<void>, size_t)
+    {}
+
+    template<bool potentially_out_of_range = false>
+    static SNMALLOC_FAST_PATH capptr::Arena<void> amplify(capptr::Alloc<void> c)
+    {
+      return capptr::Arena<void>::unsafe_from(c.unsafe_ptr());
+    }
+  };
+
+  /**
+   * Wrap a concrete Pagemap to store Arena pointers, and use these when
+   * amplifying a pointer.
+   */
+  template<typename ConcreteMap>
+  struct BasicAuthmap
+  {
+    static_assert(
+      std::is_same_v<capptr::Arena<void>, typename ConcreteMap::EntryType>,
+      "BasicAuthmap's ConcreteMap must have capptr::Arena<void> element type!");
+
+  private:
+    SNMALLOC_REQUIRE_CONSTINIT
+    static inline ConcreteMap concreteAuthmap;
+
+  public:
+    static SNMALLOC_FAST_PATH void init()
+    {
+      concreteAuthmap.template init</* randomize_location */ false>();
+    }
+
+    static SNMALLOC_FAST_PATH void
+    register_range(capptr::Arena<void> base, size_t size)
+    {
+      concreteAuthmap.register_range(address_cast(base), size);
+
+      address_t base_addr = address_cast(base);
+      for (address_t a = base_addr; a < base_addr + size;
+           a += ConcreteMap::GRANULARITY)
+      {
+        concreteAuthmap.set(a, base);
+      }
+    }
+
+    template<bool potentially_out_of_range = false>
+    static SNMALLOC_FAST_PATH capptr::Arena<void> amplify(capptr::Alloc<void> c)
+    {
+      return Aal::capptr_rebound(
+        concreteAuthmap.template get<potentially_out_of_range>(address_cast(c)),
+        c);
+    }
+  };
+
+  /**
+   * Pick between the two above implementations based on StrictProvenance
+   */
+  template<typename CA>
+  using DefaultAuthmap = std::conditional_t<
+    aal_supports<StrictProvenance>,
+    BasicAuthmap<CA>,
+    DummyAuthmap>;
+
+} // namespace snmalloc
diff --git a/src/snmalloc/backend_helpers/backend_helpers.h b/src/snmalloc/backend_helpers/backend_helpers.h
index a84a97f..32b10a4 100644
--- a/src/snmalloc/backend_helpers/backend_helpers.h
+++ b/src/snmalloc/backend_helpers/backend_helpers.h
@@ -1,4 +1,5 @@
 #include "../mem/mem.h"
+#include "authmap.h"
 #include "buddy.h"
 #include "commitrange.h"
 #include "commonconfig.h"