diff --git a/src/api/blob/controller.go b/src/api/blob/controller.go index 217e5df51..b688afa80 100644 --- a/src/api/blob/controller.go +++ b/src/api/blob/controller.go @@ -181,14 +181,14 @@ func (c *controller) FindMissingAssociationsForProject(ctx context.Context, proj return nil, nil } - associated := make(map[int64]bool, len(associatedBlobs)) + associated := make(map[string]bool, len(associatedBlobs)) for _, blob := range associatedBlobs { - associated[blob.ID] = true + associated[blob.Digest] = true } var results []*models.Blob for _, blob := range blobs { - if !associated[blob.ID] { + if !associated[blob.Digest] { results = append(results, blob) } } diff --git a/src/api/blob/controller_test.go b/src/api/blob/controller_test.go index e1b9ed46c..2732ecee4 100644 --- a/src/api/blob/controller_test.go +++ b/src/api/blob/controller_test.go @@ -144,7 +144,7 @@ func (suite *ControllerTestSuite) TestFindMissingAssociationsForProjectByArtifac suite.Len(blobs, 0) } - blobs := []*blob.Blob{{ID: 1}, {ID: 2}, {ID: 3}} + blobs := []*blob.Blob{{Digest: "1"}, {Digest: "2"}, {Digest: "3"}} { mock.OnAnything(blobMgr, "List").Return(nil, nil).Once() @@ -161,7 +161,7 @@ func (suite *ControllerTestSuite) TestFindMissingAssociationsForProjectByArtifac } { - associated := []*blob.Blob{{ID: 1}} + associated := []*blob.Blob{{Digest: "1"}} mock.OnAnything(blobMgr, "List").Return(associated, nil).Once() missing, err := ctl.FindMissingAssociationsForProject(ctx, projectID, blobs) suite.Nil(err) diff --git a/src/server/middleware/blob/put_manifest.go b/src/server/middleware/blob/put_manifest.go index 1cfd2634a..e0645ef2a 100644 --- a/src/server/middleware/blob/put_manifest.go +++ b/src/server/middleware/blob/put_manifest.go @@ -19,7 +19,6 @@ import ( "io/ioutil" "net/http" - "github.com/docker/distribution/manifest/schema2" "github.com/goharbor/harbor/src/common/utils/log" "github.com/goharbor/harbor/src/pkg/distribution" "github.com/goharbor/harbor/src/server/middleware" @@ -67,23 +66,22 @@ func PutManifestMiddleware() func(http.Handler) http.Handler { return err } - for _, digest := range findForeignBlobDigests(manifest) { - if err := blobController.AssociateWithProjectByDigest(ctx, digest, p.ProjectID); err != nil { + // NOTE: associate all blobs with project because the already exist associations may cleanup by others + for _, reference := range manifest.References() { + if err := blobController.AssociateWithProjectByDigest(ctx, reference.Digest.String(), p.ProjectID); err != nil { return err } } - artifactDigest := descriptor.Digest.String() - // ensure Blob for the manifest - blobID, err := blobController.Ensure(ctx, artifactDigest, contentType, descriptor.Size) + blobID, err := blobController.Ensure(ctx, descriptor.Digest.String(), contentType, descriptor.Size) if err != nil { - log.Errorf("%s: ensure blob %s failed, error: %v", logPrefix, descriptor.Digest, err) + log.Errorf("%s: ensure blob %s failed, error: %v", logPrefix, descriptor.Digest.String(), err) return err } if err := blobController.AssociateWithProjectByID(ctx, blobID, p.ProjectID); err != nil { - log.Errorf("%s: associate manifest with artifact %s failed, error: %v", logPrefix, descriptor.Digest, err) + log.Errorf("%s: associate manifest with artifact %s failed, error: %v", logPrefix, descriptor.Digest.String(), err) return err } @@ -93,8 +91,8 @@ func PutManifestMiddleware() func(http.Handler) http.Handler { } // associate blobs of the manifest with artifact - if err := blobController.AssociateWithArtifact(ctx, blobDigests, artifactDigest); err != nil { - log.Errorf("%s: associate blobs with artifact %s failed, error: %v", logPrefix, descriptor.Digest, err) + if err := blobController.AssociateWithArtifact(ctx, blobDigests, descriptor.Digest.String()); err != nil { + log.Errorf("%s: associate blobs with artifact %s failed, error: %v", logPrefix, descriptor.Digest.String(), err) return err } @@ -105,17 +103,3 @@ func PutManifestMiddleware() func(http.Handler) http.Handler { return alice.New(before, after).Then(next) } } - -func isForeign(d *distribution.Descriptor) bool { - return d.MediaType == schema2.MediaTypeForeignLayer -} - -func findForeignBlobDigests(manifest distribution.Manifest) []string { - var digests []string - for _, reference := range manifest.References() { - if isForeign(&reference) { - digests = append(digests, reference.Digest.String()) - } - } - return digests -} diff --git a/src/server/middleware/blob/put_manifest_test.go b/src/server/middleware/blob/put_manifest_test.go index be6088370..fba86ed07 100644 --- a/src/server/middleware/blob/put_manifest_test.go +++ b/src/server/middleware/blob/put_manifest_test.go @@ -96,12 +96,6 @@ func (suite *PutManifestMiddlewareTestSuite) TestMiddleware() { suite.WithProject(func(projectID int64, projectName string) { name := fmt.Sprintf("%s/redis", projectName) - for _, reference := range manifest.References() { - if !isForeign(&reference) { - suite.pushBlob(name, reference.Digest.String(), reference.Size) - } - } - req := suite.NewRequest(http.MethodPut, fmt.Sprintf("/v2/%s/manifests/%s", name, descriptor.Digest.String()), strings.NewReader(body)) req.Header.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json") res := httptest.NewRecorder() diff --git a/src/server/middleware/quota/put_manifest.go b/src/server/middleware/quota/put_manifest.go index e1f4e7c93..03dd0d26e 100644 --- a/src/server/middleware/quota/put_manifest.go +++ b/src/server/middleware/quota/put_manifest.go @@ -23,6 +23,7 @@ import ( "github.com/goharbor/harbor/src/api/blob" "github.com/goharbor/harbor/src/common/utils/log" "github.com/goharbor/harbor/src/internal" + "github.com/goharbor/harbor/src/pkg/blob/models" "github.com/goharbor/harbor/src/pkg/distribution" "github.com/goharbor/harbor/src/pkg/types" ) @@ -36,22 +37,16 @@ func PutManifestMiddleware() func(http.Handler) http.Handler { } var ( - parseManifestDigestAndSize = func(r *http.Request) (string, int64, error) { + unmarshalManifest = func(r *http.Request) (distribution.Manifest, distribution.Descriptor, error) { internal.NopCloseRequest(r) body, err := ioutil.ReadAll(r.Body) if err != nil { - return "", 0, err + return nil, distribution.Descriptor{}, err } contentType := r.Header.Get("Content-Type") - _, descriptor, err := distribution.UnmarshalManifest(contentType, body) - if err != nil { - - return "", 0, err - } - - return descriptor.Digest.String(), descriptor.Size, nil + return distribution.UnmarshalManifest(contentType, body) } ) @@ -60,15 +55,15 @@ func putManifestResources(r *http.Request, reference, referenceID string) (types projectID, _ := strconv.ParseInt(referenceID, 10, 64) - digest, size, err := parseManifestDigestAndSize(r) + manifest, descriptor, err := unmarshalManifest(r) if err != nil { log.Errorf("%s: unmarshal manifest failed, error: %v", logPrefix, err) return nil, err } - exist, err := blobController.Exist(r.Context(), digest, blob.IsAssociatedWithProject(projectID)) + exist, err := blobController.Exist(r.Context(), descriptor.Digest.String(), blob.IsAssociatedWithProject(projectID)) if err != nil { - log.Errorf("%s: check manifest %s is associated with project failed, error: %v", logPrefix, digest, err) + log.Errorf("%s: check manifest %s is associated with project failed, error: %v", logPrefix, descriptor.Digest.String(), err) return nil, err } @@ -76,5 +71,27 @@ func putManifestResources(r *http.Request, reference, referenceID string) (types return nil, nil } + size := descriptor.Size + + var blobs []*models.Blob + for _, reference := range manifest.References() { + blobs = append(blobs, &models.Blob{ + Digest: reference.Digest.String(), + Size: reference.Size, + ContentType: reference.MediaType, + }) + } + + missing, err := blobController.FindMissingAssociationsForProject(r.Context(), projectID, blobs) + if err != nil { + return nil, err + } + + for _, m := range missing { + if !m.IsForeignLayer() { + size += m.Size + } + } + return types.ResourceList{types.ResourceCount: 1, types.ResourceStorage: size}, nil } diff --git a/src/server/middleware/quota/put_manifest_test.go b/src/server/middleware/quota/put_manifest_test.go index d167f94fe..f8ce45e06 100644 --- a/src/server/middleware/quota/put_manifest_test.go +++ b/src/server/middleware/quota/put_manifest_test.go @@ -15,29 +15,52 @@ package quota import ( + "context" "net/http" "net/http/httptest" "testing" + "github.com/docker/distribution/manifest/schema2" + "github.com/goharbor/harbor/src/pkg/blob/models" + "github.com/goharbor/harbor/src/pkg/distribution" "github.com/goharbor/harbor/src/pkg/types" "github.com/goharbor/harbor/src/testing/mock" + distributiontesting "github.com/goharbor/harbor/src/testing/pkg/distribution" "github.com/stretchr/testify/suite" ) type PutManifestMiddlewareTestSuite struct { RequestMiddlewareTestSuite + + unmarshalManifest func(r *http.Request) (distribution.Manifest, distribution.Descriptor, error) + manifest distribution.Manifest +} + +func (suite *PutManifestMiddlewareTestSuite) SetupTest() { + suite.RequestMiddlewareTestSuite.SetupTest() + + suite.unmarshalManifest = unmarshalManifest + suite.manifest = &distributiontesting.Manifest{} + + mock.OnAnything(suite.manifest, "References").Return([]distribution.Descriptor{ + {Digest: "blob1", Size: 10, MediaType: schema2.MediaTypeLayer}, + {Digest: "blob2", Size: 20, MediaType: schema2.MediaTypeLayer}, + {Digest: "blob3", Size: 30, MediaType: schema2.MediaTypeForeignLayer}, + {Digest: "blob4", Size: 40, MediaType: schema2.MediaTypeForeignLayer}, + }) + + unmarshalManifest = func(r *http.Request) (distribution.Manifest, distribution.Descriptor, error) { + return suite.manifest, distribution.Descriptor{Digest: "digest", Size: 100}, nil + } +} + +func (suite *PutManifestMiddlewareTestSuite) TearDownTest() { + suite.RequestMiddlewareTestSuite.TearDownTest() + + unmarshalManifest = suite.unmarshalManifest } func (suite *PutManifestMiddlewareTestSuite) TestMiddleware() { - original := parseManifestDigestAndSize - defer func() { - parseManifestDigestAndSize = original - }() - - parseManifestDigestAndSize = func(r *http.Request) (string, int64, error) { - return "digest", 100, nil - } - mock.OnAnything(suite.quotaController, "IsEnabled").Return(true, nil) next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -55,7 +78,87 @@ func (suite *PutManifestMiddlewareTestSuite) TestMiddleware() { } { + // manifest not associated with project and blobs are already associated with project mock.OnAnything(suite.blobController, "Exist").Return(false, nil).Once() + mock.OnAnything(suite.blobController, "FindMissingAssociationsForProject").Return(nil, nil).Once() + mock.OnAnything(suite.quotaController, "Request").Return(nil).Once().Run(func(args mock.Arguments) { + resources := args.Get(3).(types.ResourceList) + suite.Len(resources, 2) + suite.Equal(resources[types.ResourceStorage], int64(100)) + suite.Equal(resources[types.ResourceCount], int64(1)) + + f := args.Get(4).(func() error) + f() + }) + + req := httptest.NewRequest(http.MethodPut, "/v2/library/photon/manifests/2.0", nil) + rr := httptest.NewRecorder() + + PutManifestMiddleware()(next).ServeHTTP(rr, req) + suite.Equal(http.StatusOK, rr.Code) + } + + { + // manifest not associated with project and some blobs are not associated with project + mock.OnAnything(suite.blobController, "Exist").Return(false, nil).Once() + + missing := func(ctx context.Context, projectID int64, blobs []*models.Blob) []*models.Blob { + return blobs[:1] + } + + mock.OnAnything(suite.blobController, "FindMissingAssociationsForProject").Return(missing, nil).Once() + mock.OnAnything(suite.quotaController, "Request").Return(nil).Once().Run(func(args mock.Arguments) { + resources := args.Get(3).(types.ResourceList) + suite.Len(resources, 2) + suite.Equal(resources[types.ResourceStorage], int64(100+10)) + suite.Equal(resources[types.ResourceCount], int64(1)) + + f := args.Get(4).(func() error) + f() + }) + + req := httptest.NewRequest(http.MethodPut, "/v2/library/photon/manifests/2.0", nil) + rr := httptest.NewRecorder() + + PutManifestMiddleware()(next).ServeHTTP(rr, req) + suite.Equal(http.StatusOK, rr.Code) + } + + { + // manifest not associated with project and some blobs include foreign layers are not associated with project + mock.OnAnything(suite.blobController, "Exist").Return(false, nil).Once() + + missing := func(ctx context.Context, projectID int64, blobs []*models.Blob) []*models.Blob { + return blobs[1:] + } + + mock.OnAnything(suite.blobController, "FindMissingAssociationsForProject").Return(missing, nil).Once() + mock.OnAnything(suite.quotaController, "Request").Return(nil).Once().Run(func(args mock.Arguments) { + resources := args.Get(3).(types.ResourceList) + suite.Len(resources, 2) + suite.Equal(resources[types.ResourceStorage], int64(100+20)) + suite.Equal(resources[types.ResourceCount], int64(1)) + + f := args.Get(4).(func() error) + f() + }) + + req := httptest.NewRequest(http.MethodPut, "/v2/library/photon/manifests/2.0", nil) + rr := httptest.NewRecorder() + + PutManifestMiddleware()(next).ServeHTTP(rr, req) + suite.Equal(http.StatusOK, rr.Code) + } + + { + // manifest not associated with project and only foreign layers are not associated with project + mock.OnAnything(suite.blobController, "Exist").Return(false, nil).Once() + + missing := func(ctx context.Context, projectID int64, blobs []*models.Blob) []*models.Blob { + return blobs[2:] + } + + mock.OnAnything(suite.blobController, "FindMissingAssociationsForProject").Return(missing, nil).Once() mock.OnAnything(suite.quotaController, "Request").Return(nil).Once().Run(func(args mock.Arguments) { resources := args.Get(3).(types.ResourceList) suite.Len(resources, 2) diff --git a/src/testing/pkg/distribution/manifest.go b/src/testing/pkg/distribution/manifest.go new file mode 100644 index 000000000..8ccc683ca --- /dev/null +++ b/src/testing/pkg/distribution/manifest.go @@ -0,0 +1,59 @@ +// Code generated by mockery v1.0.0. DO NOT EDIT. + +package distribution + +import ( + distribution "github.com/docker/distribution" + mock "github.com/stretchr/testify/mock" +) + +// Manifest is an autogenerated mock type for the Manifest type +type Manifest struct { + mock.Mock +} + +// Payload provides a mock function with given fields: +func (_m *Manifest) Payload() (string, []byte, error) { + ret := _m.Called() + + var r0 string + if rf, ok := ret.Get(0).(func() string); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(string) + } + + var r1 []byte + if rf, ok := ret.Get(1).(func() []byte); ok { + r1 = rf() + } else { + if ret.Get(1) != nil { + r1 = ret.Get(1).([]byte) + } + } + + var r2 error + if rf, ok := ret.Get(2).(func() error); ok { + r2 = rf() + } else { + r2 = ret.Error(2) + } + + return r0, r1, r2 +} + +// References provides a mock function with given fields: +func (_m *Manifest) References() []distribution.Descriptor { + ret := _m.Called() + + var r0 []distribution.Descriptor + if rf, ok := ret.Get(0).(func() []distribution.Descriptor); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]distribution.Descriptor) + } + } + + return r0 +} diff --git a/src/testing/pkg/pkg.go b/src/testing/pkg/pkg.go index c622dfff7..e701ee909 100644 --- a/src/testing/pkg/pkg.go +++ b/src/testing/pkg/pkg.go @@ -15,6 +15,7 @@ package pkg //go:generate mockery -case snake -dir ../../pkg/blob -name Manager -output ./blob -outpkg blob +//go:generate mockery -case snake -dir ../../vendor/github.com/docker/distribution -name Manifest -output ./distribution -outpkg distribution //go:generate mockery -case snake -dir ../../pkg/quota -name Manager -output ./quota -outpkg quota //go:generate mockery -case snake -dir ../../pkg/quota/driver -name Driver -output ./quota/driver -outpkg driver //go:generate mockery -case snake -dir ../../pkg/scan/report -name Manager -output ./scan/report -outpkg report