Restructure how WireGuard obfuscation is configured and applied

4 files changed, 92 insertions, 72 deletions

diff --git a/talpid-core/src/tunnel/wireguard/config.rs b/talpid-core/src/tunnel/wireguard/config.rs
index 307a713272..7566d5c141 100644
--- a/talpid-core/src/tunnel/wireguard/config.rs
+++ b/talpid-core/src/tunnel/wireguard/config.rs
@@ -3,7 +3,7 @@ use std::{
     ffi::CString,
     net::{Ipv4Addr, Ipv6Addr},
 };
-use talpid_types::net::{wireguard, GenericTunnelOptions};
+use talpid_types::net::{obfuscation::ObfuscatorConfig, wireguard, GenericTunnelOptions};
 
 /// Config required to set up a single WireGuard tunnel
 pub struct Config {
@@ -26,6 +26,8 @@ pub struct Config {
     /// Temporary switch for wireguard-nt
     #[cfg(target_os = "windows")]
     pub use_wireguard_nt: bool,
+    /// Obfuscator config to be used for reaching the relay.
+    pub obfuscator_config: Option<ObfuscatorConfig>,
 }
 
 const DEFAULT_MTU: u16 = 1380;
@@ -60,6 +62,7 @@ impl Config {
             &params.connection,
             &params.options,
             &params.generic_options,
+            params.obfuscation.clone(),
         )
     }
 
@@ -70,6 +73,7 @@ impl Config {
         connection_config: &wireguard::ConnectionConfig,
         wg_options: &wireguard::TunnelOptions,
         generic_options: &GenericTunnelOptions,
+        obfuscator_config: Option<ObfuscatorConfig>,
     ) -> Result<Config, Error> {
         if peers.is_empty() {
             return Err(Error::NoPeersSuppliedError);
@@ -114,6 +118,7 @@ impl Config {
             enable_ipv6: generic_options.enable_ipv6,
             #[cfg(target_os = "windows")]
             use_wireguard_nt: wg_options.use_wireguard_nt,
+            obfuscator_config,
         })
     }
 
diff --git a/talpid-core/src/tunnel/wireguard/mod.rs b/talpid-core/src/tunnel/wireguard/mod.rs
index 263bbaeabb..59098a00ca 100644
--- a/talpid-core/src/tunnel/wireguard/mod.rs
+++ b/talpid-core/src/tunnel/wireguard/mod.rs
@@ -5,7 +5,10 @@ use super::{tun_provider::TunProvider, TunnelEvent, TunnelMetadata};
 use crate::routing::{self, RequiredRoute, RouteManagerHandle};
 #[cfg(windows)]
 use futures::{channel::mpsc, StreamExt};
-use futures::{channel::oneshot, future::abortable};
+use futures::{
+    channel::oneshot,
+    future::{abortable, AbortHandle as FutureAbortHandle},
+};
 #[cfg(target_os = "linux")]
 use lazy_static::lazy_static;
 #[cfg(target_os = "linux")]
@@ -16,14 +19,16 @@ use std::env;
 use std::io;
 use std::{
     convert::Infallible,
-    net::{IpAddr, SocketAddr},
+    net::IpAddr,
     path::Path,
     sync::{mpsc as sync_mpsc, Arc, Mutex},
 };
 #[cfg(windows)]
 use talpid_types::BoxedError;
-use talpid_types::{net::TransportProtocol, ErrorExt};
-use udp_over_tcp::{TcpOptions, Udp2Tcp};
+use talpid_types::{net::obfuscation::ObfuscatorConfig, ErrorExt};
+use tunnel_obfuscation::{
+    create_obfuscator, Error as ObfuscationError, Settings as ObfuscationSettings, Udp2TcpSettings,
+};
 
 /// WireGuard config data-types
 pub mod config;
