Exclude tunnel sockets after PSK exchange

2 files changed, 54 insertions, 2 deletions

diff --git a/talpid-wireguard/src/lib.rs b/talpid-wireguard/src/lib.rs
index c1d7682648..7830c7a072 100644
--- a/talpid-wireguard/src/lib.rs
+++ b/talpid-wireguard/src/lib.rs
@@ -281,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();
 
@@ -367,6 +369,8 @@ impl WireguardMonitor {
                     &iface_name,
                     obfuscator.clone(),
                     psk_obfs_sender,
+                    #[cfg(target_os = "android")]
+                    args.tun_provider,
                 )
                 .await?;
             }
@@ -431,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,
@@ -439,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>>)
@@ -497,6 +503,8 @@ impl WireguardMonitor {
                 entry_tun_config,
                 obfuscator.clone(),
                 close_obfs_sender,
+                #[cfg(target_os = "android")]
+                &tun_provider,
             )
             .await?;
             entry_psk = Some(
@@ -522,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,
@@ -541,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() {
@@ -548,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
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(())
         })
     }