fix: refactor tagging implementation for OCI Store

content/oci/oci.go

@@ -27,10 +27,12 @@ import (
 	"path/filepath"
 	"sync"
 
+	"github.com/opencontainers/go-digest"
 	specs "github.com/opencontainers/image-spec/specs-go"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"oras.land/oras-go/v2/content"
 	"oras.land/oras-go/v2/errdef"
+	"oras.land/oras-go/v2/internal/containers"
 	"oras.land/oras-go/v2/internal/descriptor"
 	"oras.land/oras-go/v2/internal/graph"
 	"oras.land/oras-go/v2/internal/resolver"
@@ -142,18 +144,14 @@ func (s *Store) Tag(ctx context.Context, desc ocispec.Descriptor, reference stri
 		return fmt.Errorf("%s: %s: %w", desc.Digest, desc.MediaType, errdef.ErrNotFound)
 	}
 
-	if desc.Annotations == nil {
-		desc.Annotations = map[string]string{}
-	}
-	desc.Annotations[ocispec.AnnotationRefName] = reference
 	return s.tag(ctx, desc, reference)
 }
 
 // tag tags a descriptor with a reference string.
 func (s *Store) tag(ctx context.Context, desc ocispec.Descriptor, reference string) error {
 	dgst := desc.Digest.String()
 	if reference != dgst {
-		// mark desc for deduplication in SaveIndex()
+		// also tag desc by its digest
 		if err := s.tagResolver.Tag(ctx, desc, dgst); err != nil {
 			return err
 		}
@@ -269,14 +267,35 @@ func (s *Store) SaveIndex() error {
 	defer s.indexLock.Unlock()
 
 	var manifests []ocispec.Descriptor
+	tagged := make(containers.Set[digest.Digest])
 	refMap := s.tagResolver.Map()
+
+	// 1. Add descriptors that are associated with tags
+	// Note: One descriptor can be associated with multiple tags.
+	for ref, desc := range refMap {
+		if ref != desc.Digest.String() {
+			// copy the original annotation map
+			annotations := make(map[string]string, len(desc.Annotations)+1)
+			for k, v := range desc.Annotations {
+				annotations[k] = v
+			}
+			annotations[ocispec.AnnotationRefName] = ref
+			desc.Annotations = annotations
+			manifests = append(manifests, desc)
+			// mark the digest as tagged for deduplication in step 2
+			tagged.Add(desc.Digest)
+		}
+	}
+	// 2. Add descriptors that are not associated with any tag
 	for ref, desc := range refMap {
-		if ref == desc.Digest.String() && desc.Annotations[ocispec.AnnotationRefName] != "" {
-			// skip saving desc if ref is a digest and desc is tagged
-			continue
+		if ref == desc.Digest.String() {
+			// skip tagged ones since they have been added in step 1
+			if !tagged.Contains(desc.Digest) {
+				manifests = append(manifests, desc)
+			}
 		}
-		manifests = append(manifests, desc)
 	}
+
 	s.index.Manifests = manifests
 	return s.writeIndexFile()
 }

content/oci/oci_test.go

@@ -598,7 +598,10 @@ func TestStore_RepeatTag(t *testing.T) {
 		t.Fatal("Store.Push() error =", err)
 	}
 	if got, want := len(internalResolver.Map()), 1; got != want {
-		t.Errorf("resolver.Map() = %v, want %v", got, want)
+		t.Errorf("len(resolver.Map()) = %v, want %v", got, want)
+	}
+	if got, want := len(s.index.Manifests), 1; got != want {
+		t.Errorf("len(index.Manifests) = %v, want %v", got, want)
 	}
 
 	err = s.Tag(ctx, desc, ref)
@@ -608,6 +611,9 @@ func TestStore_RepeatTag(t *testing.T) {
 	if got, want := len(internalResolver.Map()), 2; got != want {
 		t.Errorf("resolver.Map() = %v, want %v", got, want)
 	}
+	if got, want := len(s.index.Manifests), 1; got != want {
+		t.Errorf("len(index.Manifests) = %v, want %v", got, want)
+	}
 
 	gotDesc, err := s.Resolve(ctx, desc.Digest.String())
 	if err != nil {
@@ -635,6 +641,9 @@ func TestStore_RepeatTag(t *testing.T) {
 	if got, want := len(internalResolver.Map()), 3; got != want {
 		t.Errorf("resolver.Map() = %v, want %v", got, want)
 	}
+	if got, want := len(s.index.Manifests), 2; got != want {
+		t.Errorf("len(index.Manifests) = %v, want %v", got, want)
+	}
 
 	err = s.Tag(ctx, desc, ref)
 	if err != nil {
@@ -643,6 +652,9 @@ func TestStore_RepeatTag(t *testing.T) {
 	if got, want := len(internalResolver.Map()), 3; got != want {
 		t.Errorf("resolver.Map() = %v, want %v", got, want)
 	}
+	if got, want := len(s.index.Manifests), 2; got != want {
+		t.Errorf("len(index.Manifests) = %v, want %v", got, want)
+	}
 
 	gotDesc, err = s.Resolve(ctx, desc.Digest.String())
 	if err != nil {
@@ -670,6 +682,9 @@ func TestStore_RepeatTag(t *testing.T) {
 	if got, want := len(internalResolver.Map()), 3; got != want {
 		t.Errorf("resolver.Map() = %v, want %v", got, want)
 	}
+	if got, want := len(s.index.Manifests), 2; got != want {
+		t.Errorf("len(index.Manifests) = %v, want %v", got, want)
+	}
 
 	err = s.Tag(ctx, desc, ref)
 	if err != nil {
@@ -678,6 +693,50 @@ func TestStore_RepeatTag(t *testing.T) {
 	if got, want := len(internalResolver.Map()), 4; got != want {
 		t.Errorf("resolver.Map() = %v, want %v", got, want)
 	}
+	if got, want := len(s.index.Manifests), 3; got != want {
+		t.Errorf("len(index.Manifests) = %v, want %v", got, want)
+	}
+
+	gotDesc, err = s.Resolve(ctx, desc.Digest.String())
+	if err != nil {
+		t.Fatal("Store.Resolve() error =", err)
+	}
+	if !reflect.DeepEqual(gotDesc, desc) {
+		t.Errorf("Store.Resolve() = %v, want %v", gotDesc, desc)
+	}
+
+	gotDesc, err = s.Resolve(ctx, ref)
+	if err != nil {
+		t.Fatal("Store.Resolve() error =", err)
+	}
+	if !reflect.DeepEqual(gotDesc, desc) {
+		t.Errorf("Store.Resolve() = %v, want %v", gotDesc, desc)
+	}
+
+	// tag another blob
+	blob = []byte("barfoo")
+	desc = content.NewDescriptorFromBytes("test", blob)
+	err = s.Push(ctx, desc, bytes.NewReader(blob))
+	if err != nil {
+		t.Fatal("Store.Push() error =", err)
+	}
+	if got, want := len(internalResolver.Map()), 4; got != want {
+		t.Errorf("resolver.Map() = %v, want %v", got, want)
+	}
+	if got, want := len(s.index.Manifests), 3; got != want {
+		t.Errorf("len(index.Manifests) = %v, want %v", got, want)
+	}
+
+	err = s.Tag(ctx, desc, ref)
+	if err != nil {
+		t.Fatal("Store.Tag() error =", err)
+	}
+	if got, want := len(internalResolver.Map()), 5; got != want {
+		t.Errorf("resolver.Map() = %v, want %v", got, want)
+	}
+	if got, want := len(s.index.Manifests), 4; got != want {
+		t.Errorf("len(index.Manifests) = %v, want %v", got, want)
+	}
 
 	gotDesc, err = s.Resolve(ctx, desc.Digest.String())
 	if err != nil {
@@ -696,6 +755,100 @@ func TestStore_RepeatTag(t *testing.T) {
 	}
 }
 
+// Related bug: https://github.com/oras-project/oras-go/issues/461
+func TestStore_TagByDigest(t *testing.T) {
+	tempDir := t.TempDir()
+	s, err := New(tempDir)
+	if err != nil {
+		t.Fatal("New() error =", err)
+	}
+	ctx := context.Background()
+
+	// get internal resolver
+	internalResolver := s.tagResolver
+
+	manifest := []byte(`{"layers":[]}`)
+	manifestDesc := content.NewDescriptorFromBytes(ocispec.MediaTypeImageManifest, manifest)
+
+	// push a manifest
+	err = s.Push(ctx, manifestDesc, bytes.NewReader(manifest))
+	if err != nil {
+		t.Fatal("Store.Push() error =", err)
+	}
+	if got, want := len(internalResolver.Map()), 1; got != want {
+		t.Errorf("len(resolver.Map()) = %v, want %v", got, want)
+	}
+	if got, want := len(s.index.Manifests), 1; got != want {
+		t.Errorf("len(index.Manifests) = %v, want %v", got, want)
+	}
+	gotDesc, err := s.Resolve(ctx, manifestDesc.Digest.String())
+	if err != nil {
+		t.Fatal("Store.Resolve() error =", err)
+	}
+	if !reflect.DeepEqual(gotDesc, manifestDesc) {
+		t.Errorf("Store.Resolve() = %v, want %v", gotDesc, manifestDesc)
+	}
+
+	// tag manifest by digest
+	err = s.Tag(ctx, manifestDesc, manifestDesc.Digest.String())
+	if err != nil {
+		t.Fatal("Store.Tag() error =", err)
+	}
+	if got, want := len(internalResolver.Map()), 1; got != want {
+		t.Errorf("len(resolver.Map()) = %v, want %v", got, want)
+	}
+	if got, want := len(s.index.Manifests), 1; got != want {
+		t.Errorf("len(index.Manifests) = %v, want %v", got, want)
+	}
+	gotDesc, err = s.Resolve(ctx, manifestDesc.Digest.String())
+	if err != nil {
+		t.Fatal("Store.Resolve() error =", err)
+	}
+	if !reflect.DeepEqual(gotDesc, manifestDesc) {
+		t.Errorf("Store.Resolve() = %v, want %v", gotDesc, manifestDesc)
+	}
+
+	// push a blob
+	blob := []byte("foobar")
+	blobDesc := content.NewDescriptorFromBytes("test", blob)
+	err = s.Push(ctx, blobDesc, bytes.NewReader(blob))
+	if err != nil {
+		t.Fatal("Store.Push() error =", err)
+	}
+	if got, want := len(internalResolver.Map()), 1; got != want {
+		t.Errorf("resolver.Map() = %v, want %v", got, want)
+	}
+	if got, want := len(s.index.Manifests), 1; got != want {
+		t.Errorf("len(index.Manifests) = %v, want %v", got, want)
+	}
+	gotDesc, err = s.Resolve(ctx, blobDesc.Digest.String())
+	if err != nil {
+		t.Fatal("Store.Resolve() error =", err)
+	}
+	if gotDesc.Digest != blobDesc.Digest || gotDesc.Size != blobDesc.Size {
+		t.Errorf("Store.Resolve() = %v, want %v", gotDesc, blobDesc)
+	}
+
+	// tag blob by digest
+	err = s.Tag(ctx, blobDesc, blobDesc.Digest.String())
+	if err != nil {
+		t.Fatal("Store.Tag() error =", err)
+	}
+	if got, want := len(internalResolver.Map()), 2; got != want {
+		t.Errorf("resolver.Map() = %v, want %v", got, want)
+	}
+	if got, want := len(s.index.Manifests), 2; got != want {
+		t.Errorf("len(index.Manifests) = %v, want %v", got, want)
+	}
+	gotDesc, err = s.Resolve(ctx, blobDesc.Digest.String())
+	if err != nil {
+		t.Fatal("Store.Resolve() error =", err)
+	}
+	if !reflect.DeepEqual(gotDesc, blobDesc) {
+		t.Errorf("Store.Resolve() = %v, want %v", gotDesc, blobDesc)
+	}
+}
+
 func TestStore_BadIndex(t *testing.T) {
 	tempDir := t.TempDir()
 	content := []byte("whatever")
@@ -922,6 +1075,11 @@ func TestStore_ExistingStore(t *testing.T) {
 	if err := s.Tag(ctx, indexRoot, tag); err != nil {
 		t.Fatal("Tag() error =", err)
 	}
+	// tag index root by digest
+	// related bug: https://github.com/oras-project/oras-go/issues/461
+	if err := s.Tag(ctx, indexRoot, indexRoot.Digest.String()); err != nil {
+		t.Fatal("Tag() error =", err)
+	}
 
 	// test with another OCI store instance to mock loading from an existing store
 	anotherS, err := New(tempDir)
@@ -1029,7 +1187,6 @@ func TestStore_ExistingStore(t *testing.T) {
 			t.Errorf("Store.Predecessors(%d) = %v, want %v", i, predecessors, want)
 		}
 	}
-
 }
 
 func TestCopy_MemoryToOCI_FullCopy(t *testing.T) {

extendedcopy.go

@@ -25,6 +25,7 @@ import (
 	"golang.org/x/sync/semaphore"
 	"oras.land/oras-go/v2/content"
 	"oras.land/oras-go/v2/internal/cas"
+	"oras.land/oras-go/v2/internal/containers"
 	"oras.land/oras-go/v2/internal/copyutil"
 	"oras.land/oras-go/v2/internal/descriptor"
 	"oras.land/oras-go/v2/internal/docker"
@@ -131,7 +132,7 @@ func ExtendedCopyGraph(ctx context.Context, src content.ReadOnlyGraphStorage, ds
 // findRoots finds the root nodes reachable from the given node through a
 // depth-first search.
 func findRoots(ctx context.Context, storage content.ReadOnlyGraphStorage, node ocispec.Descriptor, opts ExtendedCopyGraphOptions) ([]ocispec.Descriptor, error) {
-	visited := make(map[descriptor.Descriptor]bool)
+	visited := make(containers.Set[descriptor.Descriptor])
 	rootMap := make(map[descriptor.Descriptor]ocispec.Descriptor)
 	addRoot := func(key descriptor.Descriptor, val ocispec.Descriptor) {
 		if _, exists := rootMap[key]; !exists {
@@ -158,11 +159,11 @@ func findRoots(ctx context.Context, storage content.ReadOnlyGraphStorage, node o
 		currentNode := current.Node
 		currentKey := descriptor.FromOCI(currentNode)
 
-		if visited[currentKey] {
+		if visited.Contains(currentKey) {
 			// skip the current node if it has been visited
 			continue
 		}
-		visited[currentKey] = true
+		visited.Add(currentKey)
 
 		// stop finding predecessors if the target depth is reached
 		if opts.Depth > 0 && current.Depth == opts.Depth {
@@ -186,7 +187,7 @@ func findRoots(ctx context.Context, storage content.ReadOnlyGraphStorage, node o
 		// Push the predecessor nodes to the stack and keep finding from there.
 		for _, predecessor := range predecessors {
 			predecessorKey := descriptor.FromOCI(predecessor)
-			if !visited[predecessorKey] {
+			if !visited.Contains(predecessorKey) {
 				// push the predecessor node with increased depth
 				stack.Push(copyutil.NodeInfo{Node: predecessor, Depth: current.Depth + 1})
 			}

internal/containers/set.go

@@ -0,0 +1,30 @@
+/*
+Copyright The ORAS Authors.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package containers
+
+// Set represents a set data structure.
+type Set[T comparable] map[T]struct{}
+
+// Add adds key into the set s.
+func (s Set[T]) Add(key T) {
+	s[key] = struct{}{}
+}
+
+// Contains returns true if the set s contains key.
+func (s Set[T]) Contains(key T) bool {
+	_, ok := s[key]
+	return ok
+}