From 55af09a3502ad72b3909df637b03740321d7c404 Mon Sep 17 00:00:00 2001
From: Alex X <alexey.khit@gmail.com>
Date: Sun, 5 Jan 2025 11:03:44 +0300
Subject: [PATCH] Add support fix JPEG from some MJPEG sources

---
 pkg/mjpeg/helpers.go | 63 +++++++++++++++++++++++++++++++++++---------
 pkg/mjpeg/jpeg.go    | 10 +++++++
 pkg/mjpeg/rfc2435.go | 30 ++++++++++-----------
 3 files changed, 75 insertions(+), 28 deletions(-)
 create mode 100644 pkg/mjpeg/jpeg.go

diff --git a/pkg/mjpeg/helpers.go b/pkg/mjpeg/helpers.go
index 08b4408b..d1acbd45 100644
--- a/pkg/mjpeg/helpers.go
+++ b/pkg/mjpeg/helpers.go
@@ -9,24 +9,38 @@ import (
 	"github.com/pion/rtp"
 )
 
-// FixJPEG - reencode JPEG if it has wrong header
-//
-// for example, this app produce "bad" images:
-// https://github.com/jacksonliam/mjpg-streamer
-//
-// and they can't be uploaded to the Telegram servers:
-// {"ok":false,"error_code":400,"description":"Bad Request: IMAGE_PROCESS_FAILED"}
 func FixJPEG(b []byte) []byte {
 	// skip non-JPEG
-	if len(b) < 10 || b[0] != 0xFF || b[1] != 0xD8 {
-		return b
-	}
-	// skip if header OK for imghdr library
-	// https://docs.python.org/3/library/imghdr.html
-	if string(b[2:4]) == "\xFF\xDB" || string(b[6:10]) == "JFIF" || string(b[6:10]) == "Exif" {
+	if len(b) < 10 || b[0] != 0xFF || b[1] != markerSOI {
 		return b
 	}
 
+	// skip JPEG without app marker
+	if b[2] == 0xFF && b[3] == markerDQT {
+		return b
+	}
+
+	switch string(b[6:10]) {
+	case "JFIF", "Exif":
+		// skip if header OK for imghdr library
+		// - https://docs.python.org/3/library/imghdr.html
+		return b
+	case "AVI1":
+		// adds DHT tables to JPEG file before SOS marker
+		// useful when you want to save a JPEG frame from an MJPEG stream
+		// - https://github.com/image-rs/jpeg-decoder/issues/76
+		// - https://github.com/pion/mediadevices/pull/493
+		// - https://bugzilla.mozilla.org/show_bug.cgi?id=963907#c18
+		return InjectDHT(b)
+	}
+
+	// reencode JPEG if it has wrong header
+	//
+	// for example, this app produce "bad" images:
+	// https://github.com/jacksonliam/mjpg-streamer
+	//
+	// and they can't be uploaded to the Telegram servers:
+	// {"ok":false,"error_code":400,"description":"Bad Request: IMAGE_PROCESS_FAILED"}
 	img, err := jpeg.Decode(bytes.NewReader(b))
 	if err != nil {
 		return b
@@ -54,3 +68,26 @@ func Encoder(codec *core.Codec, handler core.HandlerFunc) core.HandlerFunc {
 		handler(&clone)
 	}
 }
+
+const dhtSize = 432 // known size for 4 default tables
+
+func InjectDHT(b []byte) []byte {
+	if bytes.Index(b, []byte{0xFF, markerDHT}) > 0 {
+		return b // already exist
+	}
+
+	i := bytes.Index(b, []byte{0xFF, markerSOS})
+	if i < 0 {
+		return b
+	}
+
+	dht := make([]byte, 0, dhtSize)
+	dht = MakeHuffmanHeaders(dht)
+
+	tmp := make([]byte, len(b)+dhtSize)
+	copy(tmp, b[:i])
+	copy(tmp[i:], dht)
+	copy(tmp[i+dhtSize:], b[i:])
+
+	return tmp
+}
diff --git a/pkg/mjpeg/jpeg.go b/pkg/mjpeg/jpeg.go
new file mode 100644
index 00000000..8d6d13d1
--- /dev/null
+++ b/pkg/mjpeg/jpeg.go
@@ -0,0 +1,10 @@
+package mjpeg
+
+const (
+	markerSOF = 0xC0 // Start Of Frame (Baseline Sequential)
+	markerSOI = 0xD8 // Start Of Image
+	markerEOI = 0xD9 // End Of Image
+	markerSOS = 0xDA // Start Of Scan
+	markerDQT = 0xDB // Define Quantization Table
+	markerDHT = 0xC4 // Define Huffman Table
+)
diff --git a/pkg/mjpeg/rfc2435.go b/pkg/mjpeg/rfc2435.go
index 44307896..aa34c2f1 100644
--- a/pkg/mjpeg/rfc2435.go
+++ b/pkg/mjpeg/rfc2435.go
@@ -143,9 +143,7 @@ var chm_ac_symbols = []byte{
 
 func MakeHeaders(p []byte, t byte, w, h uint16, lqt, cqt []byte) []byte {
 	// Appendix A from https://www.rfc-editor.org/rfc/rfc2435
-	p = append(p, 0xFF,
-		0xD8, // SOI
-	)
+	p = append(p, 0xFF, markerSOI)
 
 	p = MakeQuantHeader(p, lqt, 0)
 	p = MakeQuantHeader(p, cqt, 1)
@@ -156,8 +154,7 @@ func MakeHeaders(p []byte, t byte, w, h uint16, lqt, cqt []byte) []byte {
 		t = 0x22 // hsamp = 2, vsamp = 2
 	}
 
-	p = append(p, 0xFF,
-		0xC0,  // SOF
+	p = append(p, 0xFF, markerSOF,
 		0, 17, // size
 		8, // bits per component
 		byte(h>>8), byte(h&0xFF),
@@ -174,13 +171,9 @@ func MakeHeaders(p []byte, t byte, w, h uint16, lqt, cqt []byte) []byte {
 		1,    // quant table 1
 	)
 
-	p = MakeHuffmanHeader(p, lum_dc_codelens, lum_dc_symbols, 0, 0)
-	p = MakeHuffmanHeader(p, lum_ac_codelens, lum_ac_symbols, 0, 1)
-	p = MakeHuffmanHeader(p, chm_dc_codelens, chm_dc_symbols, 1, 0)
-	p = MakeHuffmanHeader(p, chm_ac_codelens, chm_ac_symbols, 1, 1)
+	p = MakeHuffmanHeaders(p)
 
-	return append(p, 0xFF,
-		0xDA,  // SOS
+	return append(p, 0xFF, markerSOS,
 		0, 12, // size
 		3,    // 3 components
 		0,    // comp 0
@@ -196,16 +189,23 @@ func MakeHeaders(p []byte, t byte, w, h uint16, lqt, cqt []byte) []byte {
 }
 
 func MakeQuantHeader(p []byte, qt []byte, tableNo byte) []byte {
-	p = append(p, 0xFF, 0xDB, 0, 67, tableNo)
+	p = append(p, 0xFF, markerDQT, 0, 67, tableNo)
 	return append(p, qt...)
 }
 
 func MakeHuffmanHeader(p, codelens, symbols []byte, tableNo, tableClass byte) []byte {
-	p = append(p,
-		0xFF, 0xC4, 0,
-		byte(3+len(codelens)+len(symbols)),
+	p = append(p, 0xFF, markerDHT,
+		0, byte(3+len(codelens)+len(symbols)), // size
 		(tableClass<<4)|tableNo,
 	)
 	p = append(p, codelens...)
 	return append(p, symbols...)
 }
+
+func MakeHuffmanHeaders(p []byte) []byte {
+	p = MakeHuffmanHeader(p, lum_dc_codelens, lum_dc_symbols, 0, 0)
+	p = MakeHuffmanHeader(p, lum_ac_codelens, lum_ac_symbols, 0, 1)
+	p = MakeHuffmanHeader(p, chm_dc_codelens, chm_dc_symbols, 1, 0)
+	p = MakeHuffmanHeader(p, chm_ac_codelens, chm_ac_symbols, 1, 1)
+	return p
+}