override add malloc_conceal api proposal

Author: devnexen

src/mem/localalloc.h

@@ -418,15 +418,18 @@ namespace snmalloc
     /**
      * Allocate memory of a dynamically known size.
      */
-    template<ZeroMem zero_mem = NoZero>
+    template<ZeroMem zero_mem = NoZero, CoreDumpMem core_dump = YesDump>
     SNMALLOC_FAST_PATH ALLOCATOR void* alloc(size_t size)
     {
+      void* result = nullptr;
 #ifdef SNMALLOC_PASS_THROUGH
       // snmalloc guarantees a lot of alignment, so we can depend on this
       // make pass through call aligned_alloc with the alignment snmalloc
       // would guarantee.
-      void* result = external_alloc::aligned_alloc(
+      result = external_alloc::aligned_alloc(
         natural_alignment(size), round_size(size));
+      if constexpr (core_dump == NoDump)
+        SharedStateHandle::Pal::nodump(result, size);
       if constexpr (zero_mem == YesZero)
         memset(result, 0, size);
       return result;
@@ -438,11 +441,17 @@ namespace snmalloc
       {
         // Small allocations are more likely. Improve
         // branch prediction by placing this case first.
-        return capptr_reveal(small_alloc<zero_mem>(size));
+        result = capptr_reveal(small_alloc<zero_mem>(size));
+        if constexpr (core_dump == NoDump)
+          SharedStateHandle::Pal::nodump(result, size);
+        return result;
       }
 
-      return capptr_reveal(alloc_not_small<zero_mem>(size));
+      result = capptr_reveal(alloc_not_small<zero_mem>(size));
+      if constexpr (core_dump == NoDump)
+        SharedStateHandle::Pal::nodump(result, size);
 #endif
+      return result;
     }
 
     /**

src/override/malloc.cc

@@ -54,6 +54,24 @@ extern "C"
     return ThreadAlloc::get().alloc<ZeroMem::YesZero>(sz);
   }
 
+#if !defined(__OpenBSD__)
+  SNMALLOC_EXPORT void* SNMALLOC_NAME_MANGLE(malloc_conceal)(size_t size)
+  {
+    return ThreadAlloc::get().alloc<NoZero, CoreDumpMem::NoDump>(size);
+  }
+
+  SNMALLOC_EXPORT void* SNMALLOC_NAME_MANGLE(calloc_conceal)(size_t nmemb, size_t size)
+  {
+    bool overflow = false;
+    size_t sz = bits::umul(size, nmemb, overflow);
+    if (SNMALLOC_UNLIKELY(overflow))
+    {
+      return SNMALLOC_NAME_MANGLE(snmalloc_set_error)();
+    }
+    return ThreadAlloc::get().alloc<ZeroMem::YesZero, CoreDumpMem::NoDump>(sz);
+  }
+#endif
+
   SNMALLOC_EXPORT
   size_t SNMALLOC_NAME_MANGLE(malloc_usable_size)(
     MALLOC_USABLE_SIZE_QUALIFIER void* ptr)

src/pal/pal_consts.h

@@ -94,6 +94,12 @@ namespace snmalloc
     YesZero
   };
 
+  enum CoreDumpMem
+  {
+    NoDump,
+    YesDump
+  };
+
   /**
    * Default Tag ID for the Apple class
    */

src/pal/pal_freebsd.h

@@ -63,6 +63,13 @@ namespace snmalloc
           ~static_cast<unsigned int>(CHERI_PERM_CHERIABI_VMMAP)));
     }
 #  endif
+
+    static void nodump(void* p, size_t size) noexcept
+    {
+      SNMALLOC_ASSERT(is_aligned_block<page_size>(p, size));
+      madvise(p, size, MADV_NOCORE)
+    }
+  };
   };
 } // namespace snmalloc
 #endif

src/pal/pal_linux.h

@@ -85,6 +85,12 @@ namespace snmalloc
         madvise(p, size, MADV_FREE);
       }
     }
+
+    static void nodump(void* p, size_t size) noexcept
+    {
+      SNMALLOC_ASSERT(is_aligned_block<page_size>(p, size));
+      madvise(p, size, MADV_DONTDUMP);
+    }
   };
 } // namespace snmalloc
 #endif

src/pal/pal_open_enclave.h

@@ -41,6 +41,12 @@ namespace snmalloc
       oe_memset_s(p, size, 0, size);
     }
 
+    static void dump(void* p, size_t size) noexcept
+    {
+      UNUSED(p);
+      UNUSED(size);
+    }
+
     /**
      * Source of Entropy
      */

src/pal/pal_posix.h

@@ -318,6 +318,12 @@ namespace snmalloc
       return nullptr;
     }
 
+    static void nodump(void* p, size_t size) noexcept
+    {
+      UNUSED(p);
+      UNUSED(size);
+    }
+
     /**
      * Source of Entropy
      *

src/pal/pal_windows.h

@@ -195,6 +195,12 @@ namespace snmalloc
       return VirtualAlloc(nullptr, size, MEM_RESERVE, PAGE_READWRITE);
     }
 
+    static void nodump(void* p, size_t size) noexcept
+    {
+      UNUSED(p);
+      UNUSED(size);
+    }
+
     /**
      * Source of Entropy
      */

src/test/func/malloc/malloc.cc

@@ -87,11 +87,15 @@ void check_result(size_t size, size_t align, void* p, int err, bool null)
   our_free(p);
 }
 
-void test_calloc(size_t nmemb, size_t size, int err, bool null)
+void test_calloc(void *(*calloc_fn)(size_t, size_t), size_t nmemb, size_t size, int err, bool null)
 {
-  printf("calloc(%zu, %zu)  combined size %zu\n", nmemb, size, nmemb * size);
+  printf("calloc");
+  if (calloc_fn == our_calloc_conceal)
+    printf("_conceal");
+
+  printf("(%zu, %zu)  combined size %zu\n", nmemb, size, nmemb * size);
   errno = SUCCESS;
-  void* p = our_calloc(nmemb, size);
+  void* p = calloc_fn(nmemb, size);
 
   if (p != nullptr)
   {
@@ -157,7 +161,7 @@ int main(int argc, char** argv)
     check_result(size + 1, 1, our_malloc(size + 1), SUCCESS, false);
   }
 
-  test_calloc(0, 0, SUCCESS, false);
+  test_calloc(our_calloc, 0, 0, SUCCESS, false);
 
   our_free(nullptr);
 
@@ -173,10 +177,10 @@ int main(int argc, char** argv)
       if (overflow)
         break;
 
-      test_calloc(n, size, SUCCESS, false);
-      test_calloc(n, 0, SUCCESS, false);
+      test_calloc(our_calloc, n, size, SUCCESS, false);
+      test_calloc(our_calloc, n, 0, SUCCESS, false);
     }
-    test_calloc(0, size, SUCCESS, false);
+    test_calloc(our_calloc, 0, size, SUCCESS, false);
   }
 
   for (smallsizeclass_t sc = 0; sc < NUM_SMALL_SIZECLASSES; sc++)
@@ -210,6 +214,34 @@ int main(int argc, char** argv)
 
   test_realloc(our_malloc(64), 4194304, SUCCESS, false);
 
+  for (smallsizeclass_t sc = 0; sc < (MAX_SMALL_SIZECLASS_BITS + 4); sc++)
+  {
+    const size_t size = bits::one_at_bit(sc);
+    printf("malloc_conceal: %zu\n", size);
+    errno = SUCCESS;
+    check_result(size, 1, our_malloc_conceal(size), SUCCESS, false);
+    errno = SUCCESS;
+    check_result(size + 1, 1, our_malloc_conceal(size + 1), SUCCESS, false);
+  }
+
+  for (smallsizeclass_t sc = 0; sc < NUM_SMALL_SIZECLASSES; sc++)
+  {
+    const size_t size = sizeclass_to_size(sc);
+
+    bool overflow = false;
+    for (size_t n = 1;
+         bits::umul(size, n, overflow) <= MAX_SMALL_SIZECLASS_SIZE;
+         n *= 5)
+    {
+      if (overflow)
+        break;
+
+      test_calloc(our_calloc_conceal, n, size, SUCCESS, false);
+      test_calloc(our_calloc_conceal, n, 0, SUCCESS, false);
+    }
+    test_calloc(our_calloc_conceal, 0, size, SUCCESS, false);
+  }
+
   test_posix_memalign(0, 0, EINVAL, true);
   test_posix_memalign(((size_t)-1) / 2, 0, EINVAL, true);
   test_posix_memalign(OS_PAGE_SIZE, sizeof(uintptr_t) / 2, EINVAL, true);