fix: bound and validate stream reassembly on the transport level

src/levels/transport.py

@@ -22,6 +22,18 @@ class Transport(Base):
     # Сколько чанков одного потока держать в полёте одновременно
     SEND_WINDOW = 8
 
+    # Сколько ждать продолжения незавершённого потока, прежде чем выбросить его
+    STREAM_TIMEOUT = ACK_TIMEOUT * ACK_RETRIES
+
+    # Границы незавершённых потоков
+    MAX_WAITING_STREAMS = 8
+    MAX_STREAM_CHUNKS = 4096
+    MAX_STREAM_BYTES = 1024 * 1024
+
+    # Пакет старше PACKET_MAX_AGE отбрасывается
+    PACKET_MAX_AGE = 300
+    CLOCK_SKEW_TOLERANCE = 60
+
 
     def __init__(self):
         super().__init__()
@@ -179,6 +191,62 @@ def already_seen(self, raw_packet_bytes):
         return False
 
 
+    # Границы заголовка потока
+    def valid_stream_header(self, packet):
+
+        if packet.chunk_count == 0 or packet.chunk_count > self.MAX_STREAM_CHUNKS:
+            self.logger.warning(f"bogus chunk count {packet.chunk_count}: packet dropped")
+            return False
+
+        if packet.chunk_id >= packet.chunk_count:
+            self.logger.warning(f"chunk id {packet.chunk_id} out of range for {packet.chunk_count} chunk(s): packet dropped")
+            return False
+
+        if packet.size > self.MAX_STREAM_BYTES:
+            self.logger.warning(f"chunk of {packet.size} bytes exceeds the stream budget: packet dropped")
+            return False
+
+        return True
+
+
+    # Выбрасывает потоки, в которые давно ничего не приходило
+    def expire_waiting_streams(self):
+
+        now = time.monotonic()
+        expired = [stream_id for stream_id, stream in self.WAITING_STREAMS.items() if stream["deadline"] <= now]
+
+        for stream_id in expired:
+            self.logger.warning(f"stream {stream_id}: no chunks for {self.STREAM_TIMEOUT}s, incomplete stream dropped")
+            del self.WAITING_STREAMS[stream_id]
+
+
+    # Поток, в который ложится чанк. Создаёт запись, если потока ещё нет
+    def stream_for(self, packet):
+
+        stream = self.WAITING_STREAMS.get(packet.stream_id)
+
+        if stream is not None and stream["count"] != packet.chunk_count:
+            self.logger.warning(f"stream {packet.stream_id}: chunk count changed {stream['count']} -> {packet.chunk_count}, stale stream dropped")
+            del self.WAITING_STREAMS[packet.stream_id]
+            stream = None
+
+        if stream is None:
+            if len(self.WAITING_STREAMS) >= self.MAX_WAITING_STREAMS:
+                oldest_id = min(self.WAITING_STREAMS, key=lambda stream_id: self.WAITING_STREAMS[stream_id]["deadline"])
+                self.logger.warning(f"stream table is full: evicting the oldest stream {oldest_id}")
+                del self.WAITING_STREAMS[oldest_id]
+
+            stream = {
+                "count": packet.chunk_count,
+                "packets": {},
+                "bytes": 0,
+                "deadline": time.monotonic() + self.STREAM_TIMEOUT,
+            }
+            self.WAITING_STREAMS[packet.stream_id] = stream
+
+        return stream
+
+
     # постоянно читает данные из PENDING_PROCESSING_BUF и обрабатывает их и отправляет выше
     def rworker(self, data):
 
@@ -196,8 +264,8 @@ def rworker(self, data):
         # Проверка на возраст пакета
         difference_seconds = int(time.time()) - packet.time
         # Если пакет старше 5 минут, отбрасываем
-        if difference_seconds >= 300:
-            self.logger.info(f"old packet. bye")
+        if difference_seconds >= self.PACKET_MAX_AGE or difference_seconds < -self.CLOCK_SKEW_TOLERANCE:
+            self.logger.info(f"packet timestamp is outside the accepted window ({difference_seconds}s). bye")
             return
 
         # Обнуляем счетчик секунд который означает сколько секунд прошло с момента получения последнего пакета
@@ -226,6 +294,10 @@ def rworker(self, data):
 
             self.logger.info(f"data packet")
 
+            # Границы заголовка проверяем раньше подтверждения
+            if not self.valid_stream_header(packet):
+                return
+
             # Подтверждение отправляем всегда: отправитель повторяет пакет
             # именно потому, что не увидел предыдущего подтверждения
             self.send_acknowledgment(data)
@@ -236,13 +308,27 @@ def rworker(self, data):
                 self.logger.info(f"duplicate packet, acknowledged and ignored")
                 return
 
