h264: fix Annex-B decoding

Some Annex-B streams use both [0 0 1] and [0 0 0 1] as delimiters. These streams are now parsed correctly.
bluenviron · Oct 8, 2022 · b3c70f5 · b3c70f5
1 parent 92cc1b5
commit b3c70f5
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 42 deletions.
diff --git a/pkg/h264/annexb.go b/pkg/h264/annexb.go
@@ -1,20 +1,19 @@
 package h264
 
 import (
-	"bytes"
 	"fmt"
 )
 
 // AnnexBUnmarshal decodes NALUs from the Annex-B stream format.
 func AnnexBUnmarshal(byts []byte) ([][]byte, error) {
 	bl := len(byts)
-	zeroCount := 0
+	initZeroCount := 0
 
 outer:
 	for i := 0; i < bl; i++ {
 		switch byts[i] {
 		case 0:
-			zeroCount++
+			initZeroCount++
 
 		case 1:
 			break outer
@@ -23,35 +22,46 @@ outer:
 			return nil, fmt.Errorf("unexpected byte: %d", byts[i])
 		}
 	}
-	if zeroCount != 2 && zeroCount != 3 {
+	if initZeroCount != 2 && initZeroCount != 3 {
 		return nil, fmt.Errorf("initial delimiter not found")
 	}
 
-	start := zeroCount + 1
+	start := initZeroCount + 1
+	zeroCount := 0
+	n := 0
+
+	for i := start; i < bl; i++ {
+		switch byts[i] {
+		case 0:
+			zeroCount++
+
+		case 1:
+			if zeroCount == 2 || zeroCount == 3 {
+				n++
+			}
+			zeroCount = 0
 
-	var n int
-	if zeroCount == 2 {
-		n = bytes.Count(byts[start:], []byte{0x00, 0x00, 0x01})
-	} else {
-		n = bytes.Count(byts[start:], []byte{0x00, 0x00, 0x00, 0x01})
+		default:
+			zeroCount = 0
+		}
 	}
 
 	ret := make([][]byte, n+1)
 	pos := 0
-
-	curZeroCount := 0
+	start = initZeroCount + 1
+	zeroCount = 0
 	delimStart := 0
 
 	for i := start; i < bl; i++ {
 		switch byts[i] {
 		case 0:
-			if curZeroCount == 0 {
+			if zeroCount == 0 {
 				delimStart = i
 			}
-			curZeroCount++
+			zeroCount++
 
 		case 1:
-			if curZeroCount == zeroCount {
+			if zeroCount == 2 || zeroCount == 3 {
 				if (delimStart - start) > MaxNALUSize {
 					return nil, fmt.Errorf("NALU size (%d) is too big (maximum is %d)", delimStart-start, MaxNALUSize)
 				}
@@ -65,10 +75,10 @@ outer:
 				pos++
 				start = i + 1
 			}
-			curZeroCount = 0
+			zeroCount = 0
 
 		default:
-			curZeroCount = 0
+			zeroCount = 0
 		}
 	}
 

diff --git a/pkg/h264/annexb_test.go b/pkg/h264/annexb_test.go
@@ -13,23 +13,15 @@ var casesAnnexB = []struct {
 	dec    [][]byte
 }{
 	{
-		"2 zeros, single",
-		[]byte{0x00, 0x00, 0x01, 0xaa, 0xbb},
-		[]byte{0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb},
-		[][]byte{
-			{0xaa, 0xbb},
-		},
-	},
-	{
-		"2 zeros, multiple",
+		"2 zeros",
 		[]byte{
 			0x00, 0x00, 0x01, 0xaa, 0xbb, 0x00, 0x00, 0x01,
 			0xcc, 0xdd, 0x00, 0x00, 0x01, 0xee, 0xff,
 		},
 		[]byte{
-			0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb, 0x00, 0x00,
-			0x00, 0x01, 0xcc, 0xdd, 0x00, 0x00, 0x00, 0x01,
-			0xee, 0xff,
+			0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb,
+			0x00, 0x00, 0x00, 0x01, 0xcc, 0xdd,
+			0x00, 0x00, 0x00, 0x01, 0xee, 0xff,
 		},
 		[][]byte{
 			{0xaa, 0xbb},
@@ -38,29 +30,46 @@ var casesAnnexB = []struct {
 		},
 	},
 	{
-		"3 zeros, single",
-		[]byte{0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb},
-		[]byte{0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb},
+		"3 zeros",
+		[]byte{
+			0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb,
+			0x00, 0x00, 0x00, 0x01, 0xcc, 0xdd,
+			0x00, 0x00, 0x00, 0x01, 0xee, 0xff,
+		},
+		[]byte{
+			0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb,
+			0x00, 0x00, 0x00, 0x01, 0xcc, 0xdd,
+			0x00, 0x00, 0x00, 0x01, 0xee, 0xff,
+		},
 		[][]byte{
 			{0xaa, 0xbb},
+			{0xcc, 0xdd},
+			{0xee, 0xff},
 		},
 	},
 	{
-		"3 zeros, multiple",
+		// used by Apple inside HLS test streams
+		"2 or 3 zeros",
 		[]byte{
-			0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb, 0x00, 0x00,
-			0x00, 0x01, 0xcc, 0xdd, 0x00, 0x00, 0x00, 0x01,
-			0xee, 0xff,
+			0, 0, 0, 1, 9, 240,
+			0, 0, 0, 1, 39, 66, 224, 21, 169, 24, 60, 23, 252, 184, 3, 80, 96, 16, 107, 108, 43, 94, 247, 192, 64,
+			0, 0, 0, 1, 40, 222, 9, 200,
+			0, 0, 1, 6, 0, 7, 131, 236, 119, 0, 0, 0, 0, 1, 3, 0, 64, 128,
+			0, 0, 1, 6, 5, 17, 3, 135, 244, 78, 205, 10, 75, 220, 161, 148, 58, 195, 212, 155, 23, 31, 0, 128,
 		},
 		[]byte{
-			0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb, 0x00, 0x00,
-			0x00, 0x01, 0xcc, 0xdd, 0x00, 0x00, 0x00, 0x01,
-			0xee, 0xff,
+			0, 0, 0, 1, 9, 240,
+			0, 0, 0, 1, 39, 66, 224, 21, 169, 24, 60, 23, 252, 184, 3, 80, 96, 16, 107, 108, 43, 94, 247, 192, 64,
+			0, 0, 0, 1, 40, 222, 9, 200,
+			0, 0, 0, 1, 6, 0, 7, 131, 236, 119, 0, 0, 0, 0, 1, 3, 0, 64, 128,
+			0, 0, 0, 1, 6, 5, 17, 3, 135, 244, 78, 205, 10, 75, 220, 161, 148, 58, 195, 212, 155, 23, 31, 0, 128,
 		},
 		[][]byte{
-			{0xaa, 0xbb},
-			{0xcc, 0xdd},
-			{0xee, 0xff},
+			{9, 240},
+			{39, 66, 224, 21, 169, 24, 60, 23, 252, 184, 3, 80, 96, 16, 107, 108, 43, 94, 247, 192, 64},
+			{40, 222, 9, 200},
+			{6, 0, 7, 131, 236, 119, 0, 0, 0, 0, 1, 3, 0, 64, 128},
+			{6, 5, 17, 3, 135, 244, 78, 205, 10, 75, 220, 161, 148, 58, 195, 212, 155, 23, 31, 0, 128},
 		},
 	},
 }