diff --git a/internal/mp4/mp4.go b/internal/mp4/mp4.go index 9144a497..654d9ae7 100644 --- a/internal/mp4/mp4.go +++ b/internal/mp4/mp4.go @@ -1,7 +1,6 @@ package mp4 import ( - "io" "net/http" "strconv" "strings" @@ -56,34 +55,25 @@ func handlerKeyframe(w http.ResponseWriter, r *http.Request) { return } - wr := &once{} // init and first frame - _, _ = cons.WriteTo(wr) + once := &core.OnceBuffer{} // init and first frame + _, _ = cons.WriteTo(once) stream.RemoveConsumer(cons) // Apple Safari won't show frame without length header := w.Header() - header.Set("Content-Length", strconv.Itoa(len(wr.buf))) + header.Set("Content-Length", strconv.Itoa(once.Len())) 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(wr.buf); err != nil { + if _, err := once.WriteTo(w); 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) diff --git a/pkg/core/writebuffer.go b/pkg/core/writebuffer.go index dce7affb..edc72e9e 100644 --- a/pkg/core/writebuffer.go +++ b/pkg/core/writebuffer.go @@ -6,6 +6,10 @@ import ( "sync" ) +// WriteBuffer by defaul Write(s) to bytes.Buffer. +// But after WriteTo to new io.Writer - calls Reset. +// Reset will flush current buffer data to new writer and starts to Write to new io.Writer +// WriteTo will be locked until Write fails or Close will be called. type WriteBuffer struct { io.Writer err error @@ -52,7 +56,7 @@ func (w *WriteBuffer) Close() error { func (w *WriteBuffer) Reset(wr io.Writer) { w.mu.Lock() w.add() - if buf, ok := wr.(*bytes.Buffer); ok { + if buf, ok := w.Writer.(*bytes.Buffer); ok && buf.Len() != 0 { if _, err := io.Copy(wr, buf); err != nil { w.err = err w.done() @@ -81,3 +85,27 @@ func (w *WriteBuffer) done() { w.wg.Done() } } + +// OnceBuffer will catch only first message +type OnceBuffer struct { + buf []byte +} + +func (o *OnceBuffer) Write(p []byte) (n int, err error) { + if o.buf == nil { + o.buf = p + } + return 0, io.EOF +} + +func (o *OnceBuffer) WriteTo(w io.Writer) (n int64, err error) { + return io.Copy(w, bytes.NewReader(o.buf)) +} + +func (o *OnceBuffer) Buffer() []byte { + return o.buf +} + +func (o *OnceBuffer) Len() int { + return len(o.buf) +}