+            # Брошенные потоки убираем прежде, чем разместить новый чанк:
+            # иначе огрызок дождётся оборота stream_id и склеится с чужими данными
+            self.expire_waiting_streams()
+
+            stream = self.stream_for(packet)
+
+            if stream["bytes"] + packet.size > self.MAX_STREAM_BYTES:
+                self.logger.warning(f"stream {packet.stream_id}: over the {self.MAX_STREAM_BYTES} byte budget, dropped")
+                del self.WAITING_STREAMS[packet.stream_id]
+                return
+
             # Чанки лежат в словаре по chunk_id: повторно присланный чанк
             # перезаписывает себя же, а не задваивает счётчик
-            stream = self.WAITING_STREAMS.setdefault(
-                packet.stream_id, {"count": packet.chunk_count, "packets": {}}
-            )
+            if packet.chunk_id not in stream["packets"]:
+                stream["bytes"] += packet.size
             stream["packets"][packet.chunk_id] = packet.payload
 
+            # Каждый новый чанк продлевает жизнь потока: медленный канал с
+            # повторами не должен ронять сборку на полпути
+            stream["deadline"] = time.monotonic() + self.STREAM_TIMEOUT
+
             self.logger.info(f"stream {packet.stream_id}: {len(stream['packets'])} packet of {stream['count']}")
 
             if stream["count"] == len(stream["packets"]):

tests/test_transport.py

