Add partial.Manifests for lazy index access (#1631)

jonjohnsonjr · web-flow · commit 65e78dc801e4 · 2023-04-06T15:43:26.000-04:00
This allows us to manipulate an index that contains an image that
contains a streaming layer.
diff --git a/cmd/crane/cmd/flatten.go b/cmd/crane/cmd/flatten.go
@@ -123,22 +123,13 @@ func push(flat partial.Describable, ref name.Reference, o crane.Options) error {
 	return fmt.Errorf("can't push %T", flat)
 }
 
-type remoteIndex interface {
-	Manifests() ([]partial.Describable, error)
-}
-
 func flattenIndex(old v1.ImageIndex, repo name.Repository, use string, o crane.Options) (partial.Describable, error) {
-	ri, ok := old.(remoteIndex)
-	if !ok {
-		return nil, fmt.Errorf("unexpected index")
-	}
-
 	m, err := old.IndexManifest()
 	if err != nil {
 		return nil, err
 	}
 
-	manifests, err := ri.Manifests()
+	manifests, err := partial.Manifests(old)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/v1/mutate/index.go b/pkg/v1/mutate/index.go
@@ -16,13 +16,15 @@ package mutate
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"sync"
 
 	"github.com/google/go-containerregistry/pkg/logs"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/match"
 	"github.com/google/go-containerregistry/pkg/v1/partial"
+	"github.com/google/go-containerregistry/pkg/v1/stream"
 	"github.com/google/go-containerregistry/pkg/v1/types"
 )
 
@@ -208,3 +210,23 @@ func (i *index) RawManifest() ([]byte, error) {
 	}
 	return json.Marshal(i.manifest)
 }
+
+func (i *index) Manifests() ([]partial.Describable, error) {
+	if err := i.compute(); errors.Is(err, stream.ErrNotComputed) {
+		// Index contains a streamable layer which has not yet been
+		// consumed. Just return the manifests we have in case the caller
+		// is going to consume the streamable layers.
+		manifests, err := partial.Manifests(i.base)
+		if err != nil {
+			return nil, err
+		}
+		for _, add := range i.adds {
+			manifests = append(manifests, add.Add)
+		}
+		return manifests, nil
+	} else if err != nil {
+		return nil, err
+	}
+
+	return partial.ComputeManifests(i)
+}
diff --git a/pkg/v1/partial/index.go b/pkg/v1/partial/index.go
@@ -19,6 +19,7 @@ import (
 
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/match"
+	"github.com/google/go-containerregistry/pkg/v1/types"
 )
 
 // FindManifests given a v1.ImageIndex, find the manifests that fit the matcher.
@@ -83,3 +84,82 @@ func FindIndexes(index v1.ImageIndex, matcher match.Matcher) ([]v1.ImageIndex, e
 	}
 	return matches, nil
 }
+
+type withManifests interface {
+	Manifests() ([]Describable, error)
+}
+
+type withLayer interface {
+	Layer(v1.Hash) (v1.Layer, error)
+}
+
+type describable struct {
+	desc v1.Descriptor
+}
+
+func (d describable) Digest() (v1.Hash, error) {
+	return d.desc.Digest, nil
+}
+
+func (d describable) Size() (int64, error) {
+	return d.desc.Size, nil
+}
+
+func (d describable) MediaType() (types.MediaType, error) {
+	return d.desc.MediaType, nil
+}
+
+func (d describable) Descriptor() (*v1.Descriptor, error) {
+	return &d.desc, nil
+}
+
+// Manifests is analogous to v1.Image.Layers in that it allows values in the
+// returned list to be lazily evaluated, which enables an index to contain
+// an image that contains a streaming layer.
+//
+// This should have been part of the v1.ImageIndex interface, but wasn't.
+// It is instead usable through this extension interface.
+func Manifests(idx v1.ImageIndex) ([]Describable, error) {
+	if wm, ok := idx.(withManifests); ok {
+		return wm.Manifests()
+	}
+
+	return ComputeManifests(idx)
+}
+
+// ComputeManifests provides a fallback implementation for Manifests.
+func ComputeManifests(idx v1.ImageIndex) ([]Describable, error) {
+	m, err := idx.IndexManifest()
+	if err != nil {
+		return nil, err
+	}
+	manifests := []Describable{}
+	for _, desc := range m.Manifests {
+		switch {
+		case desc.MediaType.IsImage():
+			img, err := idx.Image(desc.Digest)
+			if err != nil {
+				return nil, err
+			}
+			manifests = append(manifests, img)
+		case desc.MediaType.IsIndex():
+			idx, err := idx.ImageIndex(desc.Digest)
+			if err != nil {
+				return nil, err
+			}
+			manifests = append(manifests, idx)
+		default:
+			if wl, ok := idx.(withLayer); ok {
+				layer, err := wl.Layer(desc.Digest)
+				if err != nil {
+					return nil, err
+				}
+				manifests = append(manifests, layer)
+			} else {
+				manifests = append(manifests, describable{desc})
+			}
+		}
+	}
+
+	return manifests, nil
+}
diff --git a/pkg/v1/remote/index.go b/pkg/v1/remote/index.go
@@ -148,40 +148,6 @@ func (r *remoteIndex) Layer(h v1.Hash) (v1.Layer, error) {
 	return nil, fmt.Errorf("layer not found: %s", h)
 }
 
-// Experiment with a better API for v1.ImageIndex. We might want to move this
-// to partial?
-func (r *remoteIndex) Manifests() ([]partial.Describable, error) {
-	m, err := r.IndexManifest()
-	if err != nil {
-		return nil, err
-	}
-	manifests := []partial.Describable{}
-	for _, desc := range m.Manifests {
-		switch {
-		case desc.MediaType.IsImage():
-			img, err := r.Image(desc.Digest)
-			if err != nil {
-				return nil, err
-			}
-			manifests = append(manifests, img)
-		case desc.MediaType.IsIndex():
-			idx, err := r.ImageIndex(desc.Digest)
-			if err != nil {
-				return nil, err
-			}
-			manifests = append(manifests, idx)
-		default:
-			layer, err := r.Layer(desc.Digest)
-			if err != nil {
-				return nil, err
-			}
-			manifests = append(manifests, layer)
-		}
-	}
-
-	return manifests, nil
-}
-
 func (r *remoteIndex) imageByPlatform(platform v1.Platform) (v1.Image, error) {
 	desc, err := r.childByPlatform(platform)
 	if err != nil {
