From 3a10cb25bbddc3471e45fa42e29d8d419e8d9d37 Mon Sep 17 00:00:00 2001
From: Alexey Khit <alexey.khit@gmail.com>
Date: Thu, 15 Sep 2022 06:55:05 +0300
Subject: [PATCH] Fix green video from RTSP H264

---
 pkg/h264/README.md | 11 ++++++++
 pkg/h264/rtp.go    | 70 +++++++++++++++++++++++++---------------------
 2 files changed, 49 insertions(+), 32 deletions(-)

diff --git a/pkg/h264/README.md b/pkg/h264/README.md
index 8b44d478..661741d0 100644
--- a/pkg/h264/README.md
+++ b/pkg/h264/README.md
@@ -1,3 +1,14 @@
+## RTP H264
+
+Camera | NALu
+-------|-----
+EZVIZ C3S          | 7f, 8f, 28:28:28 -> 5t, 28:28:28 -> 1t, 1t, 1t, 1t
+Sonoff GK-200MP2-B | 28:28:28 -> 5t, 1t, 1t, 1t
+Dahua IPC-K42      | 7f, 8f, 28:28:28 -> 5t, 28:28:28 -> 1t, 28:28:28 -> 1t
+FFmpeg copy        | 5t, 1t, 1t, 28:28:28 -> 1t, 28:28:28 -> 1t
+FFmpeg h264        | 24 -> 6:5:5:5:5t, 24 -> 1:1:1:1t, 28:28:28 -> 5f, 28:28:28 -> 5f, 28:28:28 -> 5t
+FFmpeg resize      | 6f, 28:28:28 -> 5f, 28... -> 5t, 24 -> 1:1f, 24 -> 1:1t
+
 ## WebRTC
 
 Video codec	    | Media string | Device
diff --git a/pkg/h264/rtp.go b/pkg/h264/rtp.go
index 389cca1e..e300bf20 100644
--- a/pkg/h264/rtp.go
+++ b/pkg/h264/rtp.go
@@ -22,52 +22,62 @@ func RTPDepay(track *streamer.Track) streamer.WrapperFunc {
 		return func(packet *rtp.Packet) error {
 			//nalUnitType := packet.Payload[0] & 0x1F
 			//fmt.Printf(
-			//	"[RTP] codec: %s, nalu: %2d, size: %6d, ts: %10d, pt: %2d, ssrc: %d\n",
+			//	"[RTP] codec: %s, nalu: %2d, size: %6d, ts: %10d, pt: %2d, ssrc: %d, seq: %d\n",
 			//	track.Codec.Name, nalUnitType, len(packet.Payload), packet.Timestamp,
-			//	packet.PayloadType, packet.SSRC,
+			//	packet.PayloadType, packet.SSRC, packet.SequenceNumber,
 			//)
 
-			// NALu packets can be split in different ways:
-			// - single type 7 and type 8 packets
-			// - join type 7 and type 8 packet (type 24)
-			// - split type 5 on multiple 28 packets
-			// - split type 5 on multiple separate 28 packets
-			units, err := depack.Unmarshal(packet.Payload)
-			if len(units) == 0 || err != nil {
+			data, err := depack.Unmarshal(packet.Payload)
+			if len(data) == 0 || err != nil {
 				return nil
 			}
 
-			for len(units) > 0 {
-				i := int(binary.BigEndian.Uint32(units)) + 4
-				unit := units[:i] // NAL Unit with AVC header
-				units = units[i:]
+			for {
+				unitType := NALUType(data)
+				//fmt.Printf("[H264] nalu: %2d, size: %6d\n", unitType, len(data))
+
+				// multiple 5 and 1 in one payload is OK
+				if unitType != NALUTypeIFrame && unitType != NALUTypePFrame {
+					i := int(binary.BigEndian.Uint32(data)) + 4
+					if i < len(data) {
+						data0 := data[:i] // NAL Unit with AVC header
+						data = data[i:]
+						switch unitType {
+						case NALUTypeSPS:
+							sps = data0
+							continue
+						case NALUTypePPS:
+							pps = data0
+							continue
+						case NALUTypeSEI:
+							// some unnecessary text information
+							continue
+						}
+					}
+				}
 
-				unitType := NALUType(unit)
-				//fmt.Printf("[H264] type: %2d, size: %6d\n", unitType, i)
 				switch unitType {
 				case NALUTypeSPS:
-					//println("new SPS")
-					sps = unit
-					continue
+					sps = data
+					return nil
 				case NALUTypePPS:
-					//println("new PPS")
-					pps = unit
-					continue
+					pps = data
+					return nil
 				case NALUTypeSEI:
 					// some unnecessary text information
-					continue
+					return nil
 				}
 
 				// ffmpeg with `-tune zerolatency` enable option `-x264opts sliced-threads=1`
 				// and every NALU will be sliced to multiple NALUs
 				if !packet.Marker {
-					buffer = append(buffer, unit...)
-					continue
+					buffer = append(buffer, data...)
+					return nil
 				}
 
 				if buffer != nil {
-					buffer = append(buffer, unit...)
-					unit = buffer
+					buffer = append(buffer, data...)
+					data = buffer
 					buffer = nil
 				}
 
@@ -91,13 +101,9 @@ func RTPDepay(track *streamer.Track) streamer.WrapperFunc {
 
 				clone = *packet
 				clone.Version = RTPPacketVersionAVC
-				clone.Payload = unit
-				if err = push(&clone); err != nil {
-					return err
-				}
+				clone.Payload = data
+				return push(&clone)
 			}
-
-			return nil
 		}
 	}
 }