@@ -0,0 +1,206 @@
+import logging
+import os
+import sys
+import time
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.join(ROOT, "src"))
+
+from levels.base import Base  # noqa: E402
+from levels.packet import TransportPacket  # noqa: E402
+from levels.transport import Transport  # noqa: E402
+
+DATA_FLAG = 0x0
+ACK_FLAG = 0x1
+
+
+class CollectingUpper:
+    """Stands in for the presentation level: keeps whatever got reassembled."""
+
+    def __init__(self):
+        self.received = []
+
+    def receive(self, data):
+        self.received.append(data)
+
+
+class RecordingLower:
+    """Stands in for the transitional level: keeps whatever hit the channel."""
+
+    def __init__(self):
+        self.sent = []
+
+    def send(self, data):
+        self.sent.append(data)
+
+    def send_without_encrypt(self, data):
+        self.sent.append(data)
+
+
+def make_transport():
+    transport = Transport()
+    upper, lower = CollectingUpper(), RecordingLower()
+    transport.update_levels(upper, lower)
+    return transport, upper, lower
+
+
+def stop(transport):
+    Base.stop_event.set()
+    time.sleep(1.2)
+    Base.stop_event.clear()
+
+
+def chunk(payload, stream_id=0, chunk_count=1, chunk_id=0, age=0):
+    packet = TransportPacket(DATA_FLAG, stream_id, chunk_count, chunk_id, int(time.time()) - age, payload)
+    return packet.to_bytes()
+
+
+def acknowledgments(lower):
+    return [raw for raw in lower.sent if TransportPacket.from_bytes(raw).flags == ACK_FLAG]
+
+
+# ---------------------------------------------------------------- tests
+
+def test_whole_stream_reassembles():
+    transport, upper, lower = make_transport()
+    try:
+        for chunk_id, payload in enumerate((b"one", b"two", b"three")):
+            transport.rworker(chunk(payload, stream_id=5, chunk_count=3, chunk_id=chunk_id))
+
+        assert upper.received == [b"onetwothree"], f"unexpected payload: {upper.received}"
+        assert transport.WAITING_STREAMS == {}, "completed stream was left behind"
+        assert len(acknowledgments(lower)) == 3, "every data packet must be acknowledged"
+    finally:
+        stop(transport)
+
+
+def test_bogus_chunk_count_is_refused():
+    transport, upper, lower = make_transport()
+    try:
+        bogus = [
+            ("zero chunks", chunk(b"x", stream_id=1, chunk_count=0, chunk_id=0)),
+            ("over the cap", chunk(b"x", stream_id=2, chunk_count=Transport.MAX_STREAM_CHUNKS + 1, chunk_id=0)),
+            ("id out of range", chunk(b"x", stream_id=3, chunk_count=2, chunk_id=7)),
+        ]
+        for label, raw in bogus:
+            transport.rworker(raw)
+            assert transport.WAITING_STREAMS == {}, f"{label}: reserved memory for a bogus header"
+
+        assert upper.received == [], f"bogus headers delivered data: {upper.received}"
+        assert acknowledgments(lower) == [], "a packet we refuse must not be acknowledged"
+    finally:
+        stop(transport)
+
+
+def test_abandoned_stream_is_forgotten():
+    transport, upper, lower = make_transport()
+    try:
+        transport.STREAM_TIMEOUT = 0.2
+
+        transport.rworker(chunk(b"OLD", stream_id=9, chunk_count=3, chunk_id=0))
+        assert 9 in transport.WAITING_STREAMS, "first chunk did not open a stream"
+
+        time.sleep(0.4)
+
+        # Любой следующий пакет запускает уборку просроченных потоков
+        transport.rworker(chunk(b"UNRELATED", stream_id=1, chunk_count=1, chunk_id=0))
+        assert 9 not in transport.WAITING_STREAMS, "abandoned stream survived its deadline"
+
+        # Тот же stream_id после оборота: данные не должны склеиться со старым огрызком
+        transport.rworker(chunk(b"NEW", stream_id=9, chunk_count=1, chunk_id=0))
+        assert upper.received == [b"UNRELATED", b"NEW"], f"stale chunk leaked into a new stream: {upper.received}"
+    finally:
+        stop(transport)
+
+
+def test_stream_id_reuse_does_not_merge():
+    transport, upper, lower = make_transport()
+    try:
+        transport.rworker(chunk(b"OLD", stream_id=9, chunk_count=3, chunk_id=0))
+        transport.rworker(chunk(b"NEW", stream_id=9, chunk_count=1, chunk_id=0))
+
+        assert upper.received == [b"NEW"], f"streams were merged: {upper.received}"
+        assert transport.WAITING_STREAMS == {}, "stale stream was left behind"
+    finally:
+        stop(transport)
+
+
+def test_stream_table_stays_bounded():
+    transport, upper, lower = make_transport()
+    try:
+        opened = Transport.MAX_WAITING_STREAMS + 8
+        for stream_id in range(opened):
+            transport.rworker(chunk(b"x", stream_id=stream_id, chunk_count=2, chunk_id=0))
+
+        assert len(transport.WAITING_STREAMS) <= Transport.MAX_WAITING_STREAMS, (
+            f"table grew to {len(transport.WAITING_STREAMS)} streams, cap is {Transport.MAX_WAITING_STREAMS}"
+        )
+        assert opened - 1 in transport.WAITING_STREAMS, "the newest stream was the one evicted"
+        assert upper.received == [], "nothing was complete, yet something was delivered"
+    finally:
+        stop(transport)
+
+
+def test_stream_byte_budget_is_enforced():
+    transport, upper, lower = make_transport()
+    try:
+        transport.MAX_STREAM_BYTES = 300
+
+        for chunk_id in range(4):
+            transport.rworker(chunk(b"p" * 100, stream_id=4, chunk_count=5, chunk_id=chunk_id))
+
+        assert 4 not in transport.WAITING_STREAMS, "oversized stream stayed in memory"
+        assert upper.received == [], "oversized stream was delivered"
+    finally:
+        stop(transport)
+
+
+def test_timestamp_window_is_symmetric():
+    transport, upper, lower = make_transport()
+    try:
+        transport.rworker(chunk(b"future", stream_id=10, age=-3600))
+        transport.rworker(chunk(b"ancient", stream_id=11, age=Transport.PACKET_MAX_AGE + 10))
+
+        assert upper.received == [], f"packet outside the window was accepted: {upper.received}"
+        assert acknowledgments(lower) == [], "a packet outside the window must not be acknowledged"
+        assert transport.WAITING_STREAMS == {}, "packet outside the window reserved memory"
+
+        # Расхождение часов в пределах допуска и свежий пакет из прошлого проходят
+        transport.rworker(chunk(b"skewed", stream_id=12, age=-(Transport.CLOCK_SKEW_TOLERANCE - 10)))
+        transport.rworker(chunk(b"recent", stream_id=13, age=30))
+
+        assert upper.received == [b"skewed", b"recent"], f"legitimate packets were dropped: {upper.received}"
+    finally:
+        stop(transport)
+
+
+TESTS = [
+    test_whole_stream_reassembles,
+    test_bogus_chunk_count_is_refused,
+    test_abandoned_stream_is_forgotten,
+    test_stream_id_reuse_does_not_merge,
+    test_stream_table_stays_bounded,
+    test_stream_byte_budget_is_enforced,
+    test_timestamp_window_is_symmetric,
+]
+
+
+def main():
+    logging.disable(logging.CRITICAL)
+    failures = 0
+    for test in TESTS:
+        name = test.__name__
+        start = time.perf_counter()
+        try:
+            test()
+        except Exception as exc:
+            failures += 1
+            print(f"FAIL  {name}  ({time.perf_counter() - start:.2f}s)\n      {type(exc).__name__}: {exc}")
+        else:
+            print(f"ok    {name}  ({time.perf_counter() - start:.2f}s)")
+    print(f"\n{len(TESTS) - failures}/{len(TESTS)} passed")
+    return 1 if failures else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())