Move ephemeral negotiation to own module

5 files changed, 252 insertions, 204 deletions

diff --git a/talpid-tunnel-config-client/Cargo.toml b/talpid-tunnel-config-client/Cargo.toml
index c2cb5f6734..cec3e36d46 100644
--- a/talpid-tunnel-config-client/Cargo.toml
+++ b/talpid-tunnel-config-client/Cargo.toml
@@ -24,6 +24,8 @@ classic-mceliece-rust = { version = "2.0.0", features = [
 ] }
 pqc_kyber = { version = "0.4.0", features = ["std", "kyber1024", "zeroize"] }
 zeroize = "1.5.7"
+
+[target.'cfg(unix)'.dependencies]
 libc = "0.2"
 
 [target.'cfg(windows)'.dependencies.windows-sys]
diff --git a/talpid-tunnel-config-client/src/socket.rs b/talpid-tunnel-config-client/src/socket.rs
index 87dbe23a40..478e757445 100644
--- a/talpid-tunnel-config-client/src/socket.rs
+++ b/talpid-tunnel-config-client/src/socket.rs
@@ -12,15 +12,13 @@ mod sys {
 
     pub use libc::{setsockopt, socklen_t, IPPROTO_TCP, TCP_MAXSEG};
     pub use std::os::fd::{AsRawFd, RawFd};
-    use std::pin::Pin;
-    use std::task::{Context, Poll};
-    use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 
     /// MTU to set on the tunnel config client socket. We want a low value to prevent fragmentation.
     /// Especially on Android, we've found that the real MTU is often lower than the default MTU, and
     /// we cannot lower it further. This causes the outer packets to be dropped. Also, MTU detection
     /// will likely occur after the PQ handshake, so we cannot assume that the MTU is already
     /// correctly configured.
+    /// This is set to the lowest possible IPv4 MTU.
     const CONFIG_CLIENT_MTU: u16 = 576;
 
     pub struct TcpSocket {
diff --git a/talpid-wireguard/Cargo.toml b/talpid-wireguard/Cargo.toml
index 3fea8a17c0..f9ed018562 100644
--- a/talpid-wireguard/Cargo.toml
+++ b/talpid-wireguard/Cargo.toml
@@ -16,7 +16,6 @@ futures = "0.3.15"
 hex = "0.4"
 ipnetwork = { workspace = true }
 once_cell = { workspace = true }
-libc = "0.2.150"
 log = { workspace = true }
 parking_lot = "0.12.0"
 talpid-routing = { path = "../talpid-routing" }
@@ -44,6 +43,7 @@ tokio-stream = { version = "0.1", features = ["io-util"] }
 
 [target.'cfg(unix)'.dependencies]
 nix = "0.23"
+libc = "0.2.150"
 
 [target.'cfg(any(target_os = "linux", target_os = "macos"))'.dependencies]
 talpid-net = { path = "../talpid-net" }
diff --git a/talpid-wireguard/src/ephemeral.rs b/talpid-wireguard/src/ephemeral.rs
new file mode 100644
index 0000000000..b00babf781
--- /dev/null
+++ b/talpid-wireguard/src/ephemeral.rs
@@ -0,0 +1,244 @@
+//! This module takes care of obtaining ephemeral peers, updating the WireGuard configuration and
+//! restarting obfuscation and WG tunnels when necessary.
+
+use super::{config::Config, obfuscation::ObfuscatorHandle, CloseMsg, Error, Tunnel};
+#[cfg(target_os = "android")]
+use std::sync::Mutex;
+use std::{
+    net::IpAddr,
+    sync::{mpsc as sync_mpsc, Arc},
+    time::Duration,
+};
+#[cfg(target_os = "android")]
+use talpid_tunnel::tun_provider::TunProvider;
+
+use ipnetwork::IpNetwork;
+use talpid_types::net::wireguard::{PresharedKey, PrivateKey, PublicKey};
+use tokio::sync::Mutex as AsyncMutex;
+
+const INITIAL_PSK_EXCHANGE_TIMEOUT: Duration = Duration::from_secs(8);
+const MAX_PSK_EXCHANGE_TIMEOUT: Duration = Duration::from_secs(48);
+const PSK_EXCHANGE_TIMEOUT_MULTIPLIER: u32 = 2;
+
+#[cfg(windows)]
+pub async fn config_ephemeral_peers(
+    tunnel: &Arc<AsyncMutex<Option<Box<dyn Tunnel>>>>,
+    config: &mut Config,
+    retry_attempt: u32,
+    obfuscator: Arc<AsyncMutex<Option<ObfuscatorHandle>>>,
+    close_obfs_sender: sync_mpsc::Sender<CloseMsg>,
+) -> std::result::Result<(), CloseMsg> {
+    let iface_name = {
+        let tunnel = tunnel.lock().await;
+        let tunnel = tunnel.as_ref().unwrap();
+        tunnel.get_interface_name()
+    };
+
+    log::trace!("Temporarily lowering tunnel MTU before ephemeral peer config");
+    try_set_ipv4_mtu(&iface_name, talpid_tunnel::MIN_IPV4_MTU);
+
+    config_ephemeral_peers_inner(tunnel, config, retry_attempt, obfuscator, close_obfs_sender)
+        .await?;
+
+    log::trace!("Resetting tunnel MTU");
+    try_set_ipv4_mtu(&iface_name, config.mtu);
+
+    Ok(())
+}
+
+#[cfg(windows)]
+fn try_set_ipv4_mtu(alias: &str, mtu: u16) {
+    use talpid_windows::net::*;
+    match luid_from_alias(alias) {
+        Ok(luid) => {
+            if let Err(error) = set_mtu(u32::from(mtu), luid, AddressFamily::Ipv4) {
+                log::error!("Failed to set tunnel interface MTU: {error}");
+            }
+        }
+        Err(error) => {
+            log::error!("Failed to obtain tunnel interface LUID: {error}")
+        }
+    }
+}
+
+#[cfg(not(windows))]
+pub async fn config_ephemeral_peers(
+    tunnel: &Arc<AsyncMutex<Option<Box<dyn Tunnel>>>>,
+    config: &mut Config,
+    retry_attempt: u32,
+    obfuscator: Arc<AsyncMutex<Option<ObfuscatorHandle>>>,
+    close_obfs_sender: sync_mpsc::Sender<CloseMsg>,
+    #[cfg(target_os = "android")] tun_provider: Arc<Mutex<TunProvider>>,
+) -> std::result::Result<(), CloseMsg> {
+    config_ephemeral_peers_inner(
+        tunnel,
+        config,
+        retry_attempt,
+        obfuscator,
+        close_obfs_sender,
+        #[cfg(target_os = "android")]
+        tun_provider,
+    )
+    .await
+}
+
+async fn config_ephemeral_peers_inner(
+    tunnel: &Arc<AsyncMutex<Option<Box<dyn Tunnel>>>>,
+    config: &mut Config,
+    retry_attempt: u32,
+    obfuscator: Arc<AsyncMutex<Option<ObfuscatorHandle>>>,
+    close_obfs_sender: sync_mpsc::Sender<CloseMsg>,
+    #[cfg(target_os = "android")] tun_provider: Arc<Mutex<TunProvider>>,
+) -> std::result::Result<(), CloseMsg> {
+    let ephemeral_private_key = PrivateKey::new_from_random();
+    let close_obfs_sender = close_obfs_sender.clone();
+
+    let exit_should_have_daita = config.daita && !config.is_multihop();
+    let exit_psk = request_ephemeral_peer(
+        retry_attempt,
+        config,
+        ephemeral_private_key.public_key(),
+        config.quantum_resistant,
+        exit_should_have_daita,
+    )
+    .await?;
+
+    log::debug!("Retrieved ephemeral peer");
+
+    if config.is_multihop() {
+        // Set up tunnel to lead to entry
+        let mut entry_tun_config = config.clone();
+        entry_tun_config
+            .entry_peer
+            .allowed_ips
+            .push(IpNetwork::new(IpAddr::V4(config.ipv4_gateway), 32).unwrap());
+
+        let close_obfs_sender = close_obfs_sender.clone();
+        let entry_config = reconfigure_tunnel(
+            tunnel,
+            entry_tun_config,
+            obfuscator.clone(),
+            close_obfs_sender,
+            #[cfg(target_os = "android")]
+            &tun_provider,
+        )
+        .await?;
+        let entry_psk = request_ephemeral_peer(
+            retry_attempt,
+            &entry_config,
+            ephemeral_private_key.public_key(),
+            config.quantum_resistant,
+            config.daita,
+        )
+        .await?;
+        log::debug!("Successfully exchanged PSK with entry peer");
+
+        config.entry_peer.psk = entry_psk;
+    }
+
+    config.exit_peer_mut().psk = exit_psk;
+    #[cfg(daita)]
+    if config.daita {
+        log::trace!("Enabling constant packet size for entry peer");
+        config.entry_peer.constant_packet_size = true;
+    }
+
+    config.tunnel.private_key = ephemeral_private_key;
+
+    *config = reconfigure_tunnel(
+        tunnel,
+        config.clone(),
+        obfuscator,
+        close_obfs_sender,
+        #[cfg(target_os = "android")]
+        &tun_provider,
+    )
+    .await?;
+
+    #[cfg(daita)]
+    if config.daita {
+        // Start local DAITA machines
+        let mut tunnel = tunnel.lock().await;
+        if let Some(tunnel) = tunnel.as_mut() {
+            tunnel
+                .start_daita()
+                .map_err(Error::TunnelError)
+                .map_err(CloseMsg::SetupError)?;
+        }
+    }
+
+    Ok(())
+}
+
+/// Reconfigures the tunnel to use the provided config while potentially modifying the config
+/// and restarting the obfuscation provider. Returns the new config used by the new tunnel.
+async fn reconfigure_tunnel(
+    tunnel: &Arc<AsyncMutex<Option<Box<dyn Tunnel>>>>,
+    mut config: Config,
+    obfuscator: Arc<AsyncMutex<Option<ObfuscatorHandle>>>,
+    close_obfs_sender: sync_mpsc::Sender<CloseMsg>,
+    #[cfg(target_os = "android")] tun_provider: &Arc<Mutex<TunProvider>>,
+) -> std::result::Result<Config, CloseMsg> {
+    let mut obfs_guard = obfuscator.lock().await;
+    if let Some(obfuscator_handle) = obfs_guard.take() {
+        obfuscator_handle.abort();
+        *obfs_guard = super::obfuscation::apply_obfuscation_config(
+            &mut config,
+            close_obfs_sender,
+            #[cfg(target_os = "android")]
+            tun_provider.clone(),
+        )
+        .await
+        .map_err(CloseMsg::ObfuscatorFailed)?;
+    }
+
+    let mut tunnel = tunnel.lock().await;
+
+    let set_config_future = tunnel
+        .as_mut()
+        .map(|tunnel| tunnel.set_config(config.clone()));
+
+    if let Some(f) = set_config_future {
+        f.await
+            .map_err(Error::TunnelError)
+            .map_err(CloseMsg::SetupError)?;
+    }
+
+    Ok(config)
+}
+
+async fn request_ephemeral_peer(
+    retry_attempt: u32,
+    config: &Config,
+    wg_psk_pubkey: PublicKey,
+    enable_pq: bool,
+    enable_daita: bool,
+) -> std::result::Result<Option<PresharedKey>, CloseMsg> {
+    log::debug!("Requesting ephemeral peer");
+
+    let timeout = std::cmp::min(
+        MAX_PSK_EXCHANGE_TIMEOUT,
+        INITIAL_PSK_EXCHANGE_TIMEOUT
+            .saturating_mul(PSK_EXCHANGE_TIMEOUT_MULTIPLIER.saturating_pow(retry_attempt)),
+    );
+
+    let ephemeral = tokio::time::timeout(
+        timeout,
+        talpid_tunnel_config_client::request_ephemeral_peer(
+            config.ipv4_gateway,
+            config.tunnel.private_key.public_key(),
+            wg_psk_pubkey,
+            enable_pq,
+            enable_daita,
+        ),
+    )
+    .await
+    .map_err(|_timeout_err| {
+        log::warn!("Timeout while negotiating ephemeral peer");
+        CloseMsg::EphemeralPeerNegotiationTimeout
+    })?
+    .map_err(Error::EphemeralPeerNegotiationError)
+    .map_err(CloseMsg::SetupError)?;
+
+    Ok(ephemeral.psk)
+}
diff --git a/talpid-wireguard/src/lib.rs b/talpid-wireguard/src/lib.rs
index a021cb3959..2e6584a848 100644
--- a/talpid-wireguard/src/lib.rs
+++ b/talpid-wireguard/src/lib.rs
@@ -17,7 +17,6 @@ use std::{
     path::Path,
     pin::Pin,
     sync::{mpsc as sync_mpsc, Arc, Mutex},
-    time::Duration,
 };
 #[cfg(target_os = "linux")]
 use std::{env, sync::LazyLock};
@@ -27,12 +26,8 @@ use talpid_routing::{self, RequiredRoute};
 use talpid_tunnel::tun_provider;
 use talpid_tunnel::{tun_provider::TunProvider, TunnelArgs, TunnelEvent, TunnelMetadata};
 
-use ipnetwork::IpNetwork;
 use talpid_types::{
-    net::{
-        wireguard::{PresharedKey, PrivateKey, PublicKey},
-        AllowedTunnelTraffic, Endpoint, TransportProtocol,
-    },
+    net::{AllowedTunnelTraffic, Endpoint, TransportProtocol},
     BoxedError, ErrorExt,
 };
 use tokio::sync::Mutex as AsyncMutex;
@@ -40,6 +35,7 @@ use tokio::sync::Mutex as AsyncMutex;
 /// WireGuard config data-types
 pub mod config;
 mod connectivity_check;
+mod ephemeral;
 mod logging;
 mod obfuscation;
 mod ping_monitor;
@@ -141,10 +137,6 @@ pub struct WireguardMonitor {
     obfuscator: Arc<AsyncMutex<Option<ObfuscatorHandle>>>,
 }
 
-const INITIAL_PSK_EXCHANGE_TIMEOUT: Duration = Duration::from_secs(8);
-const MAX_PSK_EXCHANGE_TIMEOUT: Duration = Duration::from_secs(48);
-const PSK_EXCHANGE_TIMEOUT_MULTIPLIER: u32 = 2;
-
 #[cfg(target_os = "linux")]
 /// Overrides the preference for the kernel module for WireGuard.
 static FORCE_USERSPACE_WIREGUARD: LazyLock<bool> = LazyLock::new(|| {
@@ -253,13 +245,7 @@ impl WireguardMonitor {
 
             let ephemeral_obfs_sender = close_obfs_sender.clone();
             if config.quantum_resistant || config.daita {
-                #[cfg(windows)]
-                {
-                    log::trace!("Temporarily lowering tunnel MTU before ephemeral peer config");
-                    try_set_ipv4_mtu(&iface_name, talpid_tunnel::MIN_IPV4_MTU);
-                }
-
-                Self::config_ephemeral_peers(
+                ephemeral::config_ephemeral_peers(
                     &tunnel,
                     &mut config,
                     args.retry_attempt,
@@ -268,12 +254,6 @@ impl WireguardMonitor {
                 )
                 .await?;
 
-                #[cfg(windows)]
-                {
-                    log::trace!("Resetting tunnel MTU");
-                    try_set_ipv4_mtu(&iface_name, config.mtu);
-                }
-
                 let metadata = Self::tunnel_metadata(&iface_name, &config);
                 (on_event)(TunnelEvent::InterfaceUp(
                     metadata,
@@ -484,7 +464,7 @@ impl WireguardMonitor {
                 // Ping before negotiating the ephemeral peer to make sure that the tunnel works.
                 tokio::task::spawn_blocking(ping()).await.unwrap()?;
                 let ephemeral_obfs_sender = close_obfs_sender.clone();
-                Self::config_ephemeral_peers(
+                ephemeral::config_ephemeral_peers(
                     &tunnel,
                     &mut config,
                     args.retry_attempt,
@@ -569,131 +549,6 @@ impl WireguardMonitor {
         AllowedTunnelTraffic::All
     }
 
-    async fn config_ephemeral_peers(
-        tunnel: &Arc<AsyncMutex<Option<Box<dyn Tunnel>>>>,
-        config: &mut Config,
-        retry_attempt: u32,
-        obfuscator: Arc<AsyncMutex<Option<ObfuscatorHandle>>>,
-        close_obfs_sender: sync_mpsc::Sender<CloseMsg>,
-        #[cfg(target_os = "android")] tun_provider: Arc<Mutex<TunProvider>>,
-    ) -> std::result::Result<(), CloseMsg> {
-        let ephemeral_private_key = PrivateKey::new_from_random();
-        let close_obfs_sender = close_obfs_sender.clone();
-
-        let exit_should_have_daita = config.daita && !config.is_multihop();
-        let exit_psk = Self::request_ephemeral_peer(
-            retry_attempt,
-            config,
-            ephemeral_private_key.public_key(),
-            config.quantum_resistant,
-            exit_should_have_daita,
-        )
-        .await?;
-
-        log::debug!("Retrieved ephemeral peer");
-
-        if config.is_multihop() {
-            // Set up tunnel to lead to entry
-            let mut entry_tun_config = config.clone();
-            entry_tun_config
-                .entry_peer
-                .allowed_ips
-                .push(IpNetwork::new(IpAddr::V4(config.ipv4_gateway), 32).unwrap());
-
-            let close_obfs_sender = close_obfs_sender.clone();
-            let entry_config = Self::reconfigure_tunnel(
-                tunnel,
-                entry_tun_config,
-                obfuscator.clone(),
-                close_obfs_sender,
-                #[cfg(target_os = "android")]
-                &tun_provider,
-            )
-            .await?;
-            let entry_psk = Self::request_ephemeral_peer(
-                retry_attempt,
-                &entry_config,
-                ephemeral_private_key.public_key(),
-                config.quantum_resistant,
-                config.daita,
-            )
-            .await?;
-            log::debug!("Successfully exchanged PSK with entry peer");
-
-            config.entry_peer.psk = entry_psk;
-        }
-
-        config.exit_peer_mut().psk = exit_psk;
-        #[cfg(daita)]
-        if config.daita {
-            log::trace!("Enabling constant packet size for entry peer");
-            config.entry_peer.constant_packet_size = true;
-        }
-
-        config.tunnel.private_key = ephemeral_private_key;
-
-        *config = Self::reconfigure_tunnel(
-            tunnel,
-            config.clone(),
-            obfuscator,
-            close_obfs_sender,
-            #[cfg(target_os = "android")]
-            &tun_provider,
-        )
-        .await?;
-
-        #[cfg(daita)]
-        if config.daita {
-            // Start local DAITA machines
-            let mut tunnel = tunnel.lock().await;
-            if let Some(tunnel) = tunnel.as_mut() {
-                tunnel
-                    .start_daita()
-                    .map_err(Error::TunnelError)
-                    .map_err(CloseMsg::SetupError)?;
-            }
-        }
-
-        Ok(())
-    }
-
-    /// Reconfigures the tunnel to use the provided config while potentially modifying the config
-    /// and restarting the obfuscation provider. Returns the new config used by the new tunnel.
-    async fn reconfigure_tunnel(
-        tunnel: &Arc<AsyncMutex<Option<Box<dyn Tunnel>>>>,
-        mut config: Config,
-        obfuscator: Arc<AsyncMutex<Option<ObfuscatorHandle>>>,
-        close_obfs_sender: sync_mpsc::Sender<CloseMsg>,
-        #[cfg(target_os = "android")] tun_provider: &Arc<Mutex<TunProvider>>,
-    ) -> std::result::Result<Config, CloseMsg> {
-        let mut obfs_guard = obfuscator.lock().await;
-        if let Some(obfuscator_handle) = obfs_guard.take() {
-            obfuscator_handle.abort();
-            *obfs_guard = obfuscation::apply_obfuscation_config(
-                &mut config,
-                close_obfs_sender,
-                #[cfg(target_os = "android")]
-                tun_provider.clone(),
-            )
-            .await
-            .map_err(CloseMsg::ObfuscatorFailed)?;
-        }
-
-        let mut tunnel = tunnel.lock().await;
-
-        let set_config_future = tunnel
-            .as_mut()
-            .map(|tunnel| tunnel.set_config(config.clone()));
-
-        if let Some(f) = set_config_future {
-            f.await
-                .map_err(Error::TunnelError)
-                .map_err(CloseMsg::SetupError)?;
-        }
-
-        Ok(config)
-    }
-
     /// Replace `0.0.0.0/0`/`::/0` with the gateway IPs when `gateway_only` is true.
     /// Used to block traffic to other destinations while connecting on Android.
     #[cfg(target_os = "android")]
@@ -764,42 +619,6 @@ impl WireguardMonitor {
         Ok(())
     }
 
-    async fn request_ephemeral_peer(
-        retry_attempt: u32,
-        config: &Config,
-        wg_psk_pubkey: PublicKey,
-        enable_pq: bool,
-        enable_daita: bool,
-    ) -> std::result::Result<Option<PresharedKey>, CloseMsg> {
-        log::debug!("Requesting ephemeral peer");
-
-        let timeout = std::cmp::min(
-            MAX_PSK_EXCHANGE_TIMEOUT,
-            INITIAL_PSK_EXCHANGE_TIMEOUT
-                .saturating_mul(PSK_EXCHANGE_TIMEOUT_MULTIPLIER.saturating_pow(retry_attempt)),
-        );
-
-        let ephemeral = tokio::time::timeout(
-            timeout,
-            talpid_tunnel_config_client::request_ephemeral_peer(
-                config.ipv4_gateway,
-                config.tunnel.private_key.public_key(),
-                wg_psk_pubkey,
-                enable_pq,
-                enable_daita,
-            ),
-        )
-        .await
-        .map_err(|_timeout_err| {
-            log::warn!("Timeout while negotiating ephemeral peer");
-            CloseMsg::EphemeralPeerNegotiationTimeout
-        })?
-        .map_err(Error::EphemeralPeerNegotiationError)
-        .map_err(CloseMsg::SetupError)?;
-
-        Ok(ephemeral.psk)
-    }
-
     #[allow(unused_variables)]
     fn open_tunnel(
         runtime: tokio::runtime::Handle,
@@ -1232,18 +1051,3 @@ fn will_nm_manage_dns() -> bool {
         })
         .unwrap_or(false)
 }
-
-#[cfg(windows)]
-fn try_set_ipv4_mtu(alias: &str, mtu: u16) {
-    use talpid_windows::net::*;
-    match luid_from_alias(alias) {
-        Ok(luid) => {
-            if let Err(error) = set_mtu(u32::from(mtu), luid, AddressFamily::Ipv4) {
-                log::error!("Failed to set tunnel interface MTU: {error}");
-            }
-        }
-        Err(error) => {
-            log::error!("Failed to obtain tunnel interface LUID: {error}")
-        }
-    }
-}