Subtract stream ID from max packet size

Co-authored-by: Joakim Hulthe <joakim.hulthe@mullvad.net>

3 files changed, 15 insertions, 13 deletions

diff --git a/mullvad-masque-proxy/src/client/mod.rs b/mullvad-masque-proxy/src/client/mod.rs
index ac02bf96b1..df65339c03 100644
--- a/mullvad-masque-proxy/src/client/mod.rs
+++ b/mullvad-masque-proxy/src/client/mod.rs
@@ -312,6 +312,7 @@ async fn server_socket_task(
     stats: Arc<Stats>,
 ) -> Result<()> {
     let mut fragment_id = 1u16;
+    let stream_id_size = VarInt::from(stream_id).size() as u16;
 
     loop {
         let packet = select! {
@@ -335,9 +336,9 @@ async fn server_socket_task(
 
         // 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
+            max_datagram_size as u16 - stream_id_size
         } else {
-            max_udp_payload_size - QUIC_HEADER_SIZE
+            max_udp_payload_size - QUIC_HEADER_SIZE - stream_id_size
         };
 
         if packet.len() <= usize::from(maximum_packet_size) {
@@ -349,10 +350,11 @@ 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)?
             {
+                debug_assert!(fragment.len() <= maximum_packet_size as usize);
+
                 stats.tx(fragment.len(), true);
                 connection
                     .send_datagram(stream_id, fragment)
diff --git a/mullvad-masque-proxy/src/fragment.rs b/mullvad-masque-proxy/src/fragment.rs
index bc85297d70..c699ed2d78 100644
--- a/mullvad-masque-proxy/src/fragment.rs
+++ b/mullvad-masque-proxy/src/fragment.rs
@@ -124,9 +124,9 @@ pub fn fragment_packet(
         .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),
-            );
+            let mut fragment = BytesMut::with_capacity(usize::from(
+                fragment_payload_size + FRAGMENT_HEADER_SIZE_FRAGMENTED,
+            ));
             crate::HTTP_MASQUE_FRAGMENTED_DATAGRAM_CONTEXT_ID.encode(&mut fragment);
             fragment.put_u16(packet_id);
             fragment.put_u8(
diff --git a/mullvad-masque-proxy/src/server/mod.rs b/mullvad-masque-proxy/src/server/mod.rs
index 5a76ee1246..a57282d172 100644
--- a/mullvad-masque-proxy/src/server/mod.rs
+++ b/mullvad-masque-proxy/src/server/mod.rs
@@ -20,7 +20,7 @@ use tokio::{net::UdpSocket, time::interval};
 use crate::{
     compute_udp_payload_size,
     fragment::{self, Fragments},
-    FRAGMENT_HEADER_SIZE_FRAGMENTED, MIN_IPV4_MTU, MIN_IPV6_MTU, QUIC_HEADER_SIZE,
+    MIN_IPV4_MTU, MIN_IPV6_MTU, QUIC_HEADER_SIZE,
 };
 
 #[derive(Debug, thiserror::Error)]
@@ -181,6 +181,7 @@ impl Server {
         let max_udp_payload_size = compute_udp_payload_size(mtu, target_addr);
 
         let stream_id = stream.id();
+        let stream_id_size = VarInt::from(stream_id).size() as u16;
         let mut proxy_recv_buf = BytesMut::with_capacity(100 * crate::PACKET_BUFFER_SIZE);
 
         let mut fragments = Fragments::default();
@@ -216,9 +217,9 @@ impl Server {
 
                             // 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
+                               max_datagram_size as u16 - stream_id_size
                             } else {
-                                max_udp_payload_size - QUIC_HEADER_SIZE
+                                max_udp_payload_size - QUIC_HEADER_SIZE - stream_id_size
                             };
 
                             if received_packet.len() <= usize::from(maximum_packet_size) {
@@ -227,9 +228,8 @@ impl Server {
                                 }
                             } else {
                                 let _ = VarInt::decode(&mut received_packet);
-                                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; };
+                                let Ok(fragments) = fragment::fragment_packet(maximum_packet_size, &mut received_packet, fragment_id) else { continue; };
                                 fragment_id += 1;
                                 for payload in fragments {
                                     if connection.send_datagram(stream_id, payload).is_err() {