From d6c89bd1704b3603b7d2fec89430369d42b0f8bc Mon Sep 17 00:00:00 2001 From: Sebastian Holmin Date: Wed, 8 Jan 2025 18:00:07 +0100 Subject: [PATCH 1/9] Generate McEliece key pairs in separate thread --- .../src/classic_mceliece.rs | 32 ++++++++++++++++--- 1 file 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 7f7edd43a7cd..7c6db29b6874 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>> = OnceLock::new(); + +type KeyPair = (PublicKey<'static>, SecretKey<'static>); +fn spawn_keypair_worker(bufsize: usize) -> mpsc::Receiver { + 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( From 094196aa4a5356750e0ae1e13945326a1060b8d1 Mon Sep 17 00:00:00 2001 From: Sebastian Holmin Date: Wed, 8 Jan 2025 18:38:02 +0100 Subject: [PATCH 2/9] Update documentation --- talpid-core/src/tunnel_state_machine/mod.rs | 3 +++ .../src/classic_mceliece.rs | 18 +++++++++++++----- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs index 57ee4e6ecbc3..6f4dc33218d2 100644 --- a/talpid-core/src/tunnel_state_machine/mod.rs +++ b/talpid-core/src/tunnel_state_machine/mod.rs @@ -177,6 +177,9 @@ pub async fn spawn( } }); + // Spawn a worker that pre-computes McEliece key pairs for PQ tunnels + KEYPAIR_RX.get_or_init(|| tokio::sync::Mutex::new(spawn_keypair_worker(BUFSIZE))); + 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 7c6db29b6874..da6024f7136d 100644 --- a/talpid-tunnel-config-client/src/classic_mceliece.rs +++ b/talpid-tunnel-config-client/src/classic_mceliece.rs @@ -8,23 +8,31 @@ use tokio::sync::{mpsc, Mutex}; /// builds. const STACK_SIZE: usize = 2 * 1024 * 1024; -/// Number of McEliece key pairs to buffer +/// 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. 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"; -static KEYPAIR_RX: OnceLock>> = OnceLock::new(); - type KeyPair = (PublicKey<'static>, SecretKey<'static>); +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. +/// +/// As it can take upwards of 200 ms to generate McEliece key pairs, it needs to be done before we +/// start connecting to the tunnel. fn spawn_keypair_worker(bufsize: usize) -> mpsc::Receiver { + // As one of the key pairs will be buffered by the stack of the spawned thread, we reduce the + // capacity of the channel by one 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 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) From b824d79cd8293d45fc08de6cac60aa0b00096ac3 Mon Sep 17 00:00:00 2001 From: Sebastian Holmin Date: Wed, 8 Jan 2025 18:48:58 +0100 Subject: [PATCH 3/9] Spawn key pair worker on launch --- talpid-core/src/tunnel_state_machine/mod.rs | 1 + talpid-tunnel-config-client/src/classic_mceliece.rs | 6 +++--- talpid-tunnel-config-client/src/lib.rs | 2 +- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs index 6f4dc33218d2..06588392e97b 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_worker, BUFSIZE, KEYPAIR_RX}; #[cfg(target_os = "macos")] use talpid_types::ErrorExt; diff --git a/talpid-tunnel-config-client/src/classic_mceliece.rs b/talpid-tunnel-config-client/src/classic_mceliece.rs index da6024f7136d..3d29e1997c93 100644 --- a/talpid-tunnel-config-client/src/classic_mceliece.rs +++ b/talpid-tunnel-config-client/src/classic_mceliece.rs @@ -10,7 +10,7 @@ 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. -const BUFSIZE: usize = 2; +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. @@ -18,14 +18,14 @@ pub const ALGORITHM_NAME: &str = "Classic-McEliece-460896f-round3"; type KeyPair = (PublicKey<'static>, SecretKey<'static>); -static KEYPAIR_RX: OnceLock>> = OnceLock::new(); +pub 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. /// /// As it can take upwards of 200 ms to generate McEliece key pairs, it needs to be done before we /// start connecting to the tunnel. -fn spawn_keypair_worker(bufsize: usize) -> mpsc::Receiver { +pub fn spawn_keypair_worker(bufsize: usize) -> mpsc::Receiver { // As one of the key pairs will be buffered by the stack of the spawned thread, we reduce the // capacity of the channel by one let bufsize = bufsize.checked_sub(1).expect("bufsize must be at least 1"); 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; From 7f57b1f767eb79acb372fb54d840c82e55a67fd5 Mon Sep 17 00:00:00 2001 From: Sebastian Holmin Date: Thu, 9 Jan 2025 13:15:10 +0100 Subject: [PATCH 4/9] Make `KEY_PAIR` private and expose it by a fn --- talpid-core/src/tunnel_state_machine/mod.rs | 4 ++-- .../src/classic_mceliece.rs | 14 +++++++++++--- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs index 06588392e97b..d9d25d95d2a9 100644 --- a/talpid-core/src/tunnel_state_machine/mod.rs +++ b/talpid-core/src/tunnel_state_machine/mod.rs @@ -25,7 +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_worker, BUFSIZE, KEYPAIR_RX}; +use talpid_tunnel_config_client::classic_mceliece::get_or_init_keypair_receiver; #[cfg(target_os = "macos")] use talpid_types::ErrorExt; @@ -179,7 +179,7 @@ pub async fn spawn( }); // Spawn a worker that pre-computes McEliece key pairs for PQ tunnels - KEYPAIR_RX.get_or_init(|| tokio::sync::Mutex::new(spawn_keypair_worker(BUFSIZE))); + get_or_init_keypair_receiver(); Ok(TunnelStateMachineHandle { command_tx, diff --git a/talpid-tunnel-config-client/src/classic_mceliece.rs b/talpid-tunnel-config-client/src/classic_mceliece.rs index 3d29e1997c93..6aa81345c8dd 100644 --- a/talpid-tunnel-config-client/src/classic_mceliece.rs +++ b/talpid-tunnel-config-client/src/classic_mceliece.rs @@ -18,7 +18,7 @@ pub const ALGORITHM_NAME: &str = "Classic-McEliece-460896f-round3"; type KeyPair = (PublicKey<'static>, SecretKey<'static>); -pub static KEYPAIR_RX: OnceLock>> = OnceLock::new(); +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. @@ -48,8 +48,7 @@ pub fn spawn_keypair_worker(bufsize: usize) -> mpsc::Receiver { } pub async fn generate_keys() -> KeyPair { - KEYPAIR_RX - .get_or_init(|| Mutex::new(spawn_keypair_worker(BUFSIZE))) + get_or_init_keypair_receiver() .lock() .await .recv() @@ -57,6 +56,15 @@ pub async fn generate_keys() -> KeyPair { .expect("Failed to receive key pair, generating working expectedly closed.") } +/// Returns a receiver for McEliece key pairs used by PQ tunnels. These are generated in a separate +/// thread to reduce latency when connecting. +/// +/// The first call will spawn the worker which immedietly starts to compute and buffer [`BUFSIZE`] +/// of key pairs. +pub fn get_or_init_keypair_receiver<'a>() -> &'a Mutex> { + KEYPAIR_RX.get_or_init(|| Mutex::new(spawn_keypair_worker(BUFSIZE))) +} + pub fn decapsulate( secret: &SecretKey<'_>, ciphertext_slice: &[u8], From 29e173c888f539ee10d641825fcaeb6ec1e46a0b Mon Sep 17 00:00:00 2001 From: Sebastian Holmin Date: Thu, 9 Jan 2025 14:06:59 +0100 Subject: [PATCH 5/9] Fix panic on capacity=1 As `tokio::sync::mpsc` doesn't allow capacity to be zero, we cannot support buffering only one key pair if we generate it before sending. To get around this we use `reserve` to wait for capacity before generating the key. --- .../src/classic_mceliece.rs | 38 ++++++++++++------- 1 file changed, 24 insertions(+), 14 deletions(-) diff --git a/talpid-tunnel-config-client/src/classic_mceliece.rs b/talpid-tunnel-config-client/src/classic_mceliece.rs index 6aa81345c8dd..ba0888e340f9 100644 --- a/talpid-tunnel-config-client/src/classic_mceliece.rs +++ b/talpid-tunnel-config-client/src/classic_mceliece.rs @@ -8,8 +8,9 @@ use tokio::sync::{mpsc, Mutex}; /// 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. +/// 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 @@ -23,26 +24,35 @@ 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. /// -/// As it can take upwards of 200 ms to generate McEliece key pairs, it needs to be done before we +/// 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. pub fn spawn_keypair_worker(bufsize: usize) -> mpsc::Receiver { - // As one of the key pairs will be buffered by the stack of the spawned thread, we reduce the - // capacity of the channel by one - 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 || loop { - let keypair = keypair_boxed(&mut rand::thread_rng()); - if tx.blocking_send(keypair).is_err() { + 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; - } - }) - .unwrap(); + }; + 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 } From 67f0b2c53830e18f2fd8f0b18291de369db7ab65 Mon Sep 17 00:00:00 2001 From: Sebastian Holmin Date: Thu, 9 Jan 2025 14:10:20 +0100 Subject: [PATCH 6/9] Fix expect message --- talpid-tunnel-config-client/src/classic_mceliece.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/talpid-tunnel-config-client/src/classic_mceliece.rs b/talpid-tunnel-config-client/src/classic_mceliece.rs index ba0888e340f9..93de357cdd00 100644 --- a/talpid-tunnel-config-client/src/classic_mceliece.rs +++ b/talpid-tunnel-config-client/src/classic_mceliece.rs @@ -63,7 +63,7 @@ pub async fn generate_keys() -> KeyPair { .await .recv() .await - .expect("Failed to receive key pair, generating working expectedly closed.") + .expect("Expected to receive key pair, but key generator has been stopped.") } /// Returns a receiver for McEliece key pairs used by PQ tunnels. These are generated in a separate From 722536451f683319f8d4e217afd761191a7e6aac Mon Sep 17 00:00:00 2001 From: Sebastian Holmin Date: Thu, 9 Jan 2025 14:55:57 +0100 Subject: [PATCH 7/9] Add panic note to docstring --- talpid-tunnel-config-client/src/classic_mceliece.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/talpid-tunnel-config-client/src/classic_mceliece.rs b/talpid-tunnel-config-client/src/classic_mceliece.rs index 93de357cdd00..fbe69149cf04 100644 --- a/talpid-tunnel-config-client/src/classic_mceliece.rs +++ b/talpid-tunnel-config-client/src/classic_mceliece.rs @@ -26,6 +26,10 @@ static KEYPAIR_RX: OnceLock>> = OnceLock::new(); /// /// 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 { let (tx, rx) = mpsc::channel(bufsize); From 38046e07cf4b97739d86fbd50d19de095fff11f7 Mon Sep 17 00:00:00 2001 From: Sebastian Holmin Date: Thu, 9 Jan 2025 15:26:00 +0100 Subject: [PATCH 8/9] Do not expose the key pair receiver publicly --- talpid-core/src/tunnel_state_machine/mod.rs | 4 ++-- talpid-tunnel-config-client/src/classic_mceliece.rs | 13 ++++++------- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs index d9d25d95d2a9..e8bd4ed64980 100644 --- a/talpid-core/src/tunnel_state_machine/mod.rs +++ b/talpid-core/src/tunnel_state_machine/mod.rs @@ -25,7 +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::get_or_init_keypair_receiver; +use talpid_tunnel_config_client::classic_mceliece::spawn_keypair_generator; #[cfg(target_os = "macos")] use talpid_types::ErrorExt; @@ -179,7 +179,7 @@ pub async fn spawn( }); // Spawn a worker that pre-computes McEliece key pairs for PQ tunnels - get_or_init_keypair_receiver(); + spawn_keypair_generator(); Ok(TunnelStateMachineHandle { command_tx, diff --git a/talpid-tunnel-config-client/src/classic_mceliece.rs b/talpid-tunnel-config-client/src/classic_mceliece.rs index fbe69149cf04..363fbec82b5b 100644 --- a/talpid-tunnel-config-client/src/classic_mceliece.rs +++ b/talpid-tunnel-config-client/src/classic_mceliece.rs @@ -19,6 +19,8 @@ pub const ALGORITHM_NAME: &str = "Classic-McEliece-460896f-round3"; 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 @@ -62,7 +64,8 @@ pub fn spawn_keypair_worker(bufsize: usize) -> mpsc::Receiver { } pub async fn generate_keys() -> KeyPair { - get_or_init_keypair_receiver() + KEYPAIR_RX + .get_or_init(|| Mutex::new(spawn_keypair_worker(BUFSIZE))) .lock() .await .recv() @@ -70,12 +73,8 @@ pub async fn generate_keys() -> KeyPair { .expect("Expected to receive key pair, but key generator has been stopped.") } -/// Returns a receiver for McEliece key pairs used by PQ tunnels. These are generated in a separate -/// thread to reduce latency when connecting. -/// -/// The first call will spawn the worker which immedietly starts to compute and buffer [`BUFSIZE`] -/// of key pairs. -pub fn get_or_init_keypair_receiver<'a>() -> &'a Mutex> { +/// 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))) } From 60f2ef6bf1ff3a918e80029e7580bbb44bd97da5 Mon Sep 17 00:00:00 2001 From: Sebastian Holmin Date: Thu, 9 Jan 2025 15:30:22 +0100 Subject: [PATCH 9/9] Remove unnecessary pub --- talpid-tunnel-config-client/src/classic_mceliece.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/talpid-tunnel-config-client/src/classic_mceliece.rs b/talpid-tunnel-config-client/src/classic_mceliece.rs index 363fbec82b5b..7484313906a3 100644 --- a/talpid-tunnel-config-client/src/classic_mceliece.rs +++ b/talpid-tunnel-config-client/src/classic_mceliece.rs @@ -32,7 +32,7 @@ static KEYPAIR_RX: OnceLock>> = OnceLock::new(); /// # Panic /// /// Panics if the buffer capacity is 0. -pub fn spawn_keypair_worker(bufsize: usize) -> mpsc::Receiver { +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: