async_posix: Allow custom stack allocation functions to be specified for POSIX contexts

Reviewed-by: Tomas Mraz <tomas@openssl.org>
Reviewed-by: Matt Caswell <matt@openssl.org>
(Merged from https://github.com/openssl/openssl/pull/17762)

3 files changed, 80 insertions, 3 deletions

diff --git a/crypto/async/arch/async_null.c b/crypto/async/arch/async_null.c
index 675c1d35bf..85268b3235 100644
--- a/crypto/async/arch/async_null.c
+++ b/crypto/async/arch/async_null.c
@@ -16,6 +16,21 @@ int ASYNC_is_capable(void)
     return 0;
 }
 
+int ASYNC_set_mem_functions(ASYNC_stack_alloc_fn alloc_fn,
+                            ASYNC_stack_free_fn free_fn)
+{
+    return 0;
+}
+
+void ASYNC_get_mem_functions(ASYNC_stack_alloc_fn *alloc_fn,
+                             ASYNC_stack_free_fn *free_fn)
+{
+    if (alloc_fn != NULL)
+        *alloc_fn = NULL;
+    if (free_fn != NULL)
+        *free_fn = NULL;
+}
+
 void async_local_cleanup(void)
 {
 }
diff --git a/crypto/async/arch/async_posix.c b/crypto/async/arch/async_posix.c
index a0ac97783f..0a1ffbf8d0 100644
--- a/crypto/async/arch/async_posix.c
+++ b/crypto/async/arch/async_posix.c
@@ -18,6 +18,13 @@
 
 #define STACKSIZE       32768
 
+static void *async_stack_alloc(size_t *num);
+static void async_stack_free(void *addr);
+
+static int allow_customize = 1;
+static ASYNC_stack_alloc_fn stack_alloc_impl = async_stack_alloc;
+static ASYNC_stack_free_fn stack_free_impl = async_stack_free;
+
 int ASYNC_is_capable(void)
 {
     ucontext_t ctx;
@@ -29,6 +36,37 @@ int ASYNC_is_capable(void)
     return getcontext(&ctx) == 0;
 }
 
+int ASYNC_set_mem_functions(ASYNC_stack_alloc_fn alloc_fn,
+                            ASYNC_stack_free_fn free_fn)
+{
+    if (!allow_customize)
+        return 0;
+    if (alloc_fn != NULL)
+        stack_alloc_impl = alloc_fn;
+    if (free_fn != NULL)
+        stack_free_impl = free_fn;
+    return 1;
+}
+
+void ASYNC_get_mem_functions(ASYNC_stack_alloc_fn *alloc_fn,
+                             ASYNC_stack_free_fn *free_fn)
+{
+    if (alloc_fn != NULL)
+        *alloc_fn = stack_alloc_impl;
+    if (free_fn != NULL)
+        *free_fn = stack_free_impl;
+}
+
+static void *async_stack_alloc(size_t *num)
+{
+    return OPENSSL_malloc(*num);
+}
+
+static void async_stack_free(void *addr)
+{
+    OPENSSL_free(addr);
+}
+
 void async_local_cleanup(void)
 {
 }
@@ -39,9 +77,18 @@ int async_fibre_makecontext(async_fibre *fibre)
     fibre->env_init = 0;
 #endif
     if (getcontext(&fibre->fibre) == 0) {
-        fibre->fibre.uc_stack.ss_sp = OPENSSL_malloc(STACKSIZE);
+        size_t num = STACKSIZE;
+
+        /*
+         *  Disallow customisation after the first
+         *  stack is allocated.
+         */
+        if (allow_customize)
+            allow_customize = 0;
+
+        fibre->fibre.uc_stack.ss_sp = stack_alloc_impl(&num);
         if (fibre->fibre.uc_stack.ss_sp != NULL) {
-            fibre->fibre.uc_stack.ss_size = STACKSIZE;
+            fibre->fibre.uc_stack.ss_size = num;
             fibre->fibre.uc_link = NULL;
             makecontext(&fibre->fibre, async_start_func, 0);
             return 1;
@@ -55,7 +102,7 @@ int async_fibre_makecontext(async_fibre *fibre)
 
 void async_fibre_free(async_fibre *fibre)
 {
-    OPENSSL_free(fibre->fibre.uc_stack.ss_sp);
+    stack_free_impl(fibre->fibre.uc_stack.ss_sp);
     fibre->fibre.uc_stack.ss_sp = NULL;
 }
 
diff --git a/crypto/async/arch/async_win.c b/crypto/async/arch/async_win.c
index 0b276fd504..1d0d9fa028 100644
--- a/crypto/async/arch/async_win.c
+++ b/crypto/async/arch/async_win.c
@@ -20,6 +20,21 @@ int ASYNC_is_capable(void)
     return 1;
 }
 
+int ASYNC_set_mem_functions(ASYNC_stack_alloc_fn alloc_fn,
+                            ASYNC_stack_free_fn free_fn)
+{
+    return 0;
+}
+
+void ASYNC_get_mem_functions(ASYNC_stack_alloc_fn *alloc_fn,
+                             ASYNC_stack_free_fn *free_fn)
+{
+    if (alloc_fn != NULL)
+        *alloc_fn = NULL;
+    if (free_fn != NULL)
+        *free_fn = NULL;
+}
+
 void async_local_cleanup(void)
 {
     async_ctx *ctx = async_get_ctx();