refactor(h264parser): extract AnnexB & AVCC (+fuzzing) (#4)

* fix(h264parser): extract AnnexB * fix(h264parser): extract AVCC * refactor: SplitNALUs use AnnexB & AVCC
RealKeyboardWarrior · Sep 2, 2023 · 469d6d2 · 469d6d2
1 parent 1a8a93b
commit 469d6d2
Show file tree

Hide file tree

Showing 5 changed files with 198 additions and 69 deletions.
diff --git a/codec/h264parser/annexb.go b/codec/h264parser/annexb.go
@@ -0,0 +1,64 @@
+package h264parser
+
+import (
+	"fmt"
+
+	"github.com/RealKeyboardWarrior/joy4/utils/bits/pio"
+)
+
+var ErrInvalidAnnexB = fmt.Errorf("bytestream b was not AnnexB")
+
+func DecodeAnnexB(b []byte) ([][]byte, error) {
+	if len(b) < 4 {
+		return nil, ErrInvalidAnnexB
+	}
+
+	val3 := pio.U24BE(b)
+	val4 := pio.U32BE(b)
+
+	nalus := [][]byte{}
+
+	if val3 == 1 || val4 == 1 {
+		_val3 := val3
+		_val4 := val4
+		start := 0
+		pos := 0
+		for {
+			if start != pos {
+				nalus = append(nalus, b[start:pos])
+			}
+			if _val3 == 1 {
+				pos += 3
+			} else if _val4 == 1 {
+				pos += 4
+			}
+			start = pos
+			if start == len(b) {
+				break
+			}
+			_val3 = 0
+			_val4 = 0
+			for pos < len(b) {
+				if pos+2 < len(b) && b[pos] == 0 {
+					_val3 = pio.U24BE(b[pos:])
+					if _val3 == 0 {
+						if pos+3 < len(b) {
+							_val4 = uint32(b[pos+3])
+							if _val4 == 1 {
+								break
+							}
+						}
+					} else if _val3 == 1 {
+						break
+					}
+					pos++
+				} else {
+					pos++
+				}
+			}
+		}
+		return nalus, nil
+	}
+
+	return nil, ErrInvalidAnnexB
+}
diff --git a/codec/h264parser/annexb_test.go b/codec/h264parser/annexb_test.go
@@ -0,0 +1,42 @@
+package h264parser
+
+import (
+	"encoding/hex"
+	"testing"
+)
+
+func FuzzDecodeAnnexB(f *testing.F) {
+	var (
+		invalidPayload, _ = hex.DecodeString("00000101ca048008fff26423fffa5800000300140e596e979403183257cda0f047bf5755ffed1c853f637fbfbaf607d3fe5ed9073be977bf583a9db48a03a09badae0690f9c98c618844a549e0f6fada8b1cd574484f7f4c6aec7c8ef48c39d8ac8bad3767c9088a0395c8fe3d7ea3bd988194efd3a000ba79cbe3fb7b77dd16155116e30ea9d54367b56e3e2a12132556a3d0c285d6ce86f796231384bfbf6a88f81b284e2dd085f0786ffe443d7b9009a79d188f1d0000b616f0983b3484b31ba157a2f1ae4dde78dd1b12f79cc5df0d7cf215066901090e1599266b571843649e73d80fbaf3c700e7467a79563cc05c3501f0821954c762e4cdfdca87ea3a58")
+		validPayload, _   = hex.DecodeString("00000001223322330000000122332233223300000133000001000001")
+	)
+	f.Add(invalidPayload)
+	f.Add(validPayload)
+	f.Fuzz(func(t *testing.T, payload []byte) {
+		DecodeAnnexB(payload)
+	})
+}
+
+func TestDecodeAnnexB(t *testing.T) {
+	var err error
+	var nalus [][]byte
+
+	annexbFrame, _ := hex.DecodeString("00000001223322330000000122332233223300000133000001000001")
+	nalus, err = DecodeAnnexB(annexbFrame)
+	if err != nil {
+		t.Errorf("did not expect error")
+	}
+	if len(nalus) != 3 {
+		t.Errorf("expected 3 NAL units")
+	}
+
+	if nal := hex.EncodeToString(nalus[0]); nal != "22332233" {
+		t.Errorf("Expected %v", nal)
+	}
+	if nal := hex.EncodeToString(nalus[1]); nal != "223322332233" {
+		t.Errorf("Expected %v", nal)
+	}
+	if nal := hex.EncodeToString(nalus[2]); nal != "33" {
+		t.Errorf("Expected %v", nal)
+	}
+}
diff --git a/codec/h264parser/avcc.go b/codec/h264parser/avcc.go
@@ -0,0 +1,44 @@
+package h264parser
+
+import (
+	"fmt"
+
+	"github.com/RealKeyboardWarrior/joy4/utils/bits/pio"
+)
+
+var ErrInvalidAVCC = fmt.Errorf("bytestream b was not AVCC")
+
+func DecodeAVCC(b []byte) ([][]byte, error) {
+	if len(b) < 4 {
+		return nil, ErrInvalidAVCC
+	}
+
+	val4 := pio.U32BE(b)
+	if val4 > uint32(len(b)-4) {
+		return nil, ErrInvalidAVCC
+	}
+
+	_val4 := val4
+	_b := b[4:]
+	nalus := [][]byte{}
+	for {
+		nalus = append(nalus, _b[:_val4])
+		if _val4 > uint32(len(_b)) {
+			break
+		}
+		_b = _b[_val4:]
+		if len(_b) < 4 {
+			break
+		}
+		_val4 = pio.U32BE(_b)
+		_b = _b[4:]
+		if _val4 > uint32(len(_b)) {
+			break
+		}
+	}
+	if len(_b) == 0 {
+		return nalus, nil
+	}
+
+	return nil, ErrInvalidAVCC
+}
diff --git a/codec/h264parser/avcc_test.go b/codec/h264parser/avcc_test.go
@@ -0,0 +1,41 @@
+package h264parser
+
+import (
+	"encoding/hex"
+	"testing"
+)
+
+func FuzzDecodeAVCC(f *testing.F) {
+	var (
+		invalidPayload, _ = hex.DecodeString("00000101ca048008fff26423fffa5800000300140e596e979403183257cda0f047bf5755ffed1c853f637fbfbaf607d3fe5ed9073be977bf583a9db48a03a09badae0690f9c98c618844a549e0f6fada8b1cd574484f7f4c6aec7c8ef48c39d8ac8bad3767c9088a0395c8fe3d7ea3bd988194efd3a000ba79cbe3fb7b77dd16155116e30ea9d54367b56e3e2a12132556a3d0c285d6ce86f796231384bfbf6a88f81b284e2dd085f0786ffe443d7b9009a79d188f1d0000b616f0983b3484b31ba157a2f1ae4dde78dd1b12f79cc5df0d7cf215066901090e1599266b571843649e73d80fbaf3c700e7467a79563cc05c3501f0821954c762e4cdfdca87ea3a58")
+		validPayload, _   = hex.DecodeString("00000008aabbccaabbccaabb00000001aa")
+	)
+	f.Add(invalidPayload)
+	f.Add(validPayload)
+	f.Fuzz(func(t *testing.T, payload []byte) {
+		DecodeAVCC(payload)
+	})
+}
+
+func TestDecodeAVCC(t *testing.T) {
+	var err error
+	var nalus [][]byte
+
+	avccFrame, _ := hex.DecodeString(
+		"00000008aabbccaabbccaabb00000001aa",
+	)
+	nalus, err = DecodeAVCC(avccFrame)
+	if err != nil {
+		t.Errorf("did not expect error")
+	}
+
+	if len(nalus) != 2 {
+		t.Errorf("expected 2 NAL unit")
+	}
+	if nal := hex.EncodeToString(nalus[0]); nal != "aabbccaabbccaabb" {
+		t.Errorf("Expected %v", nal)
+	}
+	if nal := hex.EncodeToString(nalus[1]); nal != "aa" {
+		t.Errorf("Expected %v", nal)
+	}
+}
diff --git a/codec/h264parser/parser.go b/codec/h264parser/parser.go
@@ -214,81 +214,19 @@ const (
 	NALU_ANNEXB
 )
 
-func SplitNALUs(b []byte) (nalus [][]byte, typ int) {
+func SplitNALUs(b []byte) ([][]byte, int) {
 	if len(b) < 4 {
 		return [][]byte{b}, NALU_RAW
 	}
 
-	val3 := pio.U24BE(b)
-	val4 := pio.U32BE(b)
-
-	// maybe AVCC
-	if val4 <= uint32(len(b)) {
-		_val4 := val4
-		_b := b[4:]
-		nalus := [][]byte{}
-		for {
-			nalus = append(nalus, _b[:_val4])
-			if _val4 > uint32(len(_b)) {
-				break
-			}
-			_b = _b[_val4:]
-			if len(_b) < 4 {
-				break
-			}
-			_val4 = pio.U32BE(_b)
-			_b = _b[4:]
-			if _val4 > uint32(len(_b)) {
-				break
-			}
-		}
-		if len(_b) == 0 {
-			return nalus, NALU_AVCC
-		}
+	nalus, err := DecodeAVCC(b)
+	if nalus != nil && err == nil {
+		return nalus, NALU_AVCC
 	}
 
-	// is Annex B
-	if val3 == 1 || val4 == 1 {
-		_val3 := val3
-		_val4 := val4
-		start := 0
-		pos := 0
-		for {
-			if start != pos {
-				nalus = append(nalus, b[start:pos])
-			}
-			if _val3 == 1 {
-				pos += 3
-			} else if _val4 == 1 {
-				pos += 4
-			}
-			start = pos
-			if start == len(b) {
-				break
-			}
-			_val3 = 0
-			_val4 = 0
-			for pos < len(b) {
-				if pos+2 < len(b) && b[pos] == 0 {
-					_val3 = pio.U24BE(b[pos:])
-					if _val3 == 0 {
-						if pos+3 < len(b) {
-							_val4 = uint32(b[pos+3])
-							if _val4 == 1 {
-								break
-							}
-						}
-					} else if _val3 == 1 {
-						break
-					}
-					pos++
-				} else {
-					pos++
-				}
-			}
-		}
-		typ = NALU_ANNEXB
-		return
+	nalus, err = DecodeAnnexB(b)
+	if nalus != nil && err == nil {
+		return nalus, NALU_ANNEXB
 	}
 
 	return [][]byte{b}, NALU_RAW