Check quota limits for container uploads (#22450)

The test coverage has revealed that container packages were not checked
against the quota limits.

5 files changed, 106 insertions, 31 deletions

diff --git a/routers/api/packages/container/blob.go b/routers/api/packages/container/blob.go
index 2e4309a2eb..f0457c55e1 100644
--- a/routers/api/packages/container/blob.go
+++ b/routers/api/packages/container/blob.go
@@ -26,14 +26,18 @@ var uploadVersionMutex sync.Mutex
 
 // saveAsPackageBlob creates a package blob from an upload
 // The uploaded blob gets stored in a special upload version to link them to the package/image
-func saveAsPackageBlob(hsr packages_module.HashedSizeReader, pi *packages_service.PackageInfo) (*packages_model.PackageBlob, error) {
+func saveAsPackageBlob(hsr packages_module.HashedSizeReader, pci *packages_service.PackageCreationInfo) (*packages_model.PackageBlob, error) {
+	if err := packages_service.CheckSizeQuotaExceeded(db.DefaultContext, pci.Creator, pci.Owner, packages_model.TypeContainer, hsr.Size()); err != nil {
+		return nil, err
+	}
+
 	pb := packages_service.NewPackageBlob(hsr)
 
 	exists := false
 
 	contentStore := packages_module.NewContentStore()
 
-	uploadVersion, err := getOrCreateUploadVersion(pi)
+	uploadVersion, err := getOrCreateUploadVersion(&pci.PackageInfo)
 	if err != nil {
 		return nil, err
 	}
diff --git a/routers/api/packages/container/container.go b/routers/api/packages/container/container.go
index 8b2c4e6bb2..c22cfb5009 100644
--- a/routers/api/packages/container/container.go
+++ b/routers/api/packages/container/container.go
@@ -227,8 +227,22 @@ func InitiateUploadBlob(ctx *context.Context) {
 			return
 		}
 
-		if _, err := saveAsPackageBlob(buf, &packages_service.PackageInfo{Owner: ctx.Package.Owner, Name: image}); err != nil {
-			apiError(ctx, http.StatusInternalServerError, err)
+		if _, err := saveAsPackageBlob(
+			buf,
+			&packages_service.PackageCreationInfo{
+				PackageInfo: packages_service.PackageInfo{
+					Owner: ctx.Package.Owner,
+					Name:  image,
+				},
+				Creator: ctx.Doer,
+			},
+		); err != nil {
+			switch err {
+			case packages_service.ErrQuotaTotalCount, packages_service.ErrQuotaTypeSize, packages_service.ErrQuotaTotalSize:
+				apiError(ctx, http.StatusForbidden, err)
+			default:
+				apiError(ctx, http.StatusInternalServerError, err)
+			}
 			return
 		}
 
@@ -358,8 +372,22 @@ func EndUploadBlob(ctx *context.Context) {
 		return
 	}
 
-	if _, err := saveAsPackageBlob(uploader, &packages_service.PackageInfo{Owner: ctx.Package.Owner, Name: image}); err != nil {
-		apiError(ctx, http.StatusInternalServerError, err)
+	if _, err := saveAsPackageBlob(
+		uploader,
+		&packages_service.PackageCreationInfo{
+			PackageInfo: packages_service.PackageInfo{
+				Owner: ctx.Package.Owner,
+				Name:  image,
+			},
+			Creator: ctx.Doer,
+		},
+	); err != nil {
+		switch err {
+		case packages_service.ErrQuotaTotalCount, packages_service.ErrQuotaTypeSize, packages_service.ErrQuotaTotalSize:
+			apiError(ctx, http.StatusForbidden, err)
+		default:
+			apiError(ctx, http.StatusInternalServerError, err)
+		}
 		return
 	}
 
@@ -526,7 +554,12 @@ func UploadManifest(ctx *context.Context) {
 		} else if errors.Is(err, container_model.ErrContainerBlobNotExist) {
 			apiErrorDefined(ctx, errBlobUnknown)
 		} else {
-			apiError(ctx, http.StatusInternalServerError, err)
+			switch err {
+			case packages_service.ErrQuotaTotalCount, packages_service.ErrQuotaTypeSize, packages_service.ErrQuotaTotalSize:
+				apiError(ctx, http.StatusForbidden, err)
+			default:
+				apiError(ctx, http.StatusInternalServerError, err)
+			}
 		}
 		return
 	}
diff --git a/routers/api/packages/container/manifest.go b/routers/api/packages/container/manifest.go
index 350933f3d2..491fb70639 100644
--- a/routers/api/packages/container/manifest.go
+++ b/routers/api/packages/container/manifest.go
@@ -327,6 +327,10 @@ func createPackageAndVersion(ctx context.Context, mci *manifestCreationInfo, met
 		}
 	}
 
+	if err := packages_service.CheckCountQuotaExceeded(ctx, mci.Creator, mci.Owner); err != nil {
+		return nil, err
+	}
+
 	if mci.IsTagged {
 		if _, err := packages_model.InsertProperty(ctx, packages_model.PropertyTypeVersion, pv.ID, container_module.PropertyManifestTagged, ""); err != nil {
 			log.Error("Error setting package version property: %v", err)
diff --git a/services/packages/packages.go b/services/packages/packages.go
index 410e73c048..754dfa7110 100644
--- a/services/packages/packages.go
+++ b/services/packages/packages.go
@@ -173,7 +173,7 @@ func createPackageAndVersion(ctx context.Context, pvci *PackageCreationInfo, all
 	}
 
 	if versionCreated {
-		if err := checkCountQuotaExceeded(ctx, pvci.Creator, pvci.Owner); err != nil {
+		if err := CheckCountQuotaExceeded(ctx, pvci.Creator, pvci.Owner); err != nil {
 			return nil, false, err
 		}
 
@@ -240,7 +240,7 @@ func NewPackageBlob(hsr packages_module.HashedSizeReader) *packages_model.Packag
 func addFileToPackageVersion(ctx context.Context, pv *packages_model.PackageVersion, pvi *PackageInfo, pfci *PackageFileCreationInfo) (*packages_model.PackageFile, *packages_model.PackageBlob, bool, error) {
 	log.Trace("Adding package file: %v, %s", pv.ID, pfci.Filename)
 
-	if err := checkSizeQuotaExceeded(ctx, pfci.Creator, pvi.Owner, pvi.PackageType, pfci.Data.Size()); err != nil {
+	if err := CheckSizeQuotaExceeded(ctx, pfci.Creator, pvi.Owner, pvi.PackageType, pfci.Data.Size()); err != nil {
 		return nil, nil, false, err
 	}
 
@@ -302,7 +302,9 @@ func addFileToPackageVersion(ctx context.Context, pv *packages_model.PackageVers
 	return pf, pb, !exists, nil
 }
 
-func checkCountQuotaExceeded(ctx context.Context, doer, owner *user_model.User) error {
+// CheckCountQuotaExceeded checks if the owner has more than the allowed packages
+// The check is skipped if the doer is an admin.
+func CheckCountQuotaExceeded(ctx context.Context, doer, owner *user_model.User) error {
 	if doer.IsAdmin {
 		return nil
 	}
@@ -324,7 +326,9 @@ func checkCountQuotaExceeded(ctx context.Context, doer, owner *user_model.User)
 	return nil
 }
 
-func checkSizeQuotaExceeded(ctx context.Context, doer, owner *user_model.User, packageType packages_model.Type, uploadSize int64) error {
+// CheckSizeQuotaExceeded checks if the upload size is bigger than the allowed size
+// The check is skipped if the doer is an admin.
+func CheckSizeQuotaExceeded(ctx context.Context, doer, owner *user_model.User, packageType packages_model.Type, uploadSize int64) error {
 	if doer.IsAdmin {
 		return nil
 	}
diff --git a/tests/integration/api_packages_test.go b/tests/integration/api_packages_test.go
index 9bca6a20ee..a6c335eb7e 100644
--- a/tests/integration/api_packages_test.go
+++ b/tests/integration/api_packages_test.go
@@ -5,8 +5,10 @@ package integration
 
 import (
 	"bytes"
+	"crypto/sha256"
 	"fmt"
 	"net/http"
+	"strings"
 	"testing"
 	"time"
 
@@ -171,34 +173,62 @@ func TestPackageAccess(t *testing.T) {
 func TestPackageQuota(t *testing.T) {
 	defer tests.PrepareTestEnv(t)()
 
-	limitTotalOwnerCount, limitTotalOwnerSize, limitSizeGeneric := setting.Packages.LimitTotalOwnerCount, setting.Packages.LimitTotalOwnerSize, setting.Packages.LimitSizeGeneric
+	limitTotalOwnerCount, limitTotalOwnerSize := setting.Packages.LimitTotalOwnerCount, setting.Packages.LimitTotalOwnerSize
 
+	// Exceeded quota result in StatusForbidden for normal users but admins are always allowed to upload.
 	admin := unittest.AssertExistsAndLoadBean(t, &user_model.User{ID: 1})
 	user := unittest.AssertExistsAndLoadBean(t, &user_model.User{ID: 10})
 
-	uploadPackage := func(doer *user_model.User, version string, expectedStatus int) {
-		url := fmt.Sprintf("/api/packages/%s/generic/test-package/%s/file.bin", user.Name, version)
-		req := NewRequestWithBody(t, "PUT", url, bytes.NewReader([]byte{1}))
-		AddBasicAuthHeader(req, doer.Name)
-		MakeRequest(t, req, expectedStatus)
-	}
+	t.Run("Common", func(t *testing.T) {
+		defer tests.PrintCurrentTest(t)()
 
-	// Exceeded quota result in StatusForbidden for normal users but admins are always allowed to upload.
+		limitSizeGeneric := setting.Packages.LimitSizeGeneric
+
+		uploadPackage := func(doer *user_model.User, version string, expectedStatus int) {
+			url := fmt.Sprintf("/api/packages/%s/generic/test-package/%s/file.bin", user.Name, version)
+			req := NewRequestWithBody(t, "PUT", url, bytes.NewReader([]byte{1}))
+			AddBasicAuthHeader(req, doer.Name)
+			MakeRequest(t, req, expectedStatus)
+		}
 
-	setting.Packages.LimitTotalOwnerCount = 0
-	uploadPackage(user, "1.0", http.StatusForbidden)
-	uploadPackage(admin, "1.0", http.StatusCreated)
-	setting.Packages.LimitTotalOwnerCount = limitTotalOwnerCount
+		setting.Packages.LimitTotalOwnerCount = 0
+		uploadPackage(user, "1.0", http.StatusForbidden)
+		uploadPackage(admin, "1.0", http.StatusCreated)
+		setting.Packages.LimitTotalOwnerCount = limitTotalOwnerCount
 
-	setting.Packages.LimitTotalOwnerSize = 0
-	uploadPackage(user, "1.1", http.StatusForbidden)
-	uploadPackage(admin, "1.1", http.StatusCreated)
-	setting.Packages.LimitTotalOwnerSize = limitTotalOwnerSize
+		setting.Packages.LimitTotalOwnerSize = 0
+		uploadPackage(user, "1.1", http.StatusForbidden)
+		uploadPackage(admin, "1.1", http.StatusCreated)
+		setting.Packages.LimitTotalOwnerSize = limitTotalOwnerSize
 
-	setting.Packages.LimitSizeGeneric = 0
-	uploadPackage(user, "1.2", http.StatusForbidden)
-	uploadPackage(admin, "1.2", http.StatusCreated)
-	setting.Packages.LimitSizeGeneric = limitSizeGeneric
+		setting.Packages.LimitSizeGeneric = 0
+		uploadPackage(user, "1.2", http.StatusForbidden)
+		uploadPackage(admin, "1.2", http.StatusCreated)
+		setting.Packages.LimitSizeGeneric = limitSizeGeneric
+	})
+
+	t.Run("Container", func(t *testing.T) {
+		defer tests.PrintCurrentTest(t)()
+
+		limitSizeContainer := setting.Packages.LimitSizeContainer
+
+		uploadBlob := func(doer *user_model.User, data string, expectedStatus int) {
+			url := fmt.Sprintf("/v2/%s/quota-test/blobs/uploads?digest=sha256:%x", user.Name, sha256.Sum256([]byte(data)))
+			req := NewRequestWithBody(t, "POST", url, strings.NewReader(data))
+			AddBasicAuthHeader(req, doer.Name)
+			MakeRequest(t, req, expectedStatus)
+		}
+
+		setting.Packages.LimitTotalOwnerSize = 0
+		uploadBlob(user, "2", http.StatusForbidden)
+		uploadBlob(admin, "2", http.StatusCreated)
+		setting.Packages.LimitTotalOwnerSize = limitTotalOwnerSize
+
+		setting.Packages.LimitSizeContainer = 0
+		uploadBlob(user, "3", http.StatusForbidden)
+		uploadBlob(admin, "3", http.StatusCreated)
+		setting.Packages.LimitSizeContainer = limitSizeContainer
+	})
 }
 
 func TestPackageCleanup(t *testing.T) {