From 0b9f69e269149f932aea0e602021bd0e7ef8fd43 Mon Sep 17 00:00:00 2001
From: Till Faelligen <tfaelligen@gmail.com>
Date: Sun, 25 Apr 2021 12:01:37 +0200
Subject: [PATCH] Check if bytesWritten is equal to the maxFileSize Add tests

Signed-off-by: Till Faelligen <tfaelligen@gmail.com>
---
 mediaapi/routing/upload.go      |   5 +-
 mediaapi/routing/upload_test.go | 117 ++++++++++++++++++++++++++++++++
 2 files changed, 120 insertions(+), 2 deletions(-)
 create mode 100644 mediaapi/routing/upload_test.go

diff --git a/mediaapi/routing/upload.go b/mediaapi/routing/upload.go
index 2c5d92525..06f62fa13 100644
--- a/mediaapi/routing/upload.go
+++ b/mediaapi/routing/upload.go
@@ -159,8 +159,9 @@ func (r *uploadRequest) doUpload(
 		}
 	}
 
-	// Check if temp file size exceeds max file size configuration
-	if bytesWritten > types.FileSizeBytes(*cfg.MaxFileSizeBytes) {
+	// Check if temp file size is greater (should not happen, LimitReader stops when the defined size is reached.)
+	// or equal to the max file size configuration.
+	if bytesWritten >= types.FileSizeBytes(*cfg.MaxFileSizeBytes) {
 		fileutils.RemoveDir(tmpDir, r.Logger) // delete temp file
 		return requestEntityTooLargeJSONResponse(*cfg.MaxFileSizeBytes)
 	}
diff --git a/mediaapi/routing/upload_test.go b/mediaapi/routing/upload_test.go
new file mode 100644
index 000000000..b1fdac908
--- /dev/null
+++ b/mediaapi/routing/upload_test.go
@@ -0,0 +1,117 @@
+package routing
+
+import (
+	"context"
+	"io"
+	"os"
+	"path/filepath"
+	"reflect"
+	"strings"
+	"testing"
+
+	"github.com/matrix-org/dendrite/mediaapi/fileutils"
+	"github.com/matrix-org/dendrite/mediaapi/storage"
+	"github.com/matrix-org/dendrite/mediaapi/types"
+	"github.com/matrix-org/dendrite/setup/config"
+	"github.com/matrix-org/util"
+	log "github.com/sirupsen/logrus"
+)
+
+func Test_uploadRequest_doUpload(t *testing.T) {
+	type fields struct {
+		MediaMetadata *types.MediaMetadata
+		Logger        *log.Entry
+	}
+	type args struct {
+		ctx                       context.Context
+		reqReader                 io.Reader
+		cfg                       *config.MediaAPI
+		db                        storage.Database
+		activeThumbnailGeneration *types.ActiveThumbnailGeneration
+	}
+
+	wd, err := os.Getwd()
+	if err != nil {
+		t.Errorf("failed to get current working directory: %v", err)
+	}
+
+	maxSize := config.FileSizeBytes(8)
+	logger := log.New().WithField("mediaapi", "test")
+	testdataPath := filepath.Join(wd, "./testdata")
+
+	cfg := &config.MediaAPI{
+		MaxFileSizeBytes:  &maxSize,
+		BasePath:          config.Path(testdataPath),
+		AbsBasePath:       config.Path(testdataPath),
+		DynamicThumbnails: false,
+	}
+
+	// create testdata folder and remove when done
+	_ = os.Mkdir(testdataPath, os.ModePerm)
+	t.Cleanup(func() {
+		fileutils.RemoveDir(types.Path(testdataPath), nil)
+	})
+
+	db, err := storage.Open(&config.DatabaseOptions{
+		ConnectionString:       "file::memory:?cache=shared",
+		MaxOpenConnections:     100,
+		MaxIdleConnections:     2,
+		ConnMaxLifetimeSeconds: -1,
+	})
+	if err != nil {
+		t.Errorf("error opening mediaapi database: %v", err)
+	}
+
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+		want   *util.JSONResponse
+	}{
+		{
+			name: "upload ok",
+			args: args{
+				ctx:       context.Background(),
+				reqReader: strings.NewReader("test"),
+				cfg:       cfg,
+				db:        db,
+			},
+			fields: fields{
+				Logger: logger,
+				MediaMetadata: &types.MediaMetadata{
+					MediaID:    "1337",
+					UploadName: "test ok",
+				},
+			},
+			want: nil,
+		},
+		{
+			name: "upload not ok",
+			args: args{
+				ctx:       context.Background(),
+				reqReader: strings.NewReader("test test test"),
+				cfg:       cfg,
+				db:        db,
+			},
+			fields: fields{
+				Logger: logger,
+				MediaMetadata: &types.MediaMetadata{
+					MediaID:    "1337",
+					UploadName: "test fail",
+				},
+			},
+			want: requestEntityTooLargeJSONResponse(maxSize),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			r := &uploadRequest{
+				MediaMetadata: tt.fields.MediaMetadata,
+				Logger:        tt.fields.Logger,
+			}
+			if got := r.doUpload(tt.args.ctx, tt.args.reqReader, tt.args.cfg, tt.args.db, tt.args.activeThumbnailGeneration); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("doUpload() = %+v, want %+v", got, tt.want)
+			}
+		})
+	}
+}