Merge pull request #272 from thejoeejoee/fix-avc-slice-header

fix(avc/slice): entire NALU consumed as slice header
Eyevinn · Sep 27, 2023 · 2527ed1 · 2527ed1
2 parents 77b91dd + 939329b
commit 2527ed1
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 5 deletions.
diff --git a/avc/slice.go b/avc/slice.go
@@ -208,6 +208,7 @@ func ParseSliceHeader(nalu []byte, spsMap map[uint32]*SPS, ppsMap map[uint32]*PP
 	if sliceType != SLICE_I && sliceType != SLICE_SI {
 		sh.RefPicListModificationL0Flag = r.ReadFlag()
 		if sh.RefPicListModificationL0Flag {
+		refPicListL0Loop:
 			for {
 				sh.ModificationOfPicNumsIDC = uint32(r.ReadExpGolomb())
 				switch sh.ModificationOfPicNumsIDC {
@@ -216,17 +217,18 @@ func ParseSliceHeader(nalu []byte, spsMap map[uint32]*SPS, ppsMap map[uint32]*PP
 				case 2:
 					sh.LongTermPicNum = uint32(r.ReadExpGolomb())
 				case 3:
-					break
+					break refPicListL0Loop
 				}
 				if r.AccError() != nil {
-					break
+					break refPicListL0Loop
 				}
 			}
 		}
 	}
 	if sliceType == SLICE_B {
 		sh.RefPicListModificationL1Flag = r.ReadFlag()
 		if sh.RefPicListModificationL1Flag {
+		refPicListL1Loop:
 			for {
 				sh.ModificationOfPicNumsIDC = uint32(r.ReadExpGolomb())
 				switch sh.ModificationOfPicNumsIDC {
@@ -235,7 +237,10 @@ func ParseSliceHeader(nalu []byte, spsMap map[uint32]*SPS, ppsMap map[uint32]*PP
 				case 2:
 					sh.LongTermPicNum = uint32(r.ReadExpGolomb())
 				case 3:
-					break
+					break refPicListL1Loop
+				}
+				if r.AccError() != nil {
+					break refPicListL1Loop
 				}
 			}
 		}
@@ -299,6 +304,7 @@ func ParseSliceHeader(nalu []byte, spsMap map[uint32]*SPS, ppsMap map[uint32]*PP
 		} else {
 			sh.AdaptiveRefPicMarkingModeFlag = r.ReadFlag()
 			if sh.AdaptiveRefPicMarkingModeFlag {
+			adaptiveRefPicLoop:
 				for {
 					memoryManagementControlOperation := r.ReadExpGolomb()
 					switch memoryManagementControlOperation {
@@ -313,7 +319,10 @@ func ParseSliceHeader(nalu []byte, spsMap map[uint32]*SPS, ppsMap map[uint32]*PP
 					case 4:
 						sh.MaxLongTermFrameIdxPlus1 = uint32(r.ReadExpGolomb())
 					case 0:
-						break
+						break adaptiveRefPicLoop
+					}
+					if r.AccError() != nil {
+						break adaptiveRefPicLoop
 					}
 				}
 			}

diff --git a/avc/slice_test.go b/avc/slice_test.go
@@ -23,7 +23,7 @@ func TestSliceTypeParser(t *testing.T) {
 	}
 }
 
-func TestParseSliceHeader(t *testing.T) {
+func TestParseSliceHeader_BlackFrame(t *testing.T) {
 	wantedHdr := SliceHeader{
 		SliceType:              7,
 		SliceQPDelta:           6,
@@ -65,3 +65,58 @@ func TestParseSliceHeader(t *testing.T) {
 		t.Error(diff)
 	}
 }
+
+func TestParseSliceHeader_TwoFrames(t *testing.T) {
+	wantedIdrHdr := SliceHeader{SliceType: SLICE_I, IDRPicID: 1, SliceQPDelta: 8, Size: 5}
+	wantedNonIdrHdr := SliceHeader{
+		SliceType: SLICE_P, FrameNum: 1, ModificationOfPicNumsIDC: 3, SliceQPDelta: 13,
+		Size: 5, NumRefIdxActiveOverrideFlag: true, RefPicListModificationL0Flag: true,
+	}
+
+	data, err := ioutil.ReadFile("testdata/two-frames.264")
+	if err != nil {
+		t.Error(err)
+	}
+	nalus, err := GetNalusFromSample(data)
+	if err != nil {
+		t.Error(err)
+	}
+	spsMap := make(map[uint32]*SPS, 1)
+	ppsMap := make(map[uint32]*PPS, 1)
+	var gotIdrHdr *SliceHeader
+	var gotNonIdrHdr *SliceHeader
+	for _, nalu := range nalus {
+		switch GetNaluType(nalu[0]) {
+		case NALU_SPS:
+			sps, err := ParseSPSNALUnit(nalu, true)
+			if err != nil {
+				t.Error(err)
+			}
+			spsMap[uint32(sps.ParameterID)] = sps
+		case NALU_PPS:
+			pps, err := ParsePPSNALUnit(nalu, spsMap)
+			if err != nil {
+				t.Error(err)
+			}
+			ppsMap[uint32(pps.PicParameterSetID)] = pps
+		case NALU_IDR:
+			gotIdrHdr, err = ParseSliceHeader(nalu, spsMap, ppsMap)
+			if err != nil {
+				t.Error(err)
+			}
+		case NALU_NON_IDR:
+			gotNonIdrHdr, err = ParseSliceHeader(nalu, spsMap, ppsMap)
+			if err != nil {
+				t.Error(err)
+			}
+		}
+	}
+	if diff := deep.Equal(wantedIdrHdr, *gotIdrHdr); diff != nil {
+		fmt.Printf("Got IDR Slice Header: %+v\n Diff is: ", *gotIdrHdr)
+		t.Error(diff)
+	}
+	if diff := deep.Equal(wantedNonIdrHdr, *gotNonIdrHdr); diff != nil {
+		fmt.Printf("Got NON_IDR Slice Header: %+v\n Diff is: ", *gotNonIdrHdr)
+		t.Error(diff)
+	}
+}
diff --git a/avc/testdata/two-frames.264 b/avc/testdata/two-frames.264