Move key generation to its own module

3 files changed, 128 insertions, 102 deletions

diff --git a/talpid-relay-config-client/src/client.rs b/talpid-relay-config-client/src/client.rs
new file mode 100644
index 0000000000..da8812b4e5
--- /dev/null
+++ b/talpid-relay-config-client/src/client.rs
@@ -0,0 +1,52 @@
+use super::{kem, Error};
+use std::net::IpAddr;
+use talpid_types::net::wireguard::{PresharedKey, PrivateKey, PublicKey};
+use tonic::transport::{Channel, Endpoint, Uri};
+
+mod types {
+    tonic::include_proto!("feature");
+}
+
+type RelayConfigService = types::post_quantum_secure_client::PostQuantumSecureClient<Channel>;
+
+const CONFIG_SERVICE_PORT: u16 = 1337;
+const ALGORITHM_NAME: &str = "Classic-McEliece-8192128f";
+
+// TODO: consider binding to the tunnel interface here, on non-windows platforms
+pub async fn push_pq_key(
+    service_address: IpAddr,
+    current_pubkey: PublicKey,
+) -> Result<(PrivateKey, PresharedKey), Error> {
+    let oqs_key = PrivateKey::new_from_random();
+    let (pubkey, secret) = kem::generate_keys().await?;
+
+    let mut client = new_client(service_address).await?;
+    let response = client
+        .psk_exchange(types::PskRequest {
+            wg_pubkey: current_pubkey.as_bytes().to_vec(),
+            wg_psk_pubkey: oqs_key.public_key().as_bytes().to_vec(),
+            oqs_pubkey: Some(types::OqsPubkey {
+                algorithm_name: ALGORITHM_NAME.to_string(),
+                key_data: pubkey.into_vec(),
+            }),
+        })
+        .await
+        .map_err(Error::GrpcError)?;
+
+    Ok((
+        oqs_key,
+        kem::decapsulate(&secret, &response.into_inner().ciphertext)?,
+    ))
+}
+
+async fn new_client(addr: IpAddr) -> Result<RelayConfigService, Error> {
+    let channel = Endpoint::from_shared(format!("tcp://{addr}:{CONFIG_SERVICE_PORT}"))
+        .expect("Failed to construct URI")
+        .connect_with_connector(tower::service_fn(move |_: Uri| {
+            tokio::net::TcpStream::connect((addr, CONFIG_SERVICE_PORT))
+        }))
+        .await
+        .map_err(Error::GrpcTransportError)?;
+
+    Ok(RelayConfigService::new(channel))
+}
diff --git a/talpid-relay-config-client/src/kem.rs b/talpid-relay-config-client/src/kem.rs
new file mode 100644
index 0000000000..3fd1bff497
--- /dev/null
+++ b/talpid-relay-config-client/src/kem.rs
@@ -0,0 +1,73 @@
+use std::fmt;
+
+use super::Error;
+
+use classic_mceliece_rust::{
+    crypto_kem_dec, crypto_kem_keypair, AesState, RNGState, CRYPTO_BYTES, CRYPTO_CIPHERTEXTBYTES,
+    CRYPTO_PUBLICKEYBYTES, CRYPTO_SECRETKEYBYTES,
+};
+use rand::RngCore;
+use talpid_types::net::wireguard::PresharedKey;
+
+const STACK_SIZE: usize = 8 * 1024 * 1024;
+
+#[derive(Debug)]
+pub struct PublicKey(Box<[u8; CRYPTO_PUBLICKEYBYTES]>);
+
+impl PublicKey {
+    pub fn into_vec(self) -> Vec<u8> {
+        (self.0 as Box<[u8]>).into_vec()
+    }
+}
+
+pub struct SecretKey(Box<[u8; CRYPTO_SECRETKEYBYTES]>);
+
+impl fmt::Debug for SecretKey {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("SecretKey").finish()
+    }
+}
+
+pub async fn generate_keys() -> Result<(PublicKey, SecretKey), Error> {
+    let (tx, rx) = tokio::sync::oneshot::channel();
+
+    let gen_key = move || {
+        let mut rng = AesState::new();
+
+        let mut entropy = [0u8; 48];
+        rand::thread_rng().fill_bytes(&mut entropy);
+        rng.randombytes_init(entropy);
+
+        let mut pubkey = Box::new([0u8; CRYPTO_PUBLICKEYBYTES]);
+        let mut secret = Box::new([0u8; CRYPTO_SECRETKEYBYTES]);
+        crypto_kem_keypair(&mut pubkey, &mut secret, &mut rng).map_err(|error| {
+            log::error!("KEM keypair generation failed: {error}");
+            Error::KeyGenerationFailed
+        })?;
+
+        Ok((PublicKey(pubkey), SecretKey(secret)))
+    };
+
+    std::thread::Builder::new()
+        .stack_size(STACK_SIZE)
+        .spawn(move || {
+            tx.send(gen_key()).unwrap();
+        })
+        .unwrap();
+
+    rx.await.unwrap()
+}
+
+pub fn decapsulate(secret: &SecretKey, ciphertext: &[u8]) -> Result<PresharedKey, Error> {
+    let ct: [u8; CRYPTO_CIPHERTEXTBYTES] = ciphertext
+        .try_into()
+        .map_err(|_| Error::InvalidCiphertext)?;
+    let mut psk = [0u8; CRYPTO_BYTES];
+
+    crypto_kem_dec(&mut psk, &ct, &secret.0).map_err(|error| {
+        log::error!("KEM decapsulation failed: {error}");
+        Error::DecapsulationError
+    })?;
+
+    Ok(PresharedKey::from(psk))
+}
diff --git a/talpid-relay-config-client/src/lib.rs b/talpid-relay-config-client/src/lib.rs
index c89fbc1e15..4d0ba382ba 100644
--- a/talpid-relay-config-client/src/lib.rs
+++ b/talpid-relay-config-client/src/lib.rs
@@ -1,22 +1,6 @@
-use std::net::IpAddr;
-
-use classic_mceliece_rust::{
-    crypto_kem_dec, crypto_kem_keypair, AesState, RNGState, CRYPTO_BYTES, CRYPTO_CIPHERTEXTBYTES,
-    CRYPTO_PUBLICKEYBYTES, CRYPTO_SECRETKEYBYTES,
-};
-use rand::RngCore;
-use talpid_types::net::wireguard::{PresharedKey, PrivateKey, PublicKey};
-use tonic::transport::{Channel, Endpoint, Uri};
-
-mod types {
-    tonic::include_proto!("feature");
-}
-
-type RelayConfigService = types::post_quantum_secure_client::PostQuantumSecureClient<Channel>;
-
-const CONFIG_SERVICE_PORT: u16 = 1337;
-const STACK_SIZE: usize = 8 * 1024 * 1024;
-const ALGORITHM_NAME: &str = "Classic-McEliece-8192128f";
+pub mod client;
+mod kem;
+pub use client::*;
 
 #[derive(Debug)]
 pub enum Error {
@@ -26,86 +10,3 @@ pub enum Error {
     DecapsulationError,
     InvalidCiphertext,
 }
-
-// TODO: consider binding to the tunnel interface here, on non-windows platforms
-pub async fn push_pq_key(
-    service_address: IpAddr,
-    current_pubkey: PublicKey,
-) -> Result<(PrivateKey, PresharedKey), Error> {
-    let oqs_key = PrivateKey::new_from_random();
-
-    let (pubkey, secret) = generate_key().await?;
-
-    let mut client = new_client(service_address).await?;
-    let response = client
-        .psk_exchange(types::PskRequest {
-            wg_pubkey: current_pubkey.as_bytes().to_vec(),
-            wg_psk_pubkey: oqs_key.public_key().as_bytes().to_vec(),
-            oqs_pubkey: Some(types::OqsPubkey {
-                algorithm_name: ALGORITHM_NAME.to_string(),
-                key_data: pubkey.to_vec(),
-            }),
-        })
-        .await
-        .map_err(Error::GrpcError)?;
-
-    let ciphertext = response.into_inner().ciphertext;
-    let ct: [u8; CRYPTO_CIPHERTEXTBYTES] = ciphertext
-        .try_into()
-        .map_err(|_| Error::InvalidCiphertext)?;
-    let mut psk = [0u8; CRYPTO_BYTES];
-
-    crypto_kem_dec(&mut psk, &ct, &secret).map_err(|error| {
-        log::error!("KEM decapsulation failed: {error}");
-        Error::DecapsulationError
-    })?;
-    Ok((oqs_key, PresharedKey::from(psk)))
-}
-
-async fn generate_key() -> Result<
-    (
-        Box<[u8; CRYPTO_PUBLICKEYBYTES]>,
-        Box<[u8; CRYPTO_SECRETKEYBYTES]>,
-    ),
-    Error,
-> {
-    let (tx, rx) = tokio::sync::oneshot::channel();
-
-    let gen_key = move || {
-        let mut rng = AesState::new();
-
-        let mut entropy = [0u8; 48];
-        rand::thread_rng().fill_bytes(&mut entropy);
-        rng.randombytes_init(entropy);
-
-        let mut pubkey = Box::new([0u8; CRYPTO_PUBLICKEYBYTES]);
-        let mut secret = Box::new([0u8; CRYPTO_SECRETKEYBYTES]);
-        crypto_kem_keypair(&mut pubkey, &mut secret, &mut rng).map_err(|error| {
-            log::error!("KEM keypair generation failed: {error}");
-            Error::KeyGenerationFailed
-        })?;
-
-        Ok((pubkey, secret))
-    };
-
-    std::thread::Builder::new()
-        .stack_size(STACK_SIZE)
-        .spawn(move || {
-            tx.send(gen_key()).unwrap();
-        })
-        .unwrap();
-
-    rx.await.unwrap()
-}
-
-async fn new_client(addr: IpAddr) -> Result<RelayConfigService, Error> {
-    let channel = Endpoint::from_shared(format!("tcp://{addr}:{CONFIG_SERVICE_PORT}"))
-        .expect("Failed to construct URI")
-        .connect_with_connector(tower::service_fn(move |_: Uri| {
-            tokio::net::TcpStream::connect((addr, CONFIG_SERVICE_PORT))
-        }))
-        .await
-        .map_err(Error::GrpcTransportError)?;
-
-    Ok(RelayConfigService::new(channel))
-}