From 8aa719eda5b7806a74ccbb9031ecbf93919f7d6f Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Sun, 9 Oct 2022 23:39:08 +0200
Subject: [PATCH] prepend SPS and PPS once when remuxing H264 streams

This avoids missing frames when reading with RTSP and the original
stream has NALU groups with multiple IDRs in it.
---
 internal/core/rtmp_server_test.go | 19 ++++++++++++-------
 internal/core/streamtrack_h264.go | 27 +++++++++++++++------------
 2 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/internal/core/rtmp_server_test.go b/internal/core/rtmp_server_test.go
index 705d0f886bf..358378365a6 100644
--- a/internal/core/rtmp_server_test.go
+++ b/internal/core/rtmp_server_test.go
@@ -116,7 +116,10 @@ func TestRTMPServerPublishRead(t *testing.T) {
 				MessageStreamID: 0x1000000,
 				IsKeyFrame:      true,
 				H264Type:        flvio.AVC_NALU,
-				Payload:         []byte{0x00, 0x00, 0x00, 0x04, 0x05, 0x02, 0x03, 0x04},
+				Payload: []byte{
+					0x00, 0x00, 0x00, 0x04, 0x05, 0x02, 0x03, 0x04, // IDR 1
+					0x00, 0x00, 0x00, 0x04, 0x05, 0x02, 0x03, 0x04, // IDR 2
+				},
 			})
 			require.NoError(t, err)
 
@@ -128,12 +131,14 @@ func TestRTMPServerPublishRead(t *testing.T) {
 				IsKeyFrame:      true,
 				H264Type:        flvio.AVC_NALU,
 				Payload: []byte{
-					0x00, 0x00, 0x00, 0x19, 0x67, 0x42, 0xc0, 0x28,
-					0xd9, 0x00, 0x78, 0x02, 0x27, 0xe5, 0x84, 0x00,
-					0x00, 0x03, 0x00, 0x04, 0x00, 0x00, 0x03, 0x00,
-					0xf0, 0x3c, 0x60, 0xc9, 0x20, 0x00, 0x00, 0x00,
-					0x04, 0x08, 0x06, 0x07, 0x08, 0x00, 0x00, 0x00,
-					0x04, 0x05, 0x02, 0x03, 0x04,
+					0x00, 0x00, 0x00, 0x19, // SPS
+					0x67, 0x42, 0xc0, 0x28, 0xd9, 0x00, 0x78, 0x02,
+					0x27, 0xe5, 0x84, 0x00, 0x00, 0x03, 0x00, 0x04,
+					0x00, 0x00, 0x03, 0x00, 0xf0, 0x3c, 0x60, 0xc9,
+					0x20,
+					0x00, 0x00, 0x00, 0x04, 0x08, 0x06, 0x07, 0x08, // PPS
+					0x00, 0x00, 0x00, 0x04, 0x05, 0x02, 0x03, 0x04, // IDR 1
+					0x00, 0x00, 0x00, 0x04, 0x05, 0x02, 0x03, 0x04, // IDR 2
 				},
 			}, msg1)
 		})
diff --git a/internal/core/streamtrack_h264.go b/internal/core/streamtrack_h264.go
index 4deb1f28148..ee02166ae3b 100644
--- a/internal/core/streamtrack_h264.go
+++ b/internal/core/streamtrack_h264.go
@@ -51,10 +51,10 @@ func (t *streamTrackH264) updateTrackParameters(nalus [][]byte) {
 	}
 }
 
-// remux is needed to
-// - fix corrupted streams
-// - make streams compatible with all protocols
+// remux is needed to fix corrupted streams and make streams
+// compatible with all protocols.
 func (t *streamTrackH264) remuxNALUs(nalus [][]byte) [][]byte {
+	addSPSPPS := false
 	n := 0
 	for _, nalu := range nalus {
 		typ := h264.NALUType(nalu[0] & 0x1F)
@@ -64,7 +64,11 @@ func (t *streamTrackH264) remuxNALUs(nalus [][]byte) [][]byte {
 		case h264.NALUTypeAccessUnitDelimiter:
 			continue
 		case h264.NALUTypeIDR:
-			n += 2
+			// prepend SPS and PPS to the group if there's at least an IDR
+			if !addSPSPPS {
+				addSPSPPS = true
+				n += 2
+			}
 		}
 		n++
 	}
@@ -76,23 +80,22 @@ func (t *streamTrackH264) remuxNALUs(nalus [][]byte) [][]byte {
 	filteredNALUs := make([][]byte, n)
 	i := 0
 
+	if addSPSPPS {
+		filteredNALUs[0] = t.track.SafeSPS()
+		filteredNALUs[1] = t.track.SafePPS()
+		i = 2
+	}
+
 	for _, nalu := range nalus {
 		typ := h264.NALUType(nalu[0] & 0x1F)
 		switch typ {
 		case h264.NALUTypeSPS, h264.NALUTypePPS:
-			// remove since they're automatically added before every IDR
+			// remove since they're automatically added
 			continue
 
 		case h264.NALUTypeAccessUnitDelimiter:
 			// remove since it is not needed
 			continue
-
-		case h264.NALUTypeIDR:
-			// add SPS and PPS before every IDR
-			filteredNALUs[i] = t.track.SafeSPS()
-			i++
-			filteredNALUs[i] = t.track.SafePPS()
-			i++
 		}
 
 		filteredNALUs[i] = nalu