add to layout replacedescriptor and removedescriptor

Signed-off-by: Avi Deitcher <avi@deitcher.net>
google · Dec 18, 2020 · 3e271bc · 3e271bc
1 parent 8b5370a
commit 3e271bc
Show file tree

Hide file tree

Showing 5 changed files with 377 additions and 0 deletions.
diff --git a/pkg/v1/layout/write.go b/pkg/v1/layout/write.go
@@ -23,6 +23,9 @@ import (
 	"path/filepath"
 
 	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/match"
+	"github.com/google/go-containerregistry/pkg/v1/mutate"
+	"github.com/google/go-containerregistry/pkg/v1/partial"
 	"github.com/google/go-containerregistry/pkg/v1/types"
 	"golang.org/x/sync/errgroup"
 )
@@ -127,6 +130,103 @@ func (l Path) AppendDescriptor(desc v1.Descriptor) error {
 	return l.WriteFile("index.json", rawIndex, os.ModePerm)
 }
 
+// ReplaceImage writes a v1.Image to the Path and updates
+// the index.json to reference it, replacing any existing one that matches matcher, if found.
+func (l Path) ReplaceImage(img v1.Image, matcher match.Matcher, options ...Option) error {
+	if err := l.WriteImage(img); err != nil {
+		return err
+	}
+
+	desc, err := partial.Descriptor(img)
+	if err != nil {
+		return err
+	}
+
+	for _, opt := range options {
+		if err := opt(desc); err != nil {
+			return err
+		}
+	}
+
+	return l.replaceDescriptor(*desc, img, matcher)
+}
+
+// ReplaceIndex writes a v1.ImageIndex to the Path and updates
+// the index.json to reference it, replacing any existing one that matches matcher, if found.
+func (l Path) ReplaceIndex(ii v1.ImageIndex, matcher match.Matcher, options ...Option) error {
+	if err := l.WriteIndex(ii); err != nil {
+		return err
+	}
+
+	desc, err := partial.Descriptor(ii)
+	if err != nil {
+		return err
+	}
+
+	for _, opt := range options {
+		if err := opt(desc); err != nil {
+			return err
+		}
+	}
+
+	return l.replaceDescriptor(*desc, ii, matcher)
+}
+
+// replaceDescriptor adds a descriptor to the index.json of the Path, replacing
+// any one matching matcher, if found.
+func (l Path) replaceDescriptor(desc v1.Descriptor, append mutate.Appendable, matcher match.Matcher) error {
+	ii, err := l.ImageIndex()
+	if err != nil {
+		return err
+	}
+
+	add := mutate.IndexAddendum{
+		Add:        append,
+		Descriptor: desc,
+	}
+	ii, err = mutate.RemoveManifests(ii, matcher)
+	if err != nil {
+		return err
+	}
+	ii = mutate.AppendManifests(ii, add)
+
+	index, err := ii.IndexManifest()
+	if err != nil {
+		return err
+	}
+
+	rawIndex, err := json.MarshalIndent(index, "", "   ")
+	if err != nil {
+		return err
+	}
+
+	return l.WriteFile("index.json", rawIndex, os.ModePerm)
+}
+
+// RemoveDescriptors removes any descriptors that match the match.Matcher from the index.json of the Path.
+func (l Path) RemoveDescriptors(matcher match.Matcher) error {
+	ii, err := l.ImageIndex()
+	if err != nil {
+		return err
+	}
+	ii, err = mutate.RemoveManifests(ii, matcher)
+	if err != nil {
+		return err
+	}
+
+	index, err := ii.IndexManifest()
+	if err != nil {
+		return err
+	}
+
+	rawIndex, err := json.MarshalIndent(index, "", "   ")
+	if err != nil {
+		return err
+	}
+
+	return l.WriteFile("index.json", rawIndex, os.ModePerm)
+}
+
 // WriteFile write a file with arbitrary data at an arbitrary location in a v1
 // layout. Used mostly internally to write files like "oci-layout" and
 // "index.json", also can be used to write other arbitrary files. Do *not* use

