Fix maximum packet size calculations

7 files changed, 216 insertions, 90 deletions

diff --git a/mullvad-masque-proxy/examples/masque-client.rs b/mullvad-masque-proxy/examples/masque-client.rs
index 780e95d044..3a6b591640 100644
--- a/mullvad-masque-proxy/examples/masque-client.rs
+++ b/mullvad-masque-proxy/examples/masque-client.rs
@@ -31,7 +31,7 @@ pub struct ClientArgs {
 
     /// Maximum packet size
     #[arg(long, short = 'S', default_value = "1280")]
-    maximum_packet_size: u16,
+    mtu: u16,
 }
 
 #[tokio::main]
@@ -42,7 +42,7 @@ async fn main() {
         root_cert_path,
         server_hostname,
         bind_port,
-        maximum_packet_size,
+        mtu,
     } = ClientArgs::parse();
 
     let tls_config = match root_cert_path {
@@ -67,7 +67,7 @@ async fn main() {
         target_addr,
         &server_hostname,
         tls_config,
-        maximum_packet_size,
+        mtu,
     )
     .await;
     if let Err(err) = &client {
diff --git a/mullvad-masque-proxy/examples/masque-server.rs b/mullvad-masque-proxy/examples/masque-server.rs
index 9c07423d9c..fe216070ef 100644
--- a/mullvad-masque-proxy/examples/masque-server.rs
+++ b/mullvad-masque-proxy/examples/masque-server.rs
@@ -28,7 +28,7 @@ pub struct ServerArgs {
 
     /// Maximum packet size
     #[arg(long, short = 'm', default_value = "1700")]
-    maximum_packet_size: u16,
+    mtu: u16,
 }
 
 #[tokio::main]
@@ -42,7 +42,7 @@ async fn main() {
         args.bind_addr,
         args.allowed_ips.iter().cloned().collect(),
         tls_config.into(),
-        args.maximum_packet_size,
+        args.mtu,
     )
     .expect("Failed to initialize server");
     println!("Listening on {}", args.bind_addr);
diff --git a/mullvad-masque-proxy/src/client/mod.rs b/mullvad-masque-proxy/src/client/mod.rs
index 4d59f0dd03..ac02bf96b1 100644
--- a/mullvad-masque-proxy/src/client/mod.rs
+++ b/mullvad-masque-proxy/src/client/mod.rs
@@ -21,8 +21,10 @@ use quinn::{
 };
 
 use crate::{
+    compute_udp_payload_size,
     fragment::{self, Fragments},
     stats::Stats,
+    MIN_IPV4_MTU, MIN_IPV6_MTU, QUIC_HEADER_SIZE,
 };
 
 const MAX_HEADER_SIZE: u64 = 8192;
@@ -34,6 +36,9 @@ const MAX_INFLIGHT_PACKETS: usize = 100;
 pub struct Client {
     client_socket: Arc<UdpSocket>,
 
+    /// QUIC endpoint
+    quinn_conn: quinn::Connection,
+
     /// QUIC connection, used to send the actual HTTP datagrams
     connection: h3::client::Connection<h3_quinn::Connection, bytes::Bytes>,
 
@@ -45,8 +50,8 @@ pub struct Client {
     /// connection is terminated
     request_stream: client::RequestStream<h3_quinn::BidiStream<bytes::Bytes>, bytes::Bytes>,
 
-    /// Maximum packet size
-    maximum_packet_size: u16,
+    /// Maximum UDP payload size (packet size including QUIC overhead)
+    max_udp_payload_size: u16,
 
     stats: Arc<Stats>,
 }
@@ -61,6 +66,8 @@ pub enum Error {
     Connect(#[from] quinn::ConnectError),
     #[error("Failed to connect to QUIC endpoint")]
     Connection(#[from] quinn::ConnectionError),
+    #[error("Invalid MTU: must be at least {min_mtu}")]
+    InvalidMtu { min_mtu: u16 },
     #[error("Invalid max_udp_payload_size")]
     InvalidMaxUdpPayload(#[source] quinn::ConfigError),
     #[error("Connection closed while sending request to initiate proxying")]
@@ -100,7 +107,7 @@ impl Client {
         local_addr: SocketAddr,
         target_addr: SocketAddr,
         server_host: &str,
-        maximum_packet_size: u16,
+        mtu: u16,
     ) -> Result<Self> {
         Self::connect_with_tls_config(
             client_socket,
@@ -109,7 +116,7 @@ impl Client {
             target_addr,
             server_host,
             default_tls_config(),
-            maximum_packet_size,
+            mtu,
         )
         .await
     }
@@ -121,7 +128,7 @@ impl Client {
         target_addr: SocketAddr,
         server_host: &str,
         tls_config: Arc<rustls::ClientConfig>,
-        maximum_packet_size: u16,
+        mtu: u16,
     ) -> Result<Self> {
         let quic_client_config = QuicClientConfig::try_from(tls_config)
             .expect("Failed to construct a valid TLS configuration");
@@ -140,7 +147,7 @@ impl Client {
             target_addr,
             server_host,
             client_config,
-            maximum_packet_size,
+            mtu,
         )
         .await
     }
@@ -152,34 +159,56 @@ impl Client {
         target_addr: SocketAddr,
         server_host: &str,
         client_config: ClientConfig,
-        maximum_packet_size: u16,
+        mtu: u16,
     ) -> Result<Self> {
-        let endpoint = Self::setup_quic_endpoint(local_addr, maximum_packet_size)?;
+        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 connecting = endpoint.connect_with(client_config, server_addr, server_host)?;
 
         let connection = connecting.await?;
 
-        let (connection, send_stream, request_stream) =
-            Self::setup_h3_connection(connection, target_addr, server_host, maximum_packet_size)
-                .await?;
+        let (h3_connection, send_stream, request_stream) = Self::setup_h3_connection(
+            connection.clone(),
+            target_addr,
+            server_host,
+            max_udp_payload_size,
+        )
+        .await?;
 
         Ok(Self {
-            connection,
+            quinn_conn: connection,
+            connection: h3_connection,
             client_socket: Arc::new(client_socket),
             request_stream,
             _send_stream: send_stream,
-            maximum_packet_size,
+            max_udp_payload_size,
             stats: Arc::default(),
         })
     }
 
-    fn setup_quic_endpoint(local_addr: SocketAddr, maximum_packet_size: u16) -> Result<Endpoint> {
+    const fn validate_mtu(mtu: u16, target_addr: SocketAddr) -> Result<()> {
+        let min_mtu = if target_addr.is_ipv4() {
+            MIN_IPV4_MTU
+        } else {
+            MIN_IPV6_MTU
+        };
+        if mtu >= min_mtu {
+            Ok(())
+        } else {
+            Err(Error::InvalidMtu { min_mtu })
+        }
+    }
+
+    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)?;
 
         let mut endpoint_config = EndpointConfig::default();
         endpoint_config
-            .max_udp_payload_size(maximum_packet_size)
+            .max_udp_payload_size(max_udp_payload_size)
             .map_err(Error::InvalidMaxUdpPayload)?;
 
         Endpoint::new(endpoint_config, None, local_socket, Arc::new(TokioRuntime))
@@ -191,7 +220,7 @@ impl Client {
         connection: quinn::Connection,
         target: SocketAddr,
         server_host: &str,
-        maximum_packet_size: u16,
+        mtu: u16,
     ) -> Result<(
         client::Connection<h3_quinn::Connection, bytes::Bytes>,
         client::SendRequest<h3_quinn::OpenStreams, bytes::Bytes>,
@@ -205,7 +234,7 @@ impl Client {
             .await
             .map_err(Error::CreateClient)?;
 
-        let request = new_connect_request(target, &server_host, maximum_packet_size)?;
+        let request = new_connect_request(target, &server_host, mtu)?;
 
         let request_future = async move {
             let mut request_stream = send_stream.send_request(request).await?;
@@ -251,7 +280,8 @@ impl Client {
 
         let mut server_socket_task = tokio::task::spawn(server_socket_task(
             stream_id,
-            self.maximum_packet_size,
+            self.max_udp_payload_size,
+            self.quinn_conn,
             self.connection,
             server_tx,
             client_rx,
@@ -274,7 +304,8 @@ impl Client {
 
 async fn server_socket_task(
     stream_id: StreamId,
-    maximum_packet_size: u16,
+    max_udp_payload_size: u16,
+    quinn_conn: quinn::Connection,
     mut connection: h3::client::Connection<h3_quinn::Connection, bytes::Bytes>,
     server_tx: mpsc::Sender<Datagram>,
     mut client_rx: mpsc::Receiver<Bytes>,
@@ -302,7 +333,14 @@ async fn server_socket_task(
 
         let Some(mut packet) = packet else { break };
 
-        if packet.len() < (Into::<usize>::into(maximum_packet_size) - 100usize) {
+        // Maximum QUIC payload (including fragmentation headers)
+        let maximum_packet_size = if let Some(max_datagram_size) = quinn_conn.max_datagram_size() {
+            max_datagram_size as u16 - 1
+        } else {
+            max_udp_payload_size - QUIC_HEADER_SIZE
+        };
+
+        if packet.len() <= usize::from(maximum_packet_size) {
             stats.tx(packet.len(), false);
             connection
                 .send_datagram(stream_id, packet)
@@ -311,8 +349,9 @@ async fn server_socket_task(
             // drop the added context ID, since packet will have to be fragmented.
             let _ = VarInt::decode(&mut packet);
 
-            for fragment in fragment::fragment_packet(maximum_packet_size, &mut packet, fragment_id)
-                .map_err(Error::PacketTooLarge)?
+            for fragment in
+                fragment::fragment_packet(maximum_packet_size, &mut packet, fragment_id)
+                    .map_err(Error::PacketTooLarge)?
             {
                 stats.tx(fragment.len(), true);
                 connection
@@ -331,7 +370,7 @@ async fn client_socket_rx_task(
     client_tx: mpsc::Sender<Bytes>,
     return_addr_tx: broadcast::Sender<SocketAddr>,
 ) -> Result<()> {
-    let mut client_read_buf = BytesMut::with_capacity(crate::PACKET_BUFFER_SIZE * 1024);
+    let mut client_read_buf = BytesMut::with_capacity(100 * crate::PACKET_BUFFER_SIZE);
     let mut return_addr = SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), 0);
 
     loop {
@@ -411,7 +450,7 @@ async fn client_socket_tx_task(
 fn new_connect_request(
     socket_addr: SocketAddr,
     authority: &dyn AsRef<str>,
-    maximum_packet_size: u16,
+    mtu: u16,
 ) -> Result<http::Request<()>> {
     let host = socket_addr.ip();
     let port = socket_addr.port();
@@ -432,7 +471,7 @@ fn new_connect_request(
         // TODO: Not needed since we set the max_udp_payload_size transport param
         .header(
             b"X-Mullvad-Uplink-Mtu".as_slice(),
-            format!("{maximum_packet_size}"),
+            format!("{mtu}"),
         )
         .body(())
         .expect("failed to construct a body");
@@ -503,9 +542,12 @@ fn read_cert_store_from_reader(reader: &mut dyn io::BufRead) -> Result<rustls::R
     Ok(cert_store)
 }
 
-#[test]
-fn test_zero_stream_id() {
-    h3::quic::StreamId::try_from(0).expect("need to be able to create stream IDs with 0, no?");
+#[cfg(test)]
+mod test {
+    #[test]
+    fn test_zero_stream_id() {
+        h3::quic::StreamId::try_from(0).expect("need to be able to create stream IDs with 0, no?");
+    }
 }
 
 #[derive(Debug)]
diff --git a/mullvad-masque-proxy/src/fragment.rs b/mullvad-masque-proxy/src/fragment.rs
index f60c3b927d..bc85297d70 100644
--- a/mullvad-masque-proxy/src/fragment.rs
+++ b/mullvad-masque-proxy/src/fragment.rs
@@ -6,6 +6,8 @@ use std::{
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use h3::proto::varint::VarInt;
 
+use crate::FRAGMENT_HEADER_SIZE_FRAGMENTED;
+
 #[derive(Default)]
 pub struct Fragments {
     fragment_map: BTreeMap<u16, Vec<Fragment>>,
@@ -102,19 +104,29 @@ struct Fragment {
     time_received: Instant,
 }
 
+/// Fragment packet using the given maximum fragment size (including headers).
+///
+/// `payload` must not contain any fragmentation headers.
+/// `maximum_packet_size` is the maximum fragment size including headers.
 pub fn fragment_packet(
     maximum_packet_size: u16,
     payload: &'_ mut Bytes,
     packet_id: u16,
 ) -> Result<impl Iterator<Item = Bytes> + '_, PacketTooLarge> {
-    let num_fragments: usize = payload.chunks(maximum_packet_size.into()).count();
+    let fragment_payload_size = maximum_packet_size - FRAGMENT_HEADER_SIZE_FRAGMENTED;
+
+    let num_fragments: usize = payload.chunks(fragment_payload_size.into()).count();
     let Ok(fragment_count): std::result::Result<u8, _> = num_fragments.try_into() else {
         return Err(PacketTooLarge(payload.len()));
     };
 
-    let iterator = payload.chunks(maximum_packet_size.into()).enumerate().map(
-        move |(fragment_index, fragment_payload)| {
-            let mut fragment = BytesMut::with_capacity((maximum_packet_size + 1).into());
+    let iterator = payload
+        .chunks(fragment_payload_size.into())
+        .enumerate()
+        .map(move |(fragment_index, fragment_payload)| {
+            let mut fragment = BytesMut::with_capacity(
+                usize::from(fragment_payload_size) + usize::from(FRAGMENT_HEADER_SIZE_FRAGMENTED),
+            );
             crate::HTTP_MASQUE_FRAGMENTED_DATAGRAM_CONTEXT_ID.encode(&mut fragment);
             fragment.put_u16(packet_id);
             fragment.put_u8(
@@ -123,36 +135,44 @@ pub fn fragment_packet(
                     .expect("fragment index must fit in an u8, since num_fragments fits is an u8"),
             );
             fragment.put_u8(fragment_count);
+
+            debug_assert!(fragment.len() == usize::from(FRAGMENT_HEADER_SIZE_FRAGMENTED));
+
             fragment.extend_from_slice(fragment_payload);
             fragment.freeze()
-        },
-    );
+        });
     Ok(iterator)
 }
 
-#[test]
-fn test_fragment_reconstruction() {
-    use rand::{seq::SliceRandom, thread_rng};
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_fragment_reconstruction() {
+        use rand::{seq::SliceRandom, thread_rng};
 
-    let payload = (0..255).collect::<Vec<u8>>();
-    let max_payload_size = 50;
-    let packet_id = 76;
+        let payload = (0..255).collect::<Vec<u8>>();
+        let max_payload_size = 50;
+        let packet_id = 76;
 
-    let mut fragments = Fragments::default();
+        let mut fragments = Fragments::default();
 
-    let mut payload_clone = Bytes::from(payload.clone());
-    let mut fragment_buf = fragment_packet(max_payload_size, &mut payload_clone, packet_id)
-        .unwrap()
-        .collect::<Vec<_>>();
+        let mut payload_clone = Bytes::from(payload.clone());
+        let mut fragment_buf = fragment_packet(max_payload_size, &mut payload_clone, packet_id)
+            .unwrap()
+            .collect::<Vec<_>>();
 
-    fragment_buf.shuffle(&mut thread_rng());
+        fragment_buf.shuffle(&mut thread_rng());
 
-    for fragment in fragment_buf {
-        if let Some(reconstructed_packet) = fragments.handle_incoming_packet(fragment).unwrap() {
-            assert_eq!(payload.as_slice(), reconstructed_packet.as_ref());
-            return;
+        for fragment in fragment_buf {
+            if let Some(reconstructed_packet) = fragments.handle_incoming_packet(fragment).unwrap()
+            {
+                assert_eq!(payload.as_slice(), reconstructed_packet.as_ref());
+                return;
+            }
         }
-    }
 
-    panic!("Failed to reconstruct packet");
+        panic!("Failed to reconstruct packet");
+    }
 }
diff --git a/mullvad-masque-proxy/src/lib.rs b/mullvad-masque-proxy/src/lib.rs
index d4a4e47812..b6e351d94b 100644
--- a/mullvad-masque-proxy/src/lib.rs
+++ b/mullvad-masque-proxy/src/lib.rs
@@ -1,10 +1,37 @@
 use h3::proto::varint::VarInt;
+use std::net::SocketAddr;
 
 pub mod client;
 mod fragment;
 pub mod server;
 mod stats;
 
-const PACKET_BUFFER_SIZE: usize = 1700;
+const PACKET_BUFFER_SIZE: usize = 64 * 1024;
 pub const HTTP_MASQUE_DATAGRAM_CONTEXT_ID: VarInt = VarInt::from_u32(0);
 pub const HTTP_MASQUE_FRAGMENTED_DATAGRAM_CONTEXT_ID: VarInt = VarInt::from_u32(1);
+
+/// Fragment headers size for fragmented packets
+const FRAGMENT_HEADER_SIZE_FRAGMENTED: u16 = 5;
+
+/// UDP header overhead
+const UDP_HEADER_SIZE: u16 = 8;
+
+/// QUIC header size. This is conservative, real overhead varies
+const QUIC_HEADER_SIZE: u16 = 41;
+
+/// This is the size of the payload that stores QUIC packets
+/// MTU - IP header - UDP header
+const fn compute_udp_payload_size(mtu: u16, target_addr: SocketAddr) -> u16 {
+    let ip_overhead = if target_addr.is_ipv4() { 20 } else { 40 };
+    mtu - ip_overhead - UDP_HEADER_SIZE
+}
+
+/// Minimum allowed MTU (IPv6) is the overhead of all headers, plus 1 byte for actual data.
+// 20 = IPv4 header (without optional fields)
+const MIN_IPV4_MTU: u16 =
+    20 + UDP_HEADER_SIZE + QUIC_HEADER_SIZE + FRAGMENT_HEADER_SIZE_FRAGMENTED + 1;
+
+/// Minimum allowed MTU (IPv6) is the overhead of all headers, plus 1 byte for actual data.
+// 40 = IPv6 header (without optional fields)
+const MIN_IPV6_MTU: u16 =
+    40 + UDP_HEADER_SIZE + QUIC_HEADER_SIZE + FRAGMENT_HEADER_SIZE_FRAGMENTED + 1;
diff --git a/mullvad-masque-proxy/src/server/mod.rs b/mullvad-masque-proxy/src/server/mod.rs
index 710edaa537..5a76ee1246 100644
--- a/mullvad-masque-proxy/src/server/mod.rs
+++ b/mullvad-masque-proxy/src/server/mod.rs
@@ -17,7 +17,11 @@ use http::{Request, StatusCode};
 use quinn::{crypto::rustls::QuicServerConfig, Endpoint, Incoming};
 use tokio::{net::UdpSocket, time::interval};
 
-use crate::fragment::{self, Fragments};
+use crate::{
+    compute_udp_payload_size,
+    fragment::{self, Fragments},
+    FRAGMENT_HEADER_SIZE_FRAGMENTED, MIN_IPV4_MTU, MIN_IPV6_MTU, QUIC_HEADER_SIZE,
+};
 
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
@@ -27,6 +31,8 @@ pub enum Error {
     BindSocket(#[source] io::Error),
     #[error("Failed to send negotiation response")]
     SendNegotiationResponse(#[source] h3::Error),
+    #[error("Invalid MTU: must be at least {min_mtu}")]
+    InvalidMtu { min_mtu: u16 },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -36,7 +42,7 @@ const MASQUE_WELL_KNOWN_PATH: &str = "/.well-known/masque/udp/";
 pub struct Server {
     endpoint: Endpoint,
     allowed_hosts: AllowedIps,
-    maximum_packet_size: u16,
+    mtu: u16,
 }
 
 #[derive(Clone)]
@@ -55,8 +61,10 @@ impl Server {
         bind_addr: SocketAddr,
         allowed_hosts: HashSet<IpAddr>,
         tls_config: Arc<rustls::ServerConfig>,
-        maximum_packet_size: u16,
+        mtu: u16,
     ) -> Result<Self> {
+        Self::validate_mtu(mtu, bind_addr)?;
+
         let server_config = quinn::ServerConfig::with_crypto(Arc::new(
             QuicServerConfig::try_from(tls_config).map_err(Error::BadTlsConfig)?,
         ));
@@ -68,10 +76,23 @@ impl Server {
             allowed_hosts: AllowedIps {
                 hosts: Arc::new(allowed_hosts),
             },
-            maximum_packet_size,
+            mtu,
         })
     }
 
+    const fn validate_mtu(mtu: u16, bind_addr: SocketAddr) -> Result<()> {
+        let min_mtu = if bind_addr.is_ipv4() {
+            MIN_IPV4_MTU
+        } else {
+            MIN_IPV6_MTU
+        };
+        if mtu >= min_mtu {
+            Ok(())
+        } else {
+            Err(Error::InvalidMtu { min_mtu })
+        }
+    }
+
     pub fn local_addr(&self) -> io::Result<SocketAddr> {
         self.endpoint.local_addr()
     }
@@ -81,21 +102,19 @@ impl Server {
             tokio::spawn(Self::handle_incoming_connection(
                 new_connection,
                 self.allowed_hosts.clone(),
-                self.maximum_packet_size,
+                self.mtu,
             ));
         }
         Ok(())
     }
 
-    async fn handle_incoming_connection(
-        connection: Incoming,
-        allowed_hosts: AllowedIps,
-        maximum_packet_size: u16,
-    ) {
+    async fn handle_incoming_connection(connection: Incoming, allowed_hosts: AllowedIps, mtu: u16) {
         match connection.await {
             Ok(conn) => {
                 println!("new connection established");
 
+                let quinn_conn = conn.clone();
+
                 let Ok(mut connection) = server::builder()
                     .enable_datagram(true)
                     .build(h3_quinn::Connection::new(conn))
@@ -109,10 +128,11 @@ impl Server {
                     Ok(Some((req, stream))) => {
                         tokio::spawn(Self::handle_proxy_request(
                             connection,
+                            quinn_conn,
                             req,
                             stream,
                             allowed_hosts.clone(),
-                            maximum_packet_size,
+                            mtu,
                         ));
                     }
 
@@ -132,10 +152,11 @@ impl Server {
 
     async fn handle_proxy_request<T: BidiStream<Bytes>>(
         mut connection: Connection<h3_quinn::Connection, Bytes>,
+        quinn_conn: quinn::Connection,
         request: Request<()>,
         mut stream: RequestStream<T, Bytes>,
         allowed_hosts: AllowedIps,
-        maximum_packet_size: u16,
+        mtu: u16,
     ) {
         let Some(target_addr) = get_target_socketaddr(request.uri().path()) else {
             return;
@@ -157,8 +178,10 @@ impl Server {
             return;
         }
 
+        let max_udp_payload_size = compute_udp_payload_size(mtu, target_addr);
+
         let stream_id = stream.id();
-        let mut proxy_recv_buf = BytesMut::with_capacity(crate::PACKET_BUFFER_SIZE);
+        let mut proxy_recv_buf = BytesMut::with_capacity(100 * crate::PACKET_BUFFER_SIZE);
 
         let mut fragments = Fragments::default();
         let mut fragment_id = 0u16;
@@ -191,13 +214,22 @@ impl Server {
 
                             let mut received_packet = proxy_recv_buf.split().freeze();
 
-                            if received_packet.len() < maximum_packet_size.into() {
+                            // Maximum QUIC payload (including fragmentation headers)
+                            let maximum_packet_size = if let Some(max_datagram_size) = quinn_conn.max_datagram_size() {
+                                max_datagram_size as u16 - 1
+                            } else {
+                                max_udp_payload_size - QUIC_HEADER_SIZE
+                            };
+
+                            if received_packet.len() <= usize::from(maximum_packet_size) {
                                 if connection.send_datagram(stream_id, received_packet).is_err() {
                                     return;
                                 }
                             } else {
                                 let _ = VarInt::decode(&mut received_packet);
-                                let Ok(fragments) = fragment::fragment_packet(maximum_packet_size, &mut received_packet, fragment_id) else { continue; };
+                                let fragment_payload_size = maximum_packet_size - FRAGMENT_HEADER_SIZE_FRAGMENTED;
+
+                                let Ok(fragments) = fragment::fragment_packet(fragment_payload_size, &mut received_packet, fragment_id) else { continue; };
                                 fragment_id += 1;
                                 for payload in fragments {
                                     if connection.send_datagram(stream_id, payload).is_err() {
@@ -296,21 +328,26 @@ fn unspecified_addr(addr: IpAddr) -> IpAddr {
     }
 }
 
-#[test]
-fn test_get_good_slashy_ocketaddr() {
-    let addr: IpAddr = "192.168.1.1".parse().unwrap();
-    let port: u16 = 7979;
-    let expected_addr = SocketAddr::new(addr, port);
-    let good_path = format!("{MASQUE_WELL_KNOWN_PATH}///{addr}/{port}////");
+#[cfg(test)]
+mod test {
+    use super::*;
 
-    assert_eq!(get_target_socketaddr(&good_path).unwrap(), expected_addr)
-}
+    #[test]
+    fn test_get_good_slashy_ocketaddr() {
+        let addr: IpAddr = "192.168.1.1".parse().unwrap();
+        let port: u16 = 7979;
+        let expected_addr = SocketAddr::new(addr, port);
+        let good_path = format!("{MASQUE_WELL_KNOWN_PATH}///{addr}/{port}////");
+
+        assert_eq!(get_target_socketaddr(&good_path).unwrap(), expected_addr)
+    }
 
-#[test]
-fn test_get_bad_socketaddr() {
-    let addr: IpAddr = "192.168.1.1".parse().unwrap();
-    let port: u16 = 7979;
-    let good_path = format!("{MASQUE_WELL_KNOWN_PATH}{addr}adsfasd/asdfasdf/{port}");
+    #[test]
+    fn test_get_bad_socketaddr() {
+        let addr: IpAddr = "192.168.1.1".parse().unwrap();
+        let port: u16 = 7979;
+        let good_path = format!("{MASQUE_WELL_KNOWN_PATH}{addr}adsfasd/asdfasdf/{port}");
 
-    assert_eq!(get_target_socketaddr(&good_path), None)
+        assert_eq!(get_target_socketaddr(&good_path), None)
+    }
 }
diff --git a/mullvad-masque-proxy/tests/proxy.rs b/mullvad-masque-proxy/tests/proxy.rs
index 6825f7d75d..094bce85ff 100644
--- a/mullvad-masque-proxy/tests/proxy.rs
+++ b/mullvad-masque-proxy/tests/proxy.rs
@@ -12,7 +12,7 @@ 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 MTU: u16 = 1500;
     const HOST: &str = "test.test";
 
     let any_localhost_addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
@@ -29,7 +29,7 @@ async fn test_server_and_client_forwarding() -> anyhow::Result<()> {
         any_localhost_addr,
         Default::default(),
         Arc::new(server_tls_config),
-        MAXIMUM_PACKET_SIZE,
+        MTU,
     )
     .context("Failed to start MASQUE server")?;
 
@@ -51,7 +51,7 @@ async fn test_server_and_client_forwarding() -> anyhow::Result<()> {
         target_udp_addr,
         HOST,
         client::default_tls_config(),
-        MAXIMUM_PACKET_SIZE,
+        MTU,
     )
     .await
     .context("Failed to start MASQUE client")?;