diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs index 57ee4e6ecbc3..e8bd4ed64980 100644 --- a/talpid-core/src/tunnel_state_machine/mod.rs +++ b/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, diff --git a/talpid-tunnel-config-client/src/classic_mceliece.rs b/talpid-tunnel-config-client/src/classic_mceliece.rs index 7f7edd43a7cd..7484313906a3 100644 --- a/talpid-tunnel-config-client/src/classic_mceliece.rs +++ b/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>> = 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. +fn spawn_keypair_worker(bufsize: usize) -> mpsc::Receiver { + 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> { + KEYPAIR_RX.get_or_init(|| Mutex::new(spawn_keypair_worker(BUFSIZE))) } pub fn decapsulate( diff --git a/talpid-tunnel-config-client/src/lib.rs b/talpid-tunnel-config-client/src/lib.rs index bfa3deb29277..381bc65a5365 100644 --- a/talpid-tunnel-config-client/src/lib.rs +++ b/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;