Merge branch 'add-masque-test'

8 files changed, 221 insertions, 46 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 4fbf93a399..8a377950dd 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2871,6 +2871,7 @@ dependencies = [
 name = "mullvad-masque-proxy"
 version = "0.1.0"
 dependencies = [
+ "anyhow",
  "bytes",
  "clap",
  "h3",
@@ -2881,6 +2882,7 @@ dependencies = [
  "rand 0.8.5",
  "rustls 0.23.18",
  "rustls-pemfile 2.1.3",
+ "thiserror 2.0.9",
  "tokio",
 ]
 
diff --git a/mullvad-masque-proxy/Cargo.toml b/mullvad-masque-proxy/Cargo.toml
index 644d52755c..27b1a9ca1f 100644
--- a/mullvad-masque-proxy/Cargo.toml
+++ b/mullvad-masque-proxy/Cargo.toml
@@ -10,17 +10,19 @@ description = "A limited functionality UDP over HTTP3 proxy"
 
 [dependencies]
 quinn = { version = "0.11", default-features = false, features = ["log", "runtime-tokio", "rustls-ring"] }
-tokio = { workspace = true, features = [ "macros", "io-util" ] }
+thiserror = { workspace = true }
+tokio = { workspace = true, features = ["fs",  "macros", "io-util"] }
 h3 = "0.0.7"
 h3-datagram = "0.0.1"
-h3-quinn = { version = "0.0.9", features = [ "datagram" ]}
+h3-quinn = { version = "0.0.9", features = ["datagram"] }
 http = "1"
 rustls = { version = "0.23", default-features = false }
 rustls-pemfile = "2.1.3"
 bytes = "1"
 
 [dev-dependencies]
-tokio = { workspace = true, features = [ "macros", "io-util", "rt-multi-thread" ] }
+anyhow = { workspace = true }
+tokio = { workspace = true, features = ["fs", "macros", "io-util", "rt-multi-thread"] }
 clap = { workspace = true }
 rand = "0.8.5"
 
diff --git a/mullvad-masque-proxy/src/client/mod.rs b/mullvad-masque-proxy/src/client/mod.rs
index 4fd7ee5729..182ec580ba 100644
--- a/mullvad-masque-proxy/src/client/mod.rs
+++ b/mullvad-masque-proxy/src/client/mod.rs
@@ -38,39 +38,42 @@ pub struct Client {
 
 pub type Result<T> = std::result::Result<T, Error>;
 
-#[derive(Debug)]
+#[derive(Debug, thiserror::Error)]
 pub enum Error {
-    Bind(io::Error),
-    Connect(quinn::ConnectError),
-    Connection(quinn::ConnectionError),
-    /// Connection closed while sending request to initiate proxying
+    #[error("Failed to bind local socket")]
+    Bind(#[source] io::Error),
+    #[error("Failed to begin connecting to QUIC endpoint")]
+    Connect(#[from] quinn::ConnectError),
+    #[error("Failed to connect to QUIC endpoint")]
+    Connection(#[from] quinn::ConnectionError),
+    #[error("Connection closed while sending request to initiate proxying")]
     ConnectionClosedPrematurely,
-    /// QUIC connection failed while sending request to initiate proxying
-    ConnectionFailed(h3::Error),
-    /// Request failed to illicit a response.
-    RequestError(h3::Error),
-    /// Received response was not a 200.
+    #[error("QUIC connection failed while sending request to initiate proxying")]
+    ConnectionFailed(#[source] h3::Error),
+    #[error("Request failed to illicit a response.")]
+    RequestError(#[source] h3::Error),
+    #[error("Received response was not a 200: {}", .0)]
     UnexpectedStatus(http::StatusCode),
-    /// Failed to receive data from client socket
-    ClientRead(io::Error),
-    /// Failed to send data to client socket
-    ClientWrite(io::Error),
-    /// Failed to receive data from server socket
-    ServerRead(h3::Error),
-    /// Failed to create a client
-    CreateClient(h3::Error),
-    /// Failed to receive good response from proxy
-    ProxyResponse(h3::Error),
-    /// Failed to construct a URI
-    Uri(http::Error),
-    /// Failed to send datagram to proxy
-    SendDatagram(h3::Error),
-    /// Failed to read certificates
-    ReadCerts(io::Error),
-    /// Failed to parse certificates
+    #[error("Failed to receive data from client socket")]
+    ClientRead(#[source] io::Error),
+    #[error("Failed to send data to client socket")]
+    ClientWrite(#[source] io::Error),
+    #[error("Failed to receive data from server socket")]
+    ServerRead(#[source] h3::Error),
+    #[error("Failed to create a client")]
+    CreateClient(#[source] h3::Error),
+    #[error("Failed to receive good response from proxy")]
+    ProxyResponse(#[source] h3::Error),
+    #[error("Failed to construct a URI")]
+    Uri(#[source] http::Error),
+    #[error("Failed to send datagram to proxy")]
+    SendDatagram(#[source] h3::Error),
+    #[error("Failed to read certificates")]
+    ReadCerts(#[source] io::Error),
+    #[error("Failed to parse certificates")]
     ParseCerts,
-    /// Failed to fragment a packet - it is too large
-    PacketTooLarge(fragment::PacketTooLarge),
+    #[error("Failed to fragment a packet - it is too large")]
+    PacketTooLarge(#[from] fragment::PacketTooLarge),
 }
 
 impl Client {
@@ -137,11 +140,9 @@ impl Client {
         // TODO: Set EndpointConfig::max_udp_payload_size instead of using X-Mullvad-Uplink-Mtu
         let endpoint = Endpoint::client(local_addr).map_err(Error::Bind)?;
 
-        let connecting = endpoint
-            .connect_with(client_config, server_addr, server_host)
-            .map_err(Error::Connect)?;
+        let connecting = endpoint.connect_with(client_config, server_addr, server_host)?;
 
-        let connection = connecting.await.map_err(Error::Connection)?;
+        let connection = connecting.await?;
 
         let (connection, send_stream, request_stream) =
             Self::setup_h3_connection(connection, target_addr, server_host, maximum_packet_size)
@@ -229,7 +230,6 @@ impl Client {
                                 self.maximum_packet_size,
                                 &mut send_buf,
                                 fragment_id)
-                            .map_err(Error::PacketTooLarge)
                             ? {
                                 self.connection.send_datagram(stream_id, fragment).map_err(Error::SendDatagram)?;
                             }
diff --git a/mullvad-masque-proxy/src/fragment.rs b/mullvad-masque-proxy/src/fragment.rs
index 6dbee00a34..f60c3b927d 100644
--- a/mullvad-masque-proxy/src/fragment.rs
+++ b/mullvad-masque-proxy/src/fragment.rs
@@ -12,15 +12,19 @@ pub struct Fragments {
 }
 
 // When a packet that arrives is too small to be decoded.
-#[derive(Debug)]
+#[derive(Debug, thiserror::Error)]
 pub enum DefragError {
+    #[error("Bad context id: {:?}", .0)]
     #[allow(dead_code)] // TODO: use this error or remove it.
     BadContextId(Result<VarInt, h3::proto::coding::UnexpectedEnd>),
+
+    #[error("Payload is too small")]
     PayloadTooSmall,
 }
 
 // When a packet is larger than u16::MAX, it can't be fragmented.
-#[derive(Debug)]
+#[derive(Debug, thiserror::Error)]
+#[error("Packet is too large to fragment")]
 pub struct PacketTooLarge(pub usize);
 
 impl Fragments {
diff --git a/mullvad-masque-proxy/src/server/mod.rs b/mullvad-masque-proxy/src/server/mod.rs
index 4abc810a4d..710edaa537 100644
--- a/mullvad-masque-proxy/src/server/mod.rs
+++ b/mullvad-masque-proxy/src/server/mod.rs
@@ -19,11 +19,14 @@ use tokio::{net::UdpSocket, time::interval};
 
 use crate::fragment::{self, Fragments};
 
-#[derive(Debug)]
+#[derive(Debug, thiserror::Error)]
 pub enum Error {
-    BadTlsConfig(quinn::crypto::rustls::NoInitialCipherSuite),
-    BindSocket(io::Error),
-    SendNegotiationResponse(h3::Error),
+    #[error("Bad TLS config")]
+    BadTlsConfig(#[source] quinn::crypto::rustls::NoInitialCipherSuite),
+    #[error("Failed to bind server socket")]
+    BindSocket(#[source] io::Error),
+    #[error("Failed to send negotiation response")]
+    SendNegotiationResponse(#[source] h3::Error),
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -69,6 +72,10 @@ impl Server {
         })
     }
 
+    pub fn local_addr(&self) -> io::Result<SocketAddr> {
+        self.endpoint.local_addr()
+    }
+
     pub async fn run(self) -> Result<()> {
         while let Some(new_connection) = self.endpoint.accept().await {
             tokio::spawn(Self::handle_incoming_connection(
@@ -164,7 +171,7 @@ impl Server {
                 client_send = connection.read_datagram() => {
                     match client_send {
                             Ok(Some(received_packet)) => {
-                                handle_client_packet(received_packet, stream_id, &mut fragments, &udp_socket, target_addr).await;
+                                handle_client_packet(received_packet, stream_id, &mut fragments, &udp_socket).await;
                             },
                             Ok(None) => {
                                 return;
@@ -224,7 +231,6 @@ async fn handle_client_packet(
     stream_id: StreamId,
     fragments: &mut Fragments,
     proxy_socket: &UdpSocket,
-    target_addr: SocketAddr,
 ) {
     if received_packet.stream_id() != stream_id {
         // log::trace!("Received unexpected stream ID from server");
@@ -232,7 +238,7 @@ async fn handle_client_packet(
     }
 
     if let Ok(Some(payload)) = fragments.handle_incoming_packet(received_packet.into_payload()) {
-        let _ = proxy_socket.send_to(&payload, target_addr).await;
+        let _ = proxy_socket.send(&payload).await;
     }
 }
 
diff --git a/mullvad-masque-proxy/tests/proxy.rs b/mullvad-masque-proxy/tests/proxy.rs
new file mode 100644
index 0000000000..6825f7d75d
--- /dev/null
+++ b/mullvad-masque-proxy/tests/proxy.rs
@@ -0,0 +1,111 @@
+use std::net::SocketAddr;
+use std::sync::Arc;
+
+use anyhow::Context;
+use bytes::BytesMut;
+use tokio::fs;
+
+use mullvad_masque_proxy::client;
+use mullvad_masque_proxy::server;
+use tokio::net::UdpSocket;
+
+/// Set up a MASQUE proxy and test that it can be used to communicate with some UDP destination
+#[tokio::test]
+async fn test_server_and_client_forwarding() -> anyhow::Result<()> {
+    const MAXIMUM_PACKET_SIZE: u16 = 1700;
+    const HOST: &str = "test.test";
+
+    let any_localhost_addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
+
+    // Set up destination UDP server
+    let target_udp_server = UdpSocket::bind(any_localhost_addr).await?;
+    let target_udp_addr = target_udp_server
+        .local_addr()
+        .context("Retrieve dest UDP server addr")?;
+
+    // Set up MASQUE server
+    let server_tls_config = load_server_test_cert().await?;
+    let server = server::Server::bind(
+        any_localhost_addr,
+        Default::default(),
+        Arc::new(server_tls_config),
+        MAXIMUM_PACKET_SIZE,
+    )
+    .context("Failed to start MASQUE server")?;
+
+    let masque_server_addr = server.local_addr()?;
+
+    tokio::spawn(server.run());
+
+    // Set up MASQUE client
+    let local_socket = UdpSocket::bind(any_localhost_addr)
+        .await
+        .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(),
+        MAXIMUM_PACKET_SIZE,
+    )
+    .await
+    .context("Failed to start MASQUE client")?;
+
+    tokio::spawn(client.run());
+
+    // Connect to local UDP socket
+    let proxy_client = UdpSocket::bind(any_localhost_addr).await?;
+    proxy_client
+        .connect(masque_client_addr)
+        .await
+        .context("Failed to connect to local UDP server")?;
+
+    // Proxy client -> destination
+    let mut rx_buf = BytesMut::with_capacity(128);
+    proxy_client.send(b"abc").await?;
+    let (_, proxy_addr) = target_udp_server
+        .recv_buf_from(&mut rx_buf)
+        .await
+        .context("Expected to receive message")?;
+    assert_eq!(&*rx_buf, b"abc", "Expected to receive message from client");
+
+    // Destination -> proxy client
+    let mut rx_buf = BytesMut::with_capacity(128);
+    target_udp_server.send_to(b"def", proxy_addr).await?;
+    proxy_client
+        .recv_buf(&mut rx_buf)
+        .await
+        .context("Expected to receive message")?;
+    assert_eq!(&*rx_buf, b"def", "Expected to receive message from server");
+
+    Ok(())
+}
+
+async fn load_server_test_cert() -> anyhow::Result<rustls::ServerConfig> {
+    let key = fs::read("tests/test.key").await.context("Read test key")?;
+    let key = rustls_pemfile::private_key(&mut &*key)?.context("Invalid test key")?;
+
+    let cert_chain = fs::read("tests/test.crt")
+        .await
+        .context("Read test certificate")?;
+    let cert_chain = rustls_pemfile::certs(&mut &*cert_chain)
+        .collect::<Result<_, _>>()
+        .context("Invalid test certificate")?;
+
+    let mut tls_config = rustls::ServerConfig::builder_with_provider(Arc::new(
+        rustls::crypto::ring::default_provider(),
+    ))
+    .with_protocol_versions(&[&rustls::version::TLS13])?
+    .with_no_client_auth()
+    .with_single_cert(cert_chain, key)?;
+
+    tls_config.max_early_data_size = u32::MAX;
+    tls_config.alpn_protocols = vec![b"h3".into()];
+
+    Ok(tls_config)
+}
diff --git a/mullvad-masque-proxy/tests/test.crt b/mullvad-masque-proxy/tests/test.crt
new file mode 100644
index 0000000000..2cc992fd24
--- /dev/null
+++ b/mullvad-masque-proxy/tests/test.crt
@@ -0,0 +1,22 @@
+-----BEGIN CERTIFICATE-----
+MIIDkzCCAnugAwIBAgIUC2y+OPgQ2HbWRXTZwjvZOFlGRcIwDQYJKoZIhvcNAQEL
+BQAwWTELMAkGA1UEBhMCQVUxEzARBgNVBAgMClNvbWUtU3RhdGUxITAfBgNVBAoM
+GEludGVybmV0IFdpZGdpdHMgUHR5IEx0ZDESMBAGA1UEAwwJdGVzdC50ZXN0MB4X
+DTI1MDQwMzExNDUxNVoXDTI1MDUwMzExNDUxNVowWTELMAkGA1UEBhMCQVUxEzAR
+BgNVBAgMClNvbWUtU3RhdGUxITAfBgNVBAoMGEludGVybmV0IFdpZGdpdHMgUHR5
+IEx0ZDESMBAGA1UEAwwJdGVzdC50ZXN0MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8A
+MIIBCgKCAQEAufg598upax+qDkcQCvy4Gv9RSD04JMMnYIa6wgGSZcKuur4r9FRr
+b+pnmiwwoZPlmtE5GUnvVVXB+UGr889IatQ00HxfTr8A/McHh/bSg43uEVnLBWDy
+S9ULW6gJX6xwCvZINRCyw6B6eEBOs2T1MCsaA2/x4ba0lYT8lT9ApRCZE1HFh/Mb
+c1/N/VUaTTxtsE/mhg8sYc2ig1EcZMDy/b7pPUz83oUa4zHruqNQUb2bFQ0TrcKY
+C456aJXCb6t5rWiUJmByzjFOwVGokE550km1q7dpv3yjK5Z023O2jnBVfpVCeEAA
+mnsPQAgqFpuXXTHPqBvuyzHzjsWkZN5wXwIDAQABo1MwUTAdBgNVHQ4EFgQUEEaP
+f5wUEnKp5WZTiMdoydTPq5kwHwYDVR0jBBgwFoAUEEaPf5wUEnKp5WZTiMdoydTP
+q5kwDwYDVR0TAQH/BAUwAwEB/zANBgkqhkiG9w0BAQsFAAOCAQEAPrJfHriKFeLn
+hV9lrH+EN/Q62q1ZIxQagDadcBjhfMALfsD2nr5ENL3uVvFz+XrB9pVrMlfBTnx/
++rbzlzfTpIgKHXsICnXQbY8cvA54JyXysHf9jP3yvZP4vlCjTYX8JuWlykm4StMZ
+suPWyCYprCfJd0n8T2heFyWl6Q9MeyrQolpkrhhR6JuifH/ySc4dl3yWEMrX7lEe
+435llV5WolhiyQ060tKFfgrfPu245DO7Mci9QmafKGK4WSYil95Cwy1uYevVtiNa
+oZQ9mjWJJvVab17qTe9lWukMHFDr7qCxtwTWxynDhWvLrvlPP/2DKo0nEGZU/cVV
+ozh88D0xZw==
+-----END CERTIFICATE-----
diff --git a/mullvad-masque-proxy/tests/test.key b/mullvad-masque-proxy/tests/test.key
new file mode 100644
index 0000000000..aff1bc2dae
--- /dev/null
+++ b/mullvad-masque-proxy/tests/test.key
@@ -0,0 +1,28 @@
+-----BEGIN PRIVATE KEY-----
+MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQC5+Dn3y6lrH6oO
+RxAK/Lga/1FIPTgkwydghrrCAZJlwq66viv0VGtv6meaLDChk+Wa0TkZSe9VVcH5
+Qavzz0hq1DTQfF9OvwD8xweH9tKDje4RWcsFYPJL1QtbqAlfrHAK9kg1ELLDoHp4
+QE6zZPUwKxoDb/HhtrSVhPyVP0ClEJkTUcWH8xtzX839VRpNPG2wT+aGDyxhzaKD
+URxkwPL9vuk9TPzehRrjMeu6o1BRvZsVDROtwpgLjnpolcJvq3mtaJQmYHLOMU7B
+UaiQTnnSSbWrt2m/fKMrlnTbc7aOcFV+lUJ4QACaew9ACCoWm5ddMc+oG+7LMfOO
+xaRk3nBfAgMBAAECggEAI/G5ao6fuUfOe6H6lNUR1I4CrN7ASkK6Cqsfz720CR0e
+3pNBNaFXfrMkwSTHZYOLfmfwDFZA/xJrQn0R+jbXPWa0qpNPbI34Z+MkLoBjYe/9
+0razSd/aFRQhdN6+qRJQOZ4uiKsoki0jXri3PW9HAL9j8MQjUUgaEUg59bLbEMwX
+IR98n+aKLSFBJeeUKE5tvLPt9yvegNI1/6nELT7My3SnIOMlVN9njQq8wqqIch2X
+ky7NwlfsCZ1hF1sD8xCUBloUo/LKfCxWm+2Kt0YQ0/lYfwp8WzFXt4XAL50AKW+A
+QrzVpNMwIp/k/FeKcMF/6rbz3iqQTKW2RrK9rX08WQKBgQDbluIh36xLbSFDjLNx
+j9dCiiCI3cQZVBc7FTfStc9/k4O1NBq5YkUZlCJzjbFsxakdRcb4O2Kwf/v8sBq3
+C9PmoJVV/gqzJQ8TjLJF0MPNPzynm7B0YAD4RXUODRymkqFyHBa3uXT7WG6CmdEa
+Ggh8y5xlKZ2c19J6WpQi6EwM8wKBgQDYzkD/7ad5ttSdVUXNr57Oe9JONwUDUPT6
+UaCthUAbeeKM6JeSZ658DikKALMDZ73grncjSQgLPtOSvo/LQxda+AG515ZZ7J5T
+Nauj44UYFz22Ck7guZCq5tQQVx4YUlBRtM7El9oOvYLjBaLTl9lmmCeQGj4+ugl2
+FU/qDGh55QKBgAlARQyaSL7wvQsEfXbWUYJLIW3CsgVDJqtljHGDGVfNlinnJQ0U
+V8bpF754hLYJacOC8gv5LII1Eh+mJ6n4hJfdwgzaZAcCE62GKuiIEAewl1SUWY29
+kazj+Dd8U+2slcKh7k8VMBl6s0UrR8TqvdrMFS2p4CsAaKyg7ka+NJ4DAoGAMZy5
+KRekLGkXLE24JIJcr9mL3ZQflIuxE5scTrjgW6k/m4kaLkmFlyPSZlSUomHaBJFH
++A4dRh2BYuIym4vly05XbsSTxk4sSNROS7mj2khvOboQJMKyBTm/K2IUI/KqKJhc
+fIZXQupBClxez1a/TAfjfclTlx0RTzE/UUq3mbUCgYBGqeq9cgs9qxt0WS0+Un3o
+/cgGXLknk4IsVclMgPPiDhujTs3I/dJu94zwY3dch906k5KKctYX7cS3YUhAMzBX
+e2DYRq/Z1JC9OXCr7HLTF60IjFx7KlGRL/Kmx1Dd98IZmhS2RSE+dJPovlsYqZ1n
+dGHHcZoEs4iiPe6m+umVyw==
+-----END PRIVATE KEY-----