Add error for fragments with < 2 items

1 files changed, 9 insertions, 2 deletions

diff --git a/mullvad-masque-proxy/src/fragment.rs b/mullvad-masque-proxy/src/fragment.rs
index 9a47295fb3..6f4a84b0d9 100644
--- a/mullvad-masque-proxy/src/fragment.rs
+++ b/mullvad-masque-proxy/src/fragment.rs
@@ -34,6 +34,9 @@ pub enum DefragError {
 
     #[error("Payload is too small")]
     PayloadTooSmall,
+
+    #[error("Too few fragments in fragmented packet")]
+    TooFewFragments,
 }
 
 // When a packet is larger than u16::MAX, it can't be fragmented.
@@ -84,6 +87,10 @@ impl Fragments {
         let fragment_count = payload
             .try_get_u8()
             .map_err(|_| DefragError::PayloadTooSmall)?;
+        if fragment_count < 2 {
+            // Packets with only one fragment should be sent as non-fragmented packets.
+            return Err(DefragError::TooFewFragments);
+        }
         let fragment = Fragment { index, payload };
 
         // ensure that the fifo has capacity before pushing the new fragment id
@@ -203,9 +210,9 @@ mod test {
     fn test_fragment_reconstruction() {
         let mut fragments = Fragments::default();
 
-        'outer: for packet_id in 1..255u16 {
+        let max_payload_size = 50;
+        'outer: for packet_id in max_payload_size..255u16 {
             let payload = (0..packet_id as u8).collect::<Vec<u8>>();
-            let max_payload_size = 50;
 
             let mut payload_clone = Bytes::from(payload.clone());
             let mut fragment_buf = fragment_packet(max_payload_size, &mut payload_clone, packet_id)