Skip to content

Commit 433e5ab

Browse files
Auto-allocate stack in runtime
Signed-off-by: Marc Brevoort <[email protected]>
1 parent 07e5690 commit 433e5ab

File tree

1 file changed

+39
-8
lines changed

1 file changed

+39
-8
lines changed

oqs/src/kem.rs

Lines changed: 39 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -321,15 +321,46 @@ impl Kem {
321321
let mut sk = SecretKey {
322322
bytes: Vec::with_capacity(kem.length_secret_key),
323323
};
324-
let status = unsafe { func(pk.bytes.as_mut_ptr(), sk.bytes.as_mut_ptr()) };
325-
status_to_result(status)?;
326-
// update the lengths of the vecs
327-
// this is safe to do, as we have initialised them now.
328-
unsafe {
329-
pk.bytes.set_len(kem.length_public_key);
330-
sk.bytes.set_len(kem.length_secret_key);
324+
325+
let pklen = kem.length_public_key;
326+
let sklen = kem.length_secret_key;
327+
328+
#[cfg(feature = "std")]
329+
{
330+
// Particularly the classic McEliece kem is very stack-heavy.
331+
// The reason we're using a thread is that it allows
332+
// to increase stack space, on demand, in run time.
333+
// Not only does this eliminate the need to set compiler options for each build;
334+
// it also means we won't be holding on to memory unnecessarily in runtime.
335+
let handle = std::thread::Builder::new()
336+
.stack_size(16_000_000)
337+
.spawn(move || {
338+
let status = unsafe { func(pk.bytes.as_mut_ptr(), sk.bytes.as_mut_ptr()) };
339+
status_to_result(status)?;
340+
// update the lengths of the vecs
341+
// this is safe to do, as we have initialised them now.
342+
unsafe {
343+
pk.bytes.set_len(pklen);
344+
sk.bytes.set_len(sklen);
345+
}
346+
Ok((pk, sk))
347+
})
348+
.unwrap();
349+
handle.join().unwrap()
350+
}
351+
#[cfg(not(feature = "std"))]
352+
{
353+
// For embedded, something different needs to be done to spawn a new task. For now, do it inline.
354+
let status = unsafe { func(pk.bytes.as_mut_ptr(), sk.bytes.as_mut_ptr()) };
355+
status_to_result(status)?;
356+
// update the lengths of the vecs
357+
// this is safe to do, as we have initialised them now.
358+
unsafe {
359+
pk.bytes.set_len(pklen);
360+
sk.bytes.set_len(sklen);
361+
}
362+
Ok((pk, sk))
331363
}
332-
Ok((pk, sk))
333364
}
334365

335366
/// Encapsulate to the provided public key

0 commit comments

Comments
 (0)