fix(rtc): return VideoFrame planes as views of the frame's buffer

livekit-rtc/livekit/rtc/video_frame.py

@@ -146,13 +146,16 @@ def get_plane(self, plane_nth: int) -> Optional[memoryview]:
             Optional[memoryview]: A memoryview of the specified plane's data, or None if
             the index is out of bounds for the format.
         """
-        plane_infos = _get_plane_infos(get_address(self.data), self.type, self.width, self.height)
-        if plane_nth >= len(plane_infos):
+        # with a base address of 0 the plane pointers are offsets into the frame's data
+        plane_infos = _get_plane_infos(0, self.type, self.width, self.height)
+        if not 0 <= plane_nth < len(plane_infos):
             return None
 
+        # a slice of the frame's own buffer keeps that buffer alive, stays read-only for
+        # bytes, and cannot run past its end
         plane_info = plane_infos[plane_nth]
-        cdata = (ctypes.c_uint8 * plane_info.size).from_address(plane_info.data_ptr)
-        return memoryview(cdata)
+        start = plane_info.data_ptr
+        return memoryview(self._data).cast("B")[start : start + plane_info.size]
 
     def convert(
         self, type: proto_video.VideoBufferType.ValueType, *, flip_y: bool = False

tests/rtc/test_video_frame.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
+import array
 import ctypes
+import gc
+import weakref
 
 import pytest
 
@@ -63,3 +66,29 @@ def test_get_plane_returns_each_i422_plane() -> None:
     frame = rtc.VideoFrame(w, h, T.I422, y + u + v)
 
     assert [bytes(frame.get_plane(i)) for i in range(3)] == [y, u, v]
+
+
+@pytest.mark.parametrize("plane_nth", [-1, -3, -4, 3])
+def test_get_plane_returns_none_outside_the_planes(plane_nth: int) -> None:
+    frame = rtc.VideoFrame(4, 2, T.I420, bytes(8 + 2 + 2))
+
+    assert frame.get_plane(plane_nth) is None
+
+
+def test_a_plane_keeps_the_frame_buffer_alive() -> None:
+    buffer = array.array("B", [1] * 8 + [2] * 2 + [3] * 2)
+    buffer_ref = weakref.ref(buffer)
+
+    plane = rtc.VideoFrame(4, 2, T.I420, memoryview(buffer)).get_plane(2)
+    del buffer
+    gc.collect()
+
+    # the frame is gone; the plane is the only thing left pointing at the buffer
+    assert buffer_ref() is not None
+    assert plane is not None and bytes(plane) == b"\x03\x03"
+
+
+def test_a_plane_of_a_bytes_frame_is_read_only() -> None:
+    plane = rtc.VideoFrame(4, 2, T.I420, bytes(8 + 2 + 2)).get_plane(0)
+
+    assert plane is not None and plane.readonly