Add IPv6 support to `mullvad-masque-proxy`

6 files changed, 100 insertions, 39 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 482657c481..741615496b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -6139,6 +6139,7 @@ dependencies = [
  "shadowsocks",
  "thiserror 2.0.9",
  "tokio",
+ "tokio-util 0.7.10",
  "udp-over-tcp",
 ]
 
diff --git a/Cargo.toml b/Cargo.toml
index e6c6d1eb04..41e1089912 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -84,6 +84,7 @@ implicit_clone = "warn"
 
 [workspace.dependencies]
 tokio = { version = "1.44" }
+tokio-util = "0.7"
 parity-tokio-ipc = "0.9"
 futures = "0.3.15"
 vec1 = "1.12"
diff --git a/mullvad-masque-proxy/src/client/mod.rs b/mullvad-masque-proxy/src/client/mod.rs
index b5aaf81cec..c52e3f3341 100644
--- a/mullvad-masque-proxy/src/client/mod.rs
+++ b/mullvad-masque-proxy/src/client/mod.rs
@@ -67,6 +67,8 @@ pub type Result<T> = std::result::Result<T, Error>;
 pub enum Error {
     #[error("Failed to bind local socket")]
     Bind(#[source] io::Error),
+    #[error("Failed to setup a QUIC endpoint")]
+    Endpoint(#[source] io::Error),
     #[cfg(target_os = "linux")]
     #[error("Failed to set fwmark on remote socket")]
     Fwmark(#[source] io::Error),
@@ -112,7 +114,7 @@ pub enum Error {
     InvalidHttpRedirect(#[source] anyhow::Error),
 }
 
-#[derive(TypedBuilder)]
+#[derive(TypedBuilder, Debug)]
 pub struct ClientConfig {
     /// Socket that accepts proxy clients
     pub client_socket: UdpSocket,
@@ -130,7 +132,9 @@ pub struct ClientConfig {
     pub server_host: String,
 
     /// MTU (includes IP header)
-    #[builder(default = 1500)]
+    // TODO: this can't be 1500. Or alteast we need to account for ipv6
+    // overhead.
+    #[builder(default = 1350)]
     pub mtu: u16,
 
     /// QUIC TLS config
@@ -226,32 +230,39 @@ impl Client {
         max_udp_payload_size: u16,
         #[cfg(target_os = "linux")] fwmark: Option<u32>,
     ) -> Result<Endpoint> {
-        let local_socket = socket2::Socket::new(
-            socket2::Domain::IPV4,
-            socket2::Type::DGRAM,
-            Some(socket2::Protocol::UDP),
-        )
-        .map_err(Error::Bind)?;
-
-        #[cfg(target_os = "linux")]
-        if let Some(fwmark) = fwmark {
-            local_socket.set_mark(fwmark).map_err(Error::Fwmark)?;
-        }
-
-        local_socket.bind(&local_addr.into()).map_err(Error::Bind)?;
+        // Create a UDP socket which quinn will read/write from/to.
+        let local_socket = {
+            // family
+            let domain = match &local_addr {
+                SocketAddr::V4(_) => socket2::Domain::IPV4,
+                SocketAddr::V6(_) => socket2::Domain::IPV6,
+            };
+            let ty = socket2::Type::DGRAM;
+            let protocol = Some(socket2::Protocol::UDP);
+            let socket = socket2::Socket::new(domain, ty, protocol).map_err(Error::Bind)?;
+            #[cfg(target_os = "linux")]
+            if let Some(fwmark) = fwmark {
+                socket.set_mark(fwmark).map_err(Error::Fwmark)?;
+            }
+            socket.bind(&local_addr.into()).map_err(Error::Bind)?;
+            socket
+        };
 
-        let mut endpoint_config = EndpointConfig::default();
-        endpoint_config
-            .max_udp_payload_size(max_udp_payload_size)
-            .map_err(Error::InvalidMaxUdpPayload)?;
+        let endpoint_config = {
+            let mut endpoint_config = EndpointConfig::default();
+            endpoint_config
+                .max_udp_payload_size(max_udp_payload_size)
+                .map_err(Error::InvalidMaxUdpPayload)?;
+            endpoint_config
+        };
 
         Endpoint::new(
             endpoint_config,
             None,
-            local_socket.into(),
+            std::net::UdpSocket::from(local_socket),
             Arc::new(TokioRuntime),
         )
-        .map_err(Error::Bind)
+        .map_err(Error::Endpoint)
     }
 
     /// Returns an h3 connection that is ready to be used for sending UDP datagrams.
diff --git a/talpid-wireguard/src/obfuscation.rs b/talpid-wireguard/src/obfuscation.rs
index 1b81dd190a..6085a9fcd9 100644
--- a/talpid-wireguard/src/obfuscation.rs
+++ b/talpid-wireguard/src/obfuscation.rs
@@ -102,6 +102,7 @@ fn settings_from_config(
             auth_token,
         } => ObfuscationSettings::Quic(quic::Settings {
             quic_endpoint: *endpoint,
+            // TODO: Explain why this may always be an IPv4 address
             wireguard_endpoint: SocketAddr::from((Ipv4Addr::LOCALHOST, 51820)),
             hostname: hostname.to_owned(),
             token: auth_token.to_owned(),
diff --git a/tunnel-obfuscation/Cargo.toml b/tunnel-obfuscation/Cargo.toml
index 1bee10199b..452b9f9eed 100644
--- a/tunnel-obfuscation/Cargo.toml
+++ b/tunnel-obfuscation/Cargo.toml
@@ -15,6 +15,7 @@ log = { workspace = true }
 async-trait = "0.1"
 thiserror = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "io-util"] }
+tokio-util = { workspace = true }
 udp-over-tcp = { git = "https://github.com/mullvad/udp-over-tcp", rev = "87936ac29b68b902565955f138ab02294bcc8593" }
 shadowsocks = { workspace = true }
 mullvad-masque-proxy = { path = "../mullvad-masque-proxy" }
diff --git a/tunnel-obfuscation/src/quic.rs b/tunnel-obfuscation/src/quic.rs
index 382d998020..703d692f30 100644
--- a/tunnel-obfuscation/src/quic.rs
+++ b/tunnel-obfuscation/src/quic.rs
@@ -4,9 +4,10 @@ use async_trait::async_trait;
 use mullvad_masque_proxy::client::{Client, ClientConfig};
 use std::{
     io,
-    net::{Ipv4Addr, SocketAddr},
+    net::{Ipv4Addr, Ipv6Addr, SocketAddr},
 };
 use tokio::net::UdpSocket;
+use tokio_util::sync::{CancellationToken, DropGuard};
 
 use crate::Obfuscator;
 
@@ -20,9 +21,11 @@ pub enum Error {
     MasqueProxyError(#[source] mullvad_masque_proxy::client::Error),
 }
 
+#[derive(Debug)]
 pub struct Quic {
     local_endpoint: SocketAddr,
     task: tokio::task::JoinHandle<Result<()>>,
+    _shutdown: DropGuard,
 }
 
 #[derive(Debug)]
@@ -52,15 +55,20 @@ impl Settings {
 
 impl Quic {
     pub(crate) async fn new(settings: &Settings) -> Result<Self> {
-        let local_socket = UdpSocket::bind(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
-            .await
-            .map_err(Error::BindError)?;
-
-        let local_endpoint = local_socket.local_addr().unwrap();
-
+        let (local_socket, local_udp_client_addr) =
+            Quic::create_local_udp_socket(settings.quic_endpoint.is_ipv4()).await?;
+        // The address family of the local QUIC client socket has to match the address family
+        // of the endpoint we're connecting to. The address itself is not important to consumers wanting
+        // to obfuscate traffic. It is solely used by the local proxy client to know where the QUIC
+        // obfuscator is running.
+        let quic_client_local_addr = if settings.quic_endpoint.is_ipv4() {
+            SocketAddr::from((Ipv4Addr::UNSPECIFIED, 0))
+        } else {
+            SocketAddr::from((Ipv6Addr::UNSPECIFIED, 0))
+        };
         let config_builder = ClientConfig::builder()
             .client_socket(local_socket)
-            .local_addr((Ipv4Addr::UNSPECIFIED, 0).into())
+            .local_addr(quic_client_local_addr)
             .server_addr(settings.quic_endpoint)
             .server_host(settings.hostname.clone())
             .target_addr(settings.wireguard_endpoint)
@@ -69,17 +77,55 @@ impl Quic {
         #[cfg(target_os = "linux")]
         let config_builder = config_builder.fwmark(settings.fwmark);
 
-        let task = tokio::spawn(async move {
-            let client = Client::connect(config_builder.build())
-                .await
-                .map_err(Error::MasqueProxyError)?;
-            client.run().await.map_err(Error::MasqueProxyError)
-        });
+        let client = Client::connect(config_builder.build())
+            .await
+            .map_err(Error::MasqueProxyError)?;
+
+        let token = CancellationToken::new();
+
+        let local_proxy = tokio::spawn(Quic::run_forwarding(client, token.child_token()));
+
+        let quic = Quic {
+            local_endpoint: local_udp_client_addr,
+            task: local_proxy,
+            _shutdown: token.drop_guard(),
+        };
+
+        Ok(quic)
+    }
+
+    async fn run_forwarding(
+        masque_proxy_client: Client,
+        cancel_token: CancellationToken,
+    ) -> Result<()> {
+        log::trace!("Spawning QUIC client ..");
+        let mut client = tokio::spawn(masque_proxy_client.run());
+        log::trace!("QUIC client is running! QUIC Obfuscator is serving traffic 🎉");
+        tokio::select! {
+            _ = cancel_token.cancelled() => log::trace!("Stopping QUIC obfuscation"),
+            _result = &mut client => log::trace!("QUIC client closed"),
+        };
+
+        client.abort();
+        Ok(())
+    }
+
+    /// Create a local proxy client.
+    ///
+    /// The resulting UdpSocket/the SocketAddr where programs that want to obfuscate their
+    /// traffic with QUIC will write to.
+    async fn create_local_udp_socket(ipv4: bool) -> Result<(UdpSocket, SocketAddr)> {
+        let random_bind_addr = if ipv4 {
+            SocketAddr::from((Ipv4Addr::LOCALHOST, 0))
+        } else {
+            SocketAddr::from((Ipv6Addr::LOCALHOST, 0))
+        };
+        let local_udp_socket = UdpSocket::bind(random_bind_addr)
+            .await
+            .map_err(Error::BindError)?;
+        let udp_client_addr = local_udp_socket.local_addr().unwrap();
 
-        Ok(Quic {
-            local_endpoint,
-            task,
-        })
+        Ok((local_udp_socket, udp_client_addr))
     }
 }