Rewrite MP4, HLS, MPEG-TS consumers

internal/mp4/mp4.go

@@ -1,6 +1,7 @@
 package mp4
 
 import (
+	"io"
 	"net/http"
 	"strconv"
 	"strings"
@@ -47,18 +48,7 @@ func handlerKeyframe(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	exit := make(chan []byte, 1)
-
-	cons := &mp4.Segment{OnlyKeyframe: true}
-	cons.Listen(func(msg any) {
-		if data, ok := msg.([]byte); ok && exit != nil {
-			select {
-			case exit <- data:
-			default:
-			}
-			exit = nil
-		}
-	})
+	cons := mp4.NewKeyframe(nil)
 
 	if err := stream.AddConsumer(cons); err != nil {
 		log.Error().Err(err).Caller().Send()
@@ -66,24 +56,34 @@ func handlerKeyframe(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	data := <-exit
+	wr := &once{} // init and first frame
+	_, _ = cons.WriteTo(wr)
 
 	stream.RemoveConsumer(cons)
 
 	// Apple Safari won't show frame without length
 	header := w.Header()
-	header.Set("Content-Length", strconv.Itoa(len(data)))
-	header.Set("Content-Type", cons.MimeType)
+	header.Set("Content-Length", strconv.Itoa(len(wr.buf)))
+	header.Set("Content-Type", mp4.ContentType(cons.Codecs()))
 
 	if filename := query.Get("filename"); filename != "" {
 		header.Set("Content-Disposition", `attachment; filename="`+filename+`"`)
 	}
 
-	if _, err := w.Write(data); err != nil {
+	if _, err := w.Write(wr.buf); err != nil {
 		log.Error().Err(err).Caller().Send()
 	}
 }
 
+type once struct {
+	buf []byte
+}
+
+func (o *once) Write(p []byte) (n int, err error) {
+	o.buf = p
+	return 0, io.EOF
+}
+
 func handlerMP4(w http.ResponseWriter, r *http.Request) {
 	log.Trace().Msgf("[mp4] %s %+v", r.Method, r.Header)
 
@@ -108,29 +108,11 @@ func handlerMP4(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	exit := make(chan error, 1) // Add buffer to prevent blocking
-
-	cons := &mp4.Consumer{
-		Desc:       "MP4/HTTP",
-		RemoteAddr: tcp.RemoteAddr(r),
-		UserAgent:  r.UserAgent(),
-		Medias:     mp4.ParseQuery(r.URL.Query()),
-	}
-
-	cons.Listen(func(msg any) {
-		if exit == nil {
-			return
-		}
-		if data, ok := msg.([]byte); ok {
-			if _, err := w.Write(data); err != nil {
-				select {
-				case exit <- err:
-				default:
-				}
-				exit = nil
-			}
-		}
-	})
+	medias := mp4.ParseQuery(r.URL.Query())
+	cons := mp4.NewConsumer(medias)
+	cons.Type = "MP4/HTTP active consumer"
+	cons.RemoteAddr = tcp.RemoteAddr(r)
+	cons.UserAgent = r.UserAgent()
 
 	if err := stream.AddConsumer(cons); err != nil {
 		log.Error().Err(err).Caller().Send()
@@ -140,57 +122,34 @@ func handlerMP4(w http.ResponseWriter, r *http.Request) {
 
 	defer stream.RemoveConsumer(cons)
 
-	data, err := cons.Init()
-	if err != nil {
-		log.Error().Err(err).Caller().Send()
-		http.Error(w, err.Error(), http.StatusInternalServerError)
-		return
-	}
-
 	header := w.Header()
-	header.Set("Content-Type", cons.MimeType())
+	header.Set("Content-Type", mp4.ContentType(cons.Codecs()))
 
 	if filename := query.Get("filename"); filename != "" {
 		header.Set("Content-Disposition", `attachment; filename="`+filename+`"`)
 	}
 
 	if rotate := query.Get("rotate"); rotate != "" {
-		mp4.PatchVideoRotate(data, core.Atoi(rotate))
+		cons.Rotate = core.Atoi(rotate)
 	}
 
 	if scale := query.Get("scale"); scale != "" {
 		if sx, sy, ok := strings.Cut(scale, ":"); ok {
-			mp4.PatchVideoScale(data, core.Atoi(sx), core.Atoi(sy))
+			cons.ScaleX = core.Atoi(sx)
+			cons.ScaleY = core.Atoi(sy)
 		}
 	}
 
-	if _, err = w.Write(data); err != nil {
-		log.Error().Err(err).Caller().Send()
-		http.Error(w, err.Error(), http.StatusInternalServerError)
-		return
-	}
-
-	cons.Start()
-
 	var duration *time.Timer
 	if s := query.Get("duration"); s != "" {
 		if i, _ := strconv.Atoi(s); i > 0 {
 			duration = time.AfterFunc(time.Second*time.Duration(i), func() {
-				if exit != nil {
-					select {
-					case exit <- nil:
-					default:
-					}
-					exit = nil
-				}
+				_ = cons.Stop()
 			})
 		}
 	}
 
-	err = <-exit
-	exit = nil
-
-	log.Trace().Err(err).Caller().Send()
+	_, _ = cons.WriteTo(w)
 
 	if duration != nil {
 		duration.Stop()

internal/mp4/ws.go

@@ -6,6 +6,7 @@ import (
 	"github.com/AlexxIT/go2rtc/internal/api"
 	"github.com/AlexxIT/go2rtc/internal/api/ws"
 	"github.com/AlexxIT/go2rtc/internal/streams"
+	"github.com/AlexxIT/go2rtc/pkg/core"
 	"github.com/AlexxIT/go2rtc/pkg/mp4"
 	"github.com/AlexxIT/go2rtc/pkg/tcp"
 )
@@ -16,44 +17,30 @@ func handlerWSMSE(tr *ws.Transport, msg *ws.Message) error {
 		return errors.New(api.StreamNotFound)
 	}
 
-	cons := &mp4.Consumer{
-		Desc:       "MSE/WebSocket",
-		RemoteAddr: tcp.RemoteAddr(tr.Request),
-		UserAgent:  tr.Request.UserAgent(),
-	}
-
+	var medias []*core.Media
 	if codecs := msg.String(); codecs != "" {
 		log.Trace().Str("codecs", codecs).Msgf("[mp4] new WS/MSE consumer")
-		cons.Medias = mp4.ParseCodecs(codecs, true)
+		medias = mp4.ParseCodecs(codecs, true)
 	}
 
-	cons.Listen(func(msg any) {
-		if data, ok := msg.([]byte); ok {
-			tr.Write(data)
-		}
-	})
+	cons := mp4.NewConsumer(medias)
+	cons.Type = "MSE/WebSocket active consumer"
+	cons.RemoteAddr = tcp.RemoteAddr(tr.Request)
+	cons.UserAgent = tr.Request.UserAgent()
 
 	if err := stream.AddConsumer(cons); err != nil {
 		log.Debug().Err(err).Msg("[mp4] add consumer")
 		return err
 	}
 
+	tr.Write(&ws.Message{Type: "mse", Value: mp4.ContentType(cons.Codecs())})
+
+	go cons.WriteTo(tr.Writer())
+
 	tr.OnClose(func() {
 		stream.RemoveConsumer(cons)
 	})
 
-	tr.Write(&ws.Message{Type: "mse", Value: cons.MimeType()})
-
-	data, err := cons.Init()
-	if err != nil {
-		log.Warn().Err(err).Caller().Send()
-		return err
-	}
-
-	tr.Write(data)
-
-	cons.Start()
-
 	return nil
 }
 
@@ -63,29 +50,25 @@ func handlerWSMP4(tr *ws.Transport, msg *ws.Message) error {
 		return errors.New(api.StreamNotFound)
 	}
 
-	cons := &mp4.Segment{
-		RemoteAddr:   tcp.RemoteAddr(tr.Request),
-		UserAgent:    tr.Request.UserAgent(),
-		OnlyKeyframe: true,
-	}
-
+	var medias []*core.Media
 	if codecs := msg.String(); codecs != "" {
 		log.Trace().Str("codecs", codecs).Msgf("[mp4] new WS/MP4 consumer")
-		cons.Medias = mp4.ParseCodecs(codecs, false)
+		medias = mp4.ParseCodecs(codecs, false)
 	}
 
-	cons.Listen(func(msg any) {
-		if data, ok := msg.([]byte); ok {
-			tr.Write(data)
-		}
-	})
+	cons := mp4.NewKeyframe(medias)
+	cons.Type = "MP4/WebSocket active consumer"
+	cons.RemoteAddr = tcp.RemoteAddr(tr.Request)
+	cons.UserAgent = tr.Request.UserAgent()
 
 	if err := stream.AddConsumer(cons); err != nil {
 		log.Error().Err(err).Caller().Send()
 		return err
 	}
 
-	tr.Write(&ws.Message{Type: "mp4", Value: cons.MimeType})
+	tr.Write(&ws.Message{Type: "mse", Value: mp4.ContentType(cons.Codecs())})
+
+	go cons.WriteTo(tr.Writer())
 
 	tr.OnClose(func() {
 		stream.RemoveConsumer(cons)