fix: reassemble RTP fragments (#503)

* fix: reassemble RTP fragments * test: set marker bit in test case so that it is not treated as a fragment * test: add tests for fragmented packets --------- Co-authored-by: swoga <[email protected]>
Kane610 · Oct 22, 2024 · f34d13e · f34d13e
1 parent 60b812a
commit f34d13e
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 5 deletions.
diff --git a/axis/rtsp.py b/axis/rtsp.py
@@ -35,6 +35,7 @@ class State(enum.StrEnum):
 
 
 TIME_OUT_LIMIT = 5
+RTP_HEADER_SIZE = 12
 
 
 class RTSPClient(asyncio.Protocol):
@@ -191,6 +192,7 @@ def __init__(self, callback: Callable[[Signal], None] | None) -> None:
             self.callback = callback
             self.data: deque[bytes] = deque()
             self.transport: asyncio.BaseTransport | None = None
+            self.fragment: bool = False
 
         def connection_made(self, transport: asyncio.BaseTransport) -> None:
             """Execute when port is up and listening.
@@ -207,8 +209,19 @@ def connection_lost(self, exc: Exception | None) -> None:
         def datagram_received(self, data: bytes, addr: Any) -> None:
             """Signals when new data is available."""
             if self.callback:
-                self.data.append(data[12:])
-                self.callback(Signal.DATA)
+                payload = data[RTP_HEADER_SIZE:]
+
+                # if the previous packet was a fragment, then merge it
+                if self.fragment:
+                    previous = self.data.pop()
+                    self.data.append(previous + payload)
+                else:
+                    self.data.append(payload)
+
+                # check whether the RTP marker bit is set, if not it is a fragment
+                self.fragment = (data[1] & 0b1 << 7) == 0
+                if not self.fragment:
+                    self.callback(Signal.DATA)
 
 
 class RTSPSession:

diff --git a/tests/packet_fixtures.py b/tests/packet_fixtures.py
@@ -0,0 +1,11 @@
+"""Packet fixtures."""
+
+RTP_PACKET1_FULL = bytes.fromhex(
+    "80e20ed98cfce83a1f93c2ba3c3f786d6c2076657273696f6e3d22312e302220656e636f64696e673d225554462d38223f3e0a3c74743a4d6574616461746153747265616d20786d6c6e733a74743d22687474703a2f2f7777772e6f6e7669662e6f72672f76657231302f736368656d61223e0a3c74743a4576656e743e3c77736e743a4e6f74696669636174696f6e4d65737361676520786d6c6e733a746e73313d22687474703a2f2f7777772e6f6e7669662e6f72672f76657231302f746f706963732220786d6c6e733a746e73617869733d22687474703a2f2f7777772e617869732e636f6d2f323030392f6576656e742f746f706963732220786d6c6e733a77736e743d22687474703a2f2f646f63732e6f617369732d6f70656e2e6f72672f77736e2f622d322220786d6c6e733a777361353d22687474703a2f2f7777772e77332e6f72672f323030352f30382f61646472657373696e67223e3c77736e743a546f706963204469616c6563743d22687474703a2f2f646f63732e6f617369732d6f70656e2e6f72672f77736e2f742d312f546f70696345787072657373696f6e2f53696d706c65223e746e73617869733a53746f726167652f5265636f7264696e673c2f77736e743a546f7069633e3c77736e743a50726f64756365725265666572656e63653e3c777361353a416464726573733e7572693a2f2f63666266393666322d396463372d343337662d616131322d3135356161636262306238662f50726f64756365725265666572656e63653c2f777361353a416464726573733e3c2f77736e743a50726f64756365725265666572656e63653e3c77736e743a4d6573736167653e3c74743a4d6573736167652055746354696d653d22323032342d31302d30375431393a35323a30392e3538313936305a222050726f70657274794f7065726174696f6e3d224368616e676564223e3c74743a536f757263653e3c2f74743a536f757263653e3c74743a4b65793e3c2f74743a4b65793e3c74743a446174613e3c74743a53696d706c654974656d204e616d653d227265636f7264696e67222056616c75653d2231222f3e3c2f74743a446174613e3c2f74743a4d6573736167653e3c2f77736e743a4d6573736167653e3c2f77736e743a4e6f74696669636174696f6e4d6573736167653e3c2f74743a4576656e743e3c2f74743a4d6574616461746153747265616d3e0a"
+)
+RTP_PACKET2_FRAGMENT1 = bytes.fromhex(
+    "80620eda8cfd03a61f93c2ba3c3f786d6c2076657273696f6e3d22312e302220656e636f64696e673d225554462d38223f3e0a3c74743a4d6574616461746153747265616d20786d6c6e733a74743d22687474703a2f2f7777772e6f6e7669662e6f72672f76657231302f736368656d61223e0a3c74743a4576656e743e3c77736e743a4e6f74696669636174696f6e4d65737361676520786d6c6e733a746e73313d22687474703a2f2f7777772e6f6e7669662e6f72672f76657231302f746f706963732220786d6c6e733a746e73617869733d22687474703a2f2f7777772e617869732e636f6d2f323030392f6576656e742f746f706963732220786d6c6e733a77736e743d22687474703a2f2f646f63732e6f617369732d6f70656e2e6f72672f77736e2f622d322220786d6c6e733a777361353d22687474703a2f2f7777772e77332e6f72672f323030352f30382f61646472657373696e67223e3c77736e743a546f706963204469616c6563743d22687474703a2f2f646f63732e6f617369732d6f70656e2e6f72672f77736e2f742d312f546f70696345787072657373696f6e2f53696d706c65223e746e73617869733a43616d6572614170706c69636174696f6e506c6174666f726d2f4f626a656374416e616c79746963732f78696e7465726e616c5f646174613c2f77736e743a546f7069633e3c77736e743a50726f64756365725265666572656e63653e3c777361353a416464726573733e7572693a2f2f63666266393666322d396463372d343337662d616131322d3135356161636262306238662f50726f64756365725265666572656e63653c2f777361353a416464726573733e3c2f77736e743a50726f64756365725265666572656e63653e3c77736e743a4d6573736167653e3c74743a4d6573736167652055746354696d653d22323032342d31302d30375431393a35323a30392e3330333939305a223e3c74743a536f757263653e3c2f74743a536f757263653e3c74743a446174613e3c74743a53696d706c654974656d204e616d653d227376676672616d65222056616c75653d22266c743b73766720786d6c6e733d2671756f743b687474703a2f2f7777772e77332e6f72672f323030302f7376672671756f743b2076657273696f6e3d2671756f743b312e312671756f743b206261736550726f66696c653d2671756f743b66756c6c2671756f743b200a786d6c6e733a786c696e6b3d2671756f743b687474703a2f2f7777772e77332e6f72672f313939392f786c696e6b2671756f743b200a77696474683d2671756f743b313030252671756f743b206865696768743d2671756f743b313030252671756f743b200a76696577426f783d2671756f743b2d31202d3120322032202671756f743b200a7072657365727665417370656374526174696f3d2671756f743b6e6f6e652671756f743b202667743b0a266c743b67207472616e73666f726d3d2671756f743b6d6174726978283120302030202d312030203029206d617472697828312e3030303020302e3030303020302e3030303020312e3030303020302e3030303020302e3030303020292671756f743b202667743b266c743b706f6c79676f6e20636c6173733d2671756f743b2671756f743b20706f696e74733d2671756f743b302e3937302c2d302e393730202d302e3438302c2d302e393637202d302e3934392c2d302e363738202d302e3937362c2d302e313939202d302e3836362c2d302e333736202d302e3735322c2d302e333536202d302e3937302c302e393730202671756f743b207374796c653d2671756f743b66696c6c3a234646303030303b66696c6c2d6f7061636974793a302e323b7374726f6b652d6f7061636974793a313b7374726f6b653a234646303030303b7374726f6b652d7769647468"
+)
+RTP_PACKET2_FRAGMENT2 = bytes.fromhex(
+    "80e20edb8cfd03a61f93c2ba3a322e3570783b766563746f722d6566666563743a6e6f6e2d7363616c696e672d7374726f6b652671756f743b202f2667743b0a266c743b2f672667743b266c743b2f7376672667743b222f3e3c2f74743a446174613e3c2f74743a4d6573736167653e3c2f77736e743a4d6573736167653e3c2f77736e743a4e6f74696669636174696f6e4d6573736167653e3c2f74743a4576656e743e3c2f74743a4d6574616461746153747265616d3e0a"
+)
diff --git a/tests/test_rtsp.py b/tests/test_rtsp.py
@@ -9,9 +9,14 @@
 
 import pytest
 
-from axis.rtsp import RTSPClient, Signal, State
+from axis.rtsp import RTP_HEADER_SIZE, RTSPClient, Signal, State
 
 from .conftest import HOST, RTSP_PORT
+from .packet_fixtures import (
+    RTP_PACKET1_FULL,
+    RTP_PACKET2_FRAGMENT1,
+    RTP_PACKET2_FRAGMENT2,
+)
 
 LOGGER = logging.getLogger(__name__)
 
@@ -514,14 +519,33 @@ def test_rtp_client(rtsp_client, caplog):
     assert "Stream recepient offline" in caplog.text
 
     with patch.object(rtp_client.client, "callback") as mock_callback:
-        rtp_client.client.datagram_received("0123456789ABCDEF", "addr")
+        rtp_client.client.datagram_received(
+            bytes.fromhex("008000000000000000000000AABBCCDD"), "addr"
+        )
         mock_callback.assert_called_with(Signal.DATA)
-        assert rtp_client.data == "CDEF"
+        assert rtp_client.data == bytes.fromhex("AABBCCDD")
 
     rtsp_client.stop()
     mock_transport.close.assert_called()
 
 
+@pytest.mark.parametrize(
+    ("packets"),
+    [([RTP_PACKET1_FULL]), ([RTP_PACKET2_FRAGMENT1, RTP_PACKET2_FRAGMENT2])],
+)
+def test_rtp_fragment(rtsp_client, packets: list[bytes]):
+    """Verify RTP fragment handling."""
+    rtp_client = rtsp_client.rtp
+
+    with patch.object(rtp_client.client, "callback") as mock_callback:
+        payload = b""
+        for packet in packets:
+            rtp_client.client.datagram_received(packet, "addr")
+            payload += packet[RTP_HEADER_SIZE:]
+        mock_callback.assert_called_with(Signal.DATA)
+        assert rtp_client.data == payload
+
+
 def test_methods(rtsp_client):
     """Verify method attributes."""
     method = rtsp_client.method