aboutsummaryrefslogtreecommitdiffstats
path: root/packages/meshbay-node/tests/test_sctp_fix.py
blob: c3fe62219a46644f1a7c4cc451205563cfaf4dad (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
"""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
sent again, then a cumulative ack covers all three. Nothing is outstanding, so
nothing may be counted in flight; before the fix one chunk stayed counted, and
enough of those stop the node from sending anything on that connection.
"""

import asyncio
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:
    state = "connected"

    def _register_data_receiver(self, *a):
        pass


def _sack(cumulative_tsn: int, gaps: list) -> SackChunk:
    sack = SackChunk()
    sack.cumulative_tsn = cumulative_tsn % 2**32
    sack.gaps = gaps
    sack.duplicates = []
    sack.advertised_rwnd = 1 << 20
    return sack


async def _replay() -> RTCSctpTransport:
    s = RTCSctpTransport(_Dtls())
    s._association_state = s.State.ESTABLISHED

    async def sent(*_a, **_k):
        pass

    s._send_chunk = sent
    for name in ("_t3_start", "_t3_restart", "_t3_cancel"):
        setattr(s, name, lambda: None)
    s._cwnd = 100_000

    for _ in range(3):
        await s._send(1, 51, b"x" * 1000)
    tsns = [c.tsn for c in s._sent_queue]

    await s._receive_sack_chunk(_sack(tsns[0] - 1, [(1, 2)]))
    s._t3_expired()
    await s._transmit()
    await s._receive_sack_chunk(_sack(tsns[-1], []))
    return s


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]