From d59d7c48b491e43010dd5e392d4a5550b63ccd53 Mon Sep 17 00:00:00 2001
From: Alessandro Ros <aler9.dev@gmail.com>
Date: Thu, 27 Mar 2025 11:56:51 +0100
Subject: [PATCH] improve tests (#742)

---
 pkg/format/rtpav1/decoder_test.go  |  2 +-
 pkg/format/rtph264/decoder_test.go | 38 ++++++++++++++++++++++++++++-
 pkg/format/rtph265/decoder_test.go | 39 +++++++++++++++++++++++++++++-
 3 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/pkg/format/rtpav1/decoder_test.go b/pkg/format/rtpav1/decoder_test.go
index 4c43b263..a16bee8d 100644
--- a/pkg/format/rtpav1/decoder_test.go
+++ b/pkg/format/rtpav1/decoder_test.go
@@ -33,7 +33,7 @@ func TestDecode(t *testing.T) {
 	}
 }
 
-func TestDecoderErrorLimit(t *testing.T) {
+func TestDecoderErrorOBUCount(t *testing.T) {
 	d := &Decoder{}
 	err := d.Init()
 	require.NoError(t, err)
diff --git a/pkg/format/rtph264/decoder_test.go b/pkg/format/rtph264/decoder_test.go
index b4dd2b99..93a79016 100644
--- a/pkg/format/rtph264/decoder_test.go
+++ b/pkg/format/rtph264/decoder_test.go
@@ -222,7 +222,43 @@ func TestDecodeAccessUnit(t *testing.T) {
 	require.Equal(t, [][]byte{{0x01, 0x02}, {0x01, 0x02}}, nalus)
 }
 
-func TestDecoderErrorLimit(t *testing.T) {
+func TestDecoderErrorNALUSize(t *testing.T) {
+	d := &Decoder{}
+	err := d.Init()
+	require.NoError(t, err)
+
+	size := 0
+	i := uint16(0)
+
+	for size < h264.MaxAccessUnitSize {
+		flags := byte(0)
+		if size == 0 {
+			flags = 0b10000000
+		}
+
+		_, err = d.Decode(&rtp.Packet{
+			Header: rtp.Header{
+				Version:        2,
+				Marker:         false,
+				PayloadType:    96,
+				SequenceNumber: 17645 + i,
+				Timestamp:      2289527317,
+				SSRC:           0x9dbb7812,
+			},
+			Payload: append(
+				[]byte{byte(h264.NALUTypeFUA), flags},
+				bytes.Repeat([]byte{1, 2, 3, 4}, 1400/4)...,
+			),
+		})
+
+		size += 1400
+		i++
+	}
+
+	require.EqualError(t, err, "NALU size (8388801) is too big, maximum is 8388608")
+}
+
+func TestDecoderErrorNALUCount(t *testing.T) {
 	d := &Decoder{}
 	err := d.Init()
 	require.NoError(t, err)
diff --git a/pkg/format/rtph265/decoder_test.go b/pkg/format/rtph265/decoder_test.go
index c6599328..951b860f 100644
--- a/pkg/format/rtph265/decoder_test.go
+++ b/pkg/format/rtph265/decoder_test.go
@@ -1,6 +1,7 @@
 package rtph265
 
 import (
+	"bytes"
 	"errors"
 	"testing"
 
@@ -39,7 +40,43 @@ func TestDecode(t *testing.T) {
 	}
 }
 
-func TestDecoderErrorLimit(t *testing.T) {
+func TestDecoderErrorNALUSize(t *testing.T) {
+	d := &Decoder{}
+	err := d.Init()
+	require.NoError(t, err)
+
+	size := 0
+	i := uint16(0)
+
+	for size < h265.MaxAccessUnitSize {
+		flags := byte(0)
+		if size == 0 {
+			flags = 0b10000000
+		}
+
+		_, err = d.Decode(&rtp.Packet{
+			Header: rtp.Header{
+				Version:        2,
+				Marker:         false,
+				PayloadType:    96,
+				SequenceNumber: 17645 + i,
+				Timestamp:      2289527317,
+				SSRC:           0x9dbb7812,
+			},
+			Payload: append(
+				[]byte{byte(h265.NALUType_FragmentationUnit << 1), 0, flags},
+				bytes.Repeat([]byte{1, 2, 3, 4}, 1400/4)...,
+			),
+		})
+
+		size += 1400
+		i++
+	}
+
+	require.EqualError(t, err, "NALU size (8388802) is too big, maximum is 8388608")
+}
+
+func TestDecoderErrorNALUCount(t *testing.T) {
 	d := &Decoder{}
 	err := d.Init()
 	require.NoError(t, err)