Expose an API for finer control over heap allocator blocking behavior.

The current `timeout` argument of `heap_allocate*` functions determines
whether or not and how long the allocator should block for an allocation
to be fulfiled. However it does not control on what events the allocator
should block.

Currently the allocator would block for three reasons: 1) there is no
heap space available anymore, 2) the quota is exhausted, and 3) the
quarantine is holding back memory which could be used for the
allocation.

For the network stack use-case, we need a finer way to control this
blocking behavior, for example to only wait on the quarantine.

This commit introduces such an API through a new `flags` argument.

Add tests for this new API, but not for `AllocateWaitQuotaExceeded` as
it requires manipulating the quarantine to be done reliably.

Update the stack usage of some allocator functions which has
correspondingly increased.

Signed-off-by: Hugo Lefeuvre <hugo.lefeuvre@scisemi.com>
diff --git a/sdk/core/allocator/main.cc b/sdk/core/allocator/main.cc
index bb8160c..f2299d2 100644
--- a/sdk/core/allocator/main.cc
+++ b/sdk/core/allocator/main.cc
@@ -240,7 +240,8 @@
 	                      LockGuard<decltype(lock)>      &&g,
 	                      PrivateAllocatorCapabilityState *capability,
 	                      Timeout                         *timeout,
-	                      bool isSealedAllocation = false)
+	                      bool     isSealedAllocation = false,
+	                      uint32_t flags              = AllocateWaitAny)
 	{
 		check_gm();
 
@@ -254,8 +255,9 @@
 			{
 				return std::get<Capability<void>>(ret);
 			}
-			// If the timeout is 0, fail now.
-			if (!may_block(timeout))
+			// If the call is non-blocking (`flags` is
+			// `AllocateWaitNone`, or `timeout` is 0), fail now.
+			if (flags == AllocateWaitNone || !may_block(timeout))
 			{
 				return nullptr;
 			}
@@ -265,6 +267,13 @@
 			  std::get_if<MState::AllocationFailureRevocationNeeded>(&ret);
 			if (needsRevocation)
 			{
+				if (!(flags & AllocateWaitRevocationNeeded))
+				{
+					// The flags specify that we should not
+					// wait when revocation is needed.
+					return nullptr;
+				}
+
 				// If we are able to dequeue some objects from quarantine then
 				// retry immediately, otherwise yield.
 				//
@@ -303,11 +312,24 @@
 			}
 			// If the heap is full, wait for someone to free an allocation and
 			// then retry.
-			if (std::holds_alternative<MState::AllocationFailureHeapFull>(
-			      ret) ||
-			    std::holds_alternative<MState::AllocationFailureQuotaExceeded>(
-			      ret))
+			bool isHeapFullFailure =
+			  std::holds_alternative<MState::MState::AllocationFailureHeapFull>(
+			    ret);
+			bool isQuotaExceededFailure =
+			  std::holds_alternative<MState::AllocationFailureQuotaExceeded>(
+			    ret);
+			if (isHeapFullFailure || isQuotaExceededFailure)
 			{
+				if ((isHeapFullFailure && !(flags & AllocateWaitHeapFull)) ||
+				    (isQuotaExceededFailure &&
+				     !(flags & AllocateWaitQuotaExceeded)))
+				{
+					// The flags specify that we should not
+					// wait when the heap is full and/or
+					// when the quota is exceeded.
+					return nullptr;
+				}
+
 				Debug::log("Not enough free space to handle {}-byte "
 				           "allocation, sleeping",
 				           bytes);
@@ -834,11 +856,12 @@
 	}
 }
 
-__cheriot_minimum_stack(0x1f0) void *heap_allocate(Timeout *timeout,
+__cheriot_minimum_stack(0x200) void *heap_allocate(Timeout *timeout,
                                                    SObj     heapCapability,
-                                                   size_t   bytes)
+                                                   size_t   bytes,
+                                                   uint32_t flags)
 {
-	STACK_CHECK(0x1f0);
+	STACK_CHECK(0x200);
 	if (!check_timeout_pointer(timeout))
 	{
 		return nullptr;
@@ -855,7 +878,7 @@
 		return nullptr;
 	}
 	// Use the default memory space.
-	return malloc_internal(bytes, std::move(g), cap, timeout);
+	return malloc_internal(bytes, std::move(g), cap, timeout, false, flags);
 }
 
 __cheriot_minimum_stack(0x1b0) ssize_t
@@ -957,12 +980,13 @@
 	return freed;
 }
 
