Add E2E test for masque fragmentation

1 files changed, 102 insertions, 26 deletions

diff --git a/mullvad-masque-proxy/tests/proxy.rs b/mullvad-masque-proxy/tests/proxy.rs
index 094bce85ff..c27dfdfb9d 100644
--- a/mullvad-masque-proxy/tests/proxy.rs
+++ b/mullvad-masque-proxy/tests/proxy.rs
@@ -3,6 +3,7 @@ use std::sync::Arc;
 
 use anyhow::Context;
 use bytes::BytesMut;
+use rand::RngCore;
 use tokio::fs;
 
 use mullvad_masque_proxy::client;
@@ -12,14 +13,99 @@ 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 MTU: u16 = 1500;
+    const MTU: u16 = 1700;
+    let (client, server) = setup_masque(MTU).await?;
+
+    // Proxy client -> destination
+    let mut rx_buf = BytesMut::with_capacity(128);
+    client.send(b"abc").await?;
+    let (_, proxy_addr) = 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);
+    server.send_to(b"def", proxy_addr).await?;
+    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(())
+}
+
+/// End to end test with fragmentation.
+/// Note: This doesn't actually check whether fragmentation occurs, only that packets actually
+/// reach their destinations when fragmentation *should* be present.
+#[tokio::test]
+async fn test_server_and_client_fragmentation() -> anyhow::Result<()> {
+    const MTU: u16 = 1400;
+    const SEND_PACKET_SIZE: usize = 2500;
+
+    let (client, server) = setup_masque(MTU).await?;
+
+    // Proxy client -> destination
+    // Send a random packet, large enough to be fragmented
+    let mut fragment_me = vec![0u8; SEND_PACKET_SIZE];
+    rand::thread_rng().fill_bytes(&mut fragment_me);
+
+    client.send(&fragment_me).await?;
+
+    let mut rx_buf = BytesMut::with_capacity(SEND_PACKET_SIZE + 100);
+    let (_, proxy_addr) = server
+        .recv_buf_from(&mut rx_buf)
+        .await
+        .context("Expected to receive message")?;
+    let read = rx_buf.split();
+    assert_eq!(
+        &*read, &fragment_me,
+        "Expected to receive reassembled message from client"
+    );
+
+    // Destination -> proxy client
+    // Send a random packet, large enough to be fragmented
+    let mut fragment_me = vec![0u8; SEND_PACKET_SIZE];
+    rand::thread_rng().fill_bytes(&mut fragment_me);
+
+    server.send_to(&fragment_me, proxy_addr).await?;
+
+    let mut rx_buf = BytesMut::with_capacity(SEND_PACKET_SIZE + 100);
+    let blen = client
+        .recv_buf(&mut rx_buf)
+        .await
+        .context("Expected to receive message")?;
+
+    let read = rx_buf.split();
+    eprintln!(
+        "from server: {}, {}, {}",
+        fragment_me.len(),
+        read.len(),
+        blen
+    );
+    assert_eq!(
+        &*read, &fragment_me,
+        "Expected to receive reassembled message from server"
+    );
+
+    Ok(())
+}
+
+/// Set up a client and server connected by a MASQUE proxy.
+/// This returns a UDP socket that is connected to the local MASQUE client,
+/// and a UDP socket that represents the other endpoint.
+/// Note that the server socket (second returned value) is not connected,
+/// so `recv_from` must be used.
+async fn setup_masque(mtu: u16) -> anyhow::Result<(UdpSocket, UdpSocket)> {
     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
+    let destination_udp_server = UdpSocket::bind(any_localhost_addr).await?;
+    let target_udp_addr = destination_udp_server
         .local_addr()
         .context("Retrieve dest UDP server addr")?;
 
@@ -29,13 +115,17 @@ async fn test_server_and_client_forwarding() -> anyhow::Result<()> {
         any_localhost_addr,
         Default::default(),
         Arc::new(server_tls_config),
-        MTU,
+        mtu,
     )
     .context("Failed to start MASQUE server")?;
 
     let masque_server_addr = server.local_addr()?;
 
-    tokio::spawn(server.run());
+    tokio::spawn(async move {
+        if let Err(err) = server.run().await {
+            eprintln!("server.run() failed: {err}");
+        }
+    });
 
     // Set up MASQUE client
     let local_socket = UdpSocket::bind(any_localhost_addr)
@@ -51,12 +141,16 @@ async fn test_server_and_client_forwarding() -> anyhow::Result<()> {
         target_udp_addr,
         HOST,
         client::default_tls_config(),
-        MTU,
+        mtu,
     )
     .await
     .context("Failed to start MASQUE client")?;
 
-    tokio::spawn(client.run());
+    tokio::spawn(async move {
+        if let Err(err) = client.run().await {
+            eprintln!("client.run() failed: {err}");
+        }
+    });
 
     // Connect to local UDP socket
     let proxy_client = UdpSocket::bind(any_localhost_addr).await?;
@@ -65,25 +159,7 @@ async fn test_server_and_client_forwarding() -> anyhow::Result<()> {
         .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(())
+    Ok((proxy_client, destination_udp_server))
 }
 
 async fn load_server_test_cert() -> anyhow::Result<rustls::ServerConfig> {