Refactor relay config client crate

5 files changed, 95 insertions, 69 deletions

diff --git a/talpid-core/src/tunnel/wireguard/mod.rs b/talpid-core/src/tunnel/wireguard/mod.rs
index 44440c91b6..b8c1f1bd4c 100644
--- a/talpid-core/src/tunnel/wireguard/mod.rs
+++ b/talpid-core/src/tunnel/wireguard/mod.rs
@@ -81,7 +81,7 @@ pub enum Error {
 
     /// Failed to negotiate PQ PSK
     #[error(display = "Failed to negotiate PQ PSK")]
-    PskNegotiationError(talpid_relay_config_client::Error),
+    PskNegotiationError(#[error(source)] talpid_relay_config_client::Error),
 
     /// Failed to set up IP interfaces.
     #[cfg(windows)]
@@ -221,7 +221,7 @@ impl WireguardMonitor {
         let (setup_done_tx, mut setup_done_rx) = mpsc::channel(0);
 
         // Use allowed IPs to block anything but the v4 gateway, if PSK exchange is on.
-        let patched_config_ref;
+        let config_ref;
         let mut patched_config;
         if psk_negotiation.is_some() {
             patched_config = config.clone();
@@ -233,14 +233,14 @@ impl WireguardMonitor {
                     }
                 }
             }
-            patched_config_ref = &patched_config;
+            config_ref = &patched_config;
         } else {
-            patched_config_ref = &config;
+            config_ref = &config;
         }
 
         let tunnel = Self::open_tunnel(
             runtime.clone(),
-            patched_config_ref,
+            config_ref,
             log_path,
             resource_dir,
             tun_provider,
@@ -327,6 +327,8 @@ impl WireguardMonitor {
                 .map_err(CloseMsg::SetupError)?;
 
             if let Some(pubkey) = psk_negotiation {
+                log::debug!("Performing PQ-safe PSK exchange");
+
                 let timeout = std::cmp::min(
                     MAX_PSK_EXCHANGE_TIMEOUT,
                     INITIAL_PSK_EXCHANGE_TIMEOUT.saturating_mul(
diff --git a/talpid-relay-config-client/src/client.rs b/talpid-relay-config-client/src/client.rs
deleted file mode 100644
index da8812b4e5..0000000000
--- a/talpid-relay-config-client/src/client.rs
+++ /dev/null
@@ -1,52 +0,0 @@
-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
index 7b7817e5e3..35679ccddf 100644
--- a/talpid-relay-config-client/src/kem.rs
+++ b/talpid-relay-config-client/src/kem.rs
@@ -3,12 +3,12 @@ use std::fmt;
 use super::Error;
 
 use classic_mceliece_rust::{
-    crypto_kem_dec, crypto_kem_keypair, CRYPTO_BYTES, CRYPTO_CIPHERTEXTBYTES,
-    CRYPTO_PUBLICKEYBYTES, CRYPTO_SECRETKEYBYTES,
+    crypto_kem_dec, crypto_kem_keypair, CRYPTO_BYTES, CRYPTO_PUBLICKEYBYTES, CRYPTO_SECRETKEYBYTES,
 };
 use talpid_types::net::wireguard::PresharedKey;
 
 const STACK_SIZE: usize = 8 * 1024 * 1024;
+pub use classic_mceliece_rust::CRYPTO_CIPHERTEXTBYTES;
 
 #[derive(Debug)]
 pub struct PublicKey(Box<[u8; CRYPTO_PUBLICKEYBYTES]>);
@@ -45,20 +45,20 @@ pub async fn generate_keys() -> Result<(PublicKey, SecretKey), Error> {
     std::thread::Builder::new()
         .stack_size(STACK_SIZE)
         .spawn(move || {
-            tx.send(gen_key()).unwrap();
+            let _ = tx.send(gen_key());
         })
         .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)?;
+pub fn decapsulate(
+    secret: &SecretKey,
+    ciphertext: &[u8; CRYPTO_CIPHERTEXTBYTES],
+) -> Result<PresharedKey, Error> {
     let mut psk = [0u8; CRYPTO_BYTES];
 
-    crypto_kem_dec(&mut psk, &ct, &secret.0).map_err(|error| {
+    crypto_kem_dec(&mut psk, ciphertext, &secret.0).map_err(|error| {
         log::error!("KEM decapsulation failed: {error}");
         Error::DecapsulationError
     })?;
diff --git a/talpid-relay-config-client/src/lib.rs b/talpid-relay-config-client/src/lib.rs
index 4d0ba382ba..4806a68771 100644
--- a/talpid-relay-config-client/src/lib.rs
+++ b/talpid-relay-config-client/src/lib.rs
@@ -1,12 +1,84 @@
-pub mod client;
+use std::{fmt, net::IpAddr};
+use talpid_types::net::wireguard::{PresharedKey, PrivateKey, PublicKey};
+use tonic::transport::Channel;
+
 mod kem;
-pub use client::*;
+
+mod types {
+    tonic::include_proto!("feature");
+}
 
 #[derive(Debug)]
 pub enum Error {
-    GrpcTransportError(tonic::transport::Error),
+    GrpcConnectError(tonic::transport::Error),
     GrpcError(tonic::Status),
     KeyGenerationFailed,
     DecapsulationError,
     InvalidCiphertext,
 }
+
+impl std::fmt::Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use Error::*;
+        match self {
+            GrpcConnectError(_) => "Failed to connect to config service".fmt(f),
+            GrpcError(status) => write!(f, "RPC failed: {}", status),
+            KeyGenerationFailed => "Failed to generate KEM key pair".fmt(f),
+            DecapsulationError => "Failed to decapsulate secret".fmt(f),
+            InvalidCiphertext => "The service returned an invalid ciphertext".fmt(f),
+        }
+    }
+}
+
+impl std::error::Error for Error {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        match self {
+            Self::GrpcConnectError(error) => Some(error),
+            _ => None,
+        }
+    }
+}
+
+type RelayConfigService = types::post_quantum_secure_client::PostQuantumSecureClient<Channel>;
+
+const CONFIG_SERVICE_PORT: u16 = 1337;
+const ALGORITHM_NAME: &str = "Classic-McEliece-8192128f";
+
+/// Generates a new WireGuard key pair and negotiates a PSK with the relay in a PQ-safe
+/// manner. This creates a peer on the relay with the new WireGuard pubkey and PSK,
+/// which can then be used to establish a PQ-safe tunnel to the relay.
+// TODO: consider binding to the tunnel interface here, on non-windows platforms
+pub async fn push_pq_key(
+    service_address: IpAddr,
+    wg_pubkey: PublicKey,
+) -> Result<(PrivateKey, PresharedKey), Error> {
+    let wg_psk_privkey = PrivateKey::new_from_random();
+    let (kem_pubkey, kem_secret) = kem::generate_keys().await?;
+
+    let mut client = new_client(service_address).await?;
+    let response = client
+        .psk_exchange(types::PskRequest {
+            wg_pubkey: wg_pubkey.as_bytes().to_vec(),
+            wg_psk_pubkey: wg_psk_privkey.public_key().as_bytes().to_vec(),
+            oqs_pubkey: Some(types::OqsPubkey {
+                algorithm_name: ALGORITHM_NAME.to_string(),
+                key_data: kem_pubkey.into_vec(),
+            }),
+        })
+        .await
+        .map_err(Error::GrpcError)?;
+
+    let ct: [u8; kem::CRYPTO_CIPHERTEXTBYTES] = response
+        .into_inner()
+        .ciphertext
+        .try_into()
+        .map_err(|_| Error::InvalidCiphertext)?;
+
+    Ok((wg_psk_privkey, kem::decapsulate(&kem_secret, &ct)?))
+}
+
+async fn new_client(addr: IpAddr) -> Result<RelayConfigService, Error> {
+    RelayConfigService::connect(format!("tcp://{addr}:{CONFIG_SERVICE_PORT}"))
+        .await
+        .map_err(Error::GrpcConnectError)
+}
diff --git a/talpid-types/src/net/mod.rs b/talpid-types/src/net/mod.rs
index 31c872da87..a35c227ddd 100644
--- a/talpid-types/src/net/mod.rs
+++ b/talpid-types/src/net/mod.rs
@@ -147,7 +147,11 @@ pub struct TunnelEndpoint {
 
 impl fmt::Display for TunnelEndpoint {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
-        write!(f, "{} - {}", self.tunnel_type, self.endpoint)?;
+        write!(f, "{} ", self.tunnel_type)?;
+        if self.quantum_resistant {
+            write!(f, "(quantum resistant) ")?;
+        }
+        write!(f, "- {}", self.endpoint)?;
         match self.tunnel_type {
             TunnelType::OpenVpn => {
                 if let Some(ref proxy) = self.proxy {