Skip to content

Commit b7f7525

Browse files
fix(rtc): reject VideoFrame data shorter than the frame (#867)
1 parent d8d4010 commit b7f7525

2 files changed

Lines changed: 34 additions & 0 deletions

File tree

‎livekit-rtc/livekit/rtc/video_frame.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,9 +47,20 @@ def __init__(
4747
type (proto_video.VideoBufferType.ValueType): The format type of the video frame data
4848
(e.g., RGBA, BGRA, RGB24, etc.).
4949
data (Union[bytes, bytearray, memoryview]): The raw pixel data for the video frame.
50+
51+
Raises:
52+
ValueError: If `data` is smaller than a frame of this size and type.
5053
"""
5154
data = _ensure_compatible_buffer(data)
5255

56+
min_size = _get_plane_length(type, width, height)
57+
data_size = data.nbytes if isinstance(data, memoryview) else len(data)
58+
if data_size < min_size:
59+
raise ValueError(
60+
f"data length must be >= {min_size} bytes for a {width}x{height} "
61+
f"{proto_video.VideoBufferType.Name(type)} frame, got {data_size}"
62+
)
63+
5364
self._width = width
5465
self._height = height
5566
self._type = type

‎tests/rtc/test_video_frame.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,3 +63,26 @@ def test_get_plane_returns_each_i422_plane() -> None:
6363
frame = rtc.VideoFrame(w, h, T.I422, y + u + v)
6464

6565
assert [bytes(frame.get_plane(i)) for i in range(3)] == [y, u, v]
66+
67+
68+
@pytest.mark.parametrize(
69+
"name",
70+
["RGBA", "ABGR", "ARGB", "BGRA", "RGB24", "I420", "I420A", "I422", "I444", "I010", "NV12"],
71+
)
72+
def test_frame_rejects_data_shorter_than_the_format_needs(name: str) -> None:
73+
t, w, h = getattr(T, name), 5, 3
74+
length = video_frame._get_plane_length(t, w, h)
75+
76+
rtc.VideoFrame(w, h, t, bytearray(length))
77+
with pytest.raises(ValueError):
78+
rtc.VideoFrame(w, h, t, bytearray(length - 1))
79+
80+
81+
def test_frame_measures_a_memoryview_in_bytes() -> None:
82+
w, h = 4, 3
83+
pixels = memoryview(bytearray(w * h * 4))
84+
85+
rtc.VideoFrame(w, h, T.RGBA, pixels.cast("B", shape=[h, w, 4]))
86+
rtc.VideoFrame(w, h, T.RGBA, pixels.cast("I"))
87+
with pytest.raises(ValueError):
88+
rtc.VideoFrame(w, h, T.RGBA, pixels[:-4].cast("I"))

0 commit comments

Comments
 (0)