h264: filter out invalid NALUs from Annex-B

pkg/codecs/h264/annexb.go

@@ -1,33 +1,36 @@
 package h264
 
 import (
+	"errors"
 	"fmt"
 )
 
+// ErrAnnexBNoNALUs is returned by AnnexBUnmarshal when no NALUs have been decoded.
+var ErrAnnexBNoNALUs = errors.New("Annex-B unit doesn't contain any NALU")
+
 // AnnexBUnmarshal decodes an access unit from the Annex-B stream format.
 // Specification: ITU-T Rec. H.264, Annex B
 func AnnexBUnmarshal(buf []byte) ([][]byte, error) {
 	bl := len(buf)
 	initZeroCount := 0
-	start := 0
+	i := 0
 
 outer:
 	for {
-		if start >= bl || start >= 4 {
+		if i >= bl || i >= 4 {
 			return nil, fmt.Errorf("initial delimiter not found")
 		}
 
 		switch initZeroCount {
 		case 0, 1:
-			if buf[start] != 0 {
+			if buf[i] != 0 {
 				return nil, fmt.Errorf("initial delimiter not found")
 			}
 			initZeroCount++
 
 		case 2, 3:
-			switch buf[start] {
+			switch buf[i] {
 			case 1:
-				start++
 				break outer
 
 			case 0:
@@ -38,20 +41,36 @@ outer:
 			initZeroCount++
 		}
 
-		start++
+		i++
 	}
 
+	start := initZeroCount + 1
 	zeroCount := 0
 	n := 0
+	delimStart := 0
+	auSize := 0
 
 	for i := start; i < bl; i++ {
 		switch buf[i] {
 		case 0:
+			if zeroCount == 0 {
+				delimStart = i
+			}
 			zeroCount++
 
 		case 1:
 			if zeroCount == 2 || zeroCount == 3 {
-				n++
+				l := delimStart - start
+
+				if l != 0 {
+					if (auSize + l) > MaxAccessUnitSize {
+						return nil, fmt.Errorf("access unit size (%d) is too big, maximum is %d", auSize+l, MaxAccessUnitSize)
+					}
+					n++
+				}
+
+				auSize += l
+				start = i + 1
 			}
 			zeroCount = 0
 
@@ -60,17 +79,29 @@ outer:
 		}
 	}
 
-	if (n + 1) > MaxNALUsPerAccessUnit {
+	l := bl - start
+
+	if l != 0 {
+		if (auSize + l) > MaxAccessUnitSize {
+			return nil, fmt.Errorf("access unit size (%d) is too big, maximum is %d", auSize+l, MaxAccessUnitSize)
+		}
+		n++
+	}
+
+	if n == 0 {
+		return nil, ErrAnnexBNoNALUs
+	}
+
+	if n > MaxNALUsPerAccessUnit {
 		return nil, fmt.Errorf("NALU count (%d) exceeds maximum allowed (%d)",
-			n+1, MaxNALUsPerAccessUnit)
+			n, MaxNALUsPerAccessUnit)
 	}
 
-	ret := make([][]byte, n+1)
+	ret := make([][]byte, n)
 	pos := 0
 	start = initZeroCount + 1
 	zeroCount = 0
-	delimStart := 0
-	auSize := 0
+	delimStart = 0
 
 	for i := start; i < bl; i++ {
 		switch buf[i] {
@@ -84,17 +115,11 @@ outer:
 			if zeroCount == 2 || zeroCount == 3 {
 				l := delimStart - start
 
-				if l == 0 {
-					return nil, fmt.Errorf("invalid NALU")
+				if l != 0 {
+					ret[pos] = buf[start:delimStart]
+					pos++
 				}
 
-				if (auSize + l) > MaxAccessUnitSize {
-					return nil, fmt.Errorf("access unit size (%d) is too big, maximum is %d", auSize+l, MaxAccessUnitSize)
-				}
-
-				ret[pos] = buf[start:delimStart]
-				pos++
-				auSize += l
 				start = i + 1
 			}
 			zeroCount = 0
@@ -104,18 +129,12 @@ outer:
 		}
 	}
 
-	l := bl - start
+	l = bl - start
 
-	if l == 0 {
-		return nil, fmt.Errorf("invalid NALU")
+	if l != 0 {
+		ret[pos] = buf[start:bl]
 	}
 
-	if (auSize + l) > MaxAccessUnitSize {
-		return nil, fmt.Errorf("access unit size (%d) is too big, maximum is %d", auSize+l, MaxAccessUnitSize)
-	}
-
-	ret[pos] = buf[start:bl]
-
 	return ret, nil
 }
 

pkg/codecs/h264/annexb_test.go

@@ -84,6 +84,17 @@ func TestAnnexBUnmarshal(t *testing.T) {
 	}
 }
 
+func TestAnnexBUnmarshalEmpty(t *testing.T) {
+	buf := []byte{0, 0, 0, 1, 0, 0, 0, 1}
+	_, err := AnnexBUnmarshal(buf)
+	require.Equal(t, err, ErrAnnexBNoNALUs)
+
+	buf = []byte{0, 0, 0, 1, 0, 0, 0, 1, 1}
+	dec, err := AnnexBUnmarshal(buf)
+	require.NoError(t, err)
+	require.Equal(t, [][]byte{{1}}, dec)
+}
+
 func TestAnnexBMarshal(t *testing.T) {
 	for _, ca := range casesAnnexB {
 		t.Run(ca.name, func(t *testing.T) {