Better check for max file size

Author: Forceu

internal/storage/chunking/Chunking.go

@@ -202,19 +202,22 @@ func FileExists(id string) bool {
 }
 
 // NewChunk allocates the space for the new file and writes the chunk
-func NewChunk(chunkContent io.Reader, fileHeader *multipart.FileHeader, info ChunkInfo) error {
-	err := allocateFile(info)
+func NewChunk(chunkContent io.Reader, fileHeader *multipart.FileHeader, info ChunkInfo, maxAllowedSize int64) error {
+	err := allocateFile(info, maxAllowedSize)
 	if err != nil {
 		return err
 	}
 	return writeChunk(chunkContent, fileHeader, info)
 }
 
-func allocateFile(info ChunkInfo) error {
+func allocateFile(info ChunkInfo, maxAllowedSize int64) error {
+	if maxAllowedSize <= 0 {
+		return errors.New("invalid maxAllowedSize")
+	}
 	if FileExists(info.UUID) {
 		return nil
 	}
-	maxSizeBytes := int64(configuration.Get().MaxFileSizeMB) * 1024 * 1024
+	maxSizeBytes := min(int64(configuration.Get().MaxFileSizeMB)*1024*1024, maxAllowedSize)
 	if info.TotalFilesizeBytes > maxSizeBytes {
 		return errors.New("declared file size exceeds the maximum allowed size")
 	}

internal/storage/chunking/Chunking_test.go

@@ -4,18 +4,19 @@ import (
 	"bytes"
 	"crypto/sha1"
 	"encoding/hex"
-	"github.com/forceu/gokapi/internal/configuration"
-	"github.com/forceu/gokapi/internal/helper"
-	"github.com/forceu/gokapi/internal/test"
-	"github.com/forceu/gokapi/internal/test/testconfiguration"
-	"github.com/juju/ratelimit"
-	"golang.org/x/sync/errgroup"
 	"mime/multipart"
 	"net/textproto"
 	"net/url"
 	"os"
 	"strings"
 	"testing"
+
+	"github.com/forceu/gokapi/internal/configuration"
+	"github.com/forceu/gokapi/internal/helper"
+	"github.com/forceu/gokapi/internal/test"
+	"github.com/forceu/gokapi/internal/test/testconfiguration"
+	"github.com/juju/ratelimit"
+	"golang.org/x/sync/errgroup"
 )
 
 func TestMain(m *testing.M) {
@@ -286,32 +287,32 @@ func TestNewChunk(t *testing.T) {
 	header := multipart.FileHeader{
 		Size: 21,
 	}
-	err := NewChunk(strings.NewReader("This is a test content"), &header, info)
+	err := NewChunk(strings.NewReader("This is a test content"), &header, info, 100000)
 	test.IsNil(t, err)
 	test.IsEqualString(t, sha1sumFile("test/data/chunk-testuuid12345"), "a69ec3c3a031e3540d0c2a864ca931f3d54e2c13")
 
 	info.Offset = 52
 	header = multipart.FileHeader{
 		Size: 11,
 	}
-	err = NewChunk(strings.NewReader("More content"), &header, info)
+	err = NewChunk(strings.NewReader("More content"), &header, info, 100000)
 	test.IsNil(t, err)
 	test.IsEqualString(t, sha1sumFile("test/data/chunk-testuuid12345"), "8794d8352fae46b83bab83d3e613dde8f0244ded")
 
 	info.Offset = 99
-	err = NewChunk(strings.NewReader("More content"), &header, info)
+	err = NewChunk(strings.NewReader("More content"), &header, info, 100000)
 	test.IsNotNil(t, err)
 
 	err = os.Remove("test/data/chunk-testuuid12345")
 	test.IsNil(t, err)
 
 	info.TotalFilesizeBytes = -4
-	err = NewChunk(strings.NewReader("More content"), &header, info)
+	err = NewChunk(strings.NewReader("More content"), &header, info, 100000)
 	test.IsNotNil(t, err)
 
 	info.TotalFilesizeBytes = 100
 	info.UUID = "../../../../../../../../../../invalid"
-	err = NewChunk(strings.NewReader("More content"), &header, info)
+	err = NewChunk(strings.NewReader("More content"), &header, info, 100000)
 	test.IsNotNil(t, err)
 
 	// Testing simultaneous writes
@@ -341,7 +342,7 @@ func writeRateLimitedChunk(firstHalf bool) error {
 	}
 	content := []byte(helper.GenerateRandomString(500 * 1024))
 	bucket := ratelimit.NewBucketWithRate(400*1024, 400*1024)
-	err := NewChunk(ratelimit.Reader(bytes.NewReader(content), bucket), &header, info)
+	err := NewChunk(ratelimit.Reader(bytes.NewReader(content), bucket), &header, info, 2000000)
 	return err
 }
 

internal/webserver/Webserver.go

@@ -1005,7 +1005,7 @@ func uploadChunk(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	r.Body = http.MaxBytesReader(w, r.Body, maxUpload)
-	_, err := fileupload.ProcessNewChunk(w, r, false, "")
+	_, err := fileupload.ProcessNewChunk(w, r, false, "", maxUpload)
 	responseError(w, err)
 }
 

internal/webserver/api/Api.go

@@ -476,7 +476,7 @@ func processNewChunk(w http.ResponseWriter, request chunkParams, maxFileSizeMb i
 		return http.StatusBadRequest, errorcodes.FileTooLarge, storage.ErrorFileTooLarge.Error()
 	}
 	request.GetRequest().Body = http.MaxBytesReader(w, request.GetRequest().Body, maxUpload)
-	errCode, err := fileupload.ProcessNewChunk(w, request.GetRequest(), true, filerequestId)
+	errCode, err := fileupload.ProcessNewChunk(w, request.GetRequest(), true, filerequestId, maxUpload)
 	if err != nil {
 		return http.StatusBadRequest, errCode, err.Error()
 	}

internal/webserver/fileupload/FileUpload.go

@@ -65,7 +65,7 @@ func isChunkMinChunkSize(r *http.Request, offset, fileSize int64) bool {
 }
 
 // ProcessNewChunk processes a file chunk upload request
-func ProcessNewChunk(w http.ResponseWriter, r *http.Request, isApiCall bool, filerequestId string) (int, error) {
+func ProcessNewChunk(w http.ResponseWriter, r *http.Request, isApiCall bool, filerequestId string, maxFileSize int64) (int, error) {
 	err := r.ParseMultipartForm(int64(configuration.Get().MaxMemory) * 1024 * 1024)
 	if err != nil {
 		return errorcodes.CannotParse, err
@@ -83,14 +83,17 @@ func ProcessNewChunk(w http.ResponseWriter, r *http.Request, isApiCall bool, fil
 	if !isChunkMinChunkSize(r, chunkInfo.Offset, chunkInfo.TotalFilesizeBytes) {
 		return errorcodes.ChunkTooSmall, storage.ErrorChunkTooSmall
 	}
+	if chunkInfo.TotalFilesizeBytes > maxFileSize || chunkInfo.Offset > maxFileSize {
+		return errorcodes.FileTooLarge, storage.ErrorFileTooLarge
+	}
 
 	if filerequestId != "" {
 		if !chunkreservation.SetUploading(filerequestId, chunkInfo.UUID) {
 			return errorcodes.InvalidChunkReservation, errors.New("chunk reservation has expired or was not requested")
 		}
 	}
 
-	err = chunking.NewChunk(file, header, chunkInfo)
+	err = chunking.NewChunk(file, header, chunkInfo, maxFileSize)
 	defer file.Close()
 	if err != nil {
 		return errorcodes.CannotAllocateFile, err

internal/webserver/fileupload/FileUpload_test.go

@@ -99,17 +99,17 @@ func TestProcess(t *testing.T) {
 
 func TestProcessNewChunk(t *testing.T) {
 	w, r := test.GetRecorder("POST", "/uploadChunk", nil, nil, strings.NewReader("invalid§$%&%§"))
-	_, err := ProcessNewChunk(w, r, false, "")
+	_, err := ProcessNewChunk(w, r, false, "", 100*1000*1000)
 	test.IsNotNil(t, err)
 
 	w = httptest.NewRecorder()
 	r = getFileUploadRecorder(false)
-	_, err = ProcessNewChunk(w, r, false, "")
+	_, err = ProcessNewChunk(w, r, false, "", 100*1000*1000)
 	test.IsNotNil(t, err)
 
 	w = httptest.NewRecorder()
 	r = getFileUploadRecorder(true)
-	_, err = ProcessNewChunk(w, r, false, "")
+	_, err = ProcessNewChunk(w, r, false, "", 100*1000*1000)
 	test.IsNil(t, err)
 	response, err := io.ReadAll(w.Result().Body)
 	test.IsNil(t, err)