From e0d5f6a195cb1c15808229f34b6807e097253046 Mon Sep 17 00:00:00 2001
From: Knut Ahlers <knut@ahlers.me>
Date: Mon, 29 Apr 2019 13:02:23 +0200
Subject: [PATCH] Replace mjpeg reader to fix broken frames

Signed-off-by: Knut Ahlers <knut@ahlers.me>
---
 main.go  | 44 +++++++++++++++++++++++++++++++-------------
 mjpeg.go |  2 ++
 2 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/main.go b/main.go
index dc4deb5..1a4fa9d 100644
--- a/main.go
+++ b/main.go
@@ -3,7 +3,6 @@ package main
 import (
 	"bytes"
 	"fmt"
-	"io"
 	"net/http"
 	"os"
 	"os/exec"
@@ -96,27 +95,46 @@ func main() {
 
 	log.Debug("ffmpeg spawned")
 
-	buf := new(bytes.Buffer)
+	var (
+		br, bw int
+		buf    = make([]byte, 10*1024*1024) // 10MB (jpg should be smaller)
+	)
 
 	for {
-		if _, err := io.CopyN(buf, out, 1024); err != nil {
-			log.WithError(err).Error("Failed to read ffmpeg output")
-			break
+		// If buffer was read, slide the remains to the beginning
+		if br > 0 {
+			copy(buf, buf[br:bw])
+			bw -= br
+			br = 0
 		}
 
-		eoj := bytes.Index(buf.Bytes(), endOfJPEG)
-		if eoj == -1 {
+		// Fill buffer
+		n, err := out.Read(buf[bw:])
+		if err != nil {
+			log.WithError(err).Fatal("Unable to read from output")
+		}
+		bw += n
+
+		if n == 0 {
+			// Nothing read, try again
 			continue
 		}
 
-		img := buf.Next(eoj + len(endOfJPEG))
+		// Extract as many images as possible before next read
+		for eoj := bytes.Index(buf[br:bw], endOfJPEG); eoj >= 0; eoj = bytes.Index(buf[br:bw], endOfJPEG) {
+			eoj += len(endOfJPEG)
+			img := make([]byte, eoj-br)
+			copy(img, buf[br:br+eoj])
 
-		if !bytes.HasPrefix(img, beginOfJPEG) || !bytes.HasSuffix(img, endOfJPEG) {
-			log.Warn("Found invalid JPEG, skipping")
-			continue
+			br += eoj
+
+			if !bytes.HasPrefix(img, beginOfJPEG) || !bytes.HasSuffix(img, endOfJPEG) {
+				log.Warn("Found invalid JPEG, skipping")
+				continue
+			}
+
+			go sendImage(img)
 		}
-
-		go sendImage(img)
 	}
 }
 
diff --git a/mjpeg.go b/mjpeg.go
index cb74dd9..a167b76 100644
--- a/mjpeg.go
+++ b/mjpeg.go
@@ -5,6 +5,7 @@ import (
 	"mime/multipart"
 	"net/http"
 	"net/textproto"
+	"strconv"
 
 	"github.com/pkg/errors"
 	log "github.com/sirupsen/logrus"
@@ -38,6 +39,7 @@ func handleMJPEG(res http.ResponseWriter, r *http.Request, imgs chan []byte, uid
 			err := func() error {
 				partHeader := make(textproto.MIMEHeader)
 				partHeader.Add("Content-Type", "image/jpeg")
+				partHeader.Add("Content-Length", strconv.Itoa(len(img)))
 
 				partWriter, err := mimeWriter.CreatePart(partHeader)
 				if err != nil {