From 9e85fbd13fb3aefe9d710f0d71d46f6618dbfd0b Mon Sep 17 00:00:00 2001
From: Marcus Asteborg <maastebo@microsoft.com>
Date: Fri, 15 Mar 2024 15:02:13 -0700
Subject: [PATCH] Wrap around ssn to 0 and avoid panic

---
 src/queue/queue_test.rs       | 45 +++++++++++++++++++++++++++++++++++
 src/queue/reassembly_queue.rs |  4 ++--
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/src/queue/queue_test.rs b/src/queue/queue_test.rs
index a1828d2..03f8ec9 100644
--- a/src/queue/queue_test.rs
+++ b/src/queue/queue_test.rs
@@ -991,3 +991,48 @@ fn test_chunk_set_incomplete_chunk_set_no_contiguous_tsn() -> Result<()> {
     );
     Ok(())
 }
+
+#[test]
+fn test_reassembly_queue_ssn_overflow() -> Result<()> {
+    let mut rq = ReassemblyQueue::new(0);
+    let org_ppi = PayloadProtocolIdentifier::Binary;
+
+    for stream_sequence_number in 0..=u16::MAX {
+        let chunk = ChunkPayloadData {
+            payload_type: org_ppi,
+            beginning_fragment: true,
+            ending_fragment: true,
+            tsn: 10,
+            stream_sequence_number,
+            user_data: Bytes::from_static(b"123"),
+            ..Default::default()
+        };
+        assert!(rq.push(chunk));
+        assert!(rq.read().is_some());
+    }
+
+    Ok(())
+}
+
+#[test]
+fn test_reassembly_queue_ssn_overflow_in_forward_tsn_for_ordered() -> Result<()> {
+    let mut rq = ReassemblyQueue::new(0);
+    let org_ppi = PayloadProtocolIdentifier::Binary;
+
+    for stream_sequence_number in 0..u16::MAX {
+        let chunk = ChunkPayloadData {
+            payload_type: org_ppi,
+            beginning_fragment: true,
+            ending_fragment: true,
+            tsn: 10,
+            stream_sequence_number,
+            user_data: Bytes::from_static(b"123"),
+            ..Default::default()
+        };
+        assert!(rq.push(chunk));
+        assert!(rq.read().is_some());
+    }
+    rq.forward_tsn_for_ordered(u16::MAX);
+
+    Ok(())
+}
diff --git a/src/queue/reassembly_queue.rs b/src/queue/reassembly_queue.rs
index 30dad10..132e3fc 100644
--- a/src/queue/reassembly_queue.rs
+++ b/src/queue/reassembly_queue.rs
@@ -334,7 +334,7 @@ impl ReassemblyQueue {
                 return None;
             }
             if chunks.ssn == self.next_ssn {
-                self.next_ssn += 1;
+                self.next_ssn = self.next_ssn.wrapping_add(1);
             }
             self.ordered.remove(0)
         } else {
@@ -363,7 +363,7 @@ impl ReassemblyQueue {
 
         // Finally, forward next_ssn
         if sna16lte(self.next_ssn, last_ssn) {
-            self.next_ssn = last_ssn + 1;
+            self.next_ssn = last_ssn.wrapping_add(1);
         }
     }