Add fwmark option

5 files changed, 59 insertions, 5 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 8359d58614..f9e94d4959 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2884,6 +2884,7 @@ dependencies = [
  "rand 0.8.5",
  "rustls 0.23.18",
  "rustls-pemfile 2.1.3",
+ "socket2",
  "thiserror 2.0.9",
  "tokio",
 ]
diff --git a/mullvad-masque-proxy/Cargo.toml b/mullvad-masque-proxy/Cargo.toml
index 2911991296..7baa07851e 100644
--- a/mullvad-masque-proxy/Cargo.toml
+++ b/mullvad-masque-proxy/Cargo.toml
@@ -21,6 +21,7 @@ rustls-pemfile = "2.1.3"
 bytes = "1"
 anyhow = { workspace = true }
 log = { workspace = true }
+socket2 = { workspace = true }
 
 [dev-dependencies]
 env_logger = { workspace = true }
diff --git a/mullvad-masque-proxy/examples/masque-client.rs b/mullvad-masque-proxy/examples/masque-client.rs
index 73f7402381..7fb5424ea7 100644
--- a/mullvad-masque-proxy/examples/masque-client.rs
+++ b/mullvad-masque-proxy/examples/masque-client.rs
@@ -32,6 +32,11 @@ pub struct ClientArgs {
     /// Maximum packet size
     #[arg(long, short = 'S', default_value = "1280")]
     mtu: u16,
+
+    /// fwmark to use for the `server_addr` connection
+    #[cfg(target_os = "linux")]
+    #[arg(long)]
+    fwmark: Option<u16>,
 }
 
 #[tokio::main]
@@ -48,6 +53,8 @@ async fn main() {
         server_hostname,
         bind_port,
         mtu,
+        #[cfg(target_os = "linux")]
+        fwmark,
     } = ClientArgs::parse();
 
     let tls_config = match root_cert_path {
@@ -74,6 +81,8 @@ async fn main() {
         &server_hostname,
         tls_config,
         mtu,
+        #[cfg(target_os = "linux")]
+        fwmark,
     )
     .await;
     if let Err(err) = &client {
diff --git a/mullvad-masque-proxy/src/client/mod.rs b/mullvad-masque-proxy/src/client/mod.rs
index 5716d3b700..36c62a6e34 100644
--- a/mullvad-masque-proxy/src/client/mod.rs
+++ b/mullvad-masque-proxy/src/client/mod.rs
@@ -60,6 +60,9 @@ pub type Result<T> = std::result::Result<T, Error>;
 pub enum Error {
     #[error("Failed to bind local socket")]
     Bind(#[source] io::Error),
+    #[cfg(target_os = "linux")]
+    #[error("Failed to set fwmark on remote socket")]
+    Fwmark(#[source] io::Error),
     #[error("Failed to begin connecting to QUIC endpoint")]
     Connect(#[from] quinn::ConnectError),
     #[error("Failed to connect to QUIC endpoint")]
@@ -99,6 +102,7 @@ pub enum Error {
 }
 
 impl Client {
+    #[allow(clippy::too_many_arguments)]
     pub async fn connect(
         client_socket: UdpSocket,
         server_addr: SocketAddr,
@@ -106,6 +110,7 @@ impl Client {
         target_addr: SocketAddr,
         server_host: &str,
         mtu: u16,
+        #[cfg(target_os = "linux")] fwmark: Option<u16>,
     ) -> Result<Self> {
         Self::connect_with_tls_config(
             client_socket,
@@ -115,10 +120,13 @@ impl Client {
             server_host,
             default_tls_config(),
             mtu,
+            #[cfg(target_os = "linux")]
+            fwmark,
         )
         .await
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub async fn connect_with_tls_config(
         client_socket: UdpSocket,
         server_addr: SocketAddr,
@@ -127,6 +135,7 @@ impl Client {
         server_host: &str,
         tls_config: Arc<rustls::ClientConfig>,
         mtu: u16,
+        #[cfg(target_os = "linux")] fwmark: Option<u16>,
     ) -> Result<Self> {
         let quic_client_config = QuicClientConfig::try_from(tls_config)
             .expect("Failed to construct a valid TLS configuration");
@@ -146,10 +155,13 @@ impl Client {
             server_host,
             client_config,
             mtu,
+            #[cfg(target_os = "linux")]
+            fwmark,
         )
         .await
     }
 
+    #[allow(clippy::too_many_arguments)]
     async fn connect_with_local_addr(
         client_socket: UdpSocket,
         server_addr: SocketAddr,
@@ -158,12 +170,18 @@ impl Client {
         server_host: &str,
         client_config: ClientConfig,
         mtu: u16,
+        #[cfg(target_os = "linux")] fwmark: Option<u16>,
     ) -> Result<Self> {
         Self::validate_mtu(mtu, target_addr)?;
 
         let max_udp_payload_size = compute_udp_payload_size(mtu, target_addr);
 
-        let endpoint = Self::setup_quic_endpoint(local_addr, max_udp_payload_size)?;
+        let endpoint = Self::setup_quic_endpoint(
+            local_addr,
+            max_udp_payload_size,
+            #[cfg(target_os = "linux")]
+            fwmark,
+        )?;
 
         let connecting = endpoint.connect_with(client_config, server_addr, server_host)?;
 
@@ -201,16 +219,39 @@ impl Client {
         }
     }
 
-    fn setup_quic_endpoint(local_addr: SocketAddr, max_udp_payload_size: u16) -> Result<Endpoint> {
-        let local_socket = std::net::UdpSocket::bind(local_addr).map_err(Error::Bind)?;
+    fn setup_quic_endpoint(
+        local_addr: SocketAddr,
+        max_udp_payload_size: u16,
+        #[cfg(target_os = "linux")] fwmark: Option<u16>,
+    ) -> 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(u32::from(fwmark))
+                .map_err(Error::Fwmark)?;
+        }
+
+        local_socket.bind(&local_addr.into()).map_err(Error::Bind)?;
 
         let mut endpoint_config = EndpointConfig::default();
         endpoint_config
             .max_udp_payload_size(max_udp_payload_size)
             .map_err(Error::InvalidMaxUdpPayload)?;
 
-        Endpoint::new(endpoint_config, None, local_socket, Arc::new(TokioRuntime))
-            .map_err(Error::Bind)
+        Endpoint::new(
+            endpoint_config,
+            None,
+            local_socket.into(),
+            Arc::new(TokioRuntime),
+        )
+        .map_err(Error::Bind)
     }
 
     // Returns an h3 connection that is ready to be used for sending UDP datagrams.
diff --git a/mullvad-masque-proxy/tests/proxy.rs b/mullvad-masque-proxy/tests/proxy.rs
index 61a15e5136..054a77d86b 100644
--- a/mullvad-masque-proxy/tests/proxy.rs
+++ b/mullvad-masque-proxy/tests/proxy.rs
@@ -176,6 +176,8 @@ async fn setup_masque(mtu: u16) -> anyhow::Result<(UdpSocket, UdpSocket)> {
         HOST,
         client::default_tls_config(),
         mtu,
+        #[cfg(target_os = "linux")]
+        None,
     )
     .await
     .context("Failed to start MASQUE client")?;