Renet: correctly calculate small packet size for unreliable channels

lucaspoffo · May 29, 2023 · 6d65002 · 6d65002
1 parent 0c6311c
commit 6d65002
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 3 deletions.
diff --git a/renet/src/channel/reliable.rs b/renet/src/channel/reliable.rs
@@ -363,6 +363,8 @@ impl ReceiveChannelReliable {
 
 #[cfg(test)]
 mod tests {
+    use octets::OctetsMut;
+
     use super::*;
 
     #[test]
@@ -523,4 +525,30 @@ mod tests {
         let packets = send.get_packets_to_send(&mut sequence, &mut available_bytes, current_time);
         assert_eq!(packets.len(), 0);
     }
+
+    #[test]
+    fn small_packet_max_size() {
+        let mut sequence: u64 = 0;
+        let current_time: Duration = Duration::ZERO;
+        let mut available_bytes = u64::MAX;
+        let resend_time = Duration::from_millis(100);
+        let mut send = SendChannelReliable::new(0, resend_time, usize::MAX);
+
+        // 4 bytes
+        let message: Bytes = vec![0, 1, 2, 3].into();
+
+        // (4 + 1 + 2) * 300 = 2100 = 2 packets
+        for _ in 0..300 {
+            send.send_message(message.clone()).unwrap();
+        }
+
+        let packets = send.get_packets_to_send(&mut sequence, &mut available_bytes, current_time);
+        assert_eq!(packets.len(), 2);
+        let mut buffer = [0u8; 1400];
+        for packet in packets {
+            let mut oct = OctetsMut::with_slice(&mut buffer);
+            let len = packet.to_bytes(&mut oct).unwrap();
+            assert!(len < 1300);
+        }
+    }
 }
diff --git a/renet/src/channel/unreliable.rs b/renet/src/channel/unreliable.rs
@@ -79,7 +79,8 @@ impl SendChannelUnreliable {
 
                 self.sliced_message_id += 1;
             } else {
-                if small_messages_bytes + message.len() > SLICE_SIZE {
+                let serialized_size = message.len() + octets::varint_len(message.len() as u64);
+                if small_messages_bytes + serialized_size > SLICE_SIZE {
                     packets.push(Packet::SmallUnreliable {
                         sequence: *packet_sequence,
                         channel_id: self.channel_id,
@@ -89,7 +90,7 @@ impl SendChannelUnreliable {
                     small_messages_bytes = 0;
                 }
 
-                small_messages_bytes += message.len();
+                small_messages_bytes += serialized_size;
                 small_messages.push(message);
             }
         }
@@ -210,6 +211,8 @@ impl ReceiveChannelUnreliable {
 
 #[cfg(test)]
 mod tests {
+    use octets::OctetsMut;
+
     use super::*;
 
     #[test]
@@ -337,4 +340,28 @@ mod tests {
         let packets = send.get_packets_to_send(&mut sequence, &mut available_bytes);
         assert_eq!(packets.len(), 0);
     }
+
+    #[test]
+    fn small_packet_max_size() {
+        let mut sequence: u64 = 0;
+        let mut available_bytes = u64::MAX;
+        let mut send = SendChannelUnreliable::new(0, usize::MAX);
+
+        // 4 bytes
+        let message: Bytes = vec![0, 1, 2, 3].into();
+
+        // (4 + 1) * 400 = 2000 = 2 packets
+        for _ in 0..400 {
+            send.send_message(message.clone());
+        }
+
+        let packets = send.get_packets_to_send(&mut sequence, &mut available_bytes);
+        assert_eq!(packets.len(), 2);
+        let mut buffer = [0u8; 1400];
+        for packet in packets {
+            let mut oct = OctetsMut::with_slice(&mut buffer);
+            let len = packet.to_bytes(&mut oct).unwrap();
+            assert!(len < 1300);
+        }
+    }
 }
diff --git a/renet/src/packet.rs b/renet/src/packet.rs
@@ -257,7 +257,7 @@ impl Packet {
 
                 let payload = b.get_bytes_with_varint_length()?;
 
-                if payload.len() == 0 {
+                if payload.is_empty() {
                     return Err(SerializationError::EmptySlice);
                 }