h264: fix crash in Annex-B decoding, add fuzz tests

pkg/codecs/h264/annexb.go

@@ -8,25 +8,38 @@ import (
 func AnnexBUnmarshal(byts []byte) ([][]byte, error) {
 	bl := len(byts)
 	initZeroCount := 0
+	start := 0
 
 outer:
-	for i := 0; i < bl; i++ {
+	for {
-		switch byts[i] {
+		if start >= bl || start >= 4 {
-		case 0:
+			return nil, fmt.Errorf("initial delimiter not found")
+		}
+
+		switch initZeroCount {
+		case 0, 1:
+			if byts[start] != 0 {
+				return nil, fmt.Errorf("initial delimiter not found")
+			}
 			initZeroCount++
 
-		case 1:
+		case 2, 3:
-			break outer
+			switch byts[start] {
+			case 1:
+				start++
+				break outer
 
-		default:
+			case 0:
-			return nil, fmt.Errorf("unexpected byte: %d", byts[i])
+
+			default:
+				return nil, fmt.Errorf("initial delimiter not found")
+			}
+			initZeroCount++
 		}
-	}
+
-	if initZeroCount != 2 && initZeroCount != 3 {
+		start++
-		return nil, fmt.Errorf("initial delimiter not found")
 	}
 
-	start := initZeroCount + 1
 	zeroCount := 0
 	n := 0
 
@@ -67,16 +80,15 @@ outer:
 
 		case 1:
 			if zeroCount == 2 || zeroCount == 3 {
-				if (delimStart - start) > MaxNALUSize {
+				l := delimStart - start
-					return nil, fmt.Errorf("NALU size (%d) is too big (maximum is %d)", delimStart-start, MaxNALUSize)
+				if l == 0 {
+					return nil, fmt.Errorf("invalid NALU")
+				}
+				if l > MaxNALUSize {
+					return nil, fmt.Errorf("NALU size (%d) is too big (maximum is %d)", l, MaxNALUSize)
 				}
 
-				nalu := byts[start:delimStart]
+				ret[pos] = byts[start:delimStart]
-				if len(nalu) == 0 {
-					return nil, fmt.Errorf("empty NALU")
-				}
-
-				ret[pos] = nalu
 				pos++
 				start = i + 1
 			}
@@ -87,15 +99,15 @@ outer:
 		}
 	}
 
-	if (bl - start) > MaxNALUSize {
+	l := bl - start
-		return nil, fmt.Errorf("NALU size (%d) is too big (maximum is %d)", bl-start, MaxNALUSize)
+	if l == 0 {
+		return nil, fmt.Errorf("invalid NALU")
+	}
+	if l > MaxNALUSize {
+		return nil, fmt.Errorf("NALU size (%d) is too big (maximum is %d)", l, MaxNALUSize)
 	}
 
-	nalu := byts[start:bl]
+	ret[pos] = byts[start:bl]
-	if len(nalu) == 0 {
-		return nil, fmt.Errorf("empty NALU")
-	}
-	ret[pos] = nalu
 
 	return ret, nil
 }

pkg/codecs/h264/annexb_test.go

@@ -1,7 +1,6 @@
 package h264
 
 import (
-	"bytes"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -95,50 +94,6 @@ func TestAnnexBMarshal(t *testing.T) {
 	}
 }
 
-func TestAnnexBUnmarshalError(t *testing.T) {
-	for _, ca := range []struct {
-		name string
-		enc  []byte
-		err  string
-	}{
-		{
-			"empty",
-			[]byte{},
-			"initial delimiter not found",
-		},
-		{
-			"invalid initial delimiter 1",
-			[]byte{0xaa, 0xbb},
-			"unexpected byte: 170",
-		},
-		{
-			"invalid initial delimiter 2",
-			[]byte{0x00, 0x00, 0x00, 0x00, 0x01},
-			"initial delimiter not found",
-		},
-		{
-			"empty NALU 1",
-			[]byte{0x00, 0x00, 0x01, 0x00, 0x00, 0x01},
-			"empty NALU",
-		},
-		{
-			"empty NALU 2",
-			[]byte{0x00, 0x00, 0x01, 0xaa, 0x00, 0x00, 0x01},
-			"empty NALU",
-		},
-		{
-			"too many nalus",
-			bytes.Repeat([]byte{0x00, 0x00, 0x01, 0x0a}, 21),
-			"NALU count (21) exceeds maximum allowed (20)",
-		},
-	} {
-		t.Run(ca.name, func(t *testing.T) {
-			_, err := AnnexBUnmarshal(ca.enc)
-			require.EqualError(t, err, ca.err)
-		})
-	}
-}
-
 func BenchmarkAnnexBUnmarshal(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		AnnexBUnmarshal([]byte{
@@ -161,3 +116,9 @@ func BenchmarkAnnexBUnmarshal(b *testing.B) {
 		})
 	}
 }
+
+func FuzzAnnexBUnmarshal(f *testing.F) {
+	f.Fuzz(func(t *testing.T, b []byte) {
+		AnnexBUnmarshal(b)
+	})
+}

pkg/codecs/h264/avcc.go

@@ -15,15 +15,15 @@ func AVCCUnmarshal(buf []byte) ([][]byte, error) {
 			return nil, fmt.Errorf("invalid length")
 		}
 
-		le := int(uint32(buf[pos])<<24 | uint32(buf[pos+1])<<16 | uint32(buf[pos+2])<<8 | uint32(buf[pos+3]))
+		l := int(uint32(buf[pos])<<24 | uint32(buf[pos+1])<<16 | uint32(buf[pos+2])<<8 | uint32(buf[pos+3]))
 		pos += 4
 
-		if (bl - pos) < le {
+		if l == 0 {
-			return nil, fmt.Errorf("invalid length")
+			return nil, fmt.Errorf("invalid NALU")
 		}
 
-		if (bl - pos) > MaxNALUSize {
+		if l > MaxNALUSize {
-			return nil, fmt.Errorf("NALU size (%d) is too big (maximum is %d)", bl-pos, MaxNALUSize)
+			return nil, fmt.Errorf("NALU size (%d) is too big (maximum is %d)", l, MaxNALUSize)
 		}
 
 		if (len(ret) + 1) > MaxNALUsPerGroup {
@@ -31,8 +31,12 @@ func AVCCUnmarshal(buf []byte) ([][]byte, error) {
 				len(ret)+1, MaxNALUsPerGroup)
 		}
 
-		ret = append(ret, buf[pos:pos+le])
+		if (bl - pos) < l {
-		pos += le
+			return nil, fmt.Errorf("invalid length")
+		}
+
+		ret = append(ret, buf[pos:pos+l])
+		pos += l
 
 		if (bl - pos) == 0 {
 			break

pkg/codecs/h264/avcc_test.go

@@ -1,7 +1,6 @@
 package h264
 
 import (
-	"bytes"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -60,36 +59,8 @@ func TestAVCCMarshal(t *testing.T) {
 	}
 }
 
-func TestAVCCUnmarshalError(t *testing.T) {
+func FuzzAVCCUnmarshal(f *testing.F) {
-	for _, ca := range []struct {
+	f.Fuzz(func(t *testing.T, b []byte) {
-		name string
+		AVCCUnmarshal(b)
-		enc  []byte
+	})
-		err  string
-	}{
-		{
-			"empty",
-			[]byte{},
-			"invalid length",
-		},
-		{
-			"invalid length",
-			[]byte{0x01},
-			"invalid length",
-		},
-		{
-			"invalid length",
-			[]byte{0x00, 0x00, 0x00, 0x03},
-			"invalid length",
-		},
-		{
-			"too many nalus",
-			bytes.Repeat([]byte{0x00, 0x00, 0x00, 0x01, 0x0a}, 21),
-			"NALU count (21) exceeds maximum allowed (20)",
-		},
-	} {
-		t.Run(ca.name, func(t *testing.T) {
-			_, err := AVCCUnmarshal(ca.enc)
-			require.EqualError(t, err, ca.err)
-		})
-	}
 }

pkg/codecs/h264/emulationprevention.go

@@ -1,6 +1,6 @@
 package h264
 
-// EmulationPreventionRemove removes emlation prevention bytes from a NALU.
+// EmulationPreventionRemove removes emulation prevention bytes from a NALU.
 func EmulationPreventionRemove(nalu []byte) []byte {
 	// 0x00 0x00 0x03 0x00 -> 0x00 0x00 0x00
 	// 0x00 0x00 0x03 0x01 -> 0x00 0x00 0x01

pkg/codecs/h264/emulationprevention_test.go

@@ -54,3 +54,9 @@ func TestEmulationPreventionRemove(t *testing.T) {
 		})
 	}
 }
+
+func FuzzEmulationPreventionRemove(f *testing.F) {
+	f.Fuzz(func(t *testing.T, b []byte) {
+		EmulationPreventionRemove(b)
+	})
+}

pkg/codecs/h264/testdata/fuzz/FuzzAVCCUnmarshal/002e276d0123e6698e29b86a6db8644720a23ad912d7be68bf441e186095f7fe

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("\x00\x00\x00\x00")

pkg/codecs/h264/testdata/fuzz/FuzzAVCCUnmarshal/06ba4bdb19de593e669c642987e270fe2488d4d58ecd712db136a3e011071253

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0000")

pkg/codecs/h264/testdata/fuzz/FuzzAVCCUnmarshal/582528ddfad69eb57775199a43e0f9fd5c94bba343ce7bb6724d4ebafe311ed4

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0")

pkg/codecs/h264/testdata/fuzz/FuzzAVCCUnmarshal/dad1340efc3fed8006eb91c7ec4be5cffca97bbe1c447ed47eda11adfd7ab0cc

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("\x000\x00\x00")

pkg/codecs/h264/testdata/fuzz/FuzzAnnexBUnmarshal/002e276d0123e6698e29b86a6db8644720a23ad912d7be68bf441e186095f7fe

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("\x00\x00\x00\x00")

pkg/codecs/h264/testdata/fuzz/FuzzAnnexBUnmarshal/2ce2df4de457e581f3404dedd305e7f281b04eeb75cc2339d58c64d6618e8702

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("\x00\x000")

pkg/codecs/h264/testdata/fuzz/FuzzAnnexBUnmarshal/582528ddfad69eb57775199a43e0f9fd5c94bba343ce7bb6724d4ebafe311ed4

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0")

pkg/codecs/h264/testdata/fuzz/FuzzAnnexBUnmarshal/f0b7708a3a98029bd20ee500be4f0a4b1969c5386f513788f118019a7e9e32b1

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("\x00\x00\x01")

pkg/codecs/h264/testdata/fuzz/FuzzAnnexBUnmarshal/fb2cff69a3eec4b1a50ea017aa5604d9b3910ddfa989bdb6bc7323714fda66b3

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("\x00\x00\x01\x00\x00\x01")

pkg/codecs/h264/testdata/fuzz/FuzzAnnexBUnmarshal/fb5317ecce56275a254428a8267b56951b5b6ec2fa924f1e8d8fcab394abc2d3

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("\x00\x00")