From 1bbfc023f1b7dd4c7d6516869c08ea7b74cace09 Mon Sep 17 00:00:00 2001 From: He Weiwei Date: Sun, 18 Aug 2019 16:14:36 +0000 Subject: [PATCH] fix(quota): fix computeResources method of qutoa interceptor Signed-off-by: He Weiwei --- src/core/middlewares/chart/builder.go | 27 ++-- src/core/middlewares/chart/handler_test.go | 137 ++++++++++++++++++ .../middlewares/interceptor/quota/quota.go | 3 +- src/testing/suite.go | 93 ++++++++++++ 4 files changed, 247 insertions(+), 13 deletions(-) create mode 100644 src/core/middlewares/chart/handler_test.go create mode 100644 src/testing/suite.go diff --git a/src/core/middlewares/chart/builder.go b/src/core/middlewares/chart/builder.go index 19827e2d3..ba54cd2de 100644 --- a/src/core/middlewares/chart/builder.go +++ b/src/core/middlewares/chart/builder.go @@ -103,20 +103,23 @@ func (*chartVersionCreationBuilder) Build(req *http.Request) (interceptor.Interc return nil, fmt.Errorf("project %s not found", namespace) } - chart, err := parseChart(req) - if err != nil { - return nil, fmt.Errorf("failed to parse chart from body, error: %v", err) - } - chartName, version := chart.Metadata.Name, chart.Metadata.Version + info, ok := util.ChartVersionInfoFromContext(req.Context()) + if !ok { + chart, err := parseChart(req) + if err != nil { + return nil, fmt.Errorf("failed to parse chart from body, error: %v", err) + } + chartName, version := chart.Metadata.Name, chart.Metadata.Version - info := &util.ChartVersionInfo{ - ProjectID: project.ProjectID, - Namespace: namespace, - ChartName: chartName, - Version: version, + info = &util.ChartVersionInfo{ + ProjectID: project.ProjectID, + Namespace: namespace, + ChartName: chartName, + Version: version, + } + // Chart version info will be used by computeQuotaForUpload + *req = *req.WithContext(util.NewChartVersionInfoContext(req.Context(), info)) } - // Chart version info will be used by computeQuotaForUpload - *req = *req.WithContext(util.NewChartVersionInfoContext(req.Context(), info)) opts := []quota.Option{ quota.EnforceResources(config.QuotaPerProjectEnable()), diff --git a/src/core/middlewares/chart/handler_test.go b/src/core/middlewares/chart/handler_test.go new file mode 100644 index 000000000..aedf1218e --- /dev/null +++ b/src/core/middlewares/chart/handler_test.go @@ -0,0 +1,137 @@ +// Copyright Project Harbor Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package chart + +import ( + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "testing" + + "github.com/goharbor/harbor/src/chartserver" + "github.com/goharbor/harbor/src/common/dao" + "github.com/goharbor/harbor/src/core/middlewares/util" + "github.com/goharbor/harbor/src/pkg/types" + htesting "github.com/goharbor/harbor/src/testing" + "github.com/stretchr/testify/suite" +) + +func deleteChartVersion(projectName, chartName, version string) { + url := fmt.Sprintf("/api/chartrepo/%s/charts/%s/%s", projectName, chartName, version) + req, _ := http.NewRequest(http.MethodDelete, url, nil) + + next := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + rr := httptest.NewRecorder() + h := New(next) + h.ServeHTTP(util.NewCustomResponseWriter(rr), req) +} + +func uploadChartVersion(projectID int64, projectName, chartName, version string) { + url := fmt.Sprintf("/api/chartrepo/%s/charts/", projectName) + req, _ := http.NewRequest(http.MethodPost, url, nil) + + info := &util.ChartVersionInfo{ + ProjectID: projectID, + Namespace: projectName, + ChartName: chartName, + Version: version, + } + *req = *req.WithContext(util.NewChartVersionInfoContext(req.Context(), info)) + + next := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + w.WriteHeader(http.StatusCreated) + }) + + rr := httptest.NewRecorder() + h := New(next) + h.ServeHTTP(util.NewCustomResponseWriter(rr), req) +} + +func mockChartController() (*httptest.Server, *chartserver.Controller, error) { + mockServer := httptest.NewServer(htesting.MockChartRepoHandler) + + var oldController, newController *chartserver.Controller + url, err := url.Parse(mockServer.URL) + if err == nil { + newController, err = chartserver.NewController(url) + } + + if err != nil { + mockServer.Close() + return nil, nil, err + } + + chartController() // Init chart controller + + // Override current controller and keep the old one for restoring + oldController = controller + controller = newController + + return mockServer, oldController, nil +} + +type HandlerSuite struct { + htesting.Suite + oldController *chartserver.Controller + mockChartServer *httptest.Server +} + +func (suite *HandlerSuite) SetupTest() { + mockServer, oldController, err := mockChartController() + suite.Nil(err, "Mock chart controller failed") + + suite.oldController = oldController + suite.mockChartServer = mockServer +} + +func (suite *HandlerSuite) TearDownTest() { + for _, table := range []string{ + "quota", "quota_usage", + } { + dao.ClearTable(table) + } + + controller = suite.oldController + suite.mockChartServer.Close() +} + +func (suite *HandlerSuite) TestUpload() { + suite.WithProject(func(projectID int64, projectName string) { + uploadChartVersion(projectID, projectName, "harbor", "0.2.1") + suite.AssertResourceUsage(1, types.ResourceCount, projectID) + + // harbor:0.2.0 exists in repo1, upload it again + uploadChartVersion(projectID, projectName, "harbor", "0.2.0") + suite.AssertResourceUsage(1, types.ResourceCount, projectID) + }, "repo1") +} + +func (suite *HandlerSuite) TestDelete() { + suite.WithProject(func(projectID int64, projectName string) { + uploadChartVersion(projectID, projectName, "harbor", "0.2.1") + suite.AssertResourceUsage(1, types.ResourceCount, projectID) + + deleteChartVersion(projectName, "harbor", "0.2.1") + suite.AssertResourceUsage(0, types.ResourceCount, projectID) + }, "repo1") +} + +func TestRunHandlerSuite(t *testing.T) { + suite.Run(t, new(HandlerSuite)) +} diff --git a/src/core/middlewares/interceptor/quota/quota.go b/src/core/middlewares/interceptor/quota/quota.go index 2914af8ee..607f58dde 100644 --- a/src/core/middlewares/interceptor/quota/quota.go +++ b/src/core/middlewares/interceptor/quota/quota.go @@ -135,7 +135,8 @@ func (qi *quotaInterceptor) computeResources(req *http.Request) error { return nil } - if len(qi.opts.Resources) == 0 && qi.opts.OnResources != nil { + qi.resources = qi.opts.Resources + if len(qi.resources) == 0 && qi.opts.OnResources != nil { resources, err := qi.opts.OnResources(req) if err != nil { return fmt.Errorf("failed to compute the resources for quota, error: %v", err) diff --git a/src/testing/suite.go b/src/testing/suite.go new file mode 100644 index 000000000..679de182a --- /dev/null +++ b/src/testing/suite.go @@ -0,0 +1,93 @@ +// Copyright Project Harbor Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package testing + +import ( + "fmt" + "math/rand" + "strconv" + "time" + + "github.com/goharbor/harbor/src/common/dao" + "github.com/goharbor/harbor/src/common/models" + "github.com/goharbor/harbor/src/core/config" + "github.com/goharbor/harbor/src/pkg/types" + "github.com/stretchr/testify/suite" +) + +func init() { + rand.Seed(time.Now().UnixNano()) +} + +// Suite ... +type Suite struct { + suite.Suite +} + +// SetupSuite ... +func (suite *Suite) SetupSuite() { + config.Init() + dao.PrepareTestForPostgresSQL() +} + +// RandString ... +func (suite *Suite) RandString(n int, letters ...string) string { + if len(letters) == 0 || len(letters[0]) == 0 { + letters = []string{"abcdefghijklmnopqrstuvwxyz"} + } + + letterBytes := []byte(letters[0]) + + b := make([]byte, n) + for i := range b { + b[i] = letterBytes[rand.Intn(len(letterBytes))] + } + return string(b) +} + +// WithProject ... +func (suite *Suite) WithProject(f func(int64, string), projectNames ...string) { + var projectName string + if len(projectNames) > 0 { + projectName = projectNames[0] + } else { + projectName = suite.RandString(5) + } + + projectID, err := dao.AddProject(models.Project{ + Name: projectName, + OwnerID: 1, + }) + if err != nil { + panic(err) + } + + defer func() { + dao.DeleteProject(projectID) + }() + + f(projectID, projectName) +} + +// AssertResourceUsage ... +func (suite *Suite) AssertResourceUsage(expected int64, resource types.ResourceName, projectID int64) { + usage := models.QuotaUsage{Reference: "project", ReferenceID: strconv.FormatInt(projectID, 10)} + err := dao.GetOrmer().Read(&usage, "reference", "reference_id") + suite.Nil(err, fmt.Sprintf("Failed to get resource %s usage of project %d, error: %v", resource, projectID, err)) + + used, err := types.NewResourceList(usage.Used) + suite.Nil(err, "Bad resource usage of project %d", projectID) + suite.Equal(expected, used[resource]) +}