Merge branch 'minor-cleanup-in-pq-code-des-1335'

4 files changed, 72 insertions, 61 deletions

diff --git a/talpid-tunnel-config-client/src/classic_mceliece.rs b/talpid-tunnel-config-client/src/classic_mceliece.rs
index 2036bc3fc7..7f7edd43a7 100644
--- a/talpid-tunnel-config-client/src/classic_mceliece.rs
+++ b/talpid-tunnel-config-client/src/classic_mceliece.rs
@@ -2,9 +2,7 @@ use classic_mceliece_rust::{keypair_boxed, Ciphertext, CRYPTO_CIPHERTEXTBYTES};
 pub use classic_mceliece_rust::{PublicKey, SecretKey, SharedSecret};
 
 /// The `keypair_boxed` function needs just under 1 MiB of stack in debug
-/// builds. Even though it probably works to run it directly on the main
-/// thread on all OSes, we take this precaution and always generate the huge
-/// keys on a separate thread with a large enough stack.
+/// builds.
 const STACK_SIZE: usize = 2 * 1024 * 1024;
 
 /// Use the smallest CME variant with NIST security level 3. This variant has significantly smaller
@@ -14,6 +12,10 @@ 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();
 
+    // 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 || {
diff --git a/talpid-tunnel-config-client/src/lib.rs b/talpid-tunnel-config-client/src/lib.rs
index ab47c13be7..2a6163b61d 100644
--- a/talpid-tunnel-config-client/src/lib.rs
+++ b/talpid-tunnel-config-client/src/lib.rs
@@ -85,24 +85,49 @@ pub struct EphemeralPeer {
     pub psk: Option<PresharedKey>,
 }
 
