From 0b26b3673780b428a390e4eecc6388be96f3a5c7 Mon Sep 17 00:00:00 2001 From: He Weiwei Date: Wed, 8 Apr 2020 19:17:32 +0000 Subject: [PATCH] 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 --- src/controller/quota/controller_test.go | 166 ++++++++++++------ src/server/middleware/quota/quota.go | 22 ++- src/server/middleware/quota/quota_test.go | 70 ++++++++ .../middleware/quota/refresh_project.go | 3 +- 4 files changed, 205 insertions(+), 56 deletions(-) diff --git a/src/controller/quota/controller_test.go b/src/controller/quota/controller_test.go index c8acbb611..ecf4937a8 100644 --- a/src/controller/quota/controller_test.go +++ b/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 = "atesting.Manager{} + suite.ctl = &controller{quotaMgr: suite.quotaMgr, reservedExpiration: defaultReservedExpiration} + + hardLimits := types.ResourceList{types.ResourceStorage: 100} + suite.quota = "a.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 := "atesting.Manager{} - - hardLimits := types.ResourceList{types.ResourceStorage: 100} - - mock.OnAnything(quotaMgr, "GetByRefForUpdate").Return("a.Quota{Hard: hardLimits.String(), Used: types.Zero(hardLimits).String()}, nil) - - ctl := &controller{quotaMgr: quotaMgr, reservedExpiration: defaultReservedExpiration} + mock.OnAnything(suite.quotaMgr, "GetByRefForUpdate").Return(suite.quota, nil) 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 := "atesting.Manager{} - - hardLimits := types.ResourceList{types.ResourceStorage: 100} - - mock.OnAnything(quotaMgr, "GetByRefForUpdate").Return("a.Quota{Hard: hardLimits.String(), Used: types.Zero(hardLimits).String()}, nil) - - ctl := &controller{quotaMgr: quotaMgr, reservedExpiration: defaultReservedExpiration} + mock.OnAnything(suite.quotaMgr, "GetByRefForUpdate").Return(suite.quota, nil) 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() { - quotaMgr := "atesting.Manager{} - - hardLimits := types.ResourceList{types.ResourceStorage: 100} - - q := "a.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} +func (suite *ControllerTestSuite) TestRefresh() { + suite.PrepareForUpdate(suite.quota, types.ResourceList{types.ResourceStorage: 0}) 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(ctl.Request(ctx, reference, referenceID, resources, func() error { return nil })) - } + suite.Nil(suite.ctl.Request(ctx, suite.reference, referenceID, resources, func() error { return nil })) +} - { - suite.Error(ctl.Request(ctx, reference, referenceID, resources, func() error { return nil })) - } +func (suite *ControllerTestSuite) TestRequestExceed() { + 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) { diff --git a/src/server/middleware/quota/quota.go b/src/server/middleware/quota/quota.go index dbf5c1388..2d16d160f 100644 --- a/src/server/middleware/quota/quota.go +++ b/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 } diff --git a/src/server/middleware/quota/quota_test.go b/src/server/middleware/quota/quota_test.go index adcedef9e..681d9d1ef 100644 --- a/src/server/middleware/quota/quota_test.go +++ b/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) diff --git a/src/server/middleware/quota/refresh_project.go b/src/server/middleware/quota/refresh_project.go index 5ac0d8877..c5abd3477 100644 --- a/src/server/middleware/quota/refresh_project.go +++ b/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...) }