-__cheriot_minimum_stack(0x1f0) void *heap_allocate_array(Timeout *timeout,
+__cheriot_minimum_stack(0x200) void *heap_allocate_array(Timeout *timeout,
                                                          SObj   heapCapability,
                                                          size_t nElements,
-                                                         size_t elemSize)
+                                                         size_t elemSize,
+                                                         uint32_t flags)
 {
-	STACK_CHECK(0x1f0);
+	STACK_CHECK(0x200);
 	if (!check_timeout_pointer(timeout))
 	{
 		return nullptr;
@@ -984,7 +1008,7 @@
 	{
 		return nullptr;
 	}
-	return malloc_internal(req, std::move(g), cap, timeout);
+	return malloc_internal(req, std::move(g), cap, timeout, false, flags);
 }
 
 namespace
@@ -1103,14 +1127,14 @@
 	return nullptr;
 }
 
-__cheriot_minimum_stack(0x250) SObj
+__cheriot_minimum_stack(0x260) SObj
   token_sealed_unsealed_alloc(Timeout *timeout,
                               SObj     heapCapability,
                               SKey     key,
                               size_t   sz,
                               void   **unsealed)
 {
-	STACK_CHECK(0x250);
+	STACK_CHECK(0x260);
 	if (!check_timeout_pointer(timeout))
 	{
 		return INVALID_SOBJ;
diff --git a/sdk/include/stdlib.h b/sdk/include/stdlib.h
index d60e4e4..de007dd 100644
--- a/sdk/include/stdlib.h
+++ b/sdk/include/stdlib.h
@@ -95,15 +95,53 @@
 	}
 }
 
