Rework RTSP and RTMP processing

pkg/h264/rtp.go

@@ -1,7 +1,6 @@
 package h264
 
 import (
-	"encoding/binary"
 	"github.com/AlexxIT/go2rtc/pkg/streamer"
 	"github.com/pion/rtp"
 	"github.com/pion/rtp/codecs"
@@ -13,8 +12,7 @@ func RTPDepay(track *streamer.Track) streamer.WrapperFunc {
 	depack := &codecs.H264Packet{IsAVC: true}
 
 	sps, pps := GetParameterSet(track.Codec.FmtpLine)
-	sps = EncodeAVC(sps)
-	pps = EncodeAVC(pps)
+	ps := EncodeAVC(sps, pps)
 
 	var buffer []byte
 
@@ -27,83 +25,32 @@ func RTPDepay(track *streamer.Track) streamer.WrapperFunc {
 			//	packet.PayloadType, packet.SSRC, packet.SequenceNumber,
 			//)
 
-			data, err := depack.Unmarshal(packet.Payload)
-			if len(data) == 0 || err != nil {
+			payload, err := depack.Unmarshal(packet.Payload)
+			if len(payload) == 0 || err != nil {
 				return nil
 			}
 
-			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
-						}
-					}
-				}
-
-				switch unitType {
-				case NALUTypeSPS:
-					sps = data
-					return nil
-				case NALUTypePPS:
-					pps = data
-					return nil
-				case NALUTypeSEI:
-					// some unnecessary text information
-					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, data...)
-					return nil
-				}
-
-				if buffer != nil {
-					buffer = append(buffer, data...)
-					data = buffer
-					buffer = nil
-				}
-
-				var clone rtp.Packet
-
-				if unitType == NALUTypeIFrame {
-					clone = *packet
-					clone.Version = RTPPacketVersionAVC
-					clone.Payload = sps
-					if err = push(&clone); err != nil {
-						return err
-					}
-
-					clone = *packet
-					clone.Version = RTPPacketVersionAVC
-					clone.Payload = pps
-					if err = push(&clone); err != nil {
-						return err
-					}
-				}
-
-				clone = *packet
-				clone.Version = RTPPacketVersionAVC
-				clone.Payload = data
-				return push(&clone)
+			// 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, payload...)
+				return nil
 			}
+
+			if buffer != nil {
+				payload = append(buffer, payload...)
+				buffer = nil
+			}
+
+			switch NALUType(payload) {
+			case NALUTypeIFrame:
+				payload = Join(ps, payload)
+			}
+
+			clone := *packet
+			clone.Version = RTPPacketVersionAVC
+			clone.Payload = payload
+			return push(&clone)
 		}
 	}
 }
@@ -117,11 +64,12 @@ func RTPPay(mtu uint16) streamer.WrapperFunc {
 		return func(packet *rtp.Packet) error {
 			if packet.Version == RTPPacketVersionAVC {
 				payloads := payloader.Payload(mtu, packet.Payload)
+				last := len(payloads) - 1
 				for i, payload := range payloads {
 					clone := rtp.Packet{
 						Header: rtp.Header{
 							Version: 2,
-							Marker:  i == len(payloads)-1,
+							Marker:  i == last,
 							//PayloadType:    packet.PayloadType,
 							SequenceNumber: sequencer.NextSequenceNumber(),
 							Timestamp:      packet.Timestamp,