diff --git a/pkg/v1/layout/write_test.go b/pkg/v1/layout/write_test.go
@@ -12,6 +12,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/empty"
+	"github.com/google/go-containerregistry/pkg/v1/match"
 	"github.com/google/go-containerregistry/pkg/v1/random"
 	"github.com/google/go-containerregistry/pkg/v1/types"
 	"github.com/google/go-containerregistry/pkg/v1/validate"
@@ -235,3 +236,214 @@ func TestDeduplicatedWrites(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestRemoveDescriptor(t *testing.T) {
+	// need to set up a basic path
+	tmp, err := ioutil.TempDir("", "remove-descriptor-test")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer os.RemoveAll(tmp)
+
+	var ii v1.ImageIndex
+	ii = empty.Index
+	l, err := Write(tmp, ii)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// add two images
+	image1, err := random.Image(1024, 3)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := l.AppendImage(image1); err != nil {
+		t.Fatal(err)
+	}
+	image2, err := random.Image(1024, 3)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := l.AppendImage(image2); err != nil {
+		t.Fatal(err)
+	}
+
+	// remove one of the images by descriptor and ensure it is correct
+	digest1, err := image1.Digest()
+	if err != nil {
+		t.Fatal(err)
+	}
+	digest2, err := image2.Digest()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := l.RemoveDescriptors(match.Digests(digest1)); err != nil {
+		t.Fatal(err)
+	}
+	// ensure we only have one
+	ii, err = l.ImageIndex()
+	if err != nil {
+		t.Fatal(err)
+	}
+	manifest, err := ii.IndexManifest()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(manifest.Manifests) != 1 {
+		t.Fatalf("mismatched manifests count, had %d, expected %d", len(manifest.Manifests), 1)
+	}
+	if manifest.Manifests[0].Digest != digest2 {
+		t.Fatal("removed wrong digest")
+	}
+}
+
+func TestReplaceIndex(t *testing.T) {
+	// need to set up a basic path
+	tmp, err := ioutil.TempDir("", "replace-index-test")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer os.RemoveAll(tmp)
+
+	var ii v1.ImageIndex
+	ii = empty.Index
+	l, err := Write(tmp, ii)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// add two indexes
+	index1, err := random.Index(1024, 3, 3)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := l.AppendIndex(index1); err != nil {
+		t.Fatal(err)
+	}
+	index2, err := random.Index(1024, 3, 3)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := l.AppendIndex(index2); err != nil {
+		t.Fatal(err)
+	}
+	index3, err := random.Index(1024, 3, 3)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// remove one of the indexes by descriptor and ensure it is correct
+	digest1, err := index1.Digest()
+	if err != nil {
+		t.Fatal(err)
+	}
+	digest3, err := index3.Digest()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := l.ReplaceIndex(index3, match.Digests(digest1)); err != nil {
+		t.Fatal(err)
+	}
+	// ensure we only have one
+	ii, err = l.ImageIndex()
+	if err != nil {
+		t.Fatal(err)
+	}
+	manifest, err := ii.IndexManifest()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(manifest.Manifests) != 2 {
+		t.Fatalf("mismatched manifests count, had %d, expected %d", len(manifest.Manifests), 2)
+	}
+	// we should have digest3, and *not* have digest1
+	var have3 bool
+	for _, m := range manifest.Manifests {
+		if m.Digest == digest1 {
+			t.Fatal("found digest1 still not replaced", digest1)
+		}
+		if m.Digest == digest3 {
+			have3 = true
+		}
+	}
+	if !have3 {
+		t.Fatal("could not find digest3", digest3)
+	}
+}
+
+func TestReplaceImage(t *testing.T) {
+	// need to set up a basic path
+	tmp, err := ioutil.TempDir("", "replace-image-test")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer os.RemoveAll(tmp)
+
+	var ii v1.ImageIndex
+	ii = empty.Index
+	l, err := Write(tmp, ii)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// add two images
+	image1, err := random.Image(1024, 3)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := l.AppendImage(image1); err != nil {
+		t.Fatal(err)
+	}
+	image2, err := random.Image(1024, 3)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := l.AppendImage(image2); err != nil {
+		t.Fatal(err)
+	}
+	image3, err := random.Image(1024, 3)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// remove one of the images by descriptor and ensure it is correct
+	digest1, err := image1.Digest()
+	if err != nil {
+		t.Fatal(err)
+	}
+	digest3, err := image3.Digest()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := l.ReplaceImage(image3, match.Digests(digest1)); err != nil {
+		t.Fatal(err)
+	}
+	// ensure we only have one
+	ii, err = l.ImageIndex()
+	if err != nil {
+		t.Fatal(err)
+	}
+	manifest, err := ii.IndexManifest()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(manifest.Manifests) != 2 {
+		t.Fatalf("mismatched manifests count, had %d, expected %d", len(manifest.Manifests), 2)
+	}
+	// we should have digest3, and *not* have digest1
+	var have3 bool
+	for _, m := range manifest.Manifests {
+		if m.Digest == digest1 {
+			t.Fatal("found digest1 still not replaced", digest1)
+		}
+		if m.Digest == digest3 {
+			have3 = true
+		}
+	}
+	if !have3 {
+		t.Fatal("could not find digest3", digest3)
+	}
+}
diff --git a/pkg/v1/mutate/index.go b/pkg/v1/mutate/index.go
@@ -56,6 +56,8 @@ func computeDescriptor(ia IndexAddendum) (*v1.Descriptor, error) {
 type index struct {
 	base v1.ImageIndex
 	adds []IndexAddendum
+	// replaces are applied before adds
+	replaces []v1.Descriptor
 
 	computed  bool
 	manifest  *v1.IndexManifest
@@ -90,6 +92,10 @@ func (i *index) compute() error {
 	}
 	manifest := m.DeepCopy()
 	manifests := manifest.Manifests
+	if len(i.replaces) > 0 {
+		manifests = i.replaces
+	}
+
 	for _, add := range i.adds {
 		desc, err := computeDescriptor(add)
 		if err != nil {
@@ -105,6 +111,7 @@ func (i *index) compute() error {
 			logs.Warn.Printf("Unexpected index addendum: %T", add.Add)
 		}
 	}
+
 	manifest.Manifests = manifests
 
 	// With OCI media types, this should not be set, see discussion:

diff --git a/pkg/v1/mutate/mutate.go b/pkg/v1/mutate/mutate.go
@@ -26,6 +26,7 @@ import (
 
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/empty"
+	"github.com/google/go-containerregistry/pkg/v1/match"
 	"github.com/google/go-containerregistry/pkg/v1/tarball"
 	"github.com/google/go-containerregistry/pkg/v1/types"
 	"github.com/google/go-containerregistry/pkg/v1/v1util"
@@ -92,6 +93,26 @@ func AppendManifests(base v1.ImageIndex, adds ...IndexAddendum) v1.ImageIndex {
 	}
 }
 
+// RemoveManifests removes any descriptors that match the match.Matcher.
+func RemoveManifests(base v1.ImageIndex, matcher match.Matcher) (v1.ImageIndex, error) {
+	ii, err := base.IndexManifest()
+	if err != nil {
+		return base, err
+	}
+
+	// create a list without any existing ones
+	manifests := []v1.Descriptor{}
+	for _, manifest := range ii.Manifests {
+		if !matcher(manifest) {
+			manifests = append(manifests, manifest)
+		}
+	}
+	return &index{
+		base:     base,
+		replaces: manifests,
+	}, nil
+}
+
 // Config mutates the provided v1.Image to have the provided v1.Config
 func Config(base v1.Image, cfg v1.Config) (v1.Image, error) {
 	cf, err := base.ConfigFile()