From 4fd5ca2f407fdbe6dfe1d9cd243b22e4f2d6142b Mon Sep 17 00:00:00 2001
From: Diallo Han <diallo@hpcnt.com>
Date: Tue, 14 Nov 2023 12:18:59 +0900
Subject: [PATCH 1/3] add edgecase for DTSExtractor

---
 pkg/codecs/h264/dts_extractor.go      |  2 +-
 pkg/codecs/h264/dts_extractor_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/pkg/codecs/h264/dts_extractor.go b/pkg/codecs/h264/dts_extractor.go
index c590791..013f5e0 100644
--- a/pkg/codecs/h264/dts_extractor.go
+++ b/pkg/codecs/h264/dts_extractor.go
@@ -17,7 +17,7 @@ func getPictureOrderCount(buf []byte, sps *SPS) (uint32, error) {
 		return 0, fmt.Errorf("not enough bits")
 	}
 
-	buf = EmulationPreventionRemove(buf[:6])
+	buf = EmulationPreventionRemove(buf)
 
 	buf = buf[1:]
 	pos := 0
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()

From 1da01f537c34cb3cae1658bcd33587069f22a066 Mon Sep 17 00:00:00 2001
From: "Diallo.H" <diallo@hpcnt.com>
Date: Fri, 17 Nov 2023 19:32:16 +0900
Subject: [PATCH 2/3] apply review

---
 pkg/codecs/h264/dts_extractor.go | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/pkg/codecs/h264/dts_extractor.go b/pkg/codecs/h264/dts_extractor.go
index 013f5e0..00b2d05 100644
--- a/pkg/codecs/h264/dts_extractor.go
+++ b/pkg/codecs/h264/dts_extractor.go
@@ -3,6 +3,7 @@ package h264
 import (
 	"bytes"
 	"fmt"
+	"math"
 	"time"
 
 	"github.com/bluenviron/mediacommon/pkg/bits"
@@ -10,6 +11,10 @@ import (
 
 const (
 	maxReorderedFrames = 10
+	// (1 + 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) ) * 3 =
+	// (1 + 3 * max_size(golomb) + (max_size(Log2MaxFrameNumMinus4) + 4) / 8 + (max_size(Log2MaxPicOrderCntLsbMinus4) + 4) / 8 =
+	// (1 + 3 * 4 + 2 + 2) * 3 = 51
+	maxBytesToGetPOC = 51
 )
 
 func getPictureOrderCount(buf []byte, sps *SPS) (uint32, error) {
@@ -17,7 +22,8 @@ func getPictureOrderCount(buf []byte, sps *SPS) (uint32, error) {
 		return 0, fmt.Errorf("not enough bits")
 	}
 
-	buf = EmulationPreventionRemove(buf)
+	trimBytes := int(math.Min(float64(len(buf)), float64(maxBytesToGetPOC)))
+	buf = EmulationPreventionRemove(buf[:trimBytes])
 
 	buf = buf[1:]
 	pos := 0

From 7cf968f98e12b5a3f904ce3e9a466f2b9323c737 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Fri, 17 Nov 2023 12:23:02 +0100
Subject: [PATCH 3/3] improve performance

---
 pkg/codecs/h264/dts_extractor.go | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/pkg/codecs/h264/dts_extractor.go b/pkg/codecs/h264/dts_extractor.go
index 00b2d05..ff81261 100644
--- a/pkg/codecs/h264/dts_extractor.go
+++ b/pkg/codecs/h264/dts_extractor.go
@@ -3,7 +3,6 @@ package h264
 import (
 	"bytes"
 	"fmt"
-	"math"
 	"time"
 
 	"github.com/bluenviron/mediacommon/pkg/bits"
@@ -11,21 +10,24 @@ import (
 
 const (
 	maxReorderedFrames = 10
-	// (1 + 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) ) * 3 =
-	// (1 + 3 * max_size(golomb) + (max_size(Log2MaxFrameNumMinus4) + 4) / 8 + (max_size(Log2MaxPicOrderCntLsbMinus4) + 4) / 8 =
-	// (1 + 3 * 4 + 2 + 2) * 3 = 51
-	maxBytesToGetPOC = 51
+	/*
+		(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)
 
-	trimBytes := int(math.Min(float64(len(buf)), float64(maxBytesToGetPOC)))
-	buf = EmulationPreventionRemove(buf[:trimBytes])
+	if lb > maxBytesToGetPOC {
+		lb = maxBytesToGetPOC
+	}
 
-	buf = buf[1:]
+	buf = EmulationPreventionRemove(buf[:lb])
 	pos := 0
 
 	_, err := bits.ReadGolombUnsigned(buf, &pos) // first_mb_in_slice