Upstream/harbor
1
0
Fork 0
mirror of https://github.com/goharbor/harbor.git synced 2024-12-18 14:47:38 +01:00
Code Issues Projects Releases Wiki Activity

Merge pull request #11527 from heww/refresh-quota-ignore-limitation

Browse Source 
feat(quota): ignore limitation support for quota RefreshMiddleware
This commit is contained in:
He Weiwei 2020-04-10 15:29:41 +08:00 committed by GitHub
commit f4821e96b3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 205 additions and 56 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

166
src/controller/quota/controller_test.go
View File

@ -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{}
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}
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 := &quotatesting.Manager{}
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}
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 := &quotatesting.Manager{}
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}
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) {

22
src/server/middleware/quota/quota.go
View File

@ -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
}

70
src/server/middleware/quota/quota_test.go
View File

@ -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)

3
src/server/middleware/quota/refresh_project.go
View File

@ -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...)
}