diff --git a/src/controller/blob/controller.go b/src/controller/blob/controller.go index 7f9232d52..05b2903a4 100644 --- a/src/controller/blob/controller.go +++ b/src/controller/blob/controller.go @@ -305,6 +305,10 @@ func (c *controller) GetAcceptedBlobSize(sessionID string) (int64, error) { key := fmt.Sprintf("upload:%s:size", sessionID) size, err := redis.Int64(conn.Do("GET", key)) if err != nil { + if err == redis.ErrNil { + return 0, nil + } + return 0, err } diff --git a/src/controller/blob/controller_test.go b/src/controller/blob/controller_test.go index 2732ecee4..75ae11ec2 100644 --- a/src/controller/blob/controller_test.go +++ b/src/controller/blob/controller_test.go @@ -257,7 +257,8 @@ func (suite *ControllerTestSuite) TestGetSetAcceptedBlobSize() { sessionID := uuid.New().String() size, err := Ctl.GetAcceptedBlobSize(sessionID) - suite.NotNil(err) + suite.Nil(err) + suite.Equal(int64(0), size) suite.Nil(Ctl.SetAcceptedBlobSize(sessionID, 100)) diff --git a/src/server/middleware/quota/put_blob_upload.go b/src/server/middleware/quota/put_blob_upload.go index 4f39240c2..7dfeddd11 100644 --- a/src/server/middleware/quota/put_blob_upload.go +++ b/src/server/middleware/quota/put_blob_upload.go @@ -44,6 +44,11 @@ func putBlobUploadResources(r *http.Request, reference, referenceID string) (typ return nil, err } + if size == 0 { + logger.Debug("blob size is 0") + return nil, nil + } + projectID, _ := strconv.ParseInt(referenceID, 10, 64) digest := r.URL.Query().Get("digest") diff --git a/src/server/middleware/quota/put_blob_upload_test.go b/src/server/middleware/quota/put_blob_upload_test.go index ce7e63ced..38ecee063 100644 --- a/src/server/middleware/quota/put_blob_upload_test.go +++ b/src/server/middleware/quota/put_blob_upload_test.go @@ -15,8 +15,10 @@ package quota import ( + "fmt" "net/http" "net/http/httptest" + "strconv" "testing" "github.com/goharbor/harbor/src/pkg/types" @@ -26,46 +28,86 @@ import ( type PutBlobUploadMiddlewareTestSuite struct { RequestMiddlewareTestSuite + + handler http.Handler } -func (suite *PutBlobUploadMiddlewareTestSuite) TestMiddleware() { - mock.OnAnything(suite.quotaController, "IsEnabled").Return(true, nil) +func (suite *PutBlobUploadMiddlewareTestSuite) SetupTest() { + suite.RequestMiddlewareTestSuite.SetupTest() next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) }) + suite.handler = PutBlobUploadMiddleware()(next) + + mock.OnAnything(suite.quotaController, "IsEnabled").Return(true, nil) +} + +func (suite *PutBlobUploadMiddlewareTestSuite) makeRequest(contentLength int) *http.Request { url := "/v2/library/photon/blobs/uploads/cbabe458-28a1-4e1b-ad15-0cb0229df4e8?digest=sha256:57c2ec3bf82f09c94be2e5c5beb124b86fcbb42e76fb82c99066c054422010e4" + req := httptest.NewRequest(http.MethodPut, url, nil) + req.Header.Set("Content-Length", strconv.Itoa(contentLength)) - { - mock.OnAnything(suite.blobController, "Exist").Return(true, nil).Once() + return req +} - req := httptest.NewRequest(http.MethodPut, url, nil) - req.Header.Set("Content-Length", "100") - rr := httptest.NewRecorder() +func (suite *PutBlobUploadMiddlewareTestSuite) TestBlobSizeIsZero() { + mock.OnAnything(suite.blobController, "GetAcceptedBlobSize").Return(int64(0), nil) - PutBlobUploadMiddleware()(next).ServeHTTP(rr, req) - suite.Equal(http.StatusOK, rr.Code) - } + req := suite.makeRequest(0) + rr := httptest.NewRecorder() - { - mock.OnAnything(suite.blobController, "Exist").Return(false, nil).Once() - mock.OnAnything(suite.quotaController, "Request").Return(nil).Once().Run(func(args mock.Arguments) { - resources := args.Get(3).(types.ResourceList) - suite.Len(resources, 1) - suite.Equal(resources[types.ResourceStorage], int64(100)) + suite.handler.ServeHTTP(rr, req) + suite.Equal(http.StatusOK, rr.Code) +} - f := args.Get(4).(func() error) - f() - }) +func (suite *PutBlobUploadMiddlewareTestSuite) TestGetAcceptedBlobSizeFailed() { + mock.OnAnything(suite.blobController, "GetAcceptedBlobSize").Return(int64(0), fmt.Errorf("error")) - req := httptest.NewRequest(http.MethodPut, url, nil) - req.Header.Set("Content-Length", "100") - rr := httptest.NewRecorder() + req := suite.makeRequest(0) + rr := httptest.NewRecorder() - PutBlobUploadMiddleware()(next).ServeHTTP(rr, req) - suite.Equal(http.StatusOK, rr.Code) - } + suite.handler.ServeHTTP(rr, req) + suite.Equal(http.StatusInternalServerError, rr.Code) +} + +func (suite *PutBlobUploadMiddlewareTestSuite) TestBlobExist() { + mock.OnAnything(suite.blobController, "Exist").Return(true, nil) + + req := suite.makeRequest(100) + rr := httptest.NewRecorder() + + suite.handler.ServeHTTP(rr, req) + suite.Equal(http.StatusOK, rr.Code) +} + +func (suite *PutBlobUploadMiddlewareTestSuite) TestBlobNotExist() { + mock.OnAnything(suite.blobController, "Exist").Return(false, nil).Once() + mock.OnAnything(suite.quotaController, "Request").Return(nil).Once().Run(func(args mock.Arguments) { + resources := args.Get(3).(types.ResourceList) + suite.Len(resources, 1) + suite.Equal(resources[types.ResourceStorage], int64(100)) + + f := args.Get(4).(func() error) + f() + }) + + req := suite.makeRequest(100) + rr := httptest.NewRecorder() + + suite.handler.ServeHTTP(rr, req) + suite.Equal(http.StatusOK, rr.Code) +} + +func (suite *PutBlobUploadMiddlewareTestSuite) TestBlobExistFailed() { + mock.OnAnything(suite.blobController, "Exist").Return(false, fmt.Errorf("error")) + + req := suite.makeRequest(100) + rr := httptest.NewRecorder() + + suite.handler.ServeHTTP(rr, req) + suite.Equal(http.StatusInternalServerError, rr.Code) } func TestPutBlobUploadMiddlewareTestSuite(t *testing.T) { diff --git a/src/server/middleware/quota/quota.go b/src/server/middleware/quota/quota.go index 12dc14c10..dbce39ab8 100644 --- a/src/server/middleware/quota/quota.go +++ b/src/server/middleware/quota/quota.go @@ -83,7 +83,7 @@ func RequestMiddleware(config RequestConfig, skippers ...middleware.Skipper) fun if !enabled { // quota is disabled for the reference object, so direct to next handler - logger.Infof("quota is disabled for %s %s, so direct to next handler", reference, referenceID) + logger.Debugf("quota is disabled for %s %s, so direct to next handler", reference, referenceID) next.ServeHTTP(w, r) return } @@ -98,7 +98,7 @@ func RequestMiddleware(config RequestConfig, skippers ...middleware.Skipper) fun if len(resources) == 0 { // no resources request for this http request, so direct to next handler - logger.Info("no resources request for this http request, so direct to next handler") + logger.Debug("no resources request for this http request, so direct to next handler") next.ServeHTTP(w, r) return } @@ -206,7 +206,7 @@ func RefreshMiddleware(config RefreshConfig, skipers ...middleware.Skipper) func } if !enabled { - logger.Infof("quota is disabled for %s %s, so return directly", reference, referenceID) + logger.Debugf("quota is disabled for %s %s, so return directly", reference, referenceID) return nil }