fix(quota): fix computeResources method of qutoa interceptor

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

src/core/middlewares/chart/builder.go

@@ -103,13 +103,15 @@ func (*chartVersionCreationBuilder) Build(req *http.Request) (interceptor.Interc
 		return nil, fmt.Errorf("project %s not found", namespace)
 	}
 
+	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{
+		info = &util.ChartVersionInfo{
 			ProjectID: project.ProjectID,
 			Namespace: namespace,
 			ChartName: chartName,
@@ -117,6 +119,7 @@ func (*chartVersionCreationBuilder) Build(req *http.Request) (interceptor.Interc
 		}
 		// Chart version info will be used by computeQuotaForUpload
 		*req = *req.WithContext(util.NewChartVersionInfoContext(req.Context(), info))
+	}
 
 	opts := []quota.Option{
 		quota.EnforceResources(config.QuotaPerProjectEnable()),

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))
+}

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)

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])
+}