Add client config builder for masque client

5 files changed, 93 insertions, 98 deletions

diff --git a/Cargo.lock b/Cargo.lock
index f9e94d4959..d5cd685851 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2887,6 +2887,7 @@ dependencies = [
  "socket2",
  "thiserror 2.0.9",
  "tokio",
+ "typed-builder",
 ]
 
 [[package]]
@@ -5677,6 +5678,26 @@ dependencies = [
 ]
 
 [[package]]
+name = "typed-builder"
+version = "0.21.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ce63bcaf7e9806c206f7d7b9c1f38e0dce8bb165a80af0898161058b19248534"
+dependencies = [
+ "typed-builder-macro",
+]
+
+[[package]]
+name = "typed-builder-macro"
+version = "0.21.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "60d8d828da2a3d759d3519cdf29a5bac49c77d039ad36d0782edadbf9cd5415b"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.100",
+]
+
+[[package]]
 name = "typenum"
 version = "1.17.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/mullvad-masque-proxy/Cargo.toml b/mullvad-masque-proxy/Cargo.toml
index 7baa07851e..6f54734d1d 100644
--- a/mullvad-masque-proxy/Cargo.toml
+++ b/mullvad-masque-proxy/Cargo.toml
@@ -22,6 +22,7 @@ bytes = "1"
 anyhow = { workspace = true }
 log = { workspace = true }
 socket2 = { workspace = true }
+typed-builder = "0.21.0"
 
 [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 7fb5424ea7..62fe53f1f9 100644
--- a/mullvad-masque-proxy/examples/masque-client.rs
+++ b/mullvad-masque-proxy/examples/masque-client.rs
@@ -1,5 +1,5 @@
 use clap::Parser;
-use mullvad_masque_proxy::client::Error;
+use mullvad_masque_proxy::client::{ClientConfig, Error};
 use tokio::net::UdpSocket;
 
 use std::{
@@ -73,18 +73,19 @@ async fn main() {
 
     log::info!("Listening on {local_addr}");
 
-    let client = mullvad_masque_proxy::client::Client::connect_with_tls_config(
-        local_socket,
-        server_addr,
-        (Ipv4Addr::UNSPECIFIED, 0).into(),
-        target_addr,
-        &server_hostname,
-        tls_config,
-        mtu,
-        #[cfg(target_os = "linux")]
-        fwmark,
-    )
-    .await;
+    let config = ClientConfig::builder()
+        .client_socket(local_socket)
+        .local_addr((Ipv4Addr::UNSPECIFIED, 0).into())
+        .server_addr(server_addr)
+        .server_host(server_hostname)
+        .target_addr(target_addr)
+        .mtu(mtu)
+        .tls_config(tls_config);
+
+    #[cfg(target_os = "linux")]
+    let config = config.fwmark(fwmark);
+
+    let client = mullvad_masque_proxy::client::Client::connect(config.build()).await;
     if let Err(err) = &client {
         log::error!("ERROR: {:?}", err);
         if let Error::Connection(err) = err {
diff --git a/mullvad-masque-proxy/src/client/mod.rs b/mullvad-masque-proxy/src/client/mod.rs
index 36c62a6e34..1e2a1a9adf 100644
--- a/mullvad-masque-proxy/src/client/mod.rs
+++ b/mullvad-masque-proxy/src/client/mod.rs
@@ -11,13 +11,13 @@ use tokio::{
     select,
     sync::{broadcast, mpsc},
 };
+use typed_builder::TypedBuilder;
 
 use h3::{client, ext::Protocol, proto::varint::VarInt, quic::StreamId};
 use h3_datagram::{datagram::Datagram, datagram_traits::HandleDatagramsExt};
 use http::{header, uri::Scheme, Response, StatusCode};
 use quinn::{
-    crypto::rustls::QuicClientConfig, ClientConfig, Endpoint, EndpointConfig, TokioRuntime,
-    TransportConfig,
+    crypto::rustls::QuicClientConfig, Endpoint, EndpointConfig, TokioRuntime, TransportConfig,
 };
 
 use crate::{
@@ -101,96 +101,70 @@ pub enum Error {
     PacketTooLarge(#[from] fragment::PacketTooLarge),
 }
 
-impl Client {
-    #[allow(clippy::too_many_arguments)]
-    pub async fn connect(
-        client_socket: UdpSocket,
-        server_addr: SocketAddr,
-        local_addr: SocketAddr,
-        target_addr: SocketAddr,
-        server_host: &str,
-        mtu: u16,
-        #[cfg(target_os = "linux")] fwmark: Option<u16>,
-    ) -> Result<Self> {
-        Self::connect_with_tls_config(
-            client_socket,
-            server_addr,
-            local_addr,
-            target_addr,
-            server_host,
-            default_tls_config(),
-            mtu,
-            #[cfg(target_os = "linux")]
-            fwmark,
-        )
-        .await
-    }
+#[derive(TypedBuilder)]
+pub struct ClientConfig {
+    /// Socket that accepts proxy clients
+    pub client_socket: UdpSocket,
 
-    #[allow(clippy::too_many_arguments)]
-    pub async fn connect_with_tls_config(
-        client_socket: UdpSocket,
-        server_addr: SocketAddr,
-        local_addr: SocketAddr,
-        target_addr: SocketAddr,
-        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)
+    /// Socket address to bind the QUIC endpoint socket to
+    pub local_addr: SocketAddr,
+
+    /// Destination to which traffic is forwarded
+    pub target_addr: SocketAddr,
+
+    /// Remote QUIC endpoint address
+    pub server_addr: SocketAddr,
+
+    /// Remote QUIC endpoint hostname
+    pub server_host: String,
+
+    /// MTU (includes IP header)
+    #[builder(default = 1500)]
+    pub mtu: u16,
+
+    /// QUIC TLS config
+    #[builder(default = default_tls_config())]
+    pub tls_config: Arc<rustls::ClientConfig>,
+
+    /// Optional fwmark to set on the QUIC endpoint socket
+    #[cfg(target_os = "linux")]
+    #[builder(default)]
+    pub fwmark: Option<u16>,
+}
+
+impl Client {
+    pub async fn connect(config: ClientConfig) -> Result<Self> {
+        let quic_client_config = QuicClientConfig::try_from(config.tls_config)
             .expect("Failed to construct a valid TLS configuration");
 
-        let mut client_config = ClientConfig::new(Arc::new(quic_client_config));
+        let mut client_config = quinn::ClientConfig::new(Arc::new(quic_client_config));
         let transport_config = TransportConfig::default();
         // TODO: Set datagram_receive_buffer_size  if needed
         // TODO: Set datagram_send_buffer_size if needed
         // When would it be needed? If we need to buffer more packets or buffer less packets for
         // better performance.
         client_config.transport_config(Arc::new(transport_config));
-        Self::connect_with_local_addr(
-            client_socket,
-            server_addr,
-            local_addr,
-            target_addr,
-            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,
-        local_addr: SocketAddr,
-        target_addr: SocketAddr,
-        server_host: &str,
-        client_config: ClientConfig,
-        mtu: u16,
-        #[cfg(target_os = "linux")] fwmark: Option<u16>,
-    ) -> Result<Self> {
-        Self::validate_mtu(mtu, target_addr)?;
+        Self::validate_mtu(config.mtu, config.target_addr)?;
 
-        let max_udp_payload_size = compute_udp_payload_size(mtu, target_addr);
+        let max_udp_payload_size = compute_udp_payload_size(config.mtu, config.target_addr);
 
         let endpoint = Self::setup_quic_endpoint(
-            local_addr,
+            config.local_addr,
             max_udp_payload_size,
             #[cfg(target_os = "linux")]
-            fwmark,
+            config.fwmark,
         )?;
 
-        let connecting = endpoint.connect_with(client_config, server_addr, server_host)?;
+        let connecting =
+            endpoint.connect_with(client_config, config.server_addr, &config.server_host)?;
 
         let connection = connecting.await?;
 
         let (h3_connection, send_stream, request_stream) = Self::setup_h3_connection(
             connection.clone(),
-            target_addr,
-            server_host,
+            config.target_addr,
+            &config.server_host,
             max_udp_payload_size,
         )
         .await?;
@@ -198,7 +172,7 @@ impl Client {
         Ok(Self {
             quinn_conn: connection,
             connection: h3_connection,
-            client_socket: Arc::new(client_socket),
+            client_socket: Arc::new(config.client_socket),
             request_stream,
             _send_stream: send_stream,
             max_udp_payload_size,
diff --git a/mullvad-masque-proxy/tests/proxy.rs b/mullvad-masque-proxy/tests/proxy.rs
index 054a77d86b..40231c02ef 100644
--- a/mullvad-masque-proxy/tests/proxy.rs
+++ b/mullvad-masque-proxy/tests/proxy.rs
@@ -167,20 +167,18 @@ async fn setup_masque(mtu: u16) -> anyhow::Result<(UdpSocket, UdpSocket)> {
         .context("Failed to bind address")?;
     let masque_client_addr = local_socket.local_addr().unwrap();
 
-    let client = client::Client::connect_with_tls_config(
-        local_socket,
-        masque_server_addr,
-        // Local QUIC address
-        any_localhost_addr,
-        target_udp_addr,
-        HOST,
-        client::default_tls_config(),
-        mtu,
-        #[cfg(target_os = "linux")]
-        None,
-    )
-    .await
-    .context("Failed to start MASQUE client")?;
+    let client_config = client::ClientConfig::builder()
+        .client_socket(local_socket)
+        .local_addr(any_localhost_addr)
+        .server_addr(masque_server_addr)
+        .server_host(HOST.to_owned())
+        .target_addr(target_udp_addr)
+        .mtu(mtu)
+        .build();
+
+    let client = client::Client::connect(client_config)
+        .await
+        .context("Failed to start MASQUE client")?;
 
     tokio::spawn(async move {
         if let Err(err) = client.run().await {