addressing PR reviews

pkg/frame/compressed.go

@@ -17,14 +17,44 @@ var (
 func decodeMJPEG(frame []byte, width, height int) (image.Image, func(), error) {
 	img, err := jpeg.Decode(bytes.NewReader(frame))
 
-	if err != nil && err.Error() == "invalid JPEG format: uninitialized Huffman table" {
+	if err == nil || (err != nil && err.Error() != "invalid JPEG format: uninitialized Huffman table") {
-		img, err = jpeg.Decode(bytes.NewReader(addMotionDht(frame)))
+		return img, func() {}, err
 	}
 
+	img, err = jpeg.Decode(bytes.NewReader(addMotionDht(frame)))
 	return img, func() {}, err
 }
 
 func addMotionDht(frame []byte) []byte {
 	jpegParts := bytes.Split(frame, sosMarker)
-	return append(jpegParts[0], append(dhtMarker, append(dht, append(sosMarker, jpegParts[1]...)...)...)...)
+	if len(jpegParts) != 2 {
+		return frame
+	}
+	correctedFrame := make([]byte, len(jpegParts[0])+len(dhtMarker)+len(dht)+len(sosMarker)+len(jpegParts[1]))
+	correctedFrameOffset := 0
+
+	for indx, item := range jpegParts[0] {
+		correctedFrame[indx] = item
+	}
+	correctedFrameOffset += len(jpegParts[0])
+
+	for indx, item := range dhtMarker {
+		correctedFrame[indx+correctedFrameOffset] = item
+	}
+	correctedFrameOffset += len(dhtMarker)
+
+	for indx, item := range dht {
+		correctedFrame[indx+correctedFrameOffset] = item
+	}
+	correctedFrameOffset += len(dht)
+
+	for indx, item := range sosMarker {
+		correctedFrame[indx+correctedFrameOffset] = item
+	}
+	correctedFrameOffset += len(sosMarker)
+
+	for indx, item := range jpegParts[1] {
+		correctedFrame[indx+correctedFrameOffset] = item
+	}
+	return correctedFrame
 }

pkg/frame/compressed_test.go

@@ -0,0 +1,44 @@
+package frame
+
+import (
+	"bytes"
+	"image/jpeg"
+	"testing"
+)
+
+func TestAddMotionDht(t *testing.T) {
+	uninitializedHuffmanTableFrame, err := jpeg.Decode(bytes.NewReader(UninitializedHuffmanTable))
+
+	// Decode fails with an uninitialized Huffman table error for sample input
+	expectedErrorMessage := "invalid JPEG format: uninitialized Huffman table"
+	if err.Error() != expectedErrorMessage {
+		t.Errorf("Wrong decode error result,\nexpected:\n%+v\ngot:\n%+v", expectedErrorMessage, err)
+	}
+
+	// Decode passes after adding default Huffman table to
+	defaultHuffmanTableFrame, err := jpeg.Decode(bytes.NewReader(addMotionDht(UninitializedHuffmanTable)))
+	if err != nil {
+		t.Errorf("Expected decode function to pass after adding default Huffman table. Failed with %v\n", err)
+	}
+
+	// Adding default Huffman table to a valid frame without a Huffman table changes the table
+	if uninitializedHuffmanTableFrame == defaultHuffmanTableFrame {
+		t.Errorf("Expected addMotionDht to update frame. Instead returned original frame")
+	}
+
+	// Check that an improperly constructed frame does not get updated by addMotionDht
+	randomBytes := []byte{1, 2, 3, 4}
+	frame1, err := jpeg.Decode(bytes.NewReader(randomBytes))
+	if err == nil {
+		t.Errorf("Expected decode function to fail with random bytes but passed.")
+	}
+
+	frame2, err := jpeg.Decode(bytes.NewReader(addMotionDht(randomBytes)))
+	if err == nil {
+		t.Errorf("Expected decode function to fail with random bytes but passed.")
+	}
+
+	if frame1 != frame2 {
+		t.Errorf("addMotionDht updated the frame despite being improperly constructed")
+	}
+}