Rework RTSP and RTMP processing

pkg/h264/avc.go

@@ -12,49 +12,33 @@ func IsAVC(codec *streamer.Codec) bool {
 	return codec.PayloadType == PayloadTypeAVC
 }
 
-func EncodeAVC(raw []byte) (avc []byte) {
-	avc = make([]byte, len(raw)+4)
-	binary.BigEndian.PutUint32(avc, uint32(len(raw)))
-	copy(avc[4:], raw)
+func EncodeAVC(nals ...[]byte) (avc []byte) {
+	n := 4 * len(nals)
+	for _, nal := range nals {
+		n += len(nal)
+	}
+
+	avc = make([]byte, n)
+
+	var i int
+	for _, nal := range nals {
+		binary.BigEndian.PutUint32(avc[i:], uint32(len(nal)))
+		i += 4 + copy(avc[i+4:], nal)
+	}
+
 	return
 }
 
 func RepairAVC(track *streamer.Track) streamer.WrapperFunc {
 	sps, pps := GetParameterSet(track.Codec.FmtpLine)
-	sps = EncodeAVC(sps)
-	pps = EncodeAVC(pps)
+	ps := EncodeAVC(sps, pps)
 
 	return func(push streamer.WriterFunc) streamer.WriterFunc {
 		return func(packet *rtp.Packet) (err error) {
-			naluType := NALUType(packet.Payload)
-			switch naluType {
-			case NALUTypeSPS:
-				sps = packet.Payload
-				return
-			case NALUTypePPS:
-				pps = packet.Payload
-				return
+			if NALUType(packet.Payload) == NALUTypeIFrame {
+				packet.Payload = Join(ps, packet.Payload)
 			}
-
-			var clone rtp.Packet
-
-			if naluType == NALUTypeIFrame {
-				clone = *packet
-				clone.Payload = sps
-				if err = push(&clone); err != nil {
-					return
-				}
-
-				clone = *packet
-				clone.Payload = pps
-				if err = push(&clone); err != nil {
-					return
-				}
-			}
-
-			clone = *packet
-			clone.Payload = packet.Payload
-			return push(&clone)
+			return push(packet)
 		}
 	}
 }