+enum [[clang::flag_enum]] AllocateWaitFlags{
+  /**
+   * Non-blocking mode. This is equivalent to passing a timeout with no time
+   * remaining.
+   */
+  AllocateWaitNone = 0,
+  /**
+   * If there is enough memory in the quarantine to fulfil the allocation, wait
+   * for the revoker to free objects from the quarantine.
+   */
+  AllocateWaitRevocationNeeded = (1 << 0),
+  /**
+   * If the quota of the passed heap capability is exceeded, wait for other
+   * threads to free allocations.
+   */
+  AllocateWaitQuotaExceeded = (1 << 1),
+  /**
+   * If the heap memory is exhausted, wait for any other thread of the system
+   * to free allocations.
+   */
+  AllocateWaitHeapFull = (1 << 2),
+  /**
+   * Block on any of the above reasons. This is the default behavior.
+   */
+  AllocateWaitAny = (AllocateWaitRevocationNeeded | AllocateWaitQuotaExceeded |
+                     AllocateWaitHeapFull),
+};
+
 /**
  * Non-standard allocation API.  Allocates `size` bytes.  Blocking behaviour is
- * controlled by the `timeout` parameter.
+ * controlled by the `flags` and the `timeout` parameters.
  *
- * The non-blocking mode will return a successful allocation if one can be
- * created immediately, or `nullptr` otherwise.
- * The blocking versions of this may return `nullptr` if the timeout has expired
- * or if the allocation cannot be satisfied under any circumstances (for example
- * if `size` is larger than the total heap size).
+ * Specifically, the `flags` parameter defines on which conditions to wait, and
+ * the `timeout` parameter how long to wait.
+ *
+ * The non-blocking mode (`AllocateWaitNone`, or `timeout` with no time
+ * remaining) will return a successful allocation if one can be created
+ * immediately, or `nullptr` otherwise.
+ *
+ * The blocking modes may return `nullptr` if the condition to wait is not
+ * fulfiled, if the timeout has expired, or if the allocation cannot be
+ * satisfied under any circumstances (for example if `size` is larger than the
+ * total heap size).
+ *
+ * This means that calling this with `AllocateWaitAny` and `UnlimitedTimeout`
+ * will only ever return `nullptr` if the allocation cannot be satisfied under
+ * any circumstances.
  *
  * In both blocking and non-blocking cases, `-ENOTENOUGHSTACK` may be returned
  * if the stack is insufficiently large to safely run the function. This means
@@ -115,24 +153,19 @@
 void *__cheri_compartment("alloc")
   heap_allocate(Timeout           *timeout,
                 struct SObjStruct *heapCapability,
-                size_t             size);
+                size_t             size,
+                uint32_t flags     __if_cxx(= AllocateWaitAny));
 
 /**
  * Non-standard allocation API.  Allocates `size` * `nmemb` bytes of memory,
- * checking for arithmetic overflow.  Blocking behaviour is controlled by the
- * `timeout` parameter:
+ * checking for arithmetic overflow. Similarly to `heap_allocate`, blocking
+ * behaviour is controlled by the `flags` and the `timeout` parameters.
  *
- *  - 0 indicates that this call may not block.
- *  - The maximum value of the type indicates that this may block indefinitely.
- *  - Any other value indicates that this may block for, at most, that many
- *    ticks.
- *
- * The non-blocking mode will return a successful allocation if one can be
- * created immediately, or `nullptr` otherwise.
- * The blocking versions of this may return `nullptr` if the timeout has expired
- * or if the allocation cannot be satisfied under any circumstances (for example
- * if `nmemb` * `size` is larger than the total heap size, or if `nmemb` *
- * `size` overflows).
+ * See `heap_allocate` for more information on the blocking behavior.  One
+ * difference between this and `heap_allocate` is the definition of when the
+ * allocation cannot be satisfied under any circumstances, which is here if
+ * `nmemb` * `size` is larger than the total heap size, or if `nmemb` * `size`
+ * overflows.
  *
  * Similarly to `heap_allocate`, `-ENOTENOUGHSTACK` may be returned if the
  * stack is insufficiently large to run the function. See `heap_allocate`.
@@ -143,7 +176,8 @@
   heap_allocate_array(Timeout           *timeout,
                       struct SObjStruct *heapCapability,
                       size_t             nmemb,
-                      size_t             size);
+                      size_t             size,
+                      uint32_t flags     __if_cxx(= AllocateWaitAny));
 
 /**
  * Add a claim to an allocation.  The object will be counted against the quota
@@ -255,7 +289,7 @@
 static inline void *malloc(size_t size)
 {
 	Timeout t   = {0, 0};
-	void   *ptr = heap_allocate(&t, MALLOC_CAPABILITY, size);
+	void   *ptr = heap_allocate(&t, MALLOC_CAPABILITY, size, AllocateWaitNone);
 	if (!__builtin_cheri_tag_get(ptr))
 	{
 		ptr = NULL;
@@ -264,8 +298,9 @@
 }
 static inline void *calloc(size_t nmemb, size_t size)
 {
-	Timeout t   = {0, 0};
-	void   *ptr = heap_allocate_array(&t, MALLOC_CAPABILITY, nmemb, size);
+	Timeout t = {0, 0};
+	void   *ptr =
+	  heap_allocate_array(&t, MALLOC_CAPABILITY, nmemb, size, AllocateWaitNone);
 	if (!__builtin_cheri_tag_get(ptr))
 	{
 		ptr = NULL;
diff --git a/tests/allocator-test.cc b/tests/allocator-test.cc
index c206f35..084a29a 100644
--- a/tests/allocator-test.cc
+++ b/tests/allocator-test.cc
@@ -24,6 +24,9 @@
 using namespace CHERI;
 #define SECOND_HEAP STATIC_SEALED_VALUE(secondHeap)
 
+DECLARE_AND_DEFINE_ALLOCATOR_CAPABILITY(emptyHeap, 0);
+#define EMPTY_HEAP STATIC_SEALED_VALUE(emptyHeap)
+
 namespace
 {
 	/**
@@ -155,8 +158,50 @@
 		TEST(heap_allocate(&noWait, MALLOC_CAPABILITY, BigAllocSize) == nullptr,
 		     "Non-blocking heap allocation did not return failure with memory "
 		     "exhausted");
-		debug_log("Trying a huge allocation");
+		debug_log("Checking that the 'heap full' flag works");
 		Timeout forever{UnlimitedTimeout};
+		TEST(heap_allocate(&forever,
+		                   MALLOC_CAPABILITY,
+		                   BigAllocSize,
+		                   AllocateWaitRevocationNeeded |
+		                     AllocateWaitQuotaExceeded) == nullptr,
+		     "Blocking heap allocation with the heap full flag unset did not "
+		     "return failure with memory "
+		     "exhausted");
+		Timeout thirtyticks{30};
+		TEST(heap_allocate(&thirtyticks,
+		                   MALLOC_CAPABILITY,
+		                   BigAllocSize,
+		                   AllocateWaitHeapFull) == nullptr,
+		     "Time-limited blocking allocation did not return failure with "
+		     "memory exhausted");
+		TEST(thirtyticks.remaining == 0,
+		     "Allocation with heap full wait flag set did not wait on memory "
+		     "exhausted");
+		debug_log("Checking that the 'quota exhausted' flag works");
+		TEST(heap_allocate(&forever,
+		                   EMPTY_HEAP,
+		                   BigAllocSize,
+		                   AllocateWaitRevocationNeeded) == nullptr,
+		     "Blocking heap allocation with the quota exhausted flag unset did "
+		     "not "
+		     "return failure with memory "
+		     "exhausted");
+		thirtyticks = Timeout{30};
+		TEST(heap_allocate(&thirtyticks,
+		                   EMPTY_HEAP,
+		                   BigAllocSize,
+		                   AllocateWaitQuotaExceeded) == nullptr,
+		     "Time-limited blocking allocation did not return failure with "
+		     "memory exhausted");
+		TEST(
+		  thirtyticks.remaining == 0,
+		  "Allocation with quota exhausted wait flag set did not wait on quota "
+		  "exhausted");
+		// Note: we do not test the functioning of
+		// `AllocateWaitQuotaExceeded` as this would require to be able
+		// to manipulate the quarantine to be reliably done.
+		debug_log("Trying a huge allocation");
 		// nullptr check because we explicitly want to check for OOM
 		TEST(heap_allocate(&forever, MALLOC_CAPABILITY, 1024 * 1024 * 1024) ==
 		       nullptr,