From 95e0090bc26a7d43d4e2dd68a415fcedc4edef85 Mon Sep 17 00:00:00 2001
From: ahrav <ahravdutta02@gmail.com>
Date: Tue, 31 Oct 2023 16:42:42 -0700
Subject: [PATCH] [chore] - correctly handle input shorter than 512 bytes
 (#2077)

* correctly handle input shorter than 512 bytes

* add tests

* reorder tests

* add another test case

* update test

* address comment
---
 pkg/handlers/archive.go      |  4 +-
 pkg/handlers/archive_test.go | 77 +++++++++++++++++++++++++++++++++++-
 pkg/sources/chunker_test.go  |  2 +-
 3 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/pkg/handlers/archive.go b/pkg/handlers/archive.go
index 6b4062556..942e51328 100644
--- a/pkg/handlers/archive.go
+++ b/pkg/handlers/archive.go
@@ -414,9 +414,11 @@ func (a *Archive) handleNestedFileMIME(ctx logContext.Context, tempEnv tempEnv,
 // determineMimeType reads from the provided reader to detect the MIME type.
 // It returns the detected MIME type and a new reader that includes the read portion.
 func determineMimeType(reader io.Reader) (mimeType, io.Reader, error) {
+	// A buffer of 512 bytes is used since many file formats store their magic numbers within the first 512 bytes.
+	// If fewer bytes are read, MIME type detection may still succeed.
 	buffer := make([]byte, 512)
 	n, err := reader.Read(buffer)
-	if err != nil {
+	if err != nil && !errors.Is(err, io.EOF) {
 		return "", nil, fmt.Errorf("unable to read file for MIME type detection: %w", err)
 	}
 
diff --git a/pkg/handlers/archive_test.go b/pkg/handlers/archive_test.go
index 009446c62..3a2d9e5e4 100644
--- a/pkg/handlers/archive_test.go
+++ b/pkg/handlers/archive_test.go
@@ -11,8 +11,9 @@ import (
 	"testing"
 	"time"
 
-	diskbufferreader "github.com/trufflesecurity/disk-buffer-reader"
+	"github.com/h2non/filetype"
 	"github.com/stretchr/testify/assert"
+	diskbufferreader "github.com/trufflesecurity/disk-buffer-reader"
 
 	logContext "github.com/trufflesecurity/trufflehog/v3/pkg/context"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
@@ -271,3 +272,77 @@ func TestNestedDirArchive(t *testing.T) {
 	}
 	assert.Equal(t, want, count)
 }
+
+func TestDetermineMimeType(t *testing.T) {
+	filetype.AddMatcher(filetype.NewType("txt", "text/plain"), func(buf []byte) bool {
+		return strings.HasPrefix(string(buf), "text:")
+	})
+
+	pngBytes := []byte("\x89PNG\r\n\x1a\n")
+	jpegBytes := []byte{0xFF, 0xD8, 0xFF}
+	textBytes := []byte("text: This is a plain text")
+	rpmBytes := []byte("\xed\xab\xee\xdb")
+
+	tests := []struct {
+		name       string
+		input      io.Reader
+		expected   mimeType
+		shouldFail bool
+	}{
+		{
+			name:       "PNG file",
+			input:      bytes.NewReader(pngBytes),
+			expected:   mimeType("image/png"),
+			shouldFail: false,
+		},
+		{
+			name:       "JPEG file",
+			input:      bytes.NewReader(jpegBytes),
+			expected:   mimeType("image/jpeg"),
+			shouldFail: false,
+		},
+		{
+			name:       "Text file",
+			input:      bytes.NewReader(textBytes),
+			expected:   mimeType("text/plain"),
+			shouldFail: false,
+		},
+		{
+			name:       "RPM file",
+			input:      bytes.NewReader(rpmBytes),
+			expected:   rpmMimeType,
+			shouldFail: false,
+		},
+		{
+			name:       "Truncated JPEG file",
+			input:      io.LimitReader(bytes.NewReader(jpegBytes), 2),
+			expected:   mimeType("unknown"),
+			shouldFail: true,
+		},
+		{
+			name:       "Empty reader",
+			input:      bytes.NewReader([]byte{}),
+			shouldFail: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			originalData, _ := io.ReadAll(io.TeeReader(tt.input, &bytes.Buffer{}))
+			tt.input = bytes.NewReader(originalData) // Reset the reader
+
+			mime, reader, err := determineMimeType(tt.input)
+			if err != nil && !tt.shouldFail {
+				t.Fatalf("unexpected error: %v", err)
+			}
+
+			if !tt.shouldFail {
+				assert.Equal(t, tt.expected, mime)
+			}
+
+			// Ensure the reader still contains all the original data.
+			data, _ := io.ReadAll(reader)
+			assert.Equal(t, originalData, data)
+		})
+	}
+}
diff --git a/pkg/sources/chunker_test.go b/pkg/sources/chunker_test.go
index 0a6711707..aa350f7f0 100644
--- a/pkg/sources/chunker_test.go
+++ b/pkg/sources/chunker_test.go
@@ -7,8 +7,8 @@ import (
 	"testing"
 	"testing/iotest"
 
-	diskbufferreader "github.com/trufflesecurity/disk-buffer-reader"
 	"github.com/stretchr/testify/assert"
+	diskbufferreader "github.com/trufflesecurity/disk-buffer-reader"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 )