@@ -56,13 +61,13 @@ pub enum Error {
     #[error(display = "Tunnel failed")]
     TunnelError(#[error(source)] TunnelError),
 
-    /// Failed to set up Udp2Tcp
-    #[error(display = "Failed to start UDP-over-TCP proxy")]
-    Udp2TcpError(#[error(source)] udp_over_tcp::udp2tcp::ConnectError),
+    /// Failed to create tunnel obfuscator
+    #[error(display = "Failed to create tunnel obfuscator")]
+    CreateObfuscatorError(#[error(source)] ObfuscationError),
 
-    /// Failed to obtain the local UDP socket address
-    #[error(display = "Failed obtain local address for the UDP socket in Udp2Tcp")]
-    GetLocalUdpAddress(#[error(source)] std::io::Error),
+    /// Failed to run tunnel obfuscator
+    #[error(display = "Tunnel obfuscator failed")]
+    ObfuscatorError(#[error(source)] ObfuscationError),
 
     /// Failed to set up connectivity monitor
     #[error(display = "Connectivity monitor failed")]
@@ -93,7 +98,24 @@ pub struct WireguardMonitor {
     >,
     close_msg_receiver: sync_mpsc::Receiver<CloseMsg>,
     pinger_stop_sender: sync_mpsc::Sender<()>,
-    _tcp_proxies: Vec<TcpProxy>,
+    _obfuscator: Option<ObfuscatorHandle>,
+}
+
+/// Simple wrapper that automatically cancels the future which runs an obfuscator.
+struct ObfuscatorHandle {
+    abort_handle: FutureAbortHandle,
+}
+
+impl ObfuscatorHandle {
+    pub fn new(abort_handle: FutureAbortHandle) -> Self {
+        Self { abort_handle }
+    }
+}
+
+impl Drop for ObfuscatorHandle {
+    fn drop(&mut self) {
+        self.abort_handle.abort();
+    }
 }
 
 #[cfg(target_os = "linux")]
@@ -108,52 +130,51 @@ lazy_static! {
         .unwrap_or(false);
 }
 
-struct TcpProxy {
-    local_addr: SocketAddr,
-    abort_handle: futures::future::AbortHandle,
-}
-
-impl TcpProxy {
-    pub fn new(runtime: &tokio::runtime::Handle, endpoint: SocketAddr) -> Result<Self> {
-        let listen_addr = if endpoint.is_ipv4() {
-            SocketAddr::new("127.0.0.1".parse().unwrap(), 0)
-        } else {
-            SocketAddr::new("::1".parse().unwrap(), 0)
-        };
+fn maybe_create_obfuscator(
+    runtime: &tokio::runtime::Handle,
+    config: &mut Config,
+    close_msg_sender: sync_mpsc::Sender<CloseMsg>,
+) -> Result<Option<ObfuscatorHandle>> {
+    // There are one or two peers.
+    // The first one is always the entry relay.
+    let mut first_peer = config.peers.get_mut(0).expect("missing peer");
 
-        let udp2tcp = runtime
-            .block_on(Udp2Tcp::new(
-                listen_addr,
-                endpoint,
-                TcpOptions {
+    if let Some(ref obfuscator_config) = config.obfuscator_config {
+        match obfuscator_config {
+            ObfuscatorConfig::Udp2Tcp { endpoint } => {
+                log::trace!("Connecting to Udp2Tcp endpoint {:?}", *endpoint);
+                let settings = Udp2TcpSettings {
+                    peer: *endpoint,
                     #[cfg(target_os = "linux")]
                     fwmark: Some(crate::linux::TUNNEL_FW_MARK),
-                    ..TcpOptions::default()
-                },
-            ))
-            .map_err(Error::Udp2TcpError)?;
-        let local_addr = udp2tcp
-            .local_udp_addr()
-            .map_err(Error::GetLocalUdpAddress)?;
-
-        let (udp2tcp_future, abort_handle) = abortable(udp2tcp.run());
-        runtime.spawn(udp2tcp_future);
-
-        Ok(Self {
-            local_addr,
-            abort_handle,
-        })
-    }
-
-    pub fn local_udp_addr(&self) -> SocketAddr {
-        self.local_addr
-    }
-}
-
-impl Drop for TcpProxy {
-    fn drop(&mut self) {
-        self.abort_handle.abort();
+                };
+                let obfuscator = runtime
+                    .block_on(create_obfuscator(&ObfuscationSettings::Udp2Tcp(settings)))
+                    .map_err(Error::CreateObfuscatorError)?;
+                let endpoint = obfuscator.endpoint();
+                log::trace!("Patching first WireGuard peer to become {:?}", endpoint);
+                first_peer.endpoint = endpoint;
+                let (runner, abort_handle) = abortable(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)));
+                        }
+                    }
+                });
+                runtime.spawn(runner);
+                return Ok(Some(ObfuscatorHandle::new(abort_handle)));
+            }
+        }
     }
+    Ok(None)
 }
 
 impl WireguardMonitor {
@@ -175,19 +196,11 @@ impl WireguardMonitor {
         retry_attempt: u32,
         tunnel_close_rx: oneshot::Receiver<()>,
     ) -> Result<WireguardMonitor> {
-        let mut tcp_proxies = vec![];
-        let mut endpoint_addrs = vec![];
-
-        for peer in &mut config.peers {
-            endpoint_addrs.push(peer.endpoint.ip());
-            if peer.protocol == TransportProtocol::Tcp {
-                let udp2tcp = TcpProxy::new(&runtime, peer.endpoint.clone())?;
+        let endpoint_addrs: Vec<IpAddr> =
+            config.peers.iter().map(|peer| peer.endpoint.ip()).collect();
+        let (close_msg_sender, close_msg_receiver) = sync_mpsc::channel();
 
-                // Replace remote peer with proxy
-                peer.endpoint = udp2tcp.local_udp_addr();
-                tcp_proxies.push(udp2tcp);
-            }
-        }
+        let obfuscator = maybe_create_obfuscator(&runtime, &mut config, close_msg_sender.clone())?;
 
         #[cfg(target_os = "windows")]
         let (setup_done_tx, mut setup_done_rx) = mpsc::channel(0);
@@ -203,7 +216,6 @@ impl WireguardMonitor {
         let iface_name = tunnel.get_interface_name().to_string();
 
         let event_callback = Box::new(on_event.clone());
-        let (close_msg_sender, close_msg_receiver) = sync_mpsc::channel();
         let (pinger_tx, pinger_rx) = sync_mpsc::channel();
         let monitor = WireguardMonitor {
             runtime: runtime.clone(),
@@ -211,7 +223,7 @@ impl WireguardMonitor {
             event_callback,
             close_msg_receiver,
             pinger_stop_sender: pinger_tx,
-            _tcp_proxies: tcp_proxies,
+            _obfuscator: obfuscator,
         };
 
         let gateway = config.ipv4_gateway;
@@ -413,8 +425,9 @@ impl WireguardMonitor {
     pub fn wait(mut self) -> Result<()> {
         let wait_result = match self.close_msg_receiver.recv() {
             Ok(CloseMsg::PingErr) => Err(Error::TimeoutError),
-            Ok(CloseMsg::Stop) => Ok(()),
+            Ok(CloseMsg::Stop) | Ok(CloseMsg::ObfuscatorExpired) => Ok(()),
             Ok(CloseMsg::SetupError(error)) => Err(error),
+            Ok(CloseMsg::ObfuscatorFailed(error)) => Err(error),
             Err(_) => Ok(()),
         };
 
@@ -570,6 +583,8 @@ enum CloseMsg {
     Stop,
     PingErr,
     SetupError(Error),
+    ObfuscatorExpired,
+    ObfuscatorFailed(Error),
 }
 
 pub(crate) trait Tunnel: Send {
diff --git a/talpid-core/src/tunnel/wireguard/wireguard_nt.rs b/talpid-core/src/tunnel/wireguard/wireguard_nt.rs
index 705d09892d..21c9b705ce 100644
--- a/talpid-core/src/tunnel/wireguard/wireguard_nt.rs
+++ b/talpid-core/src/tunnel/wireguard/wireguard_nt.rs
@@ -982,7 +982,7 @@ impl Tunnel for WgNtTunnel {
 mod tests {
     use super::*;
     use lazy_static::lazy_static;
-    use talpid_types::net::{wireguard, TransportProtocol};
+    use talpid_types::net::wireguard;
 
     #[derive(Debug, Eq, PartialEq, Clone, Copy)]
     #[repr(C)]
@@ -1006,12 +1006,12 @@ mod tests {
                     public_key: WG_PUBLIC_KEY.clone(),
                     allowed_ips: vec!["1.3.3.0/24".parse().unwrap()],
                     endpoint: "1.2.3.4:1234".parse().unwrap(),
-                    protocol: TransportProtocol::Udp,
                 }],
                 ipv4_gateway: "0.0.0.0".parse().unwrap(),
                 ipv6_gateway: None,
                 mtu: 0,
                 use_wireguard_nt: true,
+                obfuscator_config: None,
             }
         };
         static ref WG_STRUCT_CONFIG: Interface = Interface {
diff --git a/talpid-core/src/tunnel_state_machine/connecting_state.rs b/talpid-core/src/tunnel_state_machine/connecting_state.rs
index 950ee1a43f..88319a737e 100644
--- a/talpid-core/src/tunnel_state_machine/connecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connecting_state.rs
@@ -479,7 +479,7 @@ fn should_retry(error: &tunnel::Error, retry_attempt: u32) -> bool {
     use tunnel::wireguard::{Error, TunnelError};
 
     match error {
-        tunnel::Error::WireguardTunnelMonitoringError(Error::Udp2TcpError(_)) => true,
+        tunnel::Error::WireguardTunnelMonitoringError(Error::CreateObfuscatorError(_)) => true,
 
         #[cfg(not(windows))]
         tunnel::Error::WireguardTunnelMonitoringError(Error::TunnelError(