Generate McEliece key pairs in separate thread

1 files changed, 27 insertions, 5 deletions

diff --git a/talpid-tunnel-config-client/src/classic_mceliece.rs b/talpid-tunnel-config-client/src/classic_mceliece.rs
index 7f7edd43a7..7c6db29b68 100644
--- a/talpid-tunnel-config-client/src/classic_mceliece.rs
+++ b/talpid-tunnel-config-client/src/classic_mceliece.rs
@@ -1,30 +1,52 @@
+use std::sync::OnceLock;
+
 use classic_mceliece_rust::{keypair_boxed, Ciphertext, CRYPTO_CIPHERTEXTBYTES};
 pub use classic_mceliece_rust::{PublicKey, SecretKey, SharedSecret};
+use tokio::sync::{mpsc, Mutex};
 
 /// The `keypair_boxed` function needs just under 1 MiB of stack in debug
 /// builds.
 const STACK_SIZE: usize = 2 * 1024 * 1024;
 
+/// Number of McEliece key pairs to buffer
+const BUFSIZE: usize = 2;
+
 /// Use the smallest CME variant with NIST security level 3. This variant has significantly smaller
 /// keys than the larger variants, and is considered safe.
 pub const ALGORITHM_NAME: &str = "Classic-McEliece-460896f-round3";
 
-pub async fn generate_keys() -> (PublicKey<'static>, SecretKey<'static>) {
-    let (tx, rx) = tokio::sync::oneshot::channel();
+static KEYPAIR_RX: OnceLock<Mutex<mpsc::Receiver<KeyPair>>> = OnceLock::new();
+
+type KeyPair = (PublicKey<'static>, SecretKey<'static>);
 
+fn spawn_keypair_worker(bufsize: usize) -> mpsc::Receiver<KeyPair> {
+    let bufsize = bufsize.checked_sub(1).expect("bufsize must be at least 1");
+    let (tx, rx) = mpsc::channel(bufsize);
     // We fork off the key computation to a separate thread for two reasons:
     // * The computation uses a lot of stack, and we don't want to rely on the default
     //   stack being large enough or having enough space left.
     // * The computation takes a long time and must not block the async runtime thread.
     std::thread::Builder::new()
         .stack_size(STACK_SIZE)
-        .spawn(move || {
+        .spawn(move || loop {
             let keypair = keypair_boxed(&mut rand::thread_rng());
-            let _ = tx.send(keypair);
+            if tx.blocking_send(keypair).is_err() {
+                return;
+            }
         })
         .unwrap();
 
-    rx.await.unwrap()
+    rx
+}
+
+pub async fn generate_keys() -> KeyPair {
+    KEYPAIR_RX
+        .get_or_init(|| Mutex::new(spawn_keypair_worker(BUFSIZE)))
+        .lock()
+        .await
+        .recv()
+        .await
+        .expect("Failed to receive key pair, generating working expectedly closed.")
 }
 
 pub fn decapsulate(