[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

@@ -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{})
 }

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.
+	// Read the first 261 header bytes into buffer as much as possible.
-	if _, err := io.ReadFull(rc, hdrBuf); err != nil {
+	if _, err := rc.Read(hdrBuf); err != nil {
 		return gtserror.Newf("error reading incoming media: %w", err)
 	}