Use maximum size to reject file upload as needed

changelogs/unreleased/authorize-upload-filesize.yml

+---
+title: Reject upload when filesize exceeds MaximumSize returned by authorize endpoint
+merge_request:
+author:
+type: added

internal/api/api.go

@@ -148,6 +148,8 @@ type Response struct {
 	ProcessLsif bool
 	// Detects whether LSIF artifact will be parsed with references
 	ProcessLsifReferences bool
+	// The maximum accepted size in bytes of the upload
+	MaximumSize int64
 }
 
 // singleJoiningSlash is taken from reverseproxy.go:NewSingleHostReverseProxy

internal/artifacts/artifacts_store_test.go

@@ -260,13 +260,12 @@ func TestUploadHandlerMultipartUploadSizeLimit(t *testing.T) {
 
 	objectURL := server.URL + test.ObjectPath
 
-	partSize := int64(1)
 	uploadSize := 10
 	preauth := api.Response{
 		RemoteObject: api.RemoteObject{
 			ID: "store-id",
 			MultipartUpload: &api.MultipartUploadParams{
-				PartSize:    partSize,
+				PartSize:    1,
 				PartURLs:    []string{objectURL + "?partNumber=1"},
 				AbortURL:    objectURL, // DELETE
 				CompleteURL: objectURL, // POST
@@ -292,3 +291,47 @@ func TestUploadHandlerMultipartUploadSizeLimit(t *testing.T) {
 	assert.False(t, os.IsMultipartUpload(test.ObjectPath), "MultipartUpload should not be in progress anymore")
 	assert.Empty(t, os.GetObjectMD5(test.ObjectPath), "upload should have failed, so the object should not exists")
 }
+
+func TestUploadHandlerMultipartUploadMaximumSizeFromApi(t *testing.T) {
+	os, server := test.StartObjectStore()
+	defer server.Close()
+
+	err := os.InitiateMultipartUpload(test.ObjectPath)
+	require.NoError(t, err)
+
+	objectURL := server.URL + test.ObjectPath
+
+	uploadSize := int64(10)
+	maxSize := uploadSize - 1
+	preauth := api.Response{
+		MaximumSize: maxSize,
+		RemoteObject: api.RemoteObject{
+			ID: "store-id",
+			MultipartUpload: &api.MultipartUploadParams{
+				PartSize:    uploadSize,
+				PartURLs:    []string{objectURL + "?partNumber=1"},
+				AbortURL:    objectURL, // DELETE
+				CompleteURL: objectURL, // POST
+			},
+		},
+	}
+
+	responseProcessor := func(w http.ResponseWriter, r *http.Request) {
+		t.Fatal("it should not be called")
+	}
+
+	ts := testArtifactsUploadServer(t, preauth, responseProcessor)
+	defer ts.Close()
+
+	contentBuffer, contentType := createTestMultipartForm(t, make([]byte, uploadSize))
+	response := testUploadArtifacts(t, contentType, ts.URL+Path, &contentBuffer)
+	require.Equal(t, http.StatusRequestEntityTooLarge, response.Code)
+
+	testhelper.Retry(t, 5*time.Second, func() error {
+		if os.GetObjectMD5(test.ObjectPath) == "" {
+			return nil
+		}
+
+		return fmt.Errorf("file is still present")
+	})
+}

internal/filestore/file_handler.go

@@ -166,7 +166,14 @@ func SaveFileFromReader(ctx context.Context, reader io.Reader, size int64, opts
 		if err != nil {
 			return nil, err
 		}
+
 		writers = append(writers, remoteWriter)
+
+		defer func() {
+			if err != nil {
+				remoteWriter.CloseWithError(err)
+			}
+		}()
 	} else {
 		clientMode = "local"
 
@@ -182,12 +189,26 @@ func SaveFileFromReader(ctx context.Context, reader io.Reader, size int64, opts
 		return nil, errors.New("missing upload destination")
 	}
 
+	if opts.MaximumSize > 0 {
+		if size > opts.MaximumSize {
+			return nil, SizeError(fmt.Errorf("the upload size %d is over maximum of %d bytes", size, opts.MaximumSize))
+		}
+
+		// We allow to read an extra byte to check later if we exceed the max size
+		reader = &io.LimitedReader{R: reader, N: opts.MaximumSize + 1}
+	}
+
 	multiWriter := io.MultiWriter(writers...)
 	fh.Size, err = io.Copy(multiWriter, reader)
 	if err != nil {
 		return nil, err
 	}
 
+	if opts.MaximumSize > 0 && fh.Size > opts.MaximumSize {
+		// An extra byte was read thus exceeding the max size
+		return nil, ErrEntityTooLarge
+	}
+
 	if size != -1 && size != fh.Size {
 		return nil, SizeError(fmt.Errorf("expected %d bytes but got only %d", size, fh.Size))
 	}

internal/filestore/file_handler_test.go

@@ -69,6 +69,36 @@ func TestSaveFileWrongSize(t *testing.T) {
 	assert.Nil(t, fh)
 }
 
+func TestSaveFileWithKnownSizeExceedLimit(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	tmpFolder, err := ioutil.TempDir("", "workhorse-test-tmp")
+	require.NoError(t, err)
+	defer os.RemoveAll(tmpFolder)
+
+	opts := &filestore.SaveFileOpts{LocalTempPath: tmpFolder, TempFilePrefix: "test-file", MaximumSize: test.ObjectSize - 1}
+	fh, err := filestore.SaveFileFromReader(ctx, strings.NewReader(test.ObjectContent), test.ObjectSize, opts)
+	assert.Error(t, err)
+	_, isSizeError := err.(filestore.SizeError)
+	assert.True(t, isSizeError, "Should fail with SizeError")
+	assert.Nil(t, fh)
+}
+
+func TestSaveFileWithUnknownSizeExceedLimit(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	tmpFolder, err := ioutil.TempDir("", "workhorse-test-tmp")
+	require.NoError(t, err)
+	defer os.RemoveAll(tmpFolder)
+
+	opts := &filestore.SaveFileOpts{LocalTempPath: tmpFolder, TempFilePrefix: "test-file", MaximumSize: test.ObjectSize - 1}
+	fh, err := filestore.SaveFileFromReader(ctx, strings.NewReader(test.ObjectContent), -1, opts)
+	assert.Equal(t, err, filestore.ErrEntityTooLarge)
+	assert.Nil(t, fh)
+}
+
 func TestSaveFromDiskNotExistingFile(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()

internal/filestore/save_file_opts.go

@@ -51,6 +51,8 @@ type SaveFileOpts struct {
 	ObjectStorageConfig ObjectStorageConfig
 	// Deadline it the S3 operation deadline, the upload will be aborted if not completed in time
 	Deadline time.Time
+	// The maximum accepted size in bytes of the upload
+	MaximumSize int64
 
 	//MultipartUpload parameters
 	// PartSize is the exact size of each uploaded part. Only the last one can be smaller
@@ -95,6 +97,7 @@ func GetOpts(apiResponse *api.Response) (*SaveFileOpts, error) {
 		UseWorkhorseClient: apiResponse.RemoteObject.UseWorkhorseClient,
 		RemoteTempObjectID: apiResponse.RemoteObject.RemoteTempObjectID,
 		Deadline:           time.Now().Add(timeout),
+		MaximumSize:        apiResponse.MaximumSize,
 	}
 
 	if opts.LocalTempPath != "" && opts.RemoteID != "" {

internal/objectstore/gocloud_object_test.go

@@ -49,22 +49,18 @@ func TestGoCloudObjectUpload(t *testing.T) {
 	require.Equal(t, []byte(test.ObjectContent), received)
 
 	cancel()
-	deleted := false
 
-	retry(3, time.Second, func() error {
+	testhelper.Retry(t, 5*time.Second, func() error {
 		exists, err := bucket.Exists(ctx, objectName)
 		require.NoError(t, err)
 
 		if exists {
-			return fmt.Errorf("file %s is still present, retrying", objectName)
+			return fmt.Errorf("file %s is still present", objectName)
 		} else {
-			deleted = true
 			return nil
 		}
 	})
 
-	require.True(t, deleted)
-
 	// Verify no log noise when deleting a file that already is gone
 	object.Delete()
 	entries := logHook.AllEntries()

internal/objectstore/s3_object_test.go

@@ -19,6 +19,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-workhorse/internal/config"
 	"gitlab.com/gitlab-org/gitlab-workhorse/internal/objectstore"
 	"gitlab.com/gitlab-org/gitlab-workhorse/internal/objectstore/test"
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/testhelper"
 )
 
 func TestS3ObjectUpload(t *testing.T) {
@@ -59,18 +60,14 @@ func TestS3ObjectUpload(t *testing.T) {
 			test.CheckS3Metadata(t, sess, config, objectName)
 
 			cancel()
-			deleted := false
 
-			retry(3, time.Second, func() error {
+			testhelper.Retry(t, 5*time.Second, func() error {
 				if test.S3ObjectDoesNotExist(t, sess, config, objectName) {
-					deleted = true
 					return nil
-				} else {
-					return fmt.Errorf("file is still present, retrying")
 				}
-			})
 
-			require.True(t, deleted)
+				return fmt.Errorf("file is still present")
+			})
 		})
 	}
 }
@@ -153,23 +150,3 @@ func TestS3ObjectUploadCancel(t *testing.T) {
 	_, err = io.Copy(object, strings.NewReader(test.ObjectContent))
 	require.Error(t, err)
 }
-
-func retry(attempts int, sleep time.Duration, fn func() error) error {
-	if err := fn(); err != nil {
-		if s, ok := err.(stop); ok {
-			// Return the original error for later checking
-			return s.error
-		}
-
-		if attempts--; attempts > 0 {
-			time.Sleep(sleep)
-			return retry(attempts, 2*sleep, fn)
-		}
-		return err
-	}
-	return nil
-}
-
-type stop struct {
-	error
-}

internal/objectstore/uploader.go

@@ -8,6 +8,7 @@ import (
 	"hash"
 	"io"
 	"strings"
+	"sync"
 	"time"
 
 	"gitlab.com/gitlab-org/labkit/log"
@@ -16,6 +17,7 @@ import (
 // Upload represents an upload to an ObjectStorage provider
 type Upload interface {
 	io.WriteCloser
+	CloseWithError(error) error
 	ETag() string
 }
 
@@ -28,7 +30,6 @@ type uploader struct {
 	md5 hash.Hash
 
 	w io.Writer
-	c io.Closer
 
 	// uploadError is the last error occourred during upload
 	uploadError error
@@ -39,25 +40,34 @@ type uploader struct {
 	pw       *io.PipeWriter
 	strategy uploadStrategy
 	metrics  bool
+
+	// closeOnce is used to prevent multiple calls to pw.Close
+	// which may result to Close overriding the error set by CloseWithError
+	// Bug fixed in v1.14: https://github.com/golang/go/commit/f45eb9ff3c96dfd951c65d112d033ed7b5e02432
+	closeOnce sync.Once
 }
 
 func newUploader(strategy uploadStrategy) uploader {
 	pr, pw := io.Pipe()
-	return uploader{w: pw, c: pw, pr: pr, pw: pw, strategy: strategy, metrics: true}
+	return uploader{w: pw, pr: pr, pw: pw, strategy: strategy, metrics: true}
 }
 
 func newMD5Uploader(strategy uploadStrategy, metrics bool) uploader {
 	pr, pw := io.Pipe()
 	hasher := md5.New()
 	mw := io.MultiWriter(pw, hasher)
-	return uploader{w: mw, c: pw, pr: pr, pw: pw, md5: hasher, strategy: strategy, metrics: metrics}
+	return uploader{w: mw, pr: pr, pw: pw, md5: hasher, strategy: strategy, metrics: metrics}
 }
 
 // Close implements the standard io.Closer interface: it closes the http client request.
 // This method will also wait for the connection to terminate and return any error occurred during the upload
 func (u *uploader) Close() error {
-	if err := u.c.Close(); err != nil {
-		return err
+	var closeError error
+	u.closeOnce.Do(func() {
+		closeError = u.pw.Close()
+	})
+	if closeError != nil {
+		return closeError
 	}
 
 	<-u.ctx.Done()
@@ -69,6 +79,14 @@ func (u *uploader) Close() error {
 	return u.uploadError
 }
 
+func (u *uploader) CloseWithError(err error) error {
+	u.closeOnce.Do(func() {
+		u.pw.CloseWithError(err)
+	})
+
+	return nil
+}
+
 func (u *uploader) Write(p []byte) (int, error) {
 	return u.w.Write(p)
 }

internal/testhelper/testhelper.go

@@ -161,3 +161,16 @@ func WaitForLogEvent(hook *test.Hook) bool {
 
 	return false
 }
+
+func Retry(t testing.TB, timeout time.Duration, fn func() error) {
+	t.Helper()
+	start := time.Now()
+	var err error
+	for ; time.Since(start) < timeout; time.Sleep(time.Millisecond) {
+		err = fn()
+		if err == nil {
+			return
+		}
+	}
+	t.Fatalf("test timeout after %v; last error: %v", timeout, err)
+}