From e4f565f34385d3208526af46d4af62793a25ae61 Mon Sep 17 00:00:00 2001
From: Alexey Khit <alexey.khit@gmail.com>
Date: Tue, 6 Sep 2022 07:00:22 +0300
Subject: [PATCH] Fix H264 in RTSP processing

---
 pkg/h264/rtp.go | 34 ++++++++++++++++------------------
 1 file changed, 16 insertions(+), 18 deletions(-)

diff --git a/pkg/h264/rtp.go b/pkg/h264/rtp.go
index fb1d230e..9cca7bc2 100644
--- a/pkg/h264/rtp.go
+++ b/pkg/h264/rtp.go
@@ -37,32 +37,34 @@ func RTPDepay(track *streamer.Track) streamer.WrapperFunc {
 				return nil
 			}
 
-			for {
+			for len(units) > 0 {
 				i := int(binary.BigEndian.Uint32(units)) + 4
-				unitAVC := units[:i]
+				unit := units[:i] // NAL Unit with AVC header
+				units = units[i:]
 
-				unitType := NALUType(unitAVC)
+				unitType := NALUType(unit)
+				//fmt.Printf("[H264] type: %2d, size: %6d\n", unitType, i)
 				switch unitType {
 				case NALUTypeSPS:
 					//println("new SPS")
-					sps = unitAVC
-					return nil
+					sps = unit
+					continue
 				case NALUTypePPS:
 					//println("new PPS")
-					pps = unitAVC
-					return nil
+					pps = unit
+					continue
 				}
 
 				// 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, unitAVC...)
-					return nil
+					buffer = append(buffer, unit...)
+					continue
 				}
 
 				if buffer != nil {
-					buffer = append(buffer, unitAVC...)
-					unitAVC = buffer
+					buffer = append(buffer, unit...)
+					unit = buffer
 					buffer = nil
 				}
 
@@ -86,17 +88,13 @@ func RTPDepay(track *streamer.Track) streamer.WrapperFunc {
 
 				clone = *packet
 				clone.Version = RTPPacketVersionAVC
-				clone.Payload = unitAVC
+				clone.Payload = unit
 				if err = push(&clone); err != nil {
 					return err
 				}
-
-				if len(units) == i {
-					return nil
-				}
-
-				units = units[i:]
 			}
+
+			return nil
 		}
 	}
 }