From e049a17216c997470dc967c974e394774b886101 Mon Sep 17 00:00:00 2001
From: Alexey Khit <alexey.khit@gmail.com>
Date: Tue, 13 Sep 2022 14:38:54 +0300
Subject: [PATCH] Adds error handler for mp4 init

---
 cmd/mp4/mp4.go      | 19 ++++++++++++++-----
 cmd/mp4/mse.go      |  9 ++++++++-
 pkg/mp4/consumer.go |  2 +-
 pkg/mp4/muxer.go    | 10 +++++++---
 4 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/cmd/mp4/mp4.go b/cmd/mp4/mp4.go
index c75fb5e8..a43907f8 100644
--- a/cmd/mp4/mp4.go
+++ b/cmd/mp4/mp4.go
@@ -48,13 +48,17 @@ func handlerKeyframe(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	defer stream.RemoveConsumer(cons)
+
 	w.Header().Set("Content-Type", cons.MimeType())
 
-	data := cons.Init()
+	data, err := cons.Init()
+	if err != nil {
+		log.Error().Err(err).Msg("[api.keyframe] init")
+		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)))
 
@@ -97,8 +101,13 @@ func handlerMP4(w http.ResponseWriter, r *http.Request) {
 
 	w.Header().Set("Content-Type", cons.MimeType())
 
-	data := cons.Init()
-	if _, err := w.Write(data); err != nil {
+	data, err := cons.Init()
+	if err != nil {
+		log.Error().Err(err).Msg("[api.mp4] init")
+		return
+	}
+
+	if _, err = w.Write(data); err != nil {
 		log.Error().Err(err).Msg("[api.mp4] write")
 		return
 	}
diff --git a/cmd/mp4/mse.go b/cmd/mp4/mse.go
index 47039459..d632ed12 100644
--- a/cmd/mp4/mse.go
+++ b/cmd/mp4/mse.go
@@ -41,5 +41,12 @@ func handlerWS(ctx *api.Context, msg *streamer.Message) {
 		Type: MsgTypeMSE, Value: cons.MimeType(),
 	})
 
-	ctx.Write(cons.Init())
+	data, err := cons.Init()
+	if err != nil {
+		log.Warn().Err(err).Msg("[api.mse] init")
+		ctx.Error(err)
+		return
+	}
+
+	ctx.Write(data)
 }
diff --git a/pkg/mp4/consumer.go b/pkg/mp4/consumer.go
index d51c4ebe..61a3d961 100644
--- a/pkg/mp4/consumer.go
+++ b/pkg/mp4/consumer.go
@@ -86,7 +86,7 @@ func (c *Consumer) MimeType() string {
 	return c.muxer.MimeType(c.codecs)
 }
 
-func (c *Consumer) Init() []byte {
+func (c *Consumer) Init() ([]byte, error) {
 	if c.muxer == nil {
 		c.muxer = &Muxer{}
 	}
diff --git a/pkg/mp4/muxer.go b/pkg/mp4/muxer.go
index 5629fc3e..bcd2783b 100644
--- a/pkg/mp4/muxer.go
+++ b/pkg/mp4/muxer.go
@@ -2,6 +2,7 @@ package mp4
 
 import (
 	"encoding/binary"
+	"fmt"
 	"github.com/AlexxIT/go2rtc/pkg/h264"
 	"github.com/AlexxIT/go2rtc/pkg/streamer"
 	"github.com/deepch/vdk/codec/h264parser"
@@ -32,18 +33,21 @@ func (m *Muxer) MimeType(codecs []*streamer.Codec) string {
 	return s + `"`
 }
 
-func (m *Muxer) GetInit(codecs []*streamer.Codec) []byte {
+func (m *Muxer) GetInit(codecs []*streamer.Codec) ([]byte, error) {
 	moov := MOOV()
 
 	for _, codec := range codecs {
 		switch codec.Name {
 		case streamer.CodecH264:
 			sps, pps := h264.GetParameterSet(codec.FmtpLine)
+			if sps == nil {
+				return nil, fmt.Errorf("empty SPS: %#v", codec)
+			}
 
 			// TODO: remove
 			codecData, err := h264parser.NewCodecDataFromSPSAndPPS(sps, pps)
 			if err != nil {
-				return nil
+				return nil, err
 			}
 
 			width := codecData.Width()
@@ -83,7 +87,7 @@ func (m *Muxer) GetInit(codecs []*streamer.Codec) []byte {
 	data := make([]byte, moov.Len())
 	moov.Marshal(data)
 
-	return append(FTYP(), data...)
+	return append(FTYP(), data...), nil
 }
 
 func (m *Muxer) Rewind() {