From 73e9cca7019c15a92cb4cd320034652590513198 Mon Sep 17 00:00:00 2001
From: tobi <31960611+tsmethurst@users.noreply.github.com>
Date: Mon, 21 Mar 2022 13:41:44 +0100
Subject: [PATCH] [bugfix] Close ReadClosers properly in the media package
 (#434)

* defer lock reader

* close readers when finished with them

* close the reader in the teereader when finished
---
 internal/media/processingemoji.go    | 17 ++++++++++-------
 internal/media/processingmedia.go    | 27 +++++++++++++++++++--------
 internal/processing/media/getfile.go | 12 +++++++++---
 internal/processing/media/util.go    | 14 ++++++++++++++
 4 files changed, 52 insertions(+), 18 deletions(-)

diff --git a/internal/media/processingemoji.go b/internal/media/processingemoji.go
index 2a3cd1613..c8c8d18c8 100644
--- a/internal/media/processingemoji.go
+++ b/internal/media/processingemoji.go
@@ -29,6 +29,7 @@ import (
 	"time"
 
 	"codeberg.org/gruf/go-store/kv"
+	"github.com/sirupsen/logrus"
 	"github.com/superseriousbusiness/gotosocial/internal/db"
 	"github.com/superseriousbusiness/gotosocial/internal/gtsmodel"
 	"github.com/superseriousbusiness/gotosocial/internal/uris"
@@ -169,6 +170,15 @@ func (p *ProcessingEmoji) store(ctx context.Context) error {
 		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 {
+				logrus.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 {
@@ -205,13 +215,6 @@ func (p *ProcessingEmoji) store(ctx context.Context) error {
 		return fmt.Errorf("store: error storing stream: %s", err)
 	}
 
-	// if the original reader is a readcloser, close it since we're done with it now
-	if rc, ok := reader.(io.ReadCloser); ok {
-		if err := rc.Close(); err != nil {
-			return fmt.Errorf("store: error closing readcloser: %s", err)
-		}
-	}
-
 	p.read = true
 
 	if p.postData != nil {
diff --git a/internal/media/processingmedia.go b/internal/media/processingmedia.go
index 634d4eb48..3cf4805f7 100644
--- a/internal/media/processingmedia.go
+++ b/internal/media/processingmedia.go
@@ -29,6 +29,7 @@ import (
 	"time"
 
 	"codeberg.org/gruf/go-store/kv"
+	"github.com/sirupsen/logrus"
 	terminator "github.com/superseriousbusiness/exif-terminator"
 	"github.com/superseriousbusiness/gotosocial/internal/db"
 	"github.com/superseriousbusiness/gotosocial/internal/gtsmodel"
@@ -260,6 +261,15 @@ func (p *ProcessingMedia) store(ctx context.Context) error {
 		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 {
+				logrus.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 {
@@ -305,6 +315,15 @@ func (p *ProcessingMedia) store(ctx context.Context) error {
 		return fmt.Errorf("store: couldn't process %s", extension)
 	}
 
+	// defer closing the clean reader when we're done with it
+	defer func() {
+		if rc, ok := clean.(io.ReadCloser); ok {
+			if err := rc.Close(); err != nil {
+				logrus.Errorf("store: error closing clean readcloser: %s", err)
+			}
+		}
+	}()
+
 	// now set some additional fields on the attachment since
 	// we know more about what the underlying media actually is
 	p.attachment.URL = uris.GenerateURIForAttachment(p.attachment.AccountID, string(TypeAttachment), string(SizeOriginal), p.attachment.ID, extension)
@@ -317,14 +336,6 @@ func (p *ProcessingMedia) store(ctx context.Context) error {
 		return fmt.Errorf("store: error storing stream: %s", err)
 	}
 	p.attachment.Cached = true
-
-	// if the original reader is a readcloser, close it since we're done with it now
-	if rc, ok := reader.(io.ReadCloser); ok {
-		if err := rc.Close(); err != nil {
-			return fmt.Errorf("store: error closing readcloser: %s", err)
-		}
-	}
-
 	p.read = true
 
 	if p.postData != nil {
diff --git a/internal/processing/media/getfile.go b/internal/processing/media/getfile.go
index 1faa8702f..c74951e38 100644
--- a/internal/processing/media/getfile.go
+++ b/internal/processing/media/getfile.go
@@ -180,9 +180,15 @@ func (p *processor) getAttachmentContent(ctx context.Context, requestingAccount
 				return nil, 0, err
 			}
 
-			// everything read from the readCloser by the media manager will be written into the bufferedWriter
-			teeReader := io.TeeReader(readCloser, bufferedWriter)
-			return teeReader, fileSize, nil
+			// Make a TeeReader so that everything read from the readCloser by the media manager will be written into the bufferedWriter.
+			// We wrap this in a teeReadCloser which implements io.ReadCloser, so that whoever uses the teeReader can close the readCloser
+			// when they're done with it.
+			trc := teeReadCloser{
+				teeReader: io.TeeReader(readCloser, bufferedWriter),
+				close:     readCloser.Close,
+			}
+
+			return trc, fileSize, nil
 		}
 
 		// close the pipewriter after data has been piped into it, so the reader on the other side doesn't block;
diff --git a/internal/processing/media/util.go b/internal/processing/media/util.go
index 37dc87979..9739e70b7 100644
--- a/internal/processing/media/util.go
+++ b/internal/processing/media/util.go
@@ -20,6 +20,7 @@ package media
 
 import (
 	"fmt"
+	"io"
 	"strconv"
 	"strings"
 )
@@ -61,3 +62,16 @@ func parseFocus(focus string) (focusx, focusy float32, err error) {
 	focusy = float32(fy)
 	return
 }
+
+type teeReadCloser struct {
+	teeReader io.Reader
+	close     func() error
+}
+
+func (t teeReadCloser) Read(p []byte) (n int, err error) {
+	return t.teeReader.Read(p)
+}
+
+func (t teeReadCloser) Close() error {
+	return t.close()
+}