Generate quantum resistant keypairs ahead of connecting

talpid-core/src/tunnel_state_machine/mod.rs

@@ -25,6 +25,7 @@ use talpid_routing::RouteManagerHandle;
 #[cfg(target_os = "macos")]
 use talpid_tunnel::TunnelMetadata;
 use talpid_tunnel::{tun_provider::TunProvider, TunnelEvent};
+use talpid_tunnel_config_client::classic_mceliece::spawn_keypair_generator;
 #[cfg(target_os = "macos")]
 use talpid_types::ErrorExt;
 
@@ -177,6 +178,9 @@ pub async fn spawn(
         }
     });
 
+    // Spawn a worker that pre-computes McEliece key pairs for PQ tunnels
+    spawn_keypair_generator();
+
     Ok(TunnelStateMachineHandle {
         command_tx,
         shutdown_rx,

talpid-tunnel-config-client/src/classic_mceliece.rs

@@ -1,30 +1,81 @@
+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. Note that, using the below algorithm, they take up
+/// around 537 kB each. We therefore only buffer two, which is the largest useful amount, in case of
+/// multihop.
+pub 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();
+type KeyPair = (PublicKey<'static>, SecretKey<'static>);
+
+/// Receiver for McEliece key pairs used by PQ tunnels. These are generated in a separate
+/// thread to reduce latency when connecting.
+static KEYPAIR_RX: OnceLock<Mutex<mpsc::Receiver<KeyPair>>> = OnceLock::new();
+
+/// Spawn a worker that pre computes `bufsize` McEliece key pairs in a separate thread, which can be
+/// fetched asynchronously using the returned channel.
+///
+/// It can take upwards of 200 ms to generate McEliece key pairs so it needs to be done before we
+/// start connecting to the tunnel.
+///
+/// # Panic
+///
+/// Panics if the buffer capacity is 0.
+pub fn spawn_keypair_worker(bufsize: usize) -> mpsc::Receiver<KeyPair> {
+    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 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 || {
-            let keypair = keypair_boxed(&mut rand::thread_rng());
-            let _ = tx.send(keypair);
-        })
-        .unwrap();
-
-    rx.await.unwrap()
+    tokio::spawn(async move {
+        loop {
+            // We do not want generate the key before we know it can be sent, as they take a lot of
+            // space. Note that `tokio::sync::mpsc` doesn't allow zero capacity channels,
+            // otherwise we could reduce the channel capacity by one, use `send_blocking` and simply
+            // store one of the keys in the stack of the thread.
+            let Ok(permit) = tx.reserve().await else {
+                return;
+            };
+            std::thread::scope(|s| {
+                std::thread::Builder::new()
+                    .stack_size(STACK_SIZE)
+                    .name("McEliece key pair generator".to_string())
+                    .spawn_scoped(s, || {
+                        permit.send(keypair_boxed(&mut rand::thread_rng()));
+                    })
+                    .unwrap();
+            });
+        }
+    });
+
+    rx
+}
+
+pub async fn generate_keys() -> KeyPair {
+    KEYPAIR_RX
+        .get_or_init(|| Mutex::new(spawn_keypair_worker(BUFSIZE)))
+        .lock()
+        .await
+        .recv()
+        .await
+        .expect("Expected to receive key pair, but key generator has been stopped.")
+}
+
+/// Spawn a worker which computes and buffers [`BUFSIZE`] of McEliece key pairs, used by PQ tunnels.
+pub fn spawn_keypair_generator<'a>() -> &'a Mutex<mpsc::Receiver<KeyPair>> {
+    KEYPAIR_RX.get_or_init(|| Mutex::new(spawn_keypair_worker(BUFSIZE)))
 }
 
 pub fn decapsulate(

talpid-tunnel-config-client/src/lib.rs

@@ -12,7 +12,7 @@ use tonic::transport::Endpoint;
 use tower::service_fn;
 use zeroize::Zeroize;
 
-mod classic_mceliece;
+pub mod classic_mceliece;
 mod ml_kem;
 #[cfg(not(target_os = "ios"))]
 mod socket;