fix: selective ack construction breaks on overflow (#137)

* fix: selective ack construction breaks on overflow

* fix: add max selective act generation size

src/recv.rs

@@ -1,10 +1,14 @@
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, HashSet};
 
 use crate::packet::SelectiveAck;
 use crate::seq::CircularRangeInclusive;
 
 type Bytes = Vec<u8>;
 
+// https://github.com/ethereum/utp/issues/139
+// Maximum number of selective ACKs that can fit within the length of 8 bits
+const MAX_SELECTIVE_ACK_COUNT: usize = 32 * 63;
+
 #[derive(Clone, Debug)]
 pub struct ReceiveBuffer<const N: usize> {
     buf: Box<[u8; N]>,
@@ -108,17 +112,17 @@ impl<const N: usize> ReceiveBuffer<N> {
         }
 
         // If there are pending packets, then the data for `ack_num + 1` must be missing.
-        let mut next = self.ack_num().wrapping_add(2);
-
-        let mut acked = Vec::new();
-        for seq_num in self.pending.keys() {
-            while *seq_num != next {
+        let mut last_ack = self.ack_num().wrapping_add(2);
+        let mut pending_packets = self.pending.keys().copied().collect::<HashSet<_>>();
+
+        let mut acked = vec![];
+        while !pending_packets.is_empty() && acked.len() < MAX_SELECTIVE_ACK_COUNT {
+            if pending_packets.remove(&last_ack) {
+                acked.push(true);
+            } else {
                 acked.push(false);
-                next = next.wrapping_add(1);
             }
-
-            acked.push(true);
-            next = next.wrapping_add(1);
+            last_ack = last_ack.wrapping_add(1);
         }
 
         Some(SelectiveAck::new(acked))
@@ -293,4 +297,36 @@ mod test {
         let selective_ack = buf.selective_ack();
         assert!(selective_ack.is_none());
     }
+
+    #[test]
+    fn selective_ack_overflow() {
+        let init_seq_num = u16::MAX - 2;
+        let mut buf = ReceiveBuffer::<SIZE>::new(init_seq_num);
+
+        let selective_ack = buf.selective_ack();
+        assert!(selective_ack.is_none());
+
+        const DATA_LEN: usize = 64;
+        let data = vec![0xef; DATA_LEN];
+
+        // Write out-of-order packet.
+        let seq_num = init_seq_num.wrapping_add(2);
+        buf.write(&data, seq_num);
+        // Write overflow packet, which is at seq_num 0.
+        let seq_num = init_seq_num.wrapping_add(3);
+        buf.write(&data, seq_num);
+        // Write another out of order but received packet.
+        let seq_num = init_seq_num.wrapping_add(5);
+        buf.write(&data, seq_num);
+
+        // Selective ACK should mark received packets as set.
+        // Selective ACK begins with ack_num + 2, onwards.
+        // Hence since we received packets 65535, 0, and 2, we should have 3 packets set, in the respective positions.
+        let selective_ack = buf.selective_ack().unwrap();
+        let mut acked = vec![false; 32];
+        acked[0] = true;
+        acked[1] = true;
+        acked[3] = true;
+        assert_eq!(selective_ack.acked(), acked);
+    }
 }