[bugfix] Wrap media reader in length reader to determine length if no `content-length` given (#848)

* use lengthReader 2 determine fileSize if not given

* update tests

* small fixes

* go fmt

1 files changed, 31 insertions, 0 deletions

diff --git a/internal/media/util.go b/internal/media/util.go
index aea2ad990..b89196f87 100644
--- a/internal/media/util.go
+++ b/internal/media/util.go
@@ -19,12 +19,15 @@
 package media
 
 import (
+	"context"
 	"errors"
 	"fmt"
+	"io"
 	"time"
 
 	"github.com/h2non/filetype"
 	"github.com/superseriousbusiness/gotosocial/internal/log"
+	"github.com/superseriousbusiness/gotosocial/internal/storage"
 )
 
 // AllSupportedMIMETypes just returns all media
@@ -144,3 +147,31 @@ func parseOlderThan(olderThanDays int) (time.Time, error) {
 
 	return olderThan, nil
 }
+
+// lengthReader wraps a reader and reads the length of total bytes written as it goes.
+type lengthReader struct {
+	source io.Reader
+	length int
+}
+
+func (r *lengthReader) Read(b []byte) (int, error) {
+	n, err := r.source.Read(b)
+	r.length += n
+	return n, err
+}
+
+// putStream either puts a file with a known fileSize into storage directly, and returns the
+// fileSize unchanged, or it wraps the reader with a lengthReader and returns the discovered
+// fileSize.
+func putStream(ctx context.Context, storage storage.Driver, key string, r io.Reader, fileSize int) (int, error) {
+	if fileSize > 0 {
+		return fileSize, storage.PutStream(ctx, key, r)
+	}
+
+	lr := &lengthReader{
+		source: r,
+	}
+
+	err := storage.PutStream(ctx, key, lr)
+	return lr.length, err
+}