From 63cbbcaa07bbff7c0023b74163b9940ddd6eb0ff Mon Sep 17 00:00:00 2001
From: Avi Deitcher <avi@deitcher.net>
Date: Mon, 4 Nov 2024 15:06:06 +0200
Subject: [PATCH] properly calculate bytes for S3.Push

Signed-off-by: Avi Deitcher <avi@deitcher.net>
---
 pkg/storage/s3/reader.go | 29 +++++++++++++++++++++++++++++
 pkg/storage/s3/s3.go     |  5 +++--
 2 files changed, 32 insertions(+), 2 deletions(-)
 create mode 100644 pkg/storage/s3/reader.go

diff --git a/pkg/storage/s3/reader.go b/pkg/storage/s3/reader.go
new file mode 100644
index 0000000..c77e9f8
--- /dev/null
+++ b/pkg/storage/s3/reader.go
@@ -0,0 +1,29 @@
+package s3
+
+import (
+	"io"
+)
+
+type CountingReader interface {
+	io.Reader
+	Bytes() int64
+}
+
+func NewCountingReader(r io.Reader) CountingReader {
+	return &countingReader{r: r}
+}
+
+type countingReader struct {
+	r     io.Reader
+	bytes int64
+}
+
+func (cr *countingReader) Read(p []byte) (int, error) {
+	n, err := cr.r.Read(p)
+	cr.bytes += int64(n)
+	return n, err
+}
+
+func (cr *countingReader) Bytes() int64 {
+	return cr.bytes
+}
diff --git a/pkg/storage/s3/s3.go b/pkg/storage/s3/s3.go
index e6f2793..450af65 100644
--- a/pkg/storage/s3/s3.go
+++ b/pkg/storage/s3/s3.go
@@ -113,6 +113,7 @@ func (s *S3) Push(target, source string, logger *log.Entry) (int64, error) {
 		return 0, fmt.Errorf("failed to read input file %q, %v", source, err)
 	}
 	defer f.Close()
+	countingReader := NewCountingReader(f)
 
 	// S3 always prepends a /, so if it already has one, it would become //
 	// For some services, that is ok, but for others, it causes issues.
@@ -122,12 +123,12 @@ func (s *S3) Push(target, source string, logger *log.Entry) (int64, error) {
 	_, err = uploader.Upload(context.TODO(), &s3.PutObjectInput{
 		Bucket: aws.String(bucket),
 		Key:    aws.String(key),
-		Body:   f,
+		Body:   countingReader,
 	})
 	if err != nil {
 		return 0, fmt.Errorf("failed to upload file, %v", err)
 	}
-	return 0, nil
+	return countingReader.Bytes(), nil
 }
 
 func (s *S3) Clean(filename string) string {