Allow mutate.Annotations to annotate an Image or ImageIndex (#1082)

* Allow mutate.Annotations to annotate an Image or ImageIndex

* Correctly spell annotateable

* Support arbitrary RawManifest inputs, add unit tests

cmd/crane/cmd/mutate.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/google/go-containerregistry/pkg/crane"
 	"github.com/google/go-containerregistry/pkg/name"
+	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/mutate"
 	"github.com/spf13/cobra"
 )
@@ -91,7 +92,7 @@ func NewCmdMutate(options *[]crane.Option) *cobra.Command {
 				log.Fatalf("mutating config: %v", err)
 			}
 
-			img = mutate.Annotations(img, annotations)
+			img = mutate.Annotations(img, annotations).(v1.Image)
 
 			// If the new ref isn't provided, write over the original image.
 			// If that ref was provided by digest (e.g., output from

pkg/v1/mutate/index.go

@@ -61,12 +61,13 @@ type index struct {
 	// remove is removed before adds
 	remove match.Matcher
 
-	computed  bool
-	manifest  *v1.IndexManifest
-	mediaType *types.MediaType
-	imageMap  map[v1.Hash]v1.Image
-	indexMap  map[v1.Hash]v1.ImageIndex
-	layerMap  map[v1.Hash]v1.Layer
+	computed    bool
+	manifest    *v1.IndexManifest
+	annotations map[string]string
+	mediaType   *types.MediaType
+	imageMap    map[v1.Hash]v1.Image
+	indexMap    map[v1.Hash]v1.ImageIndex
+	layerMap    map[v1.Hash]v1.Layer
 }
 
 var _ v1.ImageIndex = (*index)(nil)
@@ -137,6 +138,15 @@ func (i *index) compute() error {
 		}
 	}
 
+	if i.annotations != nil {
+		if manifest.Annotations == nil {
+			manifest.Annotations = map[string]string{}
+		}
+		for k, v := range i.annotations {
+			manifest.Annotations[k] = v
+		}
+	}
+
 	i.manifest = manifest
 	i.computed = true
 	return nil

pkg/v1/mutate/mutate.go

@@ -17,6 +17,7 @@ package mutate
 import (
 	"archive/tar"
 	"bytes"
+	"encoding/json"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -28,6 +29,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/partial"
 	"github.com/google/go-containerregistry/pkg/v1/tarball"
 	"github.com/google/go-containerregistry/pkg/v1/types"
 )
@@ -113,12 +115,70 @@ func Config(base v1.Image, cfg v1.Config) (v1.Image, error) {
 	return ConfigFile(base, cf)
 }
 
-// Annotations mutates the provided v1.Image to have the provided annotations
-func Annotations(base v1.Image, annotations map[string]string) v1.Image {
-	return &image{
-		base:        base,
-		annotations: annotations,
+// Annotatable represents a manifest that can carry annotations.
+type Annotatable interface {
+	partial.WithRawManifest
+}
+
+// Annotations mutates the annotations on an annotatable image or index manifest.
+//
+// The annotatable input is expected to be a v1.Image or v1.ImageIndex, and
+// returns the same type. You can type-assert the result like so:
+//
+//     img := Annotations(empty.Image, map[string]string{
+//         "foo": "bar",
+//     }).(v1.Image)
+//
+// Or for an index:
+//
+//     idx := Annotations(empty.Index, map[string]string{
+//         "foo": "bar",
+//     }).(v1.ImageIndex)
+//
+// If the input Annotatable is not an Image or ImageIndex, the result will
+// attempt to lazily annotate the raw manifest.
+func Annotations(f Annotatable, anns map[string]string) Annotatable {
+	if img, ok := f.(v1.Image); ok {
+		return &image{
+			base:        img,
+			annotations: anns,
+		}
+	}
+	if idx, ok := f.(v1.ImageIndex); ok {
+		return &index{
+			base:        idx,
+			annotations: anns,
+		}
+	}
+	return arbitraryRawManifest{f, anns}
+}
+
+type arbitraryRawManifest struct {
+	a    Annotatable
+	anns map[string]string
+}
+
+func (a arbitraryRawManifest) RawManifest() ([]byte, error) {
+	b, err := a.a.RawManifest()
+	if err != nil {
+		return nil, err
+	}
+	var m map[string]interface{}
+	if err := json.Unmarshal(b, &m); err != nil {
+		return nil, err
+	}
+	if ann, ok := m["annotations"]; ok {
+		if annm, ok := ann.(map[string]string); ok {
+			for k, v := range a.anns {
+				annm[k] = v
+			}
+		} else {
+			return nil, fmt.Errorf(".annotations is not a map: %T", ann)
+		}
+	} else {
+		m["annotations"] = a.anns
 	}
+	return json.Marshal(m)
 }
 
 // ConfigFile mutates the provided v1.Image to have the provided v1.ConfigFile

pkg/v1/mutate/mutate_test.go

@@ -266,21 +266,43 @@ func TestMutateConfig(t *testing.T) {
 	}
 }
 
-func TestAnnotations(t *testing.T) {
-	source := sourceImage(t)
-
-	newAnnotations := map[string]string{
-		"im.the.first.annotation": "hello world",
-	}
-
-	result := mutate.Annotations(source, newAnnotations)
-
-	if configDigestsAreEqual(t, source, result) {
-		t.Errorf("mutating the manifest annotations MUST mutate the config digest")
-	}
+type arbitrary struct {
+}
 
-	if err := validate.Image(result); err != nil {
-		t.Errorf("validate.Image() = %v", err)
+func (arbitrary) RawManifest() ([]byte, error) {
+	return []byte(`{"hello":"world"}`), nil
+}
+func TestAnnotations(t *testing.T) {
+	anns := map[string]string{
+		"foo": "bar",
+	}
+
+	for _, c := range []struct {
+		desc string
+		in   mutate.Annotatable
+		want string
+	}{{
+		desc: "image",
+		in:   empty.Image,
+		want: `{"schemaVersion":2,"mediaType":"application/vnd.docker.distribution.manifest.v2+json","config":{"mediaType":"application/vnd.docker.container.image.v1+json","size":115,"digest":"sha256:5b943e2b943f6c81dbbd4e2eca5121f4fcc39139e3d1219d6d89bd925b77d9fe"},"layers":[],"annotations":{"foo":"bar"}}`,
+	}, {
+		desc: "index",
+		in:   empty.Index,
+		want: `{"schemaVersion":2,"manifests":null,"annotations":{"foo":"bar"}}`,
+	}, {
+		desc: "arbitrary",
+		in:   arbitrary{},
+		want: `{"annotations":{"foo":"bar"},"hello":"world"}`,
+	}} {
+		t.Run(c.desc, func(t *testing.T) {
+			got, err := mutate.Annotations(c.in, anns).RawManifest()
+			if err != nil {
+				t.Fatalf("Annotations: %v", err)
+			}
+			if d := cmp.Diff(c.want, string(got)); d != "" {
+				t.Errorf("Diff(-want,+got): %s", d)
+			}
+		})
 	}
 }