Merge pull request #34 from Microsoft/pagemap_access_refactor

Pagemap access refactor

Author: davidchisnall

src/mem/alloc.h

@@ -62,20 +62,82 @@ namespace snmalloc
     Pagemap<SUPERSLAB_BITS, uint8_t, 0>>;
 
   HEADER_GLOBAL SuperslabPagemap global_pagemap;
+
+  /**
+   * Mixin used by `SuperslabMap` to directly access the pagemap via a global
+   * variable.  This should be used from within the library or program that
+   * owns the pagemap.
+   */
+  struct GlobalPagemap
+  {
+    /**
+     * Returns the pagemap.
+     */
+    SuperslabPagemap& pagemap()
+    {
+      return global_pagemap;
+    }
+  };
+
+  /**
+   * Optionally exported function that accesses the global pagemap provided by
+   * a shared library.
+   */
+  extern "C" void* snmalloc_pagemap_global_get(snmalloc::PagemapConfig const**);
+
+  /**
+   * Mixin used by `SuperslabMap` to access the global pagemap via a
+   * type-checked C interface.  This should be used when another library (e.g.
+   * your C standard library) uses snmalloc and you wish to use a different
+   * configuration in your program or library, but wish to share a pagemap so
+   * that either version can deallocate memory.
+   */
+  class ExternalGlobalPagemap
+  {
+    /**
+     * A pointer to the pagemap.
+     */
+    SuperslabPagemap* external_pagemap;
+
+  public:
+    /**
+     * Constructor.  Accesses the pagemap via the C ABI accessor and casts it to
+     * the expected type, failing in cases of ABI mismatch.
+     */
+    ExternalGlobalPagemap()
+    {
+      const snmalloc::PagemapConfig* c;
+      external_pagemap =
+        SuperslabPagemap::cast_to_pagemap(snmalloc_pagemap_global_get(&c), c);
+      // FIXME: Report an error somehow in non-debug builds.
+      assert(external_pagemap);
+    }
+
+    /**
+     * Returns the exported pagemap.
+     */
+    SuperslabPagemap& pagemap()
+    {
+      return *external_pagemap;
+    }
+  };
+
   /**
    * Class that defines an interface to the pagemap.  This is provided to
    * `Allocator` as a template argument and so can be replaced by a compatible
    * implementation (for example, to move pagemap updates to a different
    * protection domain).
    */
-  struct SuperslabMap
+  template<typename PagemapProvider = GlobalPagemap>
+  struct SuperslabMap : public PagemapProvider
   {
+    using PagemapProvider::PagemapProvider;
     /**
      * Get the pagemap entry corresponding to a specific address.
      */
     uint8_t get(void* p)
     {
-      return global_pagemap.get(p);
+      return PagemapProvider::pagemap().get(p);
     }
     /**
      * Set a pagemap entry indicating that there is a superslab at the
@@ -121,11 +183,11 @@ namespace snmalloc
       for (size_t i = 0; i < size_bits - SUPERSLAB_BITS; i++)
       {
         size_t run = 1ULL << i;
-        global_pagemap.set_range(
+        PagemapProvider::pagemap().set_range(
           (void*)ss, (uint8_t)(64 + i + SUPERSLAB_BITS), run);
         ss = (uintptr_t)ss + SUPERSLAB_SIZE * run;
       }
-      global_pagemap.set(p, (uint8_t)size_bits);
+      PagemapProvider::pagemap().set(p, (uint8_t)size_bits);
     }
     /**
      * Update the pagemap to remove a large allocation, of `size` bytes from
@@ -136,7 +198,7 @@ namespace snmalloc
       size_t rounded_size = bits::next_pow2(size);
       assert(get(p) == bits::next_pow2_bits(size));
       auto count = rounded_size >> SUPERSLAB_BITS;
-      global_pagemap.set_range((void*)p, PMNotOurs, count);
+      PagemapProvider::pagemap().set_range((void*)p, PMNotOurs, count);
     }
 
   private:
@@ -147,12 +209,12 @@ namespace snmalloc
      */
     void set(void* p, uint8_t x)
     {
-      global_pagemap.set(p, x);
+      PagemapProvider::pagemap().set(p, x);
     }
   };
 
 #ifndef SNMALLOC_DEFAULT_PAGEMAP
-#  define SNMALLOC_DEFAULT_PAGEMAP snmalloc::SuperslabMap
+#  define SNMALLOC_DEFAULT_PAGEMAP snmalloc::SuperslabMap<>
 #endif
 
   /**

src/mem/allocconfig.h

@@ -4,12 +4,6 @@
 
 namespace snmalloc
 {
-  enum ZeroMem
-  {
-    NoZero,
-    YesZero
-  };
-
   // 0 intermediate bits results in power of 2 small allocs. 1 intermediate
   // bit gives additional sizeclasses at the midpoint between each power of 2.
   // 2 intermediate bits gives 3 intermediate sizeclasses, etc.

src/mem/pagemap.h

@@ -10,6 +10,40 @@ namespace snmalloc
   static constexpr size_t PAGEMAP_NODE_BITS = 16;
   static constexpr size_t PAGEMAP_NODE_SIZE = 1ULL << PAGEMAP_NODE_BITS;
 
+  /**
+   * Structure describing the configuration of a pagemap.  When querying a
+   * pagemap from a different instantiation of snmalloc, the pagemap is exposed
+   * as a `void*`.  This structure allows the caller to check whether the
+   * pagemap is of the format that they expect.
+   */
+  struct PagemapConfig
+  {
+    /**
+     * The version of the pagemap structure.  This is always 1 in existing
+     * versions of snmalloc.  This will be incremented every time the format
+     * changes in an incompatible way.  Changes to the format may add fields to
+     * the end of this structure.
+     */
+    uint32_t version;
+    /**
+     * Is this a flat pagemap?  If this field is false, the pagemap is the
+     * hierarchical structure.
+     */
+    bool is_flat_pagemap;
+    /**
+     * Number of bytes in a pointer.
+     */
+    uint8_t sizeof_pointer;
+    /**
+     * The number of bits of the address used to index into the pagemap.
+     */
+    uint64_t pagemap_bits;
+    /**
+     * The size (in bytes) of a pagemap entry.
+     */
+    size_t size_of_entry;
+  };
+
   template<size_t GRANULARITY_BITS, typename T, T default_content>
   class Pagemap
   {
@@ -168,6 +202,32 @@ namespace snmalloc
     }
 
   public:
+    /**
+     * The pagemap configuration describing this instantiation of the template.
+     */
+    static constexpr PagemapConfig config = {
+      1, false, sizeof(void*), GRANULARITY_BITS, sizeof(T)};
+
+    /**
+     * Cast a `void*` to a pointer to this template instantiation, given a
+     * config describing the configuration.  Return null if the configuration
+     * passed does not correspond to this template instantiation.
+     *
+     * This intended to allow code that depends on the pagemap having a
+     * specific representation to fail gracefully.
+     */
+    static Pagemap* cast_to_pagemap(void* pm, const PagemapConfig* c)
+    {
+      if (
+        (c->version != 1) || (c->is_flat_pagemap) ||
+        (c->sizeof_pointer != sizeof(void*)) ||
+        (c->pagemap_bits != GRANULARITY_BITS) ||
+        (c->size_of_entry != sizeof(T)) || (!std::is_integral_v<T>))
+      {
+        return nullptr;
+      }
+      return static_cast<Pagemap*>(pm);
+    }
     /**
      * Returns the index of a pagemap entry within a given page.  This is used
      * in code that propagates changes to the pagemap elsewhere.
@@ -247,6 +307,32 @@ namespace snmalloc
     std::atomic<T> top[ENTRIES];
 
   public:
+    /**
+     * The pagemap configuration describing this instantiation of the template.
+     */
+    static constexpr PagemapConfig config = {
+      1, true, sizeof(void*), GRANULARITY_BITS, sizeof(T)};
+
+    /**
+     * Cast a `void*` to a pointer to this template instantiation, given a
+     * config describing the configuration.  Return null if the configuration
+     * passed does not correspond to this template instantiation.
+     *
+     * This intended to allow code that depends on the pagemap having a
+     * specific representation to fail gracefully.
+     */
+    static FlatPagemap* cast_to_pagemap(void* pm, const PagemapConfig* c)
+    {
+      if (
+        (c->version != 1) || (!c->is_flat_pagemap) ||
+        (c->sizeof_pointer != sizeof(void*)) ||
+        (c->pagemap_bits != GRANULARITY_BITS) ||
+        (c->size_of_entry != sizeof(T)) || (!std::is_integral_v<T>))
+      {
+        return nullptr;
+      }
+      return static_cast<FlatPagemap*>(pm);
+    }
     T get(void* p)
     {
       return top[(size_t)p >> SHIFT].load(std::memory_order_relaxed);

src/override/malloc.cc

@@ -206,8 +206,24 @@ extern "C"
   }
 
 #ifdef SNMALLOC_EXPOSE_PAGEMAP
-  SNMALLOC_EXPORT void* SNMALLOC_NAME_MANGLE(snmalloc_get_global_pagemap)(void)
+  /**
+   * Export the pagemap.  The return value is a pointer to the pagemap
+   * structure.  The argument is used to return a pointer to a `PagemapConfig`
+   * structure describing the type of the pagemap.  Static methods on the
+   * concrete pagemap templates can then be used to safely cast the return from
+   * this function to the correct type.  This allows us to preserve some
+   * semblance of ABI safety via a pure C API.
+   */
+  SNMALLOC_EXPORT void* SNMALLOC_NAME_MANGLE(snmalloc_pagemap_global_get)(
+    PagemapConfig const** config)
   {
+    if (config)
+    {
+      *config = &decltype(snmalloc::global_pagemap)::config;
+      assert(
+        decltype(snmalloc::global_pagemap)::cast_to_pagemap(
+          &snmalloc::global_pagemap, *config) == &snmalloc::global_pagemap);
+    }
     return &snmalloc::global_pagemap;
   }
 #endif

src/pal/pal_consts.h

@@ -1,3 +1,5 @@
+#pragma once
+
 namespace snmalloc
 {
   /**
@@ -25,4 +27,19 @@ namespace snmalloc
      */
     AlignedAllocation = (1 << 1)
   };
+  /**
+   * Flag indicating whether requested memory should be zeroed.
+   */
+  enum ZeroMem
+  {
+    /**
+     * Memory should not be zeroed, contents are undefined.
+     */
+    NoZero,
+    /**
+     * Memory must be zeroed.  This can be lazily allocated via a copy-on-write
+     * mechanism as long as any load from the memory returns zero.
+     */
+    YesZero
+  };
 }

src/test/func/two_alloc_types/alloc1.cc

@@ -4,6 +4,7 @@
 #define USE_RESERVE_MULTIPLE 1
 #define NO_BOOTSTRAP_ALLOCATOR
 #define IS_ADDRESS_SPACE_CONSTRAINED
+#define SNMALLOC_EXPOSE_PAGEMAP
 #define SNMALLOC_NAME_MANGLE(a) enclave_##a
 // Redefine the namespace, so we can have two versions.
 #define snmalloc snmalloc_enclave

src/test/func/two_alloc_types/alloc2.cc

@@ -1,6 +1,7 @@
 #undef IS_ADDRESS_SPACE_CONSTRAINED
 #define SNMALLOC_NAME_MANGLE(a) host_##a
 #define NO_BOOTSTRAP_ALLOCATOR
+#define SNMALLOC_EXPOSE_PAGEMAP
 // Redefine the namespace, so we can have two versions.
 #define snmalloc snmalloc_host
-#include "../../../override/malloc.cc"
+#include "../../../override/malloc.cc"

src/test/func/two_alloc_types/main.cc

@@ -32,6 +32,11 @@ extern "C" void host_free(void*);
 extern "C" void* enclave_malloc(size_t);
 extern "C" void enclave_free(void*);
 
+extern "C" void*
+enclave_snmalloc_pagemap_global_get(snmalloc::PagemapConfig const**);
+extern "C" void*
+host_snmalloc_pagemap_global_get(snmalloc::PagemapConfig const**);
+
 using namespace snmalloc;
 int main()
 {
@@ -42,6 +47,11 @@ int main()
   oe_end = (uint8_t*)oe_base + size;
   std::cout << "Allocated region " << oe_base << " - " << oe_end << std::endl;
 
+  // Call these functions to trigger asserts if the cast-to-self doesn't work.
+  const PagemapConfig* c;
+  enclave_snmalloc_pagemap_global_get(&c);
+  host_snmalloc_pagemap_global_get(&c);
+
   auto a = host_malloc(128);
   auto b = enclave_malloc(128);