From a31869f2704d00531cdf53db28eb9f1605f0baf7 Mon Sep 17 00:00:00 2001
From: Alex Stan <alexandrustan96@yahoo.ro>
Date: Thu, 14 Jul 2022 15:30:24 +0300
Subject: [PATCH] fix GetReferrers function to be able to retrieve referrers of
 any specified artifactType

Signed-off-by: Alex Stan <alexandrustan96@yahoo.ro>
---
 pkg/storage/local.go      | 11 +++++----
 pkg/storage/local_test.go | 50 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/pkg/storage/local.go b/pkg/storage/local.go
index 6ebd0927..2261f4d8 100644
--- a/pkg/storage/local.go
+++ b/pkg/storage/local.go
@@ -19,7 +19,6 @@ import (
 	guuid "github.com/gofrs/uuid"
 	"github.com/minio/sha256-simd"
 	"github.com/notaryproject/notation-go-lib"
-	notreg "github.com/notaryproject/notation/pkg/registry"
 	godigest "github.com/opencontainers/go-digest"
 	ispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/opencontainers/umoci"
@@ -1443,7 +1442,7 @@ func (is *ImageStoreLocal) DeleteBlob(repo, digest string) error {
 	return nil
 }
 
-func (is *ImageStoreLocal) GetReferrers(repo, digest, mediaType string) ([]artifactspec.Descriptor, error) {
+func (is *ImageStoreLocal) GetReferrers(repo, digest, artifactType string) ([]artifactspec.Descriptor, error) {
 	var lockLatency time.Time
 
 	dir := path.Join(is.rootDir, repo)
@@ -1509,14 +1508,16 @@ func (is *ImageStoreLocal) GetReferrers(repo, digest, mediaType string) ([]artif
 			return nil, err
 		}
 
-		if mediaType != artManifest.ArtifactType || gdigest != artManifest.Subject.Digest {
+		if artifactType != artManifest.ArtifactType || gdigest != artManifest.Subject.Digest {
 			continue
 		}
 
 		result = append(result, artifactspec.Descriptor{
 			MediaType:    manifest.MediaType,
-			ArtifactType: notreg.ArtifactTypeNotation,
-			Digest:       manifest.Digest, Size: manifest.Size, Annotations: manifest.Annotations,
+			ArtifactType: artManifest.ArtifactType,
+			Digest:       manifest.Digest,
+			Size:         manifest.Size,
+			Annotations:  manifest.Annotations,
 		})
 
 		found = true
diff --git a/pkg/storage/local_test.go b/pkg/storage/local_test.go
index c0566ac2..ff8bc3c2 100644
--- a/pkg/storage/local_test.go
+++ b/pkg/storage/local_test.go
@@ -17,6 +17,7 @@ import (
 
 	godigest "github.com/opencontainers/go-digest"
 	ispec "github.com/opencontainers/image-spec/specs-go/v1"
+	artifactspec "github.com/oras-project/artifacts-spec/specs-go/v1"
 	"github.com/rs/zerolog"
 	. "github.com/smartystreets/goconvey/convey"
 	"zotregistry.io/zot/errors"
@@ -163,6 +164,55 @@ func TestStorageFSAPIs(t *testing.T) {
 	})
 }
 
+func TestGetReferrers(t *testing.T) {
+	dir := t.TempDir()
+
+	log := log.Logger{Logger: zerolog.New(os.Stdout)}
+	metrics := monitoring.NewMetricsServer(false, log)
+	imgStore := storage.NewImageStore(dir, true, storage.DefaultGCDelay, true, true, log, metrics)
+
+	Convey("Get referrers", t, func(c C) {
+		err := test.CopyFiles("../../test/data/zot-test", path.Join(dir, "zot-test"))
+		So(err, ShouldBeNil)
+		body := []byte("this is a blob")
+		digest := godigest.FromBytes(body)
+		buf := bytes.NewBuffer(body)
+		buflen := buf.Len()
+		err = ioutil.WriteFile(path.Join(imgStore.RootDir(), //nolint: gosec
+			"zot-test", "blobs", digest.Algorithm().String(), digest.Encoded()),
+			buf.Bytes(), 0o644)
+		So(err, ShouldBeNil)
+		_, n, err := imgStore.FullBlobUpload("zot-test", buf, digest.String())
+		So(err, ShouldBeNil)
+		So(n, ShouldEqual, buflen)
+
+		artifactManifest := artifactspec.Manifest{}
+		artifactManifest.ArtifactType = "signature-example"
+		artifactManifest.Subject = artifactspec.Descriptor{
+			MediaType: ispec.MediaTypeImageManifest,
+			Digest:    digest,
+			Size:      int64(buflen),
+		}
+		artifactManifest.Blobs = []artifactspec.Descriptor{}
+
+		manBuf, err := json.Marshal(artifactManifest)
+		manBufLen := len(manBuf)
+		So(err, ShouldBeNil)
+		manDigest := godigest.FromBytes(manBuf)
+		_, err = imgStore.PutImageManifest("zot-test", manDigest.Encoded(), artifactspec.MediaTypeArtifactManifest, manBuf)
+		So(err, ShouldBeNil)
+
+		So(err, ShouldBeNil)
+		descriptors, err := imgStore.GetReferrers("zot-test", digest.String(), "signature-example")
+		So(err, ShouldBeNil)
+		So(descriptors, ShouldNotBeEmpty)
+		So(descriptors[0].ArtifactType, ShouldEqual, "signature-example")
+		So(descriptors[0].MediaType, ShouldEqual, artifactspec.MediaTypeArtifactManifest)
+		So(descriptors[0].Size, ShouldEqual, manBufLen)
+		So(descriptors[0].Digest, ShouldEqual, manDigest)
+	})
+}
+
 func TestDedupeLinks(t *testing.T) {
 	dir := t.TempDir()