Force WireGuard handshake before PQ handshake

mullvad · Dec 16, 2024 · a122a93 · a122a93
1 parent 3911c21
commit a122a93
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 3 deletions.
diff --git a/talpid-wireguard/src/connectivity/mod.rs b/talpid-wireguard/src/connectivity/mod.rs
@@ -6,7 +6,6 @@ mod mock;
 mod monitor;
 mod pinger;
 
-#[cfg(target_os = "android")]
 pub use check::Cancellable;
 pub use check::Check;
 pub use error::Error;

diff --git a/talpid-wireguard/src/ephemeral.rs b/talpid-wireguard/src/ephemeral.rs
@@ -1,6 +1,8 @@
 //! This module takes care of obtaining ephemeral peers, updating the WireGuard configuration and
 //! restarting obfuscation and WG tunnels when necessary.
 
+#[cfg(not(target_os = "android"))]
+use super::connectivity;
 #[cfg(target_os = "android")] // On Android, the Tunnel trait is not imported by default.
 use super::Tunnel;
 use super::{config::Config, obfuscation::ObfuscatorHandle, CloseMsg, Error, TunnelType};
@@ -31,6 +33,7 @@ pub async fn config_ephemeral_peers(
     retry_attempt: u32,
     obfuscator: Arc<AsyncMutex<Option<ObfuscatorHandle>>>,
     close_obfs_sender: sync_mpsc::Sender<CloseMsg>,
+    connectivity: &mut connectivity::Check<connectivity::Cancellable>,
 ) -> std::result::Result<(), CloseMsg> {
     let iface_name = {
         let tunnel = tunnel.lock().await;
@@ -44,8 +47,15 @@ pub async fn config_ephemeral_peers(
     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?;
+    config_ephemeral_peers_inner(
+        tunnel,
+        config,
+        retry_attempt,
+        obfuscator,
+        close_obfs_sender,
+        connectivity,
+    )
+    .await?;
 
     log::trace!("Resetting tunnel MTU");
     try_set_ipv4_mtu(&iface_name, config.mtu);
@@ -75,6 +85,9 @@ pub async fn config_ephemeral_peers(
     retry_attempt: u32,
     obfuscator: Arc<AsyncMutex<Option<ObfuscatorHandle>>>,
     close_obfs_sender: sync_mpsc::Sender<CloseMsg>,
+    #[cfg(not(target_os = "android"))] connectivity: &mut connectivity::Check<
+        connectivity::Cancellable,
+    >,
     #[cfg(target_os = "android")] tun_provider: Arc<Mutex<TunProvider>>,
 ) -> Result<(), CloseMsg> {
     config_ephemeral_peers_inner(
@@ -95,8 +108,14 @@ async fn config_ephemeral_peers_inner(
     retry_attempt: u32,
     obfuscator: Arc<AsyncMutex<Option<ObfuscatorHandle>>>,
     close_obfs_sender: sync_mpsc::Sender<CloseMsg>,
+    #[cfg(not(target_os = "android"))] connectivity: &mut connectivity::Check<
+        connectivity::Cancellable,
+    >,
     #[cfg(target_os = "android")] tun_provider: Arc<Mutex<TunProvider>>,
 ) -> Result<(), CloseMsg> {
+    #[cfg(not(target_os = "android"))]
+    reestablish_tunnel_connection(tunnel, connectivity).await?;
+
     let ephemeral_private_key = PrivateKey::new_from_random();
     let close_obfs_sender = close_obfs_sender.clone();
 
@@ -134,6 +153,10 @@ async fn config_ephemeral_peers_inner(
             &tun_provider,
         )
         .await?;
+
+        #[cfg(not(target_os = "android"))]
+        reestablish_tunnel_connection(tunnel, connectivity).await?;
+
         let entry_ephemeral_peer = request_ephemeral_peer(
             retry_attempt,
             &entry_config,
@@ -268,6 +291,38 @@ async fn reconfigure_tunnel(
     Ok(config)
 }
 
+/// Ensure that the WireGuard tunnel works. This is useful after updating the WireGuard config, to
+/// force a WireGuard handshake.
+/// This should reduce the number of PQ timeouts.
+#[cfg(not(target_os = "android"))]
+async fn reestablish_tunnel_connection(
+    tunnel: &Arc<AsyncMutex<Option<TunnelType>>>,
+    connectivity: &mut connectivity::Check<connectivity::Cancellable>,
+) -> Result<(), CloseMsg> {
+    use talpid_types::ErrorExt;
+
+    let mut shared_tunnel = tunnel.lock().await;
+    let tunnel = shared_tunnel.take().expect("tunnel was None");
+    let ping_result = connectivity.establish_connectivity(&tunnel);
+    *shared_tunnel = Some(tunnel);
+    drop(shared_tunnel);
+
+    match ping_result {
+        Ok(true) => Ok(()),
+        Ok(false) => {
+            log::warn!("Timeout while checking tunnel connection");
+            Err(CloseMsg::PingErr)
+        }
+        Err(error) => {
+            log::error!(
+                "{}",
+                error.display_chain_with_msg("Failed to check tunnel connection")
+            );
+            Err(CloseMsg::PingErr)
+        }
+    }
+}
+
 async fn request_ephemeral_peer(
     retry_attempt: u32,
     config: &Config,

diff --git a/talpid-wireguard/src/lib.rs b/talpid-wireguard/src/lib.rs
@@ -274,6 +274,7 @@ impl WireguardMonitor {
                     args.retry_attempt,
                     obfuscator.clone(),
                     ephemeral_obfs_sender,
+                    &mut connectivity_monitor,
                 )
                 .await?;