From 524661c76e6fbf47f48822ef41d0bcbd1eaa502a Mon Sep 17 00:00:00 2001 From: Joakim Hulthe Date: Tue, 21 Jan 2025 14:52:47 +0100 Subject: [PATCH] Restart mceliece keygen worker if it stops --- .../src/classic_mceliece.rs | 27 +++++++++++++------ 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/talpid-tunnel-config-client/src/classic_mceliece.rs b/talpid-tunnel-config-client/src/classic_mceliece.rs index 7484313906a3..19f98307f18c 100644 --- a/talpid-tunnel-config-client/src/classic_mceliece.rs +++ b/talpid-tunnel-config-client/src/classic_mceliece.rs @@ -1,4 +1,4 @@ -use std::sync::OnceLock; +use std::{mem, sync::OnceLock}; use classic_mceliece_rust::{keypair_boxed, Ciphertext, CRYPTO_CIPHERTEXTBYTES}; pub use classic_mceliece_rust::{PublicKey, SecretKey, SharedSecret}; @@ -64,18 +64,29 @@ fn spawn_keypair_worker(bufsize: usize) -> mpsc::Receiver { } pub async fn generate_keys() -> KeyPair { - KEYPAIR_RX + let mut rx = 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.") + .await; + + let max_retry_attempts = 10; + + for _ in 0..max_retry_attempts { + match rx.recv().await { + Some(keypair) => return keypair, + None => { + // The key generation worker has stopped for some reason. Try to start it again. + let _old_rx = mem::replace(&mut *rx, spawn_keypair_worker(BUFSIZE)); + } + } + } + + panic!("Failed to start key generation worker") } /// 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 spawn_keypair_generator() { + KEYPAIR_RX.get_or_init(|| Mutex::new(spawn_keypair_worker(BUFSIZE))); } pub fn decapsulate(