From ccdd6ed490874d6e7324308b40c087558addc2e6 Mon Sep 17 00:00:00 2001
From: Alexey Khit <alexey.khit@gmail.com>
Date: Sat, 2 Sep 2023 06:30:23 +0300
Subject: [PATCH] Update SPS parser

---
 pkg/h264/h264_test.go | 32 +++++++++++++++++++++++
 pkg/h264/sps.go       | 59 ++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 85 insertions(+), 6 deletions(-)

diff --git a/pkg/h264/h264_test.go b/pkg/h264/h264_test.go
index cf159c32..c366c032 100644
--- a/pkg/h264/h264_test.go
+++ b/pkg/h264/h264_test.go
@@ -38,6 +38,38 @@ func TestDecodeSPS(t *testing.T) {
 	sps = DecodeSPS(b)
 	require.Equal(t, uint16(1920), sps.Width())
 	require.Equal(t, uint16(1080), sps.Height())
+
+	s = "Z01AMqaAKAC1kAA=" // Dahua
+	b, err = base64.StdEncoding.DecodeString(s)
+	require.Nil(t, err)
+
+	sps = DecodeSPS(b)
+	require.Equal(t, uint16(2560), sps.Width())
+	require.Equal(t, uint16(1440), sps.Height())
+
+	s = "Z2QAM6wVFKAoAPGQ" // Reolink
+	b, err = base64.StdEncoding.DecodeString(s)
+	require.Nil(t, err)
+
+	sps = DecodeSPS(b)
+	require.Equal(t, uint16(2560), sps.Width())
+	require.Equal(t, uint16(1920), sps.Height())
+
+	s = "Z2QAKKwa0AoAt03AQEBQAAADABAAAAMB6PFCKg==" // TP-Link
+	b, err = base64.StdEncoding.DecodeString(s)
+	require.Nil(t, err)
+
+	sps = DecodeSPS(b)
+	require.Equal(t, uint16(1280), sps.Width())
+	require.Equal(t, uint16(720), sps.Height())
+
+	s = "Z2QAFqwa0BQF/yzcBAQFAAADAAEAAAMAHo8UIqA=" // TP-Link sub
+	b, err = base64.StdEncoding.DecodeString(s)
+	require.Nil(t, err)
+
+	sps = DecodeSPS(b)
+	require.Equal(t, uint16(640), sps.Width())
+	require.Equal(t, uint16(360), sps.Height())
 }
 
 func TestGetProfileLevelID(t *testing.T) {
diff --git a/pkg/h264/sps.go b/pkg/h264/sps.go
index f34fb8cb..bb5db648 100644
--- a/pkg/h264/sps.go
+++ b/pkg/h264/sps.go
@@ -44,9 +44,28 @@ type SPS struct {
 
 	vui_parameters_present_flag    byte
 	aspect_ratio_info_present_flag byte
-	aspect_ratio_idc               uint32
-	sar_width                      uint32
-	sar_height                     uint32
+	aspect_ratio_idc               byte
+	sar_width                      uint16
+	sar_height                     uint16
+
+	overscan_info_present_flag byte
+	overscan_appropriate_flag  byte
+
+	video_signal_type_present_flag byte
+	video_format                   uint8
+	video_full_range_flag          byte
+
+	colour_description_present_flag byte
+	colour_description              uint32
+
+	chroma_loc_info_present_flag        byte
+	chroma_sample_loc_type_top_field    uint32
+	chroma_sample_loc_type_bottom_field uint32
+
+	timing_info_present_flag byte
+	num_units_in_tick        uint32
+	time_scale               uint32
+	fixed_frame_rate_flag    byte
 }
 
 func (s *SPS) Width() uint16 {
@@ -146,13 +165,41 @@ func DecodeSPS(sps []byte) *SPS {
 	if s.vui_parameters_present_flag != 0 {
 		s.aspect_ratio_info_present_flag = r.ReadBit()
 		if s.aspect_ratio_info_present_flag != 0 {
-			s.aspect_ratio_idc = r.ReadBits(8)
+			s.aspect_ratio_idc = r.ReadByte()
 			if s.aspect_ratio_idc == 255 {
-				s.sar_width = r.ReadBits(16)
-				s.sar_height = r.ReadBits(16)
+				s.sar_width = r.ReadUint16()
+				s.sar_height = r.ReadUint16()
 			}
 		}
 
+		s.overscan_info_present_flag = r.ReadBit()
+		if s.overscan_info_present_flag != 0 {
+			s.overscan_appropriate_flag = r.ReadBit()
+		}
+
+		s.video_signal_type_present_flag = r.ReadBit()
+		if s.video_signal_type_present_flag != 0 {
+			s.video_format = r.ReadBits8(3)
+			s.video_full_range_flag = r.ReadBit()
+
+			s.colour_description_present_flag = r.ReadBit()
+			if s.colour_description_present_flag != 0 {
+				s.colour_description = r.ReadUint24()
+			}
+		}
+
+		s.chroma_loc_info_present_flag = r.ReadBit()
+		if s.chroma_loc_info_present_flag != 0 {
+			s.chroma_sample_loc_type_top_field = r.ReadUEGolomb()
+			s.chroma_sample_loc_type_bottom_field = r.ReadUEGolomb()
+		}
+
+		s.timing_info_present_flag = r.ReadBit()
+		if s.timing_info_present_flag != 0 {
+			s.num_units_in_tick = r.ReadUint32()
+			s.time_scale = r.ReadUint32()
+			s.fixed_frame_rate_flag = r.ReadBit()
+		}
 		//...
 	}