fix(quota): correct size quota for image with foreign layers

1. Sync blobs from manifest for image with foreign layers.
2. Ignore size of foreign layers when compute size quota.
3. Fix repo info of artifact when upgrade from 1.8 version.

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

src/common/dao/blob.go

@@ -5,6 +5,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/docker/distribution"
 	"github.com/goharbor/harbor/src/common/models"
 	"github.com/goharbor/harbor/src/common/utils/log"
 )
@@ -64,6 +65,99 @@ func DeleteBlob(digest string) error {
 	return err
 }
 
+// ListBlobs list blobs according to the query conditions
+func ListBlobs(query *models.BlobQuery) ([]*models.Blob, error) {
+	qs := GetOrmer().QueryTable(&models.Blob{})
+
+	if query != nil {
+		if query.Digest != "" {
+			qs = qs.Filter("Digest", query.Digest)
+		}
+
+		if query.ContentType != "" {
+			qs = qs.Filter("ContentType", query.ContentType)
+		}
+
+		if len(query.Digests) > 0 {
+			qs = qs.Filter("Digest__in", query.Digests)
+		}
+
+		if query.Size > 0 {
+			qs = qs.Limit(query.Size)
+			if query.Page > 0 {
+				qs = qs.Offset((query.Page - 1) * query.Size)
+			}
+		}
+	}
+
+	blobs := []*models.Blob{}
+	_, err := qs.All(&blobs)
+	return blobs, err
+}
+
+// SyncBlobs sync references to blobs
+func SyncBlobs(references []distribution.Descriptor) error {
+	if len(references) == 0 {
+		return nil
+	}
+
+	var digests []string
+	for _, reference := range references {
+		digests = append(digests, reference.Digest.String())
+	}
+
+	existing, err := ListBlobs(&models.BlobQuery{Digests: digests})
+	if err != nil {
+		return err
+	}
+
+	mp := make(map[string]*models.Blob, len(existing))
+	for _, blob := range existing {
+		mp[blob.Digest] = blob
+	}
+
+	var missing, updating []*models.Blob
+	for _, reference := range references {
+		if blob, found := mp[reference.Digest.String()]; found {
+			if blob.ContentType != reference.MediaType {
+				blob.ContentType = reference.MediaType
+				updating = append(updating, blob)
+			}
+
+		} else {
+			missing = append(missing, &models.Blob{
+				Digest:       reference.Digest.String(),
+				ContentType:  reference.MediaType,
+				Size:         reference.Size,
+				CreationTime: time.Now(),
+			})
+		}
+	}
+
+	o := GetOrmer()
+
+	if len(updating) > 0 {
+		for _, blob := range updating {
+			if _, err := o.Update(blob, "content_type"); err != nil {
+				log.Warningf("Failed to update blob %s, error: %v", blob.Digest, err)
+			}
+		}
+	}
+
+	if len(missing) > 0 {
+		_, err = o.InsertMulti(10, missing)
+		if err != nil {
+			if strings.Contains(err.Error(), "duplicate key value violates unique constraint") {
+				return ErrDupRows
+			}
+		}
+
+		return err
+	}
+
+	return nil
+}
+
 // GetBlobsByArtifact returns blobs of artifact
 func GetBlobsByArtifact(artifactDigest string) ([]*models.Blob, error) {
 	sql := `SELECT * FROM blob WHERE digest IN (SELECT digest_blob FROM artifact_blob WHERE digest_af = ?)`

src/common/dao/blob_test.go

@@ -18,6 +18,8 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/docker/distribution"
+	"github.com/docker/distribution/manifest/schema2"
 	"github.com/goharbor/harbor/src/common/models"
 	"github.com/goharbor/harbor/src/common/utils"
 	"github.com/opencontainers/go-digest"
@@ -69,6 +71,87 @@ func TestDeleteBlob(t *testing.T) {
 	require.Nil(t, err)
 }
 
+func TestListBlobs(t *testing.T) {
+	assert := assert.New(t)
+
+	d1 := digest.FromString(utils.GenerateRandomString())
+	d2 := digest.FromString(utils.GenerateRandomString())
+	d3 := digest.FromString(utils.GenerateRandomString())
+	d4 := digest.FromString(utils.GenerateRandomString())
+	for _, e := range []struct {
+		Digest      digest.Digest
+		ContentType string
+		Size        int64
+	}{
+		{d1, schema2.MediaTypeLayer, 1},
+		{d2, schema2.MediaTypeLayer, 2},
+		{d3, schema2.MediaTypeForeignLayer, 3},
+		{d4, schema2.MediaTypeForeignLayer, 4},
+	} {
+		blob := &models.Blob{
+			Digest:      e.Digest.String(),
+			ContentType: e.ContentType,
+			Size:        e.Size,
+		}
+		_, err := AddBlob(blob)
+		assert.Nil(err)
+	}
+
+	defer func() {
+		for _, d := range []digest.Digest{d1, d2, d3, d4} {
+			DeleteBlob(d.String())
+		}
+	}()
+
+	blobs, err := ListBlobs(&models.BlobQuery{Digest: d1.String()})
+	assert.Nil(err)
+	assert.Len(blobs, 1)
+
+	blobs, err = ListBlobs(&models.BlobQuery{ContentType: schema2.MediaTypeForeignLayer})
+	assert.Nil(err)
+	assert.Len(blobs, 2)
+
+	blobs, err = ListBlobs(&models.BlobQuery{Digests: []string{d1.String(), d2.String(), d3.String()}})
+	assert.Nil(err)
+	assert.Len(blobs, 3)
+}
+
+func TestSyncBlobs(t *testing.T) {
+	assert := assert.New(t)
+
+	d1 := digest.FromString(utils.GenerateRandomString())
+	d2 := digest.FromString(utils.GenerateRandomString())
+	d3 := digest.FromString(utils.GenerateRandomString())
+	d4 := digest.FromString(utils.GenerateRandomString())
+
+	blob := &models.Blob{
+		Digest:      d1.String(),
+		ContentType: schema2.MediaTypeLayer,
+		Size:        1,
+	}
+	_, err := AddBlob(blob)
+	assert.Nil(err)
+
+	assert.Nil(SyncBlobs([]distribution.Descriptor{}))
+
+	references := []distribution.Descriptor{
+		{MediaType: schema2.MediaTypeLayer, Digest: d1, Size: 1},
+		{MediaType: schema2.MediaTypeForeignLayer, Digest: d2, Size: 2},
+		{MediaType: schema2.MediaTypeForeignLayer, Digest: d3, Size: 3},
+		{MediaType: schema2.MediaTypeForeignLayer, Digest: d4, Size: 4},
+	}
+	assert.Nil(SyncBlobs(references))
+	defer func() {
+		for _, d := range []digest.Digest{d1, d2, d3, d4} {
+			DeleteBlob(d.String())
+		}
+	}()
+
+	blobs, err := ListBlobs(&models.BlobQuery{Digests: []string{d1.String(), d2.String(), d3.String(), d4.String()}})
+	assert.Nil(err)
+	assert.Len(blobs, 4)
+}
+
 func prepareImage(projectID int64, projectName, name, tag string, layerDigests ...string) (string, error) {
 	digest := digest.FromString(strings.Join(layerDigests, ":")).String()
 	artifact := &models.Artifact{PID: projectID, Repo: projectName + "/" + name, Digest: digest, Tag: tag}

src/common/models/blob.go

@@ -2,6 +2,8 @@ package models
 
 import (
 	"time"
+
+	"github.com/docker/distribution/manifest/schema2"
 )
 
 // Blob holds the details of a blob.
@@ -17,3 +19,16 @@ type Blob struct {
 func (b *Blob) TableName() string {
 	return "blob"
 }
+
+// IsForeignLayer returns true if the blob is foreign layer
+func (b *Blob) IsForeignLayer() bool {
+	return b.ContentType == schema2.MediaTypeForeignLayer
+}
+
+// BlobQuery ...
+type BlobQuery struct {
+	Digest      string
+	ContentType string
+	Digests     []string
+	Pagination
+}

src/core/api/quota/registry/registry.go

@@ -15,9 +15,12 @@
 package registry
 
 import (
+	"strings"
+	"sync"
+	"time"
+
 	"github.com/docker/distribution/manifest/schema1"
 	"github.com/docker/distribution/manifest/schema2"
-
 	"github.com/goharbor/harbor/src/common"
 	"github.com/goharbor/harbor/src/common/dao"
 	"github.com/goharbor/harbor/src/common/models"
@@ -29,9 +32,6 @@ import (
 	"github.com/goharbor/harbor/src/core/promgr"
 	coreutils "github.com/goharbor/harbor/src/core/utils"
 	"github.com/pkg/errors"
-	"strings"
-	"sync"
-	"time"
 )
 
 // Migrator ...
@@ -423,7 +423,7 @@ func infoOfRepo(pid int64, repo string) (quota.RepoData, error) {
 		}
 		af := &models.Artifact{
 			PID:          pid,
-			Repo:         strings.Split(repo, "/")[1],
+			Repo:         repo,
 			Tag:          tag,
 			Digest:       desc.Digest.String(),
 			Kind:         "Docker-Image",

src/core/middlewares/sizequota/builder.go

@@ -119,6 +119,11 @@ func (*manifestCreationBuilder) Build(req *http.Request) (interceptor.Intercepto
 	// Replace request with manifests info context
 	*req = *req.WithContext(util.NewManifestInfoContext(req.Context(), info))
 
+	// Sync manifest layers to blobs for foreign layers not pushed and they are not in blob table
+	if err := info.SyncBlobs(); err != nil {
+		log.Warningf("Failed to sync blobs, error: %v", err)
+	}
+
 	opts := []quota.Option{
 		quota.EnforceResources(config.QuotaPerProjectEnable()),
 		quota.WithManager("project", strconv.FormatInt(info.ProjectID, 10)),

src/core/middlewares/sizequota/handler_test.go

@@ -133,6 +133,24 @@ func manifestWithAdditionalLayers(raw schema2.Manifest, layerSizes []int64) sche
 	return manifest
 }
 
+func manifestWithAdditionalForeignLayers(raw schema2.Manifest, layerSizes []int64) schema2.Manifest {
+	var manifest schema2.Manifest
+
+	manifest.Versioned = raw.Versioned
+	manifest.Config = raw.Config
+	manifest.Layers = append(manifest.Layers, raw.Layers...)
+
+	for _, size := range layerSizes {
+		manifest.Layers = append(manifest.Layers, distribution.Descriptor{
+			MediaType: schema2.MediaTypeForeignLayer,
+			Size:      size,
+			Digest:    digest.FromString(randomString(15)),
+		})
+	}
+
+	return manifest
+}
+
 func digestOfManifest(manifest schema2.Manifest) string {
 	bytes, _ := json.Marshal(manifest)
 
@@ -149,8 +167,10 @@ func sizeOfImage(manifest schema2.Manifest) int64 {
 	totalSizeOfLayers := manifest.Config.Size
 
 	for _, layer := range manifest.Layers {
+		if layer.MediaType != schema2.MediaTypeForeignLayer {
 			totalSizeOfLayers += layer.Size
 		}
+	}
 
 	return sizeOfManifest(manifest) + totalSizeOfLayers
 }
@@ -256,8 +276,10 @@ func putManifest(projectName, name, tag string, manifest schema2.Manifest) {
 func pushImage(projectName, name, tag string, manifest schema2.Manifest) {
 	putBlobUpload(projectName, name, genUUID(), manifest.Config.Digest.String(), manifest.Config.Size)
 	for _, layer := range manifest.Layers {
+		if layer.MediaType != schema2.MediaTypeForeignLayer {
 			putBlobUpload(projectName, name, genUUID(), layer.Digest.String(), layer.Size)
 		}
+	}
 
 	putManifest(projectName, name, tag, manifest)
 }
@@ -388,6 +410,25 @@ func (suite *HandlerSuite) TestDeleteManifest() {
 	})
 }
 
+func (suite *HandlerSuite) TestImageWithForeignLayers() {
+	withProject(func(projectID int64, projectName string) {
+		manifest := manifestWithAdditionalForeignLayers(makeManifest(1, []int64{2, 3, 4, 5}), []int64{6, 7})
+		size := sizeOfImage(manifest)
+
+		pushImage(projectName, "photon", "latest", manifest)
+		suite.checkStorageUsage(size, projectID)
+		suite.checkStorageUsage(sizeOfManifest(manifest)+1+2+3+4+5, projectID)
+
+		blobs, err := dao.GetBlobsByArtifact(digestOfManifest(manifest))
+		if suite.Nil(err) {
+			suite.Len(blobs, 8)
+		}
+
+		deleteManifest(projectName, "photon", digestOfManifest(manifest))
+		suite.checkStorageUsage(0, projectID)
+	})
+}
+
 func (suite *HandlerSuite) TestImageOverwrite() {
 	withProject(func(projectID int64, projectName string) {
 		manifest1 := makeManifest(1, []int64{2, 3, 4, 5})

src/core/middlewares/sizequota/util.go

@@ -274,8 +274,10 @@ func computeResourcesForManifestCreation(req *http.Request) (types.ResourceList,
 	size := info.Descriptor.Size
 
 	for _, blob := range blobs {
+		if !blob.IsForeignLayer() {
 			size += blob.Size
 		}
+	}
 
 	return types.ResourceList{types.ResourceStorage: size}, nil
 }
@@ -297,8 +299,10 @@ func computeResourcesForManifestDeletion(req *http.Request) (types.ResourceList,
 
 	var size int64
 	for _, blob := range blobs {
+		if !blob.IsForeignLayer() {
 			size = size + blob.Size
 		}
+	}
 
 	return types.ResourceList{types.ResourceStorage: size}, nil
 }

src/core/middlewares/util/util.go

@@ -174,6 +174,17 @@ func (info *ManifestInfo) BlobMutexKey(blob *models.Blob, suffix ...string) stri
 	return strings.Join(append(a, suffix...), ":")
 }
 
+// SyncBlobs sync layers of manifest to blobs
+func (info *ManifestInfo) SyncBlobs() error {
+	err := dao.SyncBlobs(info.References)
+	if err == dao.ErrDupRows {
+		log.Warning("Some blobs created by others, ignore this error")
+		return nil
+	}
+
+	return err
+}
+
 // GetBlobsNotInProject returns blobs of the manifest which not in the project
 func (info *ManifestInfo) GetBlobsNotInProject() ([]*models.Blob, error) {
 	var digests []string