From 433e5ab282ed0189a348dfca8a183afd8cf175c0 Mon Sep 17 00:00:00 2001
From: Marc Brevoort <marc@100percentit.com>
Date: Wed, 15 May 2024 10:40:28 +0100
Subject: [PATCH] Auto-allocate stack in runtime

Signed-off-by: Marc Brevoort <marc@100percentit.com>
---
 oqs/src/kem.rs | 47 +++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 39 insertions(+), 8 deletions(-)

diff --git a/oqs/src/kem.rs b/oqs/src/kem.rs
index 2f9d6e4efe..47b50ec8b0 100644
--- a/oqs/src/kem.rs
+++ b/oqs/src/kem.rs
@@ -321,15 +321,46 @@ impl Kem {
         let mut sk = SecretKey {
             bytes: Vec::with_capacity(kem.length_secret_key),
         };
-        let status = unsafe { func(pk.bytes.as_mut_ptr(), sk.bytes.as_mut_ptr()) };
-        status_to_result(status)?;
-        // update the lengths of the vecs
-        // this is safe to do, as we have initialised them now.
-        unsafe {
-            pk.bytes.set_len(kem.length_public_key);
-            sk.bytes.set_len(kem.length_secret_key);
+
+        let pklen = kem.length_public_key;
+        let sklen = kem.length_secret_key;
+
+        #[cfg(feature = "std")]
+        {
+            // Particularly the classic McEliece kem is very stack-heavy.
+            // The reason we're using a thread is that it allows
+            // to increase stack space, on demand, in run time.
+            // Not only does this eliminate the need to set compiler options for each build;
+            // it also means we won't be holding on to memory unnecessarily in runtime.
+            let handle = std::thread::Builder::new()
+                .stack_size(16_000_000)
+                .spawn(move || {
+                    let status = unsafe { func(pk.bytes.as_mut_ptr(), sk.bytes.as_mut_ptr()) };
+                    status_to_result(status)?;
+                    // update the lengths of the vecs
+                    // this is safe to do, as we have initialised them now.
+                    unsafe {
+                        pk.bytes.set_len(pklen);
+                        sk.bytes.set_len(sklen);
+                    }
+                    Ok((pk, sk))
+                })
+                .unwrap();
+            handle.join().unwrap()
+        }
+        #[cfg(not(feature = "std"))]
+        {
+            // For embedded, something different needs to be done to spawn a new task. For now, do it inline.
+            let status = unsafe { func(pk.bytes.as_mut_ptr(), sk.bytes.as_mut_ptr()) };
+            status_to_result(status)?;
+            // update the lengths of the vecs
+            // this is safe to do, as we have initialised them now.
+            unsafe {
+                pk.bytes.set_len(pklen);
+                sk.bytes.set_len(sklen);
+            }
+            Ok((pk, sk))
         }
-        Ok((pk, sk))
     }
 
     /// Encapsulate to the provided public key