[bugfix] Wrap media in read closer (#941)

* use readcloser for content.Content

* call media postdata function no matter what

* return a readcloser from data func

* tidy of logic of readertostore

* fix whoopsie

5 files changed, 61 insertions, 53 deletions

diff --git a/internal/media/manager_test.go b/internal/media/manager_test.go
index e00cdd98d..b50235054 100644
--- a/internal/media/manager_test.go
+++ b/internal/media/manager_test.go
@@ -43,13 +43,13 @@ type ManagerTestSuite struct {
 func (suite *ManagerTestSuite) TestEmojiProcessBlocking() {
 	ctx := context.Background()
 
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		// load bytes from a test image
 		b, err := os.ReadFile("./test/rainbow-original.png")
 		if err != nil {
 			panic(err)
 		}
-		return bytes.NewBuffer(b), int64(len(b)), nil
+		return io.NopCloser(bytes.NewBuffer(b)), int64(len(b)), nil
 	}
 
 	emojiID := "01GDQ9G782X42BAMFASKP64343"
@@ -114,12 +114,12 @@ func (suite *ManagerTestSuite) TestEmojiProcessBlockingRefresh() {
 	oldEmojiImagePath := emojiToUpdate.ImagePath
 	oldEmojiImageStaticPath := emojiToUpdate.ImageStaticPath
 
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		b, err := os.ReadFile("./test/gts_pixellated-original.png")
 		if err != nil {
 			panic(err)
 		}
-		return bytes.NewBuffer(b), int64(len(b)), nil
+		return io.NopCloser(bytes.NewBuffer(b)), int64(len(b)), nil
 	}
 
 	emojiID := emojiToUpdate.ID
@@ -197,13 +197,13 @@ func (suite *ManagerTestSuite) TestEmojiProcessBlockingRefresh() {
 func (suite *ManagerTestSuite) TestEmojiProcessBlockingTooLarge() {
 	ctx := context.Background()
 
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		// load bytes from a test image
 		b, err := os.ReadFile("./test/big-panda.gif")
 		if err != nil {
 			panic(err)
 		}
-		return bytes.NewBuffer(b), int64(len(b)), nil
+		return io.NopCloser(bytes.NewBuffer(b)), int64(len(b)), nil
 	}
 
 	emojiID := "01GDQ9G782X42BAMFASKP64343"
@@ -221,13 +221,13 @@ func (suite *ManagerTestSuite) TestEmojiProcessBlockingTooLarge() {
 func (suite *ManagerTestSuite) TestEmojiProcessBlockingTooLargeNoSizeGiven() {
 	ctx := context.Background()
 
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		// load bytes from a test image
 		b, err := os.ReadFile("./test/big-panda.gif")
 		if err != nil {
 			panic(err)
 		}
-		return bytes.NewBuffer(b), int64(len(b)), nil
+		return io.NopCloser(bytes.NewBuffer(b)), int64(len(b)), nil
 	}
 
 	emojiID := "01GDQ9G782X42BAMFASKP64343"
@@ -245,13 +245,13 @@ func (suite *ManagerTestSuite) TestEmojiProcessBlockingTooLargeNoSizeGiven() {
 func (suite *ManagerTestSuite) TestEmojiProcessBlockingNoFileSizeGiven() {
 	ctx := context.Background()
 
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		// load bytes from a test image
 		b, err := os.ReadFile("./test/rainbow-original.png")
 		if err != nil {
 			panic(err)
 		}
-		return bytes.NewBuffer(b), -1, nil
+		return io.NopCloser(bytes.NewBuffer(b)), -1, nil
 	}
 
 	emojiID := "01GDQ9G782X42BAMFASKP64343"
@@ -307,13 +307,13 @@ func (suite *ManagerTestSuite) TestEmojiProcessBlockingNoFileSizeGiven() {
 func (suite *ManagerTestSuite) TestSimpleJpegProcessBlocking() {
 	ctx := context.Background()
 
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		// load bytes from a test image
 		b, err := os.ReadFile("./test/test-jpeg.jpg")
 		if err != nil {
 			panic(err)
 		}
-		return bytes.NewBuffer(b), int64(len(b)), nil
+		return io.NopCloser(bytes.NewBuffer(b)), int64(len(b)), nil
 	}
 
 	accountID := "01FS1X72SK9ZPW0J1QQ68BD264"
@@ -379,14 +379,14 @@ func (suite *ManagerTestSuite) TestSimpleJpegProcessBlocking() {
 func (suite *ManagerTestSuite) TestSimpleJpegProcessBlockingNoContentLengthGiven() {
 	ctx := context.Background()
 
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		// load bytes from a test image
 		b, err := os.ReadFile("./test/test-jpeg.jpg")
 		if err != nil {
 			panic(err)
 		}
 		// give length as -1 to indicate unknown
-		return bytes.NewBuffer(b), -1, nil
+		return io.NopCloser(bytes.NewBuffer(b)), -1, nil
 	}
 
 	accountID := "01FS1X72SK9ZPW0J1QQ68BD264"
@@ -452,7 +452,7 @@ func (suite *ManagerTestSuite) TestSimpleJpegProcessBlockingNoContentLengthGiven
 func (suite *ManagerTestSuite) TestSimpleJpegProcessBlockingReadCloser() {
 	ctx := context.Background()
 
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		// open test image as a file
 		f, err := os.Open("./test/test-jpeg.jpg")
 		if err != nil {
@@ -525,13 +525,13 @@ func (suite *ManagerTestSuite) TestSimpleJpegProcessBlockingReadCloser() {
 func (suite *ManagerTestSuite) TestPngNoAlphaChannelProcessBlocking() {
 	ctx := context.Background()
 
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		// load bytes from a test image
 		b, err := os.ReadFile("./test/test-png-noalphachannel.png")
 		if err != nil {
 			panic(err)
 		}
-		return bytes.NewBuffer(b), int64(len(b)), nil
+		return io.NopCloser(bytes.NewBuffer(b)), int64(len(b)), nil
 	}
 
 	accountID := "01FS1X72SK9ZPW0J1QQ68BD264"
@@ -597,13 +597,13 @@ func (suite *ManagerTestSuite) TestPngNoAlphaChannelProcessBlocking() {
 func (suite *ManagerTestSuite) TestPngAlphaChannelProcessBlocking() {
 	ctx := context.Background()
 
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		// load bytes from a test image
 		b, err := os.ReadFile("./test/test-png-alphachannel.png")
 		if err != nil {
 			panic(err)
 		}
-		return bytes.NewBuffer(b), int64(len(b)), nil
+		return io.NopCloser(bytes.NewBuffer(b)), int64(len(b)), nil
 	}
 
 	accountID := "01FS1X72SK9ZPW0J1QQ68BD264"
@@ -669,13 +669,13 @@ func (suite *ManagerTestSuite) TestPngAlphaChannelProcessBlocking() {
 func (suite *ManagerTestSuite) TestSimpleJpegProcessBlockingWithCallback() {
 	ctx := context.Background()
 
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		// load bytes from a test image
 		b, err := os.ReadFile("./test/test-jpeg.jpg")
 		if err != nil {
 			panic(err)
 		}
-		return bytes.NewBuffer(b), int64(len(b)), nil
+		return io.NopCloser(bytes.NewBuffer(b)), int64(len(b)), nil
 	}
 
 	// test the callback function by setting a simple boolean
@@ -752,13 +752,13 @@ func (suite *ManagerTestSuite) TestSimpleJpegProcessBlockingWithCallback() {
 func (suite *ManagerTestSuite) TestSimpleJpegProcessAsync() {
 	ctx := context.Background()
 
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		// load bytes from a test image
 		b, err := os.ReadFile("./test/test-jpeg.jpg")
 		if err != nil {
 			panic(err)
 		}
-		return bytes.NewBuffer(b), int64(len(b)), nil
+		return io.NopCloser(bytes.NewBuffer(b)), int64(len(b)), nil
 	}
 
 	accountID := "01FS1X72SK9ZPW0J1QQ68BD264"
@@ -837,9 +837,9 @@ func (suite *ManagerTestSuite) TestSimpleJpegQueueSpamming() {
 		panic(err)
 	}
 
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		// load bytes from a test image
-		return bytes.NewReader(b), int64(len(b)), nil
+		return io.NopCloser(bytes.NewReader(b)), int64(len(b)), nil
 	}
 
 	accountID := "01FS1X72SK9ZPW0J1QQ68BD264"
@@ -913,13 +913,13 @@ func (suite *ManagerTestSuite) TestSimpleJpegQueueSpamming() {
 func (suite *ManagerTestSuite) TestSimpleJpegProcessBlockingWithDiskStorage() {
 	ctx := context.Background()
 
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		// load bytes from a test image
 		b, err := os.ReadFile("./test/test-jpeg.jpg")
 		if err != nil {
 			panic(err)
 		}
-		return bytes.NewBuffer(b), int64(len(b)), nil
+		return io.NopCloser(bytes.NewBuffer(b)), int64(len(b)), nil
 	}
 
 	accountID := "01FS1X72SK9ZPW0J1QQ68BD264"
diff --git a/internal/media/processingemoji.go b/internal/media/processingemoji.go
index e1c6f2efb..79bc23998 100644
--- a/internal/media/processingemoji.go
+++ b/internal/media/processingemoji.go
@@ -193,24 +193,22 @@ func (p *ProcessingEmoji) store(ctx context.Context) error {
 		return nil
 	}
 
-	// execute the data function to get the reader out of it
-	reader, fileSize, err := p.data(ctx)
+	// execute the data function to get the readcloser out of it
+	rc, fileSize, err := p.data(ctx)
 	if err != nil {
 		return fmt.Errorf("store: error executing data function: %s", err)
 	}
 
 	// defer closing the reader when we're done with it
 	defer func() {
-		if rc, ok := reader.(io.ReadCloser); ok {
-			if err := rc.Close(); err != nil {
-				log.Errorf("store: error closing readcloser: %s", err)
-			}
+		if err := rc.Close(); err != nil {
+			log.Errorf("store: error closing readcloser: %s", err)
 		}
 	}()
 
 	// extract no more than 261 bytes from the beginning of the file -- this is the header
 	firstBytes := make([]byte, maxFileHeaderBytes)
-	if _, err := reader.Read(firstBytes); err != nil {
+	if _, err := rc.Read(firstBytes); err != nil {
 		return fmt.Errorf("store: error reading initial %d bytes: %s", maxFileHeaderBytes, err)
 	}
 
@@ -242,7 +240,7 @@ func (p *ProcessingEmoji) store(ctx context.Context) error {
 	p.emoji.ImageContentType = contentType
 
 	// concatenate the first bytes with the existing bytes still in the reader (thanks Mara)
-	readerToStore := io.MultiReader(bytes.NewBuffer(firstBytes), reader)
+	readerToStore := io.MultiReader(bytes.NewBuffer(firstBytes), rc)
 
 	var maxEmojiSize int64
 	if p.emoji.Domain == "" {
diff --git a/internal/media/processingmedia.go b/internal/media/processingmedia.go
index 1f5a58b9f..573df1d0e 100644
--- a/internal/media/processingmedia.go
+++ b/internal/media/processingmedia.go
@@ -263,24 +263,31 @@ func (p *ProcessingMedia) store(ctx context.Context) error {
 		return nil
 	}
 
-	// execute the data function to get the reader out of it
-	reader, fileSize, err := p.data(ctx)
+	// execute the data function to get the readcloser out of it
+	rc, fileSize, err := p.data(ctx)
 	if err != nil {
 		return fmt.Errorf("store: error executing data function: %s", err)
 	}
 
 	// defer closing the reader when we're done with it
 	defer func() {
-		if rc, ok := reader.(io.ReadCloser); ok {
-			if err := rc.Close(); err != nil {
-				log.Errorf("store: error closing readcloser: %s", err)
+		if err := rc.Close(); err != nil {
+			log.Errorf("store: error closing readcloser: %s", err)
+		}
+	}()
+
+	// execute the postData function no matter what happens
+	defer func() {
+		if p.postData != nil {
+			if err := p.postData(ctx); err != nil {
+				log.Errorf("store: error executing postData: %s", err)
 			}
 		}
 	}()
 
 	// extract no more than 261 bytes from the beginning of the file -- this is the header
 	firstBytes := make([]byte, maxFileHeaderBytes)
-	if _, err := reader.Read(firstBytes); err != nil {
+	if _, err := rc.Read(firstBytes); err != nil {
 		return fmt.Errorf("store: error reading initial %d bytes: %s", maxFileHeaderBytes, err)
 	}
 
@@ -303,29 +310,36 @@ func (p *ProcessingMedia) store(ctx context.Context) error {
 	extension := split[1] // something like 'jpeg'
 
 	// concatenate the cleaned up first bytes with the existing bytes still in the reader (thanks Mara)
-	readerToStore := io.MultiReader(bytes.NewBuffer(firstBytes), reader)
+	multiReader := io.MultiReader(bytes.NewBuffer(firstBytes), rc)
 
 	// use the extension to derive the attachment type
 	// and, while we're in here, clean up exif data from
 	// the image if we already know the fileSize
+	var readerToStore io.Reader
 	switch extension {
 	case mimeGif:
 		p.attachment.Type = gtsmodel.FileTypeImage
+		// nothing to terminate, we can just store the multireader
+		readerToStore = multiReader
 	case mimeJpeg, mimePng:
 		p.attachment.Type = gtsmodel.FileTypeImage
 		if fileSize > 0 {
-			var err error
-			readerToStore, err = terminator.Terminate(readerToStore, int(fileSize), extension)
+			terminated, err := terminator.Terminate(multiReader, int(fileSize), extension)
 			if err != nil {
 				return fmt.Errorf("store: exif error: %s", err)
 			}
 			defer func() {
-				if rc, ok := readerToStore.(io.ReadCloser); ok {
-					if err := rc.Close(); err != nil {
+				if closer, ok := terminated.(io.Closer); ok {
+					if err := closer.Close(); err != nil {
 						log.Errorf("store: error closing terminator reader: %s", err)
 					}
 				}
 			}()
+			// store the exif-terminated version of what was in the multireader
+			readerToStore = terminated
+		} else {
+			// can't terminate if we don't know the file size, so just store the multiReader
+			readerToStore = multiReader
 		}
 	default:
 		return fmt.Errorf("store: couldn't process %s", extension)
@@ -347,10 +361,6 @@ func (p *ProcessingMedia) store(ctx context.Context) error {
 	p.attachment.File.FileSize = int(fileSize)
 	p.read = true
 
-	if p.postData != nil {
-		return p.postData(ctx)
-	}
-
 	log.Tracef("store: finished storing initial data for attachment %s", p.attachment.URL)
 	return nil
 }
diff --git a/internal/media/pruneremote_test.go b/internal/media/pruneremote_test.go
index e71d3310b..d3a01b7be 100644
--- a/internal/media/pruneremote_test.go
+++ b/internal/media/pruneremote_test.go
@@ -74,13 +74,13 @@ func (suite *PruneRemoteTestSuite) TestPruneAndRecache() {
 	suite.ErrorIs(err, storage.ErrNotFound)
 
 	// now recache the image....
-	data := func(_ context.Context) (io.Reader, int64, error) {
+	data := func(_ context.Context) (io.ReadCloser, int64, error) {
 		// load bytes from a test image
 		b, err := os.ReadFile("../../testrig/media/thoughtsofdog-original.jpeg")
 		if err != nil {
 			panic(err)
 		}
-		return bytes.NewBuffer(b), int64(len(b)), nil
+		return io.NopCloser(bytes.NewBuffer(b)), int64(len(b)), nil
 	}
 	processingRecache, err := suite.manager.RecacheMedia(ctx, data, nil, testAttachment.ID)
 	suite.NoError(err)
diff --git a/internal/media/types.go b/internal/media/types.go
index 3238916b8..763a8137f 100644
--- a/internal/media/types.go
+++ b/internal/media/types.go
@@ -118,7 +118,7 @@ type AdditionalEmojiInfo struct {
 }
 
 // DataFunc represents a function used to retrieve the raw bytes of a piece of media.
-type DataFunc func(ctx context.Context) (reader io.Reader, fileSize int64, err error)
+type DataFunc func(ctx context.Context) (reader io.ReadCloser, fileSize int64, err error)
 
 // PostDataCallbackFunc represents a function executed after the DataFunc has been executed,
 // and the returned reader has been read. It can be used to clean up any remaining resources.