Merge branch 'refactor-wg-obfs-setup'

2 files changed, 176 insertions, 145 deletions

diff --git a/talpid-wireguard/src/lib.rs b/talpid-wireguard/src/lib.rs
index 0c918e1fc7..e26f558c69 100644
--- a/talpid-wireguard/src/lib.rs
+++ b/talpid-wireguard/src/lib.rs
@@ -6,13 +6,14 @@ use self::config::Config;
 #[cfg(windows)]
 use futures::channel::mpsc;
 use futures::future::{BoxFuture, Future};
+use obfuscation::ObfuscatorHandle;
 #[cfg(target_os = "android")]
 use std::borrow::Cow;
 #[cfg(windows)]
 use std::io;
 use std::{
     convert::Infallible,
-    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
+    net::IpAddr,
     path::Path,
     pin::Pin,
     sync::{mpsc as sync_mpsc, Arc, Mutex},
@@ -29,22 +30,18 @@ use talpid_tunnel::{tun_provider::TunProvider, TunnelArgs, TunnelEvent, TunnelMe
 use ipnetwork::IpNetwork;
 use talpid_types::{
     net::{
-        obfuscation::ObfuscatorConfig,
         wireguard::{PresharedKey, PrivateKey, PublicKey},
         AllowedTunnelTraffic, Endpoint, TransportProtocol,
     },
     BoxedError, ErrorExt,
 };
 use tokio::sync::Mutex as AsyncMutex;
-use tunnel_obfuscation::{
-    create_obfuscator, shadowsocks, udp2tcp, Error as ObfuscationError,
-    Settings as ObfuscationSettings,
-};
 
 /// WireGuard config data-types
 pub mod config;
 mod connectivity_check;
 mod logging;
+mod obfuscation;
 mod ping_monitor;
 mod stats;
 #[cfg(wireguard_go)]
@@ -78,13 +75,9 @@ pub enum Error {
     #[error("Tunnel failed")]
     TunnelError(#[source] TunnelError),
 
-    /// Failed to create tunnel obfuscator
-    #[error("Failed to create tunnel obfuscator")]
-    CreateObfuscatorError(#[source] ObfuscationError),
-
-    /// Failed to run tunnel obfuscator
-    #[error("Tunnel obfuscator failed")]
-    ObfuscatorError(#[source] ObfuscationError),
+    /// Failed to run tunnel obfuscation
+    #[error("Tunnel obfuscation failed")]
+    ObfuscationError(#[source] tunnel_obfuscation::Error),
 
     /// Failed to set up connectivity monitor
     #[error("Connectivity monitor failed")]
@@ -109,8 +102,7 @@ impl Error {
     /// Return whether retrying the operation that caused this error is likely to succeed.
     pub fn is_recoverable(&self) -> bool {
         match self {
-            Error::CreateObfuscatorError(_) => true,
-            Error::ObfuscatorError(_) => true,
+            Error::ObfuscationError(_) => true,
             Error::PskNegotiationError(_) => true,
             Error::TunnelError(TunnelError::RecoverableStartWireguardError(..)) => true,
 
@@ -153,41 +145,6 @@ 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;
 
-/// Simple wrapper that automatically cancels the future which runs an obfuscator.
-struct ObfuscatorHandle {
-    obfuscation_task: tokio::task::JoinHandle<()>,
-    #[cfg(target_os = "android")]
-    remote_socket_fd: std::os::unix::io::RawFd,
-}
-
-impl ObfuscatorHandle {
-    pub fn new(
-        obfuscation_task: tokio::task::JoinHandle<()>,
-        #[cfg(target_os = "android")] remote_socket_fd: std::os::unix::io::RawFd,
-    ) -> Self {
-        Self {
-            obfuscation_task,
-            #[cfg(target_os = "android")]
-            remote_socket_fd,
-        }
-    }
-
-    #[cfg(target_os = "android")]
-    pub fn remote_socket_fd(&self) -> std::os::unix::io::RawFd {
-        self.remote_socket_fd
-    }
-
-    pub fn abort(&self) {
-        self.obfuscation_task.abort();
-    }
-}
-
-impl Drop for ObfuscatorHandle {
-    fn drop(&mut self) {
-        self.obfuscation_task.abort();
-    }
-}
-
 #[cfg(target_os = "linux")]
 /// Overrides the preference for the kernel module for WireGuard.
 static FORCE_USERSPACE_WIREGUARD: LazyLock<bool> = LazyLock::new(|| {
@@ -196,72 +153,6 @@ static FORCE_USERSPACE_WIREGUARD: LazyLock<bool> = LazyLock::new(|| {
         .unwrap_or(false)
 });
 
-async fn maybe_create_obfuscator(
-    config: &mut Config,
-    close_msg_sender: sync_mpsc::Sender<CloseMsg>,
-) -> Result<Option<ObfuscatorHandle>> {
-    if let Some(ref obfuscator_config) = config.obfuscator_config {
-        let settings = match obfuscator_config {
-            ObfuscatorConfig::Udp2Tcp { endpoint } => {
-                ObfuscationSettings::Udp2Tcp(udp2tcp::Settings {
-                    peer: *endpoint,
-                    #[cfg(target_os = "linux")]
-                    fwmark: config.fwmark,
-                })
-            }
-            ObfuscatorConfig::Shadowsocks { endpoint } => {
-                ObfuscationSettings::Shadowsocks(shadowsocks::Settings {
-                    shadowsocks_endpoint: *endpoint,
-                    // TODO: Temporary since we may different entry IPs later?
-                    wireguard_endpoint: if config.entry_peer.endpoint.is_ipv4() {
-                        SocketAddr::from((Ipv4Addr::LOCALHOST, 51820))
-                    } else {
-                        SocketAddr::from((Ipv6Addr::LOCALHOST, 51820))
-                    },
-                    //wireguard_endpoint: config.entry_peer.endpoint,
-                    #[cfg(target_os = "linux")]
-                    fwmark: config.fwmark,
-                })
-            }
-        };
-
-        log::trace!("Obfuscation settings: {settings:?}");
-
-        let obfuscator = create_obfuscator(&settings)
-            .await
-            .map_err(Error::CreateObfuscatorError)?;
-        let endpoint = obfuscator.endpoint();
-
-        log::trace!("Patching first WireGuard peer to become {endpoint}");
-        config.entry_peer.endpoint = endpoint;
-
-        #[cfg(target_os = "android")]
-        let remote_socket_fd = obfuscator.remote_socket_fd();
-
-        let obfuscation_task = tokio::spawn(async move {
-            match obfuscator.run().await {
-                Ok(_) => {
-                    let _ = close_msg_sender.send(CloseMsg::ObfuscatorExpired);
-                }
-                Err(error) => {
-                    log::error!(
-                        "{}",
-                        error.display_chain_with_msg("Obfuscation controller failed")
-                    );
-                    let _ = close_msg_sender
-                        .send(CloseMsg::ObfuscatorFailed(Error::ObfuscatorError(error)));
-                }
-            }
-        });
-        return Ok(Some(ObfuscatorHandle::new(
-            obfuscation_task,
-            #[cfg(target_os = "android")]
-            remote_socket_fd,
-        )));
-    }
-    Ok(None)
-}
-
 impl WireguardMonitor {
     /// Starts a WireGuard tunnel with the given config
     #[cfg(not(target_os = "android"))]
@@ -282,10 +173,12 @@ impl WireguardMonitor {
         let endpoint_addrs: Vec<IpAddr> = config.peers().map(|peer| peer.endpoint.ip()).collect();
 
         let (close_obfs_sender, close_obfs_listener) = sync_mpsc::channel();
-        let obfuscator = args.runtime.block_on(maybe_create_obfuscator(
-            &mut config,
-            close_obfs_sender.clone(),
-        ))?;
+        let obfuscator = args
+            .runtime
+            .block_on(obfuscation::apply_obfuscation_config(
+                &mut config,
+                close_obfs_sender.clone(),
+            ))?;
 
         #[cfg(target_os = "windows")]
         let (setup_done_tx, setup_done_rx) = mpsc::channel(0);
@@ -502,18 +395,13 @@ impl WireguardMonitor {
         )?;
 
         let (close_obfs_sender, close_obfs_listener) = sync_mpsc::channel();
-        let obfuscator = args.runtime.block_on(maybe_create_obfuscator(
-            &mut config,
-            close_obfs_sender.clone(),
-        ))?;
-
-        if let Some(remote_socket_fd) = obfuscator.as_ref().map(|obfs| obfs.remote_socket_fd()) {
-            // Exclude remote obfuscation socket or bridge
-            log::debug!("Excluding remote socket fd from the tunnel");
-            if let Err(error) = args.tun_provider.lock().unwrap().bypass(remote_socket_fd) {
-                log::error!("Failed to exclude remote socket fd: {error}");
-            }
-        }
+        let obfuscator = args
+            .runtime
+            .block_on(obfuscation::apply_obfuscation_config(
+                &mut config,
+                close_obfs_sender.clone(),
+                args.tun_provider.clone(),
+            ))?;
 
         let iface_name = tunnel.get_interface_name();
 
@@ -769,19 +657,14 @@ impl WireguardMonitor {
         let mut obfs_guard = obfuscator.lock().await;
         if let Some(obfuscator_handle) = obfs_guard.take() {
             obfuscator_handle.abort();
-            *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}");
-                }
-            }
+            *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;
diff --git a/talpid-wireguard/src/obfuscation.rs b/talpid-wireguard/src/obfuscation.rs
new file mode 100644
index 0000000000..0e1e7273e8
--- /dev/null
+++ b/talpid-wireguard/src/obfuscation.rs
@@ -0,0 +1,148 @@
+//! Glue between tunnel-obfuscation and WireGuard configurations
+
+use super::{Error, Result};
+use crate::{config::Config, CloseMsg};
+#[cfg(target_os = "android")]
+use std::sync::{Arc, Mutex};
+use std::{
+    net::{Ipv4Addr, Ipv6Addr, SocketAddr},
+    sync::mpsc as sync_mpsc,
+};
+#[cfg(target_os = "android")]
+use talpid_tunnel::tun_provider::TunProvider;
+use talpid_types::{net::obfuscation::ObfuscatorConfig, ErrorExt};
+
+use tunnel_obfuscation::{
+    create_obfuscator, shadowsocks, udp2tcp, Settings as ObfuscationSettings,
+};
+
+/// Begin running obfuscation machine, if configured. This function will patch `config`'s endpoint
+/// to point to an endpoint on localhost
+pub async fn apply_obfuscation_config(
+    config: &mut Config,
+    close_msg_sender: sync_mpsc::Sender<CloseMsg>,
+    #[cfg(target_os = "android")] tun_provider: Arc<Mutex<TunProvider>>,
+) -> Result<Option<ObfuscatorHandle>> {
+    let Some(ref obfuscator_config) = config.obfuscator_config else {
+        return Ok(None);
+    };
+
+    let settings = settings_from_config(
+        obfuscator_config,
+        #[cfg(target_os = "linux")]
+        config.fwmark,
+    );
+    apply_obfuscation_config_inner(
+        config,
+        settings,
+        close_msg_sender,
+        #[cfg(target_os = "android")]
+        tun_provider,
+    )
+    .await
+    .map(Some)
+}
+
+async fn apply_obfuscation_config_inner(
+    config: &mut Config,
+    settings: ObfuscationSettings,
+    close_msg_sender: sync_mpsc::Sender<CloseMsg>,
+    #[cfg(target_os = "android")] tun_provider: Arc<Mutex<TunProvider>>,
+) -> Result<ObfuscatorHandle> {
+    log::trace!("Obfuscation settings: {settings:?}");
+
+    let obfuscator = create_obfuscator(&settings)
+        .await
+        .map_err(Error::ObfuscationError)?;
+
+    #[cfg(target_os = "android")]
+    bypass_vpn(tun_provider, obfuscator.remote_socket_fd()).await;
+
+    patch_endpoint(config, obfuscator.endpoint());
+
+    let obfuscation_task = tokio::spawn(async move {
+        match obfuscator.run().await {
+            Ok(_) => {
+                let _ = close_msg_sender.send(CloseMsg::ObfuscatorExpired);
+            }
+            Err(error) => {
+                log::error!(
+                    "{}",
+                    error.display_chain_with_msg("Obfuscation controller failed")
+                );
+                let _ = close_msg_sender
+                    .send(CloseMsg::ObfuscatorFailed(Error::ObfuscationError(error)));
+            }
+        }
+    });
+
+    Ok(ObfuscatorHandle::new(obfuscation_task))
+}
+
+/// Patch the first peer in the WireGuard configuration to use the local proxy endpoint
+fn patch_endpoint(config: &mut Config, endpoint: SocketAddr) {
+    log::trace!("Patching first WireGuard peer to become {endpoint}");
+    config.entry_peer.endpoint = endpoint;
+}
+
+fn settings_from_config(
+    config: &ObfuscatorConfig,
+    #[cfg(target_os = "linux")] fwmark: Option<u32>,
+) -> ObfuscationSettings {
+    match config {
+        ObfuscatorConfig::Udp2Tcp { endpoint } => ObfuscationSettings::Udp2Tcp(udp2tcp::Settings {
+            peer: *endpoint,
+            #[cfg(target_os = "linux")]
+            fwmark,
+        }),
+        ObfuscatorConfig::Shadowsocks { endpoint } => {
+            ObfuscationSettings::Shadowsocks(shadowsocks::Settings {
+                shadowsocks_endpoint: *endpoint,
+                wireguard_endpoint: if endpoint.is_ipv4() {
+                    SocketAddr::from((Ipv4Addr::LOCALHOST, 51820))
+                } else {
+                    SocketAddr::from((Ipv6Addr::LOCALHOST, 51820))
+                },
+                #[cfg(target_os = "linux")]
+                fwmark,
+            })
+        }
+    }
+}
+
+/// Route socket outside of the VPN on Android
+#[cfg(target_os = "android")]
+async fn bypass_vpn(
+    tun_provider: Arc<Mutex<TunProvider>>,
+    remote_socket_fd: std::os::unix::io::RawFd,
+) {
+    // Exclude remote obfuscation socket or bridge
+    log::debug!("Excluding remote socket fd from the tunnel");
+    let _ = tokio::task::spawn_blocking(move || {
+        if let Err(error) = tun_provider.lock().unwrap().bypass(remote_socket_fd) {
+            log::error!("Failed to exclude remote socket fd: {error}");
+        }
+    })
+    .await;
+}
+
+/// Simple wrapper that automatically cancels the future which runs an obfuscator.
+pub struct ObfuscatorHandle {
+    obfuscation_task: tokio::task::JoinHandle<()>,
+}
+
+impl ObfuscatorHandle {
+    pub fn new(obfuscation_task: tokio::task::JoinHandle<()>) -> Self {
+        Self { obfuscation_task }
+    }
+
+    pub fn abort(&self) {
+        self.obfuscation_task.abort();
+    }
+}
+
+impl Drop for ObfuscatorHandle {
+    fn drop(&mut self) {
+        self.obfuscation_task.abort();
+    }
+}