From 6dceed64edef9c62cb3a680a54522bf8a20e34c3 Mon Sep 17 00:00:00 2001
From: Alexey Khit <alexey.khit@gmail.com>
Date: Thu, 18 Aug 2022 23:25:09 +0300
Subject: [PATCH] Fix RTMP to WebRTC

---
 pkg/h264/avc.go        | 60 ++++++++++++++++++++++++++++++++++++++++++
 pkg/h264/helper.go     | 14 ----------
 pkg/webrtc/streamer.go |  6 +++--
 3 files changed, 64 insertions(+), 16 deletions(-)
 create mode 100644 pkg/h264/avc.go

diff --git a/pkg/h264/avc.go b/pkg/h264/avc.go
new file mode 100644
index 00000000..cb4be366
--- /dev/null
+++ b/pkg/h264/avc.go
@@ -0,0 +1,60 @@
+package h264
+
+import (
+	"encoding/binary"
+	"github.com/AlexxIT/go2rtc/pkg/streamer"
+	"github.com/pion/rtp"
+)
+
+const PayloadTypeAVC = 255
+
+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)
+	return
+}
+
+func RepairAVC(track *streamer.Track) streamer.WrapperFunc {
+	sps, pps := GetParameterSet(track.Codec.FmtpLine)
+	sps = EncodeAVC(sps)
+	pps = EncodeAVC(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
+			}
+
+			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)
+		}
+	}
+}
diff --git a/pkg/h264/helper.go b/pkg/h264/helper.go
index e4fafd38..da63ce94 100644
--- a/pkg/h264/helper.go
+++ b/pkg/h264/helper.go
@@ -2,7 +2,6 @@ package h264
 
 import (
 	"encoding/base64"
-	"encoding/binary"
 	"github.com/AlexxIT/go2rtc/pkg/streamer"
 	"strings"
 )
@@ -12,25 +11,12 @@ const (
 	NALUTypeIFrame = 5
 	NALUTypeSPS    = 7
 	NALUTypePPS    = 8
-
-	PayloadTypeAVC = 255
 )
 
 func NALUType(b []byte) byte {
 	return b[4] & 0x1F
 }
 
-func EncodeAVC(raw []byte) (avc []byte) {
-	avc = make([]byte, len(raw)+4)
-	binary.BigEndian.PutUint32(avc, uint32(len(raw)))
-	copy(avc[4:], raw)
-	return
-}
-
-func IsAVC(codec *streamer.Codec) bool {
-	return codec.PayloadType == PayloadTypeAVC
-}
-
 func GetParameterSet(fmtp string) (sps, pps []byte) {
 	if fmtp == "" {
 		return
diff --git a/pkg/webrtc/streamer.go b/pkg/webrtc/streamer.go
index 9d7d6340..ef4881b1 100644
--- a/pkg/webrtc/streamer.go
+++ b/pkg/webrtc/streamer.go
@@ -55,10 +55,12 @@ func (c *Conn) AddTrack(media *streamer.Media, track *streamer.Track) *streamer.
 			wrapper := h264.RTPPay(1200)
 			push = wrapper(push)
 
-			if !h264.IsAVC(codec) {
+			if h264.IsAVC(codec) {
+				wrapper = h264.RepairAVC(track)
+			} else {
 				wrapper = h264.RTPDepay(track)
-				push = wrapper(push)
 			}
+			push = wrapper(push)
 		}
 
 		track = track.Bind(push)