Merge branch 'android-fix-pq-des-232'

2 files changed, 67 insertions, 12 deletions

diff --git a/talpid-wireguard/src/lib.rs b/talpid-wireguard/src/lib.rs
index cd250856df..7830c7a072 100644
--- a/talpid-wireguard/src/lib.rs
+++ b/talpid-wireguard/src/lib.rs
@@ -9,13 +9,12 @@ use futures::future::{abortable, AbortHandle as FutureAbortHandle, BoxFuture, Fu
 use futures::{channel::mpsc, StreamExt};
 #[cfg(target_os = "linux")]
 use lazy_static::lazy_static;
-#[cfg(target_os = "android")]
-use std::borrow::Cow;
 #[cfg(target_os = "linux")]
 use std::env;
 #[cfg(windows)]
 use std::io;
 use std::{
+    borrow::Cow,
     convert::Infallible,
     net::IpAddr,
     path::Path,
@@ -259,25 +258,21 @@ impl WireguardMonitor {
         // and should be a hacky fix for the problem. In the longer term this should be fixed by
         // allowing the handshake to work even if there is fragmentation and/or setting the MTU
         // properly so fragmentation does not happen.
-        #[cfg(not(target_os = "android"))]
         let init_tunnel_config = if cfg!(target_os = "macos") {
             let mut init_tunnel_config = config.clone();
             if psk_negotiation && config.peers.len() > 1 {
                 const MH_PQ_HANDSHAKE_MTU: u16 = 1280;
                 init_tunnel_config.mtu = MH_PQ_HANDSHAKE_MTU;
             }
-            init_tunnel_config
+            Cow::Owned(init_tunnel_config)
         } else {
-            config.clone()
+            Cow::Borrowed(&config)
         };
 
         #[cfg(target_os = "windows")]
         let (setup_done_tx, setup_done_rx) = mpsc::channel(0);
         let tunnel = Self::open_tunnel(
             args.runtime.clone(),
-            #[cfg(target_os = "android")]
-            &Self::patch_allowed_ips(&config, psk_negotiation),
-            #[cfg(not(target_os = "android"))]
             &init_tunnel_config,
             log_path,
             args.resource_dir,
@@ -286,6 +281,8 @@ impl WireguardMonitor {
             args.route_manager.clone(),
             #[cfg(target_os = "windows")]
             setup_done_tx,
+            #[cfg(target_os = "android")]
+            psk_negotiation,
         )?;
         let iface_name = tunnel.get_interface_name();
 
@@ -372,6 +369,8 @@ impl WireguardMonitor {
                     &iface_name,
                     obfuscator.clone(),
                     psk_obfs_sender,
+                    #[cfg(target_os = "android")]
+                    args.tun_provider,
                 )
                 .await?;
             }
@@ -436,6 +435,7 @@ impl WireguardMonitor {
         Ok(monitor)
     }
 
+    #[allow(clippy::too_many_arguments)]
     async fn psk_negotiation<F>(
         tunnel: &mut Arc<Mutex<Option<Box<dyn Tunnel>>>>,
         config: &mut Config,
@@ -444,6 +444,7 @@ impl WireguardMonitor {
         iface_name: &str,
         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>
     where
         F: (Fn(TunnelEvent) -> Pin<Box<dyn std::future::Future<Output = ()> + Send>>)
@@ -502,6 +503,8 @@ impl WireguardMonitor {
                 entry_tun_config,
                 obfuscator.clone(),
                 close_obfs_sender,
+                #[cfg(target_os = "android")]
+                &tun_provider,
             )
             .await?;
             entry_psk = Some(
@@ -527,8 +530,15 @@ impl WireguardMonitor {
             config.peers.get_mut(0).expect("peer not found").psk = Some(exit_psk);
         }
 
-        *config =
-            Self::reconfigure_tunnel(tunnel, config.clone(), obfuscator, close_obfs_sender).await?;
+        *config = Self::reconfigure_tunnel(
+            tunnel,
+            config.clone(),
+            obfuscator,
+            close_obfs_sender,
+            #[cfg(target_os = "android")]
+            &tun_provider,
+        )
+        .await?;
         let metadata = Self::tunnel_metadata(iface_name, config);
         (on_event)(TunnelEvent::InterfaceUp(
             metadata,
@@ -546,6 +556,7 @@ impl WireguardMonitor {
         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() {
@@ -553,6 +564,16 @@ impl WireguardMonitor {
             *obfs_guard = maybe_create_obfuscator(&mut config, close_obfs_sender)
                 .await
                 .map_err(CloseMsg::ObfuscatorFailed)?;
+
+            // Exclude new remote obfuscation socket or bridge
+            #[cfg(target_os = "android")]
+            if let Some(obfuscator_handle) = &*obfs_guard {
+                let remote_socket_fd = obfuscator_handle.remote_socket_fd();
+                log::debug!("Excluding remote socket fd from the tunnel");
+                if let Err(error) = tun_provider.lock().unwrap().bypass(remote_socket_fd) {
+                    log::error!("Failed to exclude remote socket fd: {error}");
+                }
+            }
         }
 
         let set_config_future = tunnel
@@ -705,6 +726,7 @@ impl WireguardMonitor {
         log_path: Option<&Path>,
         resource_dir: &Path,
         tun_provider: Arc<Mutex<TunProvider>>,
+        #[cfg(target_os = "android")] psk_negotiation: bool,
         #[cfg(windows)] route_manager_handle: crate::routing::RouteManagerHandle,
         #[cfg(windows)] setup_done_tx: mpsc::Sender<std::result::Result<(), BoxedError>>,
     ) -> Result<Box<dyn Tunnel>> {
@@ -758,16 +780,23 @@ impl WireguardMonitor {
             .map_err(Error::TunnelError);
         }
 
+        #[cfg(not(windows))]
+        let routes = Self::get_tunnel_destinations(config).flat_map(Self::replace_default_prefixes);
+
+        #[cfg(target_os = "android")]
+        let config = Self::patch_allowed_ips(config, psk_negotiation);
+
         #[cfg(any(target_os = "linux", windows))]
         log::debug!("Using userspace WireGuard implementation");
         Ok(Box::new(
             WgGoTunnel::start_tunnel(
-                config,
+                #[allow(clippy::needless_borrow)]
+                &config,
                 log_path,
                 #[cfg(not(windows))]
                 tun_provider,
                 #[cfg(not(windows))]
-                Self::get_tunnel_destinations(config).flat_map(Self::replace_default_prefixes),
+                routes,
                 #[cfg(windows)]
                 route_manager_handle,
                 #[cfg(windows)]
diff --git a/talpid-wireguard/src/wireguard_go.rs b/talpid-wireguard/src/wireguard_go.rs
index 8b9f1f696d..d12753258b 100644
--- a/talpid-wireguard/src/wireguard_go.rs
+++ b/talpid-wireguard/src/wireguard_go.rs
@@ -62,6 +62,8 @@ pub struct WgGoTunnel {
     _route_callback_handle: Option<talpid_routing::CallbackHandle>,
     #[cfg(target_os = "windows")]
     setup_handle: tokio::task::JoinHandle<()>,
+    #[cfg(target_os = "android")]
+    tun_provider: Arc<Mutex<TunProvider>>,
 }
 
 impl WgGoTunnel {
@@ -72,8 +74,12 @@ impl WgGoTunnel {
         tun_provider: Arc<Mutex<TunProvider>>,
         routes: impl Iterator<Item = IpNetwork>,
     ) -> Result<Self> {
+        #[cfg(target_os = "android")]
+        let tun_provider_clone = tun_provider.clone();
+
         #[cfg_attr(not(target_os = "android"), allow(unused_mut))]
         let (mut tunnel_device, tunnel_fd) = Self::get_tunnel(tun_provider, config, routes)?;
+
         let interface_name: String = tunnel_device.interface_name().to_string();
         let wg_config_str = config.to_userspace_format();
         let logging_context = initialize_logging(log_path)
@@ -103,6 +109,8 @@ impl WgGoTunnel {
             handle: Some(handle),
             _tunnel_device: tunnel_device,
             _logging_context: logging_context,
+            #[cfg(target_os = "android")]
+            tun_provider: tun_provider_clone,
         })
     }
 
@@ -373,11 +381,29 @@ impl Tunnel for WgGoTunnel {
     ) -> Pin<Box<dyn Future<Output = std::result::Result<(), super::TunnelError>> + Send>> {
         let wg_config_str = config.to_userspace_format();
         let handle = self.handle.unwrap();
+        #[cfg(target_os = "android")]
+        let tun_provider = self.tun_provider.clone();
         Box::pin(async move {
             let status = unsafe { wgSetConfig(handle, wg_config_str.as_ptr() as *const i8) };
             if status != 0 {
                 return Err(TunnelError::SetConfigError);
             }
+
+            // When reapplying the config, the endpoint socket may be discarded
+            // and needs to be excluded again
+            #[cfg(target_os = "android")]
+            {
+                let socket_v4 = unsafe { wgGetSocketV4(handle) };
+                let socket_v6 = unsafe { wgGetSocketV6(handle) };
+                let mut provider = tun_provider.lock().unwrap();
+                provider
+                    .bypass(socket_v4)
+                    .map_err(super::TunnelError::BypassError)?;
+                provider
+                    .bypass(socket_v6)
+                    .map_err(super::TunnelError::BypassError)?;
+            }
+
             Ok(())
         })
     }