fix(blob,quota): process blobs already in registry no but associated with project (#11071)

1. Before put manifest request, ensure that the requested size resource
include the blobs which are referenced by the manifest but not
associated with project.
2. After put manifest request, associate the blobs which are referenced
by the manifest but not associated with project.

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

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

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)

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

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

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
 }

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)

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

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