+/// Negotiate a short-lived peer with a PQ-safe PSK or with DAITA enabled.
+#[cfg(not(target_os = "ios"))]
+pub async fn request_ephemeral_peer(
+    service_address: Ipv4Addr,
+    parent_pubkey: PublicKey,
+    ephemeral_pubkey: PublicKey,
+    enable_post_quantum: bool,
+    enable_daita: bool,
+) -> Result<EphemeralPeer, Error> {
+    let client = connect_relay_config_client(service_address).await?;
+
+    request_ephemeral_peer_with(
+        client,
+        parent_pubkey,
+        ephemeral_pubkey,
+        enable_post_quantum,
+        enable_daita,
+    )
+    .await
+}
+
 pub async fn request_ephemeral_peer_with(
     mut client: RelayConfigService,
     parent_pubkey: PublicKey,
     ephemeral_pubkey: PublicKey,
-    enable_post_quantum: bool,
+    enable_quantum_resistant: bool,
     enable_daita: bool,
 ) -> Result<EphemeralPeer, Error> {
-    let (pq_request, kem_secrets) = post_quantum_secrets(enable_post_quantum).await;
-    let daita = Some(proto::DaitaRequestV1 {
-        activate_daita: enable_daita,
-    });
+    let (pq_request, kem_secrets) = if enable_quantum_resistant {
+        let (pq_request, kem_secrets) = post_quantum_secrets().await;
+        (Some(pq_request), Some(kem_secrets))
+    } else {
+        (None, None)
+    };
 
     let response = client
         .register_peer_v1(proto::EphemeralPeerRequestV1 {
             wg_parent_pubkey: parent_pubkey.as_bytes().to_vec(),
             wg_ephemeral_peer_pubkey: ephemeral_pubkey.as_bytes().to_vec(),
             post_quantum: pq_request,
-            daita,
+            daita: Some(proto::DaitaRequestV1 {
+                activate_daita: enable_daita,
+            }),
         })
         .await
         .map_err(Error::GrpcError)?;
@@ -154,55 +179,28 @@ pub async fn request_ephemeral_peer_with(
     Ok(EphemeralPeer { psk })
 }
 
-/// Negotiate a short-lived peer with a PQ-safe PSK or with DAITA enabled.
-#[cfg(not(target_os = "ios"))]
-pub async fn request_ephemeral_peer(
-    service_address: Ipv4Addr,
-    parent_pubkey: PublicKey,
-    ephemeral_pubkey: PublicKey,
-    enable_post_quantum: bool,
-    enable_daita: bool,
-) -> Result<EphemeralPeer, Error> {
-    let client = new_client(service_address).await?;
-
-    request_ephemeral_peer_with(
-        client,
-        parent_pubkey,
-        ephemeral_pubkey,
-        enable_post_quantum,
-        enable_daita,
-    )
-    .await
-}
-
-async fn post_quantum_secrets(
-    enable_post_quantum: bool,
-) -> (
-    Option<PostQuantumRequestV1>,
-    Option<(classic_mceliece_rust::SecretKey<'static>, ml_kem::Keypair)>,
+async fn post_quantum_secrets() -> (
+    PostQuantumRequestV1,
+    (classic_mceliece_rust::SecretKey<'static>, ml_kem::Keypair),
 ) {
-    if enable_post_quantum {
-        let (cme_kem_pubkey, cme_kem_secret) = classic_mceliece::generate_keys().await;
-        let ml_kem_keypair = ml_kem::keypair();
+    let (cme_kem_pubkey, cme_kem_secret) = classic_mceliece::generate_keys().await;
+    let ml_kem_keypair = ml_kem::keypair();
 
-        (
-            Some(proto::PostQuantumRequestV1 {
-                kem_pubkeys: vec![
-                    proto::KemPubkeyV1 {
-                        algorithm_name: classic_mceliece::ALGORITHM_NAME.to_owned(),
-                        key_data: cme_kem_pubkey.as_array().to_vec(),
-                    },
-                    proto::KemPubkeyV1 {
-                        algorithm_name: ml_kem::ALGORITHM_NAME.to_owned(),
-                        key_data: ml_kem_keypair.encapsulation_key(),
-                    },
-                ],
-            }),
-            Some((cme_kem_secret, ml_kem_keypair)),
-        )
-    } else {
-        (None, None)
-    }
+    (
+        proto::PostQuantumRequestV1 {
+            kem_pubkeys: vec![
+                proto::KemPubkeyV1 {
+                    algorithm_name: classic_mceliece::ALGORITHM_NAME.to_owned(),
+                    key_data: cme_kem_pubkey.as_array().to_vec(),
+                },
+                proto::KemPubkeyV1 {
+                    algorithm_name: ml_kem::ALGORITHM_NAME.to_owned(),
+                    key_data: ml_kem_keypair.encapsulation_key(),
+                },
+            ],
+        },
+        (cme_kem_secret, ml_kem_keypair),
+    )
 }
 
 /// Performs `dst = dst ^ src`.
@@ -212,22 +210,25 @@ fn xor_assign(dst: &mut [u8; 32], src: &[u8; 32]) {
     }
 }
 
+/// Create a new `RelayConfigService` connected to the given IP.
+/// On non-Windows platforms the connection is made with a socket where the MSS
+/// value has been speficically lowered, to avoid MTU issues. See the `socket` module.
 #[cfg(not(target_os = "ios"))]
-async fn new_client(addr: Ipv4Addr) -> Result<RelayConfigService, Error> {
+async fn connect_relay_config_client(ip: Ipv4Addr) -> Result<RelayConfigService, Error> {
     use futures::TryFutureExt;
 
     let endpoint = Endpoint::from_static("tcp://0.0.0.0:0");
-    let addr = IpAddr::V4(addr);
+    let addr = SocketAddr::new(IpAddr::V4(ip), CONFIG_SERVICE_PORT);
 
-    let conn = endpoint
+    let connection = endpoint
         .connect_with_connector(service_fn(move |_| async move {
             let sock = socket::TcpSocket::new()?;
-            sock.connect(SocketAddr::new(addr, CONFIG_SERVICE_PORT))
+            sock.connect(addr)
                 .map_ok(hyper_util::rt::tokio::TokioIo::new)
                 .await
         }))
         .await
         .map_err(Error::GrpcConnectError)?;
 
-    Ok(RelayConfigService::new(conn))
+    Ok(RelayConfigService::new(connection))
 }
diff --git a/talpid-tunnel-config-client/src/socket.rs b/talpid-tunnel-config-client/src/socket.rs
index 478e757445..f7e48a6f8f 100644
--- a/talpid-tunnel-config-client/src/socket.rs
+++ b/talpid-tunnel-config-client/src/socket.rs
@@ -1,5 +1,10 @@
 //! A TCP stream with a low MSS set. This prevents incorrectly configured MTU from causing
 //! fragmentation/packet loss. This is only supported on non-Windows targets.
+//!
+//! On windows this solution does not work. So on Windows we instead temporarily lower the MTU
+//! while negotiating the ephemeral peer. This code lives in `talpid-wireguard/src/ephemeral.rs`
+//! These two solutions to the same problem should probably be refactored to end up closer
+//! to each other.
 
 use std::io;
 use std::net::SocketAddr;
diff --git a/talpid-wireguard/src/ephemeral.rs b/talpid-wireguard/src/ephemeral.rs
index b00babf781..5440a142f6 100644
--- a/talpid-wireguard/src/ephemeral.rs
+++ b/talpid-wireguard/src/ephemeral.rs
@@ -34,6 +34,9 @@ pub async fn config_ephemeral_peers(
         tunnel.get_interface_name()
     };
 
+    // Lower the MTU in order to make the ephemeral peer handshake work more reliably.
+    // On unix based operating systems this is done by setting the MSS directly on the
+    // TCP socket. But that solution does not work on Windows, so we do this MTU hack instead.
     log::trace!("Temporarily lowering tunnel MTU before ephemeral peer config");
     try_set_ipv4_mtu(&iface_name, talpid_tunnel::MIN_IPV4_MTU);