From 77c453600b458f42b39797d0a0c856b9bd1873c5 Mon Sep 17 00:00:00 2001
From: Matthew Parkinson <mjp41@users.noreply.github.com>
Date: Wed, 25 Mar 2020 08:10:39 +0000
Subject: [PATCH] OE fixes (#157)

* Only compile OE PAL if required.

* OE:reserve: Fix bug in loop.

* Handle out of memory by returning nullptr.
---
 src/mem/alloc.h                            | 19 ++++++++++++-------
 src/mem/largealloc.h                       | 11 +++++++++++
 src/pal/pal.h                              |  4 +++-
 src/pal/pal_open_enclave.h                 |  5 +++--
 src/test/func/fixed_region/fixed_region.cc | 11 +++++++----
 5 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/src/mem/alloc.h b/src/mem/alloc.h
index 9d14e0a85..87b9e813b 100644
--- a/src/mem/alloc.h
+++ b/src/mem/alloc.h
@@ -877,7 +877,7 @@ namespace snmalloc
         large_allocator.template alloc<NoZero, allow_reserve>(
           0, SUPERSLAB_SIZE));
 
-      if ((allow_reserve == NoReserve) && (super == nullptr))
+      if (super == nullptr)
         return super;
 
       super->init(public_state());
@@ -939,7 +939,7 @@ namespace snmalloc
 
         super = get_superslab<allow_reserve>();
 
-        if ((allow_reserve == NoReserve) && (super == nullptr))
+        if (super == nullptr)
           return nullptr;
 
         Slab* slab = super->alloc_short_slab(sizeclass);
@@ -949,7 +949,7 @@ namespace snmalloc
 
       Superslab* super = get_superslab<allow_reserve>();
 
-      if ((allow_reserve == NoReserve) && (super == nullptr))
+      if (super == nullptr)
         return nullptr;
 
       Slab* slab = super->alloc_slab(sizeclass);
@@ -1025,6 +1025,9 @@ namespace snmalloc
         if ((allow_reserve == NoReserve) && (slab == nullptr))
           return nullptr;
 
+        if (slab == nullptr)
+          return nullptr;
+
         sl.insert_back(slab->get_link());
       }
       auto& ffl = small_fast_free_lists[sizeclass];
@@ -1156,7 +1159,7 @@ namespace snmalloc
           large_allocator.template alloc<NoZero, allow_reserve>(
             0, SUPERSLAB_SIZE));
 
-        if ((allow_reserve == NoReserve) && (slab == nullptr))
+        if (slab == nullptr)
           return nullptr;
 
         slab->init(public_state(), sizeclass, rsize);
@@ -1230,10 +1233,12 @@ namespace snmalloc
 
       void* p = large_allocator.template alloc<zero_mem, allow_reserve>(
         large_class, size);
+      if (likely(p != nullptr))
+      {
+        chunkmap().set_large_size(p, size);
 
-      chunkmap().set_large_size(p, size);
-
-      stats().large_alloc(large_class);
+        stats().large_alloc(large_class);
+      }
       return p;
     }
 
diff --git a/src/mem/largealloc.h b/src/mem/largealloc.h
index ac8b6a2f2..ed0fd36d9 100644
--- a/src/mem/largealloc.h
+++ b/src/mem/largealloc.h
@@ -125,6 +125,12 @@ namespace snmalloc
     {
       // Reserve the smallest large_class which is SUPERSLAB_SIZE
       void* r = reserve<false>(0);
+
+      if (r == nullptr)
+        Pal::error(
+          "Unrecoverable internal error: \
+          failed to allocator internal data structure.");
+
       PAL::template notify_using<NoZero>(r, OS_PAGE_SIZE);
 
       bump = r;
@@ -274,6 +280,9 @@ namespace snmalloc
         size_t request = bits::max(size * 4, SUPERSLAB_SIZE * 8);
         void* p = PAL::template reserve<false>(request);
 
+        if (p == nullptr)
+          return nullptr;
+
         address_t p0 = address_cast(p);
         address_t start = bits::align_up(p0, align);
         address_t p1 = p0 + request;
@@ -354,6 +363,8 @@ namespace snmalloc
       if (p == nullptr)
       {
         p = memory_provider.template reserve<false>(large_class);
+        if (p == nullptr)
+          return nullptr;
         memory_provider.template notify_using<zero_mem>(p, size);
       }
       else
diff --git a/src/pal/pal.h b/src/pal/pal.h
index dfadcadff..42e1b13d9 100644
--- a/src/pal/pal.h
+++ b/src/pal/pal.h
@@ -12,7 +12,9 @@
 #  include "pal_openbsd.h"
 #  include "pal_windows.h"
 #endif
-#include "pal_open_enclave.h"
+#if defined(OPEN_ENCLAVE)
+#  include "pal_open_enclave.h"
+#endif
 #include "pal_plain.h"
 
 namespace snmalloc
diff --git a/src/pal/pal_open_enclave.h b/src/pal/pal_open_enclave.h
index 88ecbd8e0..58de14628 100644
--- a/src/pal/pal_open_enclave.h
+++ b/src/pal/pal_open_enclave.h
@@ -31,6 +31,7 @@ namespace snmalloc
       if (oe_base == 0)
       {
         void* dummy = NULL;
+        // If this CAS fails then another thread has initialised this.
         oe_base.compare_exchange_strong(
           dummy, const_cast<void*>(__oe_get_heap_base()));
       }
@@ -44,9 +45,9 @@ namespace snmalloc
         next_base = pointer_offset(new_base, size);
 
         if (next_base > end)
-          error("Out of memory");
+          return nullptr;
 
-      } while (oe_base.compare_exchange_strong(old_base, next_base));
+      } while (!oe_base.compare_exchange_strong(old_base, next_base));
 
       return old_base;
     }
diff --git a/src/test/func/fixed_region/fixed_region.cc b/src/test/func/fixed_region/fixed_region.cc
index 94ad37732..6c9679780 100644
--- a/src/test/func/fixed_region/fixed_region.cc
+++ b/src/test/func/fixed_region/fixed_region.cc
@@ -24,8 +24,7 @@ extern "C" const void* __oe_get_heap_end()
 
 extern "C" void* oe_memset_s(void* p, size_t p_size, int c, size_t size)
 {
-  std::cout << "Memset " << p << " (" << p_size << ") - " << size << std::endl;
-
+  UNUSED(p_size);
   return memset(p, c, size);
 }
 
@@ -50,10 +49,14 @@ int main()
 
   auto a = ThreadAlloc::get();
 
-  for (size_t i = 0; i < 1000; i++)
+  while (true)
   {
     auto r1 = a->alloc(100);
-    std::cout << "Allocated object " << r1 << std::endl;
+
+    // Run until we exhaust the fixed region.
+    // This should return null.
+    if (r1 == nullptr)
+      return 0;
 
     if (oe_base > r1)
       abort();