1 files changed, 148 insertions, 3 deletions
diff --git a/packages/meshbay-node/tests/test_sctp_fix.py b/packages/meshbay-node/tests/test_sctp_fix.py
index 0b2536a..c3fe622 100644
--- a/packages/meshbay-node/tests/test_sctp_fix.py
+++ b/packages/meshbay-node/tests/test_sctp_fix.py
@@ -1,4 +1,5 @@
-"""aiortc's SCTP sender must not leak bytes in flight (see sctp_fix.py).
+"""aiortc's SCTP stack as sctp_fix.py leaves it: no bytes leaked in flight,
+reassembly in linear time, one SACK every second packet.
 
 Replayed against aiortc's own RTCSctpTransport: three chunks sent, the first
 lost and the other two acknowledged in a gap block, T3 expires and they are
@@ -8,10 +9,20 @@ enough of those stop the node from sending anything on that connection.
 """
 
 import asyncio
-
-from aiortc.rtcsctptransport import RTCSctpTransport, SackChunk
+import random
 
 import meshbay_node.transport.webrtc.core  # noqa: F401  (applies the fix)
+from aiortc.rtcsctptransport import (
+    SCTP_DATA_FIRST_FRAG,
+    SCTP_DATA_LAST_FRAG,
+    SCTP_DATA_UNORDERED,
+    DataChunk,
+    InboundStream,
+    RTCSctpTransport,
+    SackChunk,
+    serialize_packet,
+)
+from meshbay_node.transport.webrtc import sctp_fix
 
 
 class _Dtls:
@@ -57,3 +68,137 @@ def test_nothing_outstanding_means_nothing_in_flight():
     s = asyncio.run(_replay())
     assert not s._sent_queue
     assert s._flight_size == 0
+
+
+# ── Reassembly (see sctp_fix.py) ─────────────────────────────────────────────
+
+
+def _fragments(messages: list[int], *, unordered: bool = False, tsn: int = 1000):
+    """DATA chunks for messages of the given fragment counts, in TSN order."""
+    out = []
+    for seq, n in enumerate(messages):
+        for i in range(n):
+            c = DataChunk()
+            c.flags = ((SCTP_DATA_FIRST_FRAG if i == 0 else 0)
+                       | (SCTP_DATA_LAST_FRAG if i == n - 1 else 0)
+                       | (SCTP_DATA_UNORDERED if unordered else 0))
+            c.tsn, c.stream_id, c.stream_seq, c.protocol = tsn, 1, seq, 53
+            c.user_data = f"{seq}.{i};".encode()
+            tsn += 1
+            out.append(c)
+    return out
+
+
+class _Original:
+    """aiortc's own reassembly, untouched, for the same arrivals."""
+
+    def __init__(self):
+        self.reassembly, self.sequence_number = [], 0
+
+    def add(self, chunk):
+        sctp_fix._add_chunk(self, chunk)
+        return list(sctp_fix._pop_messages(self))
+
+
+def test_reassembly_gives_what_aiortc_gives_for_any_arrival_order():
+    rng = random.Random(7)
+    for trial in range(300):
+        messages = [rng.randint(1, 12) for _ in range(rng.randint(1, 8))]
+        chunks = _fragments(messages, unordered=trial % 3 == 0)
+        # Reordered within a short window, as a network reorders.
+        for i in range(len(chunks)):
+            j = min(len(chunks) - 1, i + rng.randint(0, 4))
+            chunks[i], chunks[j] = chunks[j], chunks[i]
+        fixed, original = InboundStream(), _Original()
+        got, want = [], []
+        for c in chunks:
+            fixed.add_chunk(c)
+            got.append(list(fixed.pop_messages()))
+            want.append(original.add(c))
+        assert got == want, (trial, messages)
+
+
+def test_a_large_message_is_walked_once_not_once_per_fragment(monkeypatch):
+    walks = []
+    original = sctp_fix._pop_messages
+
+    def counted(stream):
+        walks.append(len(stream.reassembly))
+        return original(stream)
+
+    monkeypatch.setattr(sctp_fix, "_pop_messages", counted)
+    s = InboundStream()
+    out = []
+    for c in _fragments([1000, 1000]):
+        s.add_chunk(c)
+        out += list(s.pop_messages())
+    assert len(out) == 2
+    assert walks == [1000, 1000], "one walk per message, when its last fragment is in"
+
+
+def test_pruning_a_last_fragment_does_not_stall_the_next_message():
+    s = InboundStream()
+    first, second = _fragments([3]), _fragments([2], tsn=1003)
+    second[0].stream_seq = second[1].stream_seq = 1
+    for c in first[1:]:          # the first fragment never comes
+        s.add_chunk(c)
+        assert not list(s.pop_messages())
+    s.prune_chunks(first[-1].tsn)  # abandoned (FORWARD TSN)
+    s.sequence_number = 1
+    for c in second:
+        s.add_chunk(c)
+    assert [m[2] for m in s.pop_messages()] == [b"0.0;0.1;"]
+
+
+# ── Delayed acknowledgement (see sctp_fix.py) ────────────────────────────────
+
+def _receiver():
+    s = RTCSctpTransport(_Dtls())
+    s._association_state = s.State.ESTABLISHED
+    s._last_received_tsn = 999
+    sacks = []
+
+    async def send_sack():
+        sacks.append(s._last_received_tsn)
+        s._sack_needed = False
+        s._sack_duplicates.clear()
+
+    async def deliver(*_a):
+        pass
+
+    s._send_sack = send_sack
+    s._receive = deliver
+    return s, sacks
+
+
+def _packet(s, chunk) -> bytes:
+    return serialize_packet(5000, 5000, s._local_verification_tag, chunk)
+
+
+def test_one_sack_every_second_packet_of_data():
+    async def go():
+        s, sacks = _receiver()
+        for c in _fragments([1] * 10):
+            await s._handle_data(_packet(s, c))
+        return sacks
+    assert asyncio.run(go()) == [1001, 1003, 1005, 1007, 1009]
+
+
+def test_a_gap_is_acknowledged_at_once():
+    async def go():
+        s, sacks = _receiver()
+        chunks = _fragments([1] * 4)
+        await s._handle_data(_packet(s, chunks[0]))
+        await s._handle_data(_packet(s, chunks[2]))   # 1001 missing
+        return sacks
+    assert asyncio.run(go()) == [1000], "the gap is the loss signal the sender waits for"
+
+
+def test_a_lone_packet_is_acknowledged_within_the_delay():
+    async def go():
+        s, sacks = _receiver()
+        await s._handle_data(_packet(s, _fragments([1])[0]))
+        assert sacks == []
+        await asyncio.sleep(sctp_fix.SACK_DELAY_S + 0.1)
+        return sacks
+    assert asyncio.run(go()) == [1000]