h264: prevent Annex-B from decoding more NALUs than MaxNALUsPerAccessUnit (#303)

Author: aler9

pkg/codecs/h264/annexb.go

@@ -84,6 +84,11 @@ func (a *AnnexB) Unmarshal(buf []byte) error {
 		return ErrAnnexBNoInitialDelimiter
 	}
 
+	if count > MaxNALUsPerAccessUnit {
+		return fmt.Errorf("NALU count (%d) exceeds maximum allowed (%d)",
+			count, MaxNALUsPerAccessUnit)
+	}
+
 	*a = make([][]byte, 0, count)
 	i := 0
 	start := 0

pkg/codecs/h264/annexb_test.go

@@ -121,6 +121,17 @@ func TestAnnexBUnmarshalEmpty(t *testing.T) {
 	require.Equal(t, AnnexB{{1}}, dec)
 }
 
+func TestAnnexBUnmarshalExceedsMaxNALUs(t *testing.T) {
+	buf := make([]byte, 0, 5*51)
+	for range 51 {
+		buf = append(buf, 0x00, 0x00, 0x00, 0x01, 0xAA)
+	}
+
+	var dec AnnexB
+	err := dec.Unmarshal(buf)
+	require.EqualError(t, err, "NALU count (51) exceeds maximum allowed (50)")
+}
+
 func TestAnnexBMarshal(t *testing.T) {
 	for _, ca := range casesAnnexB {
 		t.Run(ca.name, func(t *testing.T) {

pkg/codecs/h264/avcc_test.go

@@ -70,6 +70,17 @@ func TestAVCCUnmarshalEmpty(t *testing.T) {
 	}, dec)
 }
 
+func TestAVCCUnmarshalExceedsMaxNALUs(t *testing.T) {
+	buf := make([]byte, 0, 5*51)
+	for range 51 {
+		buf = append(buf, 0x00, 0x00, 0x00, 0x01, 0xAA)
+	}
+
+	var dec AVCC
+	err := dec.Unmarshal(buf)
+	require.EqualError(t, err, "NALU count (51) exceeds maximum allowed (50)")
+}
+
 func TestAVCCMarshal(t *testing.T) {
 	for _, ca := range casesAVCC {
 		t.Run(ca.name, func(t *testing.T) {