From 32651c74ab7a0dc72eb4f4a10d46248a85d1f483 Mon Sep 17 00:00:00 2001
From: Alexey Khit <alexey.khit@gmail.com>
Date: Tue, 8 Nov 2022 01:13:38 +0300
Subject: [PATCH] Fix frame.mp4 support

---
 cmd/mp4/mp4.go      | 14 +++-----
 pkg/mp4/keyframe.go | 85 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 10 deletions(-)
 create mode 100644 pkg/mp4/keyframe.go

diff --git a/cmd/mp4/mp4.go b/cmd/mp4/mp4.go
index 43272463..55535a16 100644
--- a/cmd/mp4/mp4.go
+++ b/cmd/mp4/mp4.go
@@ -35,7 +35,7 @@ func handlerKeyframe(w http.ResponseWriter, r *http.Request) {
 
 	exit := make(chan []byte)
 
-	cons := &mp4.Consumer{}
+	cons := &mp4.Keyframe{}
 	cons.Listen(func(msg interface{}) {
 		if data, ok := msg.([]byte); ok && exit != nil {
 			exit <- data
@@ -48,19 +48,13 @@ func handlerKeyframe(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	defer stream.RemoveConsumer(cons)
+	data := <-exit
 
-	w.Header().Set("Content-Type", cons.MimeType())
-
-	data, err := cons.Init()
-	if err != nil {
-		log.Error().Err(err).Caller().Send()
-		return
-	}
-	data = append(data, <-exit...)
+	stream.RemoveConsumer(cons)
 
 	// Apple Safari won't show frame without length
 	w.Header().Set("Content-Length", strconv.Itoa(len(data)))
+	w.Header().Set("Content-Type", cons.MimeType)
 
 	if _, err := w.Write(data); err != nil {
 		log.Error().Err(err).Caller().Send()
diff --git a/pkg/mp4/keyframe.go b/pkg/mp4/keyframe.go
new file mode 100644
index 00000000..fd9c97f4
--- /dev/null
+++ b/pkg/mp4/keyframe.go
@@ -0,0 +1,85 @@
+package mp4
+
+import (
+	"github.com/AlexxIT/go2rtc/pkg/h264"
+	"github.com/AlexxIT/go2rtc/pkg/h265"
+	"github.com/AlexxIT/go2rtc/pkg/streamer"
+	"github.com/pion/rtp"
+)
+
+type Keyframe struct {
+	streamer.Element
+
+	MimeType string
+}
+
+func (c *Keyframe) GetMedias() []*streamer.Media {
+	return []*streamer.Media{
+		{
+			Kind:      streamer.KindVideo,
+			Direction: streamer.DirectionRecvonly,
+			Codecs: []*streamer.Codec{
+				{Name: streamer.CodecH264},
+				{Name: streamer.CodecH265},
+			},
+		},
+	}
+}
+
+func (c *Keyframe) AddTrack(media *streamer.Media, track *streamer.Track) *streamer.Track {
+	muxer := &Muxer{}
+
+	codecs := []*streamer.Codec{track.Codec}
+
+	init, err := muxer.GetInit(codecs)
+	if err != nil {
+		return nil
+	}
+
+	c.MimeType = muxer.MimeType(codecs)
+
+	switch track.Codec.Name {
+	case streamer.CodecH264:
+		push := func(packet *rtp.Packet) error {
+			if !h264.IsKeyframe(packet.Payload) {
+				return nil
+			}
+
+			buf := muxer.Marshal(0, packet)
+			c.Fire(append(init, buf...))
+
+			return nil
+		}
+
+		var wrapper streamer.WrapperFunc
+		if track.Codec.IsMP4() {
+			wrapper = h264.RepairAVC(track)
+		} else {
+			wrapper = h264.RTPDepay(track)
+		}
+		push = wrapper(push)
+
+		return track.Bind(push)
+
+	case streamer.CodecH265:
+		push := func(packet *rtp.Packet) error {
+			if !h265.IsKeyframe(packet.Payload) {
+				return nil
+			}
+
+			buf := muxer.Marshal(0, packet)
+			c.Fire(append(init, buf...))
+
+			return nil
+		}
+
+		if !track.Codec.IsMP4() {
+			wrapper := h265.RTPDepay(track)
+			push = wrapper(push)
+		}
+
+		return track.Bind(push)
+	}
+
+	panic("unsupported codec")
+}