From 7af4363a0c21897d109cb8a0c54b986878669219 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Sun, 24 Sep 2023 11:06:00 +0200
Subject: [PATCH] h264: support AVCC with empty NALUs

(https://github.com/bluenviron/mediamtx/issues/2375)
---
 pkg/codecs/h264/avcc.go        | 33 ++++++++++++++++-----------------
 pkg/codecs/h264/avcc_test.go   | 23 +++++++++++++++++++++++
 pkg/formats/fmp4/part_track.go |  6 ++++++
 3 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/pkg/codecs/h264/avcc.go b/pkg/codecs/h264/avcc.go
index b7812b8..620e1b1 100644
--- a/pkg/codecs/h264/avcc.go
+++ b/pkg/codecs/h264/avcc.go
@@ -5,6 +5,7 @@ import (
 )
 
 // AVCCUnmarshal decodes an access unit from the AVCC stream format.
+// This can also return nil in case the AVCC unit doesn't contain any NALU.
 // Specification: ?
 func AVCCUnmarshal(buf []byte) ([][]byte, error) {
 	bl := len(buf)
@@ -21,28 +22,26 @@ func AVCCUnmarshal(buf []byte) ([][]byte, error) {
 		l := int(uint32(buf[pos])<<24 | uint32(buf[pos+1])<<16 | uint32(buf[pos+2])<<8 | uint32(buf[pos+3]))
 		pos += 4
 
-		if l == 0 {
-			return nil, fmt.Errorf("invalid NALU")
-		}
+		if l != 0 {
+			if (auSize + l) > MaxAccessUnitSize {
+				return nil, fmt.Errorf("access unit size (%d) is too big, maximum is %d", auSize+l, MaxAccessUnitSize)
+			}
 
-		if (auSize + l) > MaxAccessUnitSize {
-			return nil, fmt.Errorf("access unit size (%d) is too big, maximum is %d", auSize+l, MaxAccessUnitSize)
-		}
+			if (naluCount + 1) > MaxNALUsPerAccessUnit {
+				return nil, fmt.Errorf("NALU count (%d) exceeds maximum allowed (%d)",
+					len(ret)+1, MaxNALUsPerAccessUnit)
+			}
 
-		if (naluCount + 1) > MaxNALUsPerAccessUnit {
-			return nil, fmt.Errorf("NALU count (%d) exceeds maximum allowed (%d)",
-				len(ret)+1, MaxNALUsPerAccessUnit)
-		}
+			if (bl - pos) < l {
+				return nil, fmt.Errorf("invalid length")
+			}
 
-		if (bl - pos) < l {
-			return nil, fmt.Errorf("invalid length")
+			ret = append(ret, buf[pos:pos+l])
+			auSize += l
+			naluCount++
+			pos += l
 		}
 
-		ret = append(ret, buf[pos:pos+l])
-		auSize += l
-		naluCount++
-		pos += l
-
 		if (bl - pos) == 0 {
 			break
 		}
diff --git a/pkg/codecs/h264/avcc_test.go b/pkg/codecs/h264/avcc_test.go
index 867ddd5..60c087a 100644
--- a/pkg/codecs/h264/avcc_test.go
+++ b/pkg/codecs/h264/avcc_test.go
@@ -49,6 +49,29 @@ func TestAVCCUnmarshal(t *testing.T) {
 	}
 }
 
+// issue mediamtx/2375
+func TestAVCCUnmarshalEmpty(t *testing.T) {
+	caenc := []byte{
+		0x0, 0x0, 0x0, 0x0,
+	}
+	cadec := [][]byte(nil)
+
+	dec, err := AVCCUnmarshal(caenc)
+	require.NoError(t, err)
+	require.Equal(t, cadec, dec)
+
+	caenc = []byte{
+		0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x3, 0x1, 0x2, 0x3,
+	}
+	cadec = [][]byte{
+		{1, 2, 3},
+	}
+
+	dec, err = AVCCUnmarshal(caenc)
+	require.NoError(t, err)
+	require.Equal(t, cadec, dec)
+}
+
 func TestAVCCMarshal(t *testing.T) {
 	for _, ca := range casesAVCC {
 		t.Run(ca.name, func(t *testing.T) {
diff --git a/pkg/formats/fmp4/part_track.go b/pkg/formats/fmp4/part_track.go
index f883d04..8058581 100644
--- a/pkg/formats/fmp4/part_track.go
+++ b/pkg/formats/fmp4/part_track.go
@@ -1,6 +1,8 @@
 package fmp4
 
 import (
+	"fmt"
+
 	"github.com/abema/go-mp4"
 
 	"github.com/bluenviron/mediacommon/pkg/codecs/av1"
@@ -59,6 +61,10 @@ func (ps PartSample) GetH26x() ([][]byte, error) {
 		return nil, err
 	}
 
+	if au == nil {
+		return nil, fmt.Errorf("no valid NALUs found")
+	}
+
 	return au, nil
 }