From 151a3cb5eaa0949126d925a26eedba7a3f26f9c2 Mon Sep 17 00:00:00 2001
From: Han Gyoung-Su <hsnks100@gmail.com>
Date: Fri, 17 Nov 2023 20:31:30 +0900
Subject: [PATCH] h264: fix DTS extraction in case of long NALUs (#98)

* add edgecase for DTSExtractor

* apply review

* improve performance

---------

Co-authored-by: aler9 <46489434+aler9@users.noreply.github.com>
---
 pkg/codecs/h264/dts_extractor.go      | 18 +++++++++++++-----
 pkg/codecs/h264/dts_extractor_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/pkg/codecs/h264/dts_extractor.go b/pkg/codecs/h264/dts_extractor.go
index c590791..ff81261 100644
--- a/pkg/codecs/h264/dts_extractor.go
+++ b/pkg/codecs/h264/dts_extractor.go
@@ -10,16 +10,24 @@ import (
 
 const (
 	maxReorderedFrames = 10
+	/*
+		(max_size(first_mb_in_slice) + max_size(slice_type) + max_size(pic_parameter_set_id) +
+		max_size(frame_num) + max_size(pic_order_cnt_lsb)) * 4 / 3 =
+		(3 * max_size(golomb) + (max(Log2MaxFrameNumMinus4) + 4) / 8 + (max(Log2MaxPicOrderCntLsbMinus4) + 4) / 8) * 4 / 3 =
+		(3 * 4 + 2 + 2) * 4 / 3 = 22
+	*/
+	maxBytesToGetPOC = 22
 )
 
 func getPictureOrderCount(buf []byte, sps *SPS) (uint32, error) {
-	if len(buf) < 6 {
-		return 0, fmt.Errorf("not enough bits")
-	}
+	buf = buf[1:]
+	lb := len(buf)
 
-	buf = EmulationPreventionRemove(buf[:6])
+	if lb > maxBytesToGetPOC {
+		lb = maxBytesToGetPOC
+	}
 
-	buf = buf[1:]
+	buf = EmulationPreventionRemove(buf[:lb])
 	pos := 0
 
 	_, err := bits.ReadGolombUnsigned(buf, &pos) // first_mb_in_slice
diff --git a/pkg/codecs/h264/dts_extractor_test.go b/pkg/codecs/h264/dts_extractor_test.go
index 19e59ed..5be41e2 100644
--- a/pkg/codecs/h264/dts_extractor_test.go
+++ b/pkg/codecs/h264/dts_extractor_test.go
@@ -278,6 +278,33 @@ func TestDTSExtractor(t *testing.T) {
 				},
 			},
 		},
+		{
+			"Log2MaxPicOrderCntLsbMinus4 = 12",
+			[]sample{
+				{
+					[][]byte{
+						{ // SPS
+							0x67, 0x42, 0xc0, 0x1e, 0x8c, 0x8d, 0x40, 0x50, 0x17, 0xfc, 0xb0, 0x0f, 0x08, 0x84, 0x6a,
+						},
+						{ // PPS
+							0x68, 0xce, 0x3c, 0x80,
+						},
+						{ // IDR
+							0x65, 0x88, 0x80, 0x14, 0x3, 0xff, 0xde, 0x8, 0xe4, 0x74,
+						},
+					},
+					0,
+					0,
+				},
+				{
+					[][]byte{
+						{0x61, 0x00, 0xf0, 0xe0, 0x00, 0x40, 0x00, 0xbe, 0x47, 0x9b}, // non-IDR
+					},
+					40 * time.Millisecond,
+					40 * time.Millisecond,
+				},
+			},
+		},
 	} {
 		t.Run(ca.name, func(t *testing.T) {
 			ex := NewDTSExtractor()