Start to build up a mutate package.

internal/oci/mutate/map.go

@@ -0,0 +1,171 @@
+//
+// Copyright 2021 The Sigstore 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 mutate
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/empty"
+	"github.com/google/go-containerregistry/pkg/v1/mutate"
+	"github.com/google/go-containerregistry/pkg/v1/types"
+	"github.com/sigstore/cosign/internal/oci"
+)
+
+// Mutator is the signature of the callback supplied to Map.
+// The oci.SignedEntity is either an oci.SignedImageIndex or an oci.SignedImage.
+// This callback is called on oci.SignedImageIndex *before* its children are
+// processed with a context that returns IsBeforeChildren(ctx) == true.
+// If the images within the SignedImageIndex change after the Before pass, then
+// the Mutator will be invoked again on the new SignedImageIndex with a context
+// that returns IsAfterChildren(ctx) == true.
+// If the returned entity is nil, it is filtered from the result of Map.
+type Mutator func(context.Context, oci.SignedEntity) (oci.SignedEntity, error)
+
+// ErrSkipChildren is a special error that may be returned from a Mutator
+// to skip processing of an index's child entities.
+var ErrSkipChildren = errors.New("skip child entities")
+
+// Map calls `fn` on the signed entity and each of its constituent entities (`SignedImageIndex`
+// or `SignedImage`) transitively.
+// Any errors returned by an `fn` are returned by `Map`.
+func Map(ctx context.Context, parent oci.SignedEntity, fn Mutator) (oci.SignedEntity, error) {
+	parent, err := fn(before(ctx), parent)
+	switch {
+	case errors.Is(err, ErrSkipChildren):
+		return parent, nil
+	case err != nil:
+		return nil, err
+	case parent == nil:
+		// If the function returns nil, it filters it.
+		return nil, nil
+	}
+
+	sii, ok := parent.(oci.SignedImageIndex)
+	if !ok {
+		return parent, nil
+	}
+	im, err := sii.IndexManifest()
+	if err != nil {
+		return nil, err
+	}
+
+	// Track whether any of the child entities change.
+	changed := false
+
+	adds := []IndexAddendum{}
+	for _, desc := range im.Manifests {
+		switch desc.MediaType {
+		case types.OCIImageIndex, types.DockerManifestList:
+			x, err := sii.SignedImageIndex(desc.Digest)
+			if err != nil {
+				return nil, err
+			}
+
+			se, err := Map(ctx, x, fn)
+			if err != nil {
+				return nil, err
+			} else if se == nil {
+				// If the function returns nil, it filters it.
+				changed = true
+				continue
+			}
+
+			changed = changed || (x != se)
+			adds = append(adds, IndexAddendum{
+				Add: se.(oci.SignedImageIndex), // Must be an image index.
+				Descriptor: v1.Descriptor{
+					URLs:        desc.URLs,
+					MediaType:   desc.MediaType,
+					Annotations: desc.Annotations,
+					Platform:    desc.Platform,
+				},
+			})
+
+		case types.OCIManifestSchema1, types.DockerManifestSchema2:
+			x, err := sii.SignedImage(desc.Digest)
+			if err != nil {
+				return nil, err
+			}
+
+			se, err := fn(ctx, x)
+			if err != nil {
+				return nil, err
+			} else if se == nil {
+				// If the function returns nil, it filters it.
+				changed = true
+				continue
+			}
+
+			changed = changed || (x != se)
+			adds = append(adds, IndexAddendum{
+				Add: se.(oci.SignedImage), // Must be an image
+				Descriptor: v1.Descriptor{
+					URLs:        desc.URLs,
+					MediaType:   desc.MediaType,
+					Annotations: desc.Annotations,
+					Platform:    desc.Platform,
+				},
+			})
+
+		default:
+			return nil, fmt.Errorf("unknown mime type: %v", desc.MediaType)
+		}
+	}
+
+	if !changed {
+		return parent, nil
+	}
+
+	// Preserve the key attributes from the base IndexManifest.
+	e := mutate.IndexMediaType(empty.Index, im.MediaType)
+	e = mutate.Annotations(e, im.Annotations).(v1.ImageIndex)
+
+	// Construct a new ImageIndex from the new consituent signed images.
+	result := AppendManifests(e, adds...)
+
+	// Since the children changed, give the callback a crack at the new image index.
+	return fn(after(ctx), result)
+}
+
+// This is used to associate which pass of the Map a particular
+// callback is being invoked for.
+type mapPassKey struct{}
+
+// before decorates the context such that IsBeforeChildren(ctx) is true.
+func before(ctx context.Context) context.Context {
+	return context.WithValue(ctx, mapPassKey{}, "before")
+}
+
+// after decorates the context such that IsAfterChildren(ctx) is true.
+func after(ctx context.Context) context.Context {
+	return context.WithValue(ctx, mapPassKey{}, "after")
+}
+
+// IsBeforeChildren is true within a Mutator when it is called before the children
+// have been processed.
+func IsBeforeChildren(ctx context.Context) bool {
+	return ctx.Value(mapPassKey{}) == "before"
+}
+
+// IsAfterChildren is true within a Mutator when it is called after the children
+// have been processed; however, this call is only made if the set of children
+// changes since the Before call.
+func IsAfterChildren(ctx context.Context) bool {
+	return ctx.Value(mapPassKey{}) == "after"
+}