feat(quota): ignore limitation support for quota RefreshMiddleware

1. Ignore limitation when refresh quota for project.
2. Return 403 when quota errors occurred.
3. Add test for Refresh method of quota controller.

Closes #11512

Signed-off-by: He Weiwei <hweiwei@vmware.com>

src/controller/quota/controller_test.go

@@ -16,6 +16,7 @@ package quota
 
 import (
 	"context"
+	"fmt"
 	"strconv"
 	"testing"
 	"time"
@@ -34,6 +35,34 @@ import (
 
 type ControllerTestSuite struct {
 	suite.Suite
+
+	reference string
+	driver    *drivertesting.Driver
+	quotaMgr  *quotatesting.Manager
+	ctl       Controller
+
+	quota *quota.Quota
+}
+
+func (suite *ControllerTestSuite) SetupTest() {
+	suite.reference = "mock"
+
+	suite.driver = &drivertesting.Driver{}
+	driver.Register(suite.reference, suite.driver)
+
+	suite.quotaMgr = &quotatesting.Manager{}
+	suite.ctl = &controller{quotaMgr: suite.quotaMgr, reservedExpiration: defaultReservedExpiration}
+
+	hardLimits := types.ResourceList{types.ResourceStorage: 100}
+	suite.quota = &quota.Quota{Hard: hardLimits.String(), Used: types.Zero(hardLimits).String()}
+}
+
+func (suite *ControllerTestSuite) PrepareForUpdate(q *quota.Quota, newUsage interface{}) {
+	mock.OnAnything(suite.quotaMgr, "GetByRefForUpdate").Return(q, nil)
+
+	mock.OnAnything(suite.driver, "CalculateUsage").Return(newUsage, nil)
+
+	mock.OnAnything(suite.quotaMgr, "Update").Return(nil)
 }
 
 func (suite *ControllerTestSuite) TestGetReservedResources() {
@@ -66,80 +95,113 @@ func (suite *ControllerTestSuite) TestGetReservedResources() {
 }
 
 func (suite *ControllerTestSuite) TestReserveResources() {
-	quotaMgr := &quotatesting.Manager{}
+	mock.OnAnything(suite.quotaMgr, "GetByRefForUpdate").Return(suite.quota, nil)
-
-	hardLimits := types.ResourceList{types.ResourceStorage: 100}
-
-	mock.OnAnything(quotaMgr, "GetByRefForUpdate").Return(&quota.Quota{Hard: hardLimits.String(), Used: types.Zero(hardLimits).String()}, nil)
-
-	ctl := &controller{quotaMgr: quotaMgr, reservedExpiration: defaultReservedExpiration}
 
 	ctx := orm.NewContext(context.TODO(), &ormtesting.FakeOrmer{})
-	reference, referenceID := "reference", uuid.New().String()
+	referenceID := uuid.New().String()
 	resources := types.ResourceList{types.ResourceStorage: 100}
 
-	suite.Nil(ctl.reserveResources(ctx, reference, referenceID, resources))
+	ctl := suite.ctl.(*controller)
 
-	suite.Error(ctl.reserveResources(ctx, reference, referenceID, resources))
+	suite.Nil(ctl.reserveResources(ctx, suite.reference, referenceID, resources))
+
+	suite.Error(ctl.reserveResources(ctx, suite.reference, referenceID, resources))
 }
 
 func (suite *ControllerTestSuite) TestUnreserveResources() {
-	quotaMgr := &quotatesting.Manager{}
+	mock.OnAnything(suite.quotaMgr, "GetByRefForUpdate").Return(suite.quota, nil)
-
-	hardLimits := types.ResourceList{types.ResourceStorage: 100}
-
-	mock.OnAnything(quotaMgr, "GetByRefForUpdate").Return(&quota.Quota{Hard: hardLimits.String(), Used: types.Zero(hardLimits).String()}, nil)
-
-	ctl := &controller{quotaMgr: quotaMgr, reservedExpiration: defaultReservedExpiration}
 
 	ctx := orm.NewContext(context.TODO(), &ormtesting.FakeOrmer{})
-	reference, referenceID := "reference", uuid.New().String()
+	referenceID := uuid.New().String()
 	resources := types.ResourceList{types.ResourceStorage: 100}
 
-	suite.Nil(ctl.reserveResources(ctx, reference, referenceID, resources))
+	ctl := suite.ctl.(*controller)
 
-	suite.Error(ctl.reserveResources(ctx, reference, referenceID, resources))
+	suite.Nil(ctl.reserveResources(ctx, suite.reference, referenceID, resources))
 
-	suite.Nil(ctl.unreserveResources(ctx, reference, referenceID, resources))
+	suite.Error(ctl.reserveResources(ctx, suite.reference, referenceID, resources))
 
-	suite.Nil(ctl.reserveResources(ctx, reference, referenceID, resources))
+	suite.Nil(ctl.unreserveResources(ctx, suite.reference, referenceID, resources))
+
+	suite.Nil(ctl.reserveResources(ctx, suite.reference, referenceID, resources))
 }
 
-func (suite *ControllerTestSuite) TestRequest() {
+func (suite *ControllerTestSuite) TestRefresh() {
-	quotaMgr := &quotatesting.Manager{}
+	suite.PrepareForUpdate(suite.quota, types.ResourceList{types.ResourceStorage: 0})
-
-	hardLimits := types.ResourceList{types.ResourceStorage: 100}
-
-	q := &quota.Quota{Hard: hardLimits.String(), Used: types.Zero(hardLimits).String()}
-	used := types.ResourceList{types.ResourceStorage: 0}
-
-	mock.OnAnything(quotaMgr, "GetByRefForUpdate").Return(q, nil)
-
-	mock.OnAnything(quotaMgr, "Update").Return(nil).Run(func(mock.Arguments) {
-		q.SetUsed(used)
-	})
-
-	d := &drivertesting.Driver{}
-
-	mock.OnAnything(d, "CalculateUsage").Return(used, nil).Run(func(args mock.Arguments) {
-		used[types.ResourceStorage]++
-	})
-
-	driver.Register("mock", d)
-
-	ctl := &controller{quotaMgr: quotaMgr, reservedExpiration: defaultReservedExpiration}
 
 	ctx := orm.NewContext(context.TODO(), &ormtesting.FakeOrmer{})
-	reference, referenceID := "mock", "1"
+	referenceID := uuid.New().String()
+
+	suite.Nil(suite.ctl.Refresh(ctx, suite.reference, referenceID))
+}
+
+func (suite *ControllerTestSuite) TestRefreshDriverNotFound() {
+	ctx := orm.NewContext(context.TODO(), &ormtesting.FakeOrmer{})
+
+	suite.Error(suite.ctl.Refresh(ctx, uuid.New().String(), uuid.New().String()))
+}
+
+func (suite *ControllerTestSuite) TestRefershNegativeUsage() {
+	suite.PrepareForUpdate(suite.quota, types.ResourceList{types.ResourceStorage: -1})
+
+	ctx := orm.NewContext(context.TODO(), &ormtesting.FakeOrmer{})
+	referenceID := uuid.New().String()
+
+	suite.Error(suite.ctl.Refresh(ctx, suite.reference, referenceID))
+}
+
+func (suite *ControllerTestSuite) TestRefreshUsageExceed() {
+	suite.PrepareForUpdate(suite.quota, types.ResourceList{types.ResourceStorage: 101})
+
+	ctx := orm.NewContext(context.TODO(), &ormtesting.FakeOrmer{})
+	referenceID := uuid.New().String()
+
+	suite.Error(suite.ctl.Refresh(ctx, suite.reference, referenceID))
+}
+
+func (suite *ControllerTestSuite) TestRefreshIgnoreLimitation() {
+	suite.PrepareForUpdate(suite.quota, types.ResourceList{types.ResourceStorage: 101})
+
+	ctx := orm.NewContext(context.TODO(), &ormtesting.FakeOrmer{})
+	referenceID := uuid.New().String()
+
+	suite.Nil(suite.ctl.Refresh(ctx, suite.reference, referenceID, IgnoreLimitation(true)))
+}
+
+func (suite *ControllerTestSuite) TestNoResourcesRequest() {
+	ctx := orm.NewContext(context.TODO(), &ormtesting.FakeOrmer{})
+	referenceID := uuid.New().String()
+
+	suite.Nil(suite.ctl.Request(ctx, suite.reference, referenceID, nil, func() error { return nil }))
+}
+func (suite *ControllerTestSuite) TestRequest() {
+	suite.PrepareForUpdate(suite.quota, nil)
+
+	ctx := orm.NewContext(context.TODO(), &ormtesting.FakeOrmer{})
+	referenceID := uuid.New().String()
 	resources := types.ResourceList{types.ResourceStorage: 100}
 
-	{
+	suite.Nil(suite.ctl.Request(ctx, suite.reference, referenceID, resources, func() error { return nil }))
-		suite.Nil(ctl.Request(ctx, reference, referenceID, resources, func() error { return nil }))
+}
-	}
 
-	{
+func (suite *ControllerTestSuite) TestRequestExceed() {
-		suite.Error(ctl.Request(ctx, reference, referenceID, resources, func() error { return nil }))
+	suite.PrepareForUpdate(suite.quota, nil)
-	}
+
+	ctx := orm.NewContext(context.TODO(), &ormtesting.FakeOrmer{})
+	referenceID := uuid.New().String()
+	resources := types.ResourceList{types.ResourceStorage: 101}
+
+	suite.Error(suite.ctl.Request(ctx, suite.reference, referenceID, resources, func() error { return nil }))
+}
+
+func (suite *ControllerTestSuite) TestRequestFunctionFailed() {
+	suite.PrepareForUpdate(suite.quota, nil)
+
+	ctx := orm.NewContext(context.TODO(), &ormtesting.FakeOrmer{})
+	referenceID := uuid.New().String()
+	resources := types.ResourceList{types.ResourceStorage: 100}
+
+	suite.Error(suite.ctl.Request(ctx, suite.reference, referenceID, resources, func() error { return fmt.Errorf("error") }))
 }
 
 func TestControllerTestSuite(t *testing.T) {

src/server/middleware/quota/quota.go

@@ -15,12 +15,13 @@
 package quota
 
 import (
-	"errors"
 	"fmt"
 	"net/http"
 	"strings"
 
+	cq "github.com/goharbor/harbor/src/controller/quota"
 	"github.com/goharbor/harbor/src/lib"
+	"github.com/goharbor/harbor/src/lib/errors"
 	"github.com/goharbor/harbor/src/lib/log"
 	"github.com/goharbor/harbor/src/pkg/notification"
 	"github.com/goharbor/harbor/src/pkg/notifier/event"
@@ -167,7 +168,13 @@ func RequestMiddleware(config RequestConfig, skippers ...middleware.Skipper) fun
 			}
 
 			res.Reset()
-			serror.SendError(res, err)
+
+			var errs quota.Errors
+			if errors.As(err, &errs) {
+				serror.SendError(res, errors.DeniedError(nil).WithMessage(errs.Error()))
+			} else {
+				serror.SendError(res, err)
+			}
 		}
 
 	}, skippers...)
@@ -175,6 +182,9 @@ func RequestMiddleware(config RequestConfig, skippers ...middleware.Skipper) fun
 
 // RefreshConfig refresh quota usage middleware config
 type RefreshConfig struct {
+	// IgnoreLimitation allow quota usage exceed the limitation when it's true
+	IgnoreLimitation bool
+
 	// ReferenceObject returns reference object its quota usage will refresh by reference and reference id
 	ReferenceObject func(*http.Request) (reference string, referenceID string, err error)
 }
@@ -210,8 +220,14 @@ func RefreshMiddleware(config RefreshConfig, skipers ...middleware.Skipper) func
 			return nil
 		}
 
-		if err = quotaController.Refresh(r.Context(), reference, referenceID); err != nil {
+		if err = quotaController.Refresh(r.Context(), reference, referenceID, cq.IgnoreLimitation(config.IgnoreLimitation)); err != nil {
 			logger.Errorf("refresh quota for %s %s failed, error: %v", reference, referenceID, err)
+
+			var errs quota.Errors
+			if errors.As(err, &errs) {
+				return errors.DeniedError(nil).WithMessage(errs.Error())
+			}
+
 			return err
 		}
 

src/server/middleware/quota/quota_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/goharbor/harbor/src/controller/blob"
 	"github.com/goharbor/harbor/src/controller/project"
 	"github.com/goharbor/harbor/src/controller/quota"
+	pquota "github.com/goharbor/harbor/src/pkg/quota"
 	"github.com/goharbor/harbor/src/pkg/types"
 	artifacttesting "github.com/goharbor/harbor/src/testing/controller/artifact"
 	blobtesting "github.com/goharbor/harbor/src/testing/controller/blob"
@@ -193,6 +194,27 @@ func (suite *RequestMiddlewareTestSuite) TestResourcesRequestFailed() {
 	suite.Equal(http.StatusInternalServerError, rr.Code)
 }
 
+func (suite *RequestMiddlewareTestSuite) TestResourcesRequestDenied() {
+	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
+
+	req := httptest.NewRequest(http.MethodPost, "/url", nil)
+	rr := httptest.NewRecorder()
+
+	reference, referenceID := "project", "1"
+	resources := types.ResourceList{types.ResourceStorage: 100}
+	config := suite.makeRequestConfig(reference, referenceID, resources)
+
+	mock.OnAnything(suite.quotaController, "IsEnabled").Return(true, nil)
+	var errs pquota.Errors
+	errs = errs.Add(fmt.Errorf("Exceed"))
+	mock.OnAnything(suite.quotaController, "Request").Return(errs)
+
+	RequestMiddleware(config)(next).ServeHTTP(rr, req)
+	suite.Equal(http.StatusForbidden, rr.Code)
+}
+
 func TestRequestMiddlewareTestSuite(t *testing.T) {
 	suite.Run(t, &RequestMiddlewareTestSuite{})
 }
@@ -231,6 +253,54 @@ func (suite *RefreshMiddlewareTestSuite) TestQuotaDisabled() {
 	suite.Equal(http.StatusOK, rr.Code)
 }
 
+func (suite *RefreshMiddlewareTestSuite) TestQuotaIsEnabledFailed() {
+	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
+
+	req := httptest.NewRequest(http.MethodPost, "/url", nil)
+	rr := httptest.NewRecorder()
+
+	reference, referenceID := "project", "1"
+
+	config := RefreshConfig{
+		ReferenceObject: func(*http.Request) (string, string, error) {
+			return reference, referenceID, nil
+		},
+	}
+
+	mock.OnAnything(suite.quotaController, "IsEnabled").Return(false, fmt.Errorf("error"))
+
+	RefreshMiddleware(config)(next).ServeHTTP(rr, req)
+	suite.Equal(http.StatusInternalServerError, rr.Code)
+}
+
+func (suite *RefreshMiddlewareTestSuite) TestInvalidConfig() {
+	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
+
+	req := httptest.NewRequest(http.MethodPost, "/url", nil)
+	rr := httptest.NewRecorder()
+
+	config := RefreshConfig{}
+	RefreshMiddleware(config)(next).ServeHTTP(rr, req)
+	suite.Equal(http.StatusInternalServerError, rr.Code)
+}
+
+func (suite *RefreshMiddlewareTestSuite) TestNotSuccess() {
+	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusBadRequest)
+	})
+
+	req := httptest.NewRequest(http.MethodPost, "/url", nil)
+	rr := httptest.NewRecorder()
+
+	config := RefreshConfig{}
+	RefreshMiddleware(config)(next).ServeHTTP(rr, req)
+	suite.Equal(http.StatusBadRequest, rr.Code)
+}
+
 func (suite *RefreshMiddlewareTestSuite) TestRefershOK() {
 	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusOK)

src/server/middleware/quota/refresh_project.go

@@ -23,6 +23,7 @@ import (
 // RefreshForProjectMiddleware middleware which refresh the quota usage of project after the response success
 func RefreshForProjectMiddleware(skippers ...middleware.Skipper) func(http.Handler) http.Handler {
 	return RefreshMiddleware(RefreshConfig{
-		ReferenceObject: projectReferenceObject,
+		IgnoreLimitation: true,
+		ReferenceObject:  projectReferenceObject,
 	}, skippers...)
 }