From 27f465913916505197b06503efd7c30a79ac2016 Mon Sep 17 00:00:00 2001
From: KEINOS <github@keinos.com>
Date: Wed, 25 Oct 2023 21:58:45 +0900
Subject: [PATCH] [bugfix] allow store smaller PNG image than 261 bytes (#2263)
 (#2298)

* chore: add test of golden cases before fix of #2263

* chore: add test case to reproduce error of #2263

* [bugfix] allow store smaller PNG image than 261 bytes (#2263)
---
 internal/media/manager_test.go                |  70 ++++++++++++++++++
 internal/media/processingmedia.go             |   4 +-
 internal/media/test/test-jpeg-1x1px-white.jpg | Bin 0 -> 630 bytes
 .../test/test-png-alphachannel-1x1px.png      | Bin 0 -> 91 bytes
 4 files changed, 72 insertions(+), 2 deletions(-)
 create mode 100644 internal/media/test/test-jpeg-1x1px-white.jpg
 create mode 100644 internal/media/test/test-png-alphachannel-1x1px.png

diff --git a/internal/media/manager_test.go b/internal/media/manager_test.go
index 7b9b66147..cd0f9c24b 100644
--- a/internal/media/manager_test.go
+++ b/internal/media/manager_test.go
@@ -1248,6 +1248,76 @@ func (suite *ManagerTestSuite) TestSimpleJpegProcessBlockingWithDiskStorage() {
 	suite.Equal(processedThumbnailBytesExpected, processedThumbnailBytes)
 }
 
+func (suite *ManagerTestSuite) TestSmallSizedMediaTypeDetection_issue2263() {
+	for index, test := range []struct {
+		name     string // Test title
+		path     string // File path
+		expected string // Expected ContentType
+	}{
+		{
+			name:     "big size JPEG",
+			path:     "./test/test-jpeg.jpg",
+			expected: "image/jpeg",
+		},
+		{
+			name:     "big size PNG",
+			path:     "./test/test-png-noalphachannel.png",
+			expected: "image/png",
+		},
+		{
+			name:     "small size JPEG",
+			path:     "./test/test-jpeg-1x1px-white.jpg",
+			expected: "image/jpeg",
+		},
+		{
+			name:     "golden case PNG (big size)",
+			path:     "./test/test-png-alphachannel-1x1px.png",
+			expected: "image/png",
+		},
+	} {
+		suite.Run(test.name, func() {
+			ctx, cncl := context.WithTimeout(context.Background(), time.Second*60)
+			defer cncl()
+
+			data := func(_ context.Context) (io.ReadCloser, int64, error) {
+				// load bytes from a test image
+				b, err := os.ReadFile(test.path)
+				suite.NoError(err, "Test %d: failed during test setup", index+1)
+
+				return io.NopCloser(bytes.NewBuffer(b)), int64(len(b)), nil
+			}
+
+			accountID := "01FS1X72SK9ZPW0J1QQ68BD264"
+
+			// process the media with no additional info provided
+			processingMedia, err := suite.manager.ProcessMedia(ctx, data, accountID, nil)
+			suite.NoError(err)
+
+			// fetch the attachment id from the processing media
+			attachmentID := processingMedia.AttachmentID()
+
+			// wait for processing to complete
+			var attachment *gtsmodel.MediaAttachment
+			if !testrig.WaitFor(func() bool {
+				attachment, err = suite.db.GetAttachmentByID(ctx, attachmentID)
+				return err == nil && attachment != nil
+			}) {
+				suite.FailNow("timed out waiting for attachment to process")
+			}
+
+			// make sure it's got the stuff set on it that we expect
+			// the attachment ID and accountID we expect
+			suite.Equal(attachmentID, attachment.ID)
+			suite.Equal(accountID, attachment.AccountID)
+
+			actual := attachment.File.ContentType
+			expect := test.expected
+
+			suite.Equal(expect, actual, "Test %d: %s", index+1, test.name)
+		})
+	}
+}
+
 func TestManagerTestSuite(t *testing.T) {
 	suite.Run(t, &ManagerTestSuite{})
 }
diff --git a/internal/media/processingmedia.go b/internal/media/processingmedia.go
index acfee48f3..591cbac74 100644
--- a/internal/media/processingmedia.go
+++ b/internal/media/processingmedia.go
@@ -161,8 +161,8 @@ func (p *ProcessingMedia) store(ctx context.Context) error {
 	// and https://github.com/h2non/filetype
 	hdrBuf := make([]byte, 261)
 
-	// Read the first 261 header bytes into buffer.
-	if _, err := io.ReadFull(rc, hdrBuf); err != nil {
+	// Read the first 261 header bytes into buffer as much as possible.
+	if _, err := rc.Read(hdrBuf); err != nil {
 		return gtserror.Newf("error reading incoming media: %w", err)
 	}
 
diff --git a/internal/media/test/test-jpeg-1x1px-white.jpg b/internal/media/test/test-jpeg-1x1px-white.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..13d5a4b8c29bbb776092d6b65839ce6141c84d95
GIT binary patch
literal 630
zcmex=<NpH&0WUXCHwH#VMg|WcWcYuZ!I|Mdr2uFMBc*2iKfoZ!!NADC$jqq3z$D1X
zEXer(2!lKW10yTYwIByW0V5ML3o9Et2PYRdP{CFK1|~*kW+oP9R#p}kpz2znJOhg$
ztB|6hBb#twBD+$dh*9Ijg&fLG8xM*GUHqV8oK)1r$t5N(At|M*rmmr>WnyY(ZeeNV
z?BeR??&0Yb91<E99uXOpoRXTBo{^bVTvA$AUQt=q+|t_C-qG1LY4Vh*)27duIcw44
zB}<nrU$Ju4rp;TnZri?N=dMGCj~qRA{KUyqmo8tqdhPm+o3|c5di><+v*#~fzWVs-
z^OvvRzW@073*;|G24;x2;66k1mmttzOe`$SEbJhEF*20{F|!~GtD+&BkYgZwVxh2-
zQ6q<l)5L`v4{|CS2YnDtD!Rxereg9?^&`k@V4o4^u_m&72KN}kU$+=|m>GeQ$1KQT
N&oHS=V*&sFn*g~%I{W|t

literal 0
HcmV?d00001

diff --git a/internal/media/test/test-png-alphachannel-1x1px.png b/internal/media/test/test-png-alphachannel-1x1px.png
new file mode 100644
index 0000000000000000000000000000000000000000..6f39d4d675ea86ecdd936b38b263f10a443f8c4e
GIT binary patch
literal 91
zcmeAS@N?(olHy`uVBq!ia0vp^j3CUx1|;Q0k8}bl&H|6fVg?3oVGw3ym^DWND9G#S
k;uyj)Gx^8=|NrY5IDa$$<+%7l4k*Fk>FVdQ&MBb@0N7C#0{{R3

literal 0
HcmV?d00001