Merge branch 'masque-args'

3 files changed, 40 insertions, 9 deletions

diff --git a/mullvad-masque-proxy/examples/masque-client.rs b/mullvad-masque-proxy/examples/masque-client.rs
index 62fe53f1f9..f873b14b2f 100644
--- a/mullvad-masque-proxy/examples/masque-client.rs
+++ b/mullvad-masque-proxy/examples/masque-client.rs
@@ -1,3 +1,4 @@
+use anyhow::Context;
 use clap::Parser;
 use mullvad_masque_proxy::client::{ClientConfig, Error};
 use tokio::net::UdpSocket;
@@ -5,6 +6,7 @@ use tokio::net::UdpSocket;
 use std::{
     net::{Ipv4Addr, SocketAddr},
     path::PathBuf,
+    time::Duration,
 };
 
 #[derive(Parser, Debug)]
@@ -25,9 +27,9 @@ pub struct ClientArgs {
     #[arg(long, short = 'H')]
     server_hostname: String,
 
-    /// Local bind port
-    #[arg(long, short = 'p', default_value = "0")]
-    bind_port: u16,
+    /// Bind address
+    #[arg(long, short = 'b', default_value = "127.0.0.1:0")]
+    bind_addr: SocketAddr,
 
     /// Maximum packet size
     #[arg(long, short = 'S', default_value = "1280")]
@@ -37,6 +39,17 @@ pub struct ClientArgs {
     #[cfg(target_os = "linux")]
     #[arg(long)]
     fwmark: Option<u16>,
+
+    /// Maximum duration of inactivity (in seconds) until the tunnel times out.
+    /// Inactivity happens when no data is sent over the proxy.
+    #[arg(long, short = 'i', value_parser = duration_from_seconds)]
+    idle_timeout: Option<Duration>,
+}
+
+/// Parse a duration from a decimal number of seconds
+fn duration_from_seconds(s: &str) -> anyhow::Result<Duration> {
+    let seconds: f64 = s.parse().context("Expected a decimal number, e.g. 1.0")?;
+    Ok(Duration::from_secs_f64(seconds))
 }
 
 #[tokio::main]
@@ -51,10 +64,11 @@ async fn main() {
         target_addr,
         root_cert_path,
         server_hostname,
-        bind_port,
+        bind_addr,
         mtu,
         #[cfg(target_os = "linux")]
         fwmark,
+        idle_timeout,
     } = ClientArgs::parse();
 
     let tls_config = match root_cert_path {
@@ -65,8 +79,7 @@ async fn main() {
 
     let _keylog = rustls::KeyLogFile::new();
 
-    let unbound_local_addr: SocketAddr = (Ipv4Addr::UNSPECIFIED, bind_port).into();
-    let local_socket = UdpSocket::bind(unbound_local_addr)
+    let local_socket = UdpSocket::bind(bind_addr)
         .await
         .expect("Failed to bind address");
     let local_addr = local_socket.local_addr().unwrap();
@@ -80,7 +93,8 @@ async fn main() {
         .server_host(server_hostname)
         .target_addr(target_addr)
         .mtu(mtu)
-        .tls_config(tls_config);
+        .tls_config(tls_config)
+        .idle_timeout(idle_timeout);
 
     #[cfg(target_os = "linux")]
     let config = config.fwmark(fwmark);
diff --git a/mullvad-masque-proxy/src/client/mod.rs b/mullvad-masque-proxy/src/client/mod.rs
index 1e2a1a9adf..f913318595 100644
--- a/mullvad-masque-proxy/src/client/mod.rs
+++ b/mullvad-masque-proxy/src/client/mod.rs
@@ -5,6 +5,7 @@ use std::{
     net::{Ipv4Addr, SocketAddr},
     path::Path,
     sync::{Arc, LazyLock},
+    time::Duration,
 };
 use tokio::{
     net::UdpSocket,
@@ -17,7 +18,8 @@ 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, Endpoint, EndpointConfig, TokioRuntime, TransportConfig,
+    crypto::rustls::QuicClientConfig, Endpoint, EndpointConfig, IdleTimeout, TokioRuntime,
+    TransportConfig,
 };
 
 use crate::{
@@ -99,6 +101,8 @@ pub enum Error {
     ParseCerts,
     #[error("Failed to fragment a packet - it is too large")]
     PacketTooLarge(#[from] fragment::PacketTooLarge),
+    #[error("The provided idle timeout was invalid")]
+    InvalidIdleTimeout(quinn::VarIntBoundsExceeded),
 }
 
 #[derive(TypedBuilder)]
@@ -130,6 +134,10 @@ pub struct ClientConfig {
     #[cfg(target_os = "linux")]
     #[builder(default)]
     pub fwmark: Option<u16>,
+
+    /// Optional timeout when no data is sent in the proxy.
+    #[builder(default)]
+    pub idle_timeout: Option<Duration>,
 }
 
 impl Client {
@@ -138,7 +146,15 @@ impl Client {
             .expect("Failed to construct a valid TLS configuration");
 
         let mut client_config = quinn::ClientConfig::new(Arc::new(quic_client_config));
-        let transport_config = TransportConfig::default();
+        let mut transport_config = TransportConfig::default();
+        transport_config.max_idle_timeout(
+            config
+                .idle_timeout
+                .map(IdleTimeout::try_from)
+                .transpose()
+                .map_err(Error::InvalidIdleTimeout)?,
+        );
+
         // 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
diff --git a/mullvad-masque-proxy/tests/proxy.rs b/mullvad-masque-proxy/tests/proxy.rs
index 40231c02ef..04eacfe78f 100644
--- a/mullvad-masque-proxy/tests/proxy.rs
+++ b/mullvad-masque-proxy/tests/proxy.rs
@@ -174,6 +174,7 @@ async fn setup_masque(mtu: u16) -> anyhow::Result<(UdpSocket, UdpSocket)> {
         .server_host(HOST.to_owned())
         .target_addr(target_udp_addr)
         .mtu(mtu)
+        .idle_timeout(Some(Duration::from_secs(10)))
         .build();
 
     let client = client::Client::connect(client_config)