From d9bfbcb99e526b2a9417160e209b816e1b1fb6bd Mon Sep 17 00:00:00 2001
From: jonjohnsonjr <jonjohnson@google.com>
Date: Wed, 5 Jan 2022 14:06:05 -0800
Subject: [PATCH] crane pull: support pulling index to OCI Layout (#1215)

* crane pull: support pulling index to OCI Layout

Prior to this, we'd also resolve an index to an image when using crane
pull. For OCI Layouts, we can actually pull the whole index, so when
--format=oci and --platform is unset, just do that.

* Address feedback

* Add platform test
---
 .github/workflows/e2e.yaml | 18 +++++++++++++++
 cmd/crane/cmd/pull.go      | 45 +++++++++++++++++++++++++++++++++++---
 pkg/crane/pull.go          |  7 ++++--
 3 files changed, 65 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/e2e.yaml b/.github/workflows/e2e.yaml
index 6c3cdf47e..6b0419aa4 100644
--- a/.github/workflows/e2e.yaml
+++ b/.github/workflows/e2e.yaml
@@ -63,3 +63,21 @@ jobs:
         ./app/crane pull --format=oci $img $layout
         ./app/crane push --image-refs=foo.images $layout $dst
         diff <(./app/crane manifest $img) <(./app/crane manifest $(cat foo.images))
+
+        # Make sure we can roundtrip an index (distroless).
+        distroless=$(mktemp -d)
+        remote="gcr.io/distroless/static"
+        local="localhost:1338/distroless:static"
+
+        ./app/crane pull --format=oci $remote $distroless
+        ./app/crane push $distroless $local
+        diff <(./app/crane manifest $remote) <(./app/crane manifest $local)
+
+        # And that it works for a single platform (pulling from what we just pushed).
+        distroless=$(mktemp -d)
+        remote="$local"
+        local="localhost:1338/distroless/platform:static"
+
+        ./app/crane pull --platform=linux/arm64 --format=oci $remote $distroless
+        ./app/crane push $distroless $local
+        diff <(./app/crane manifest --platform linux/arm64 $remote) <(./app/crane manifest $local)
diff --git a/cmd/crane/cmd/pull.go b/cmd/crane/cmd/pull.go
index b3c772743..8bb242cf8 100644
--- a/cmd/crane/cmd/pull.go
+++ b/cmd/crane/cmd/pull.go
@@ -18,8 +18,11 @@ import (
 	"fmt"
 
 	"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/cache"
+	"github.com/google/go-containerregistry/pkg/v1/layout"
+	"github.com/google/go-containerregistry/pkg/v1/remote"
 	"github.com/spf13/cobra"
 )
 
@@ -33,16 +36,38 @@ func NewCmdPull(options *[]crane.Option) *cobra.Command {
 		Args:  cobra.MinimumNArgs(2),
 		RunE: func(_ *cobra.Command, args []string) error {
 			imageMap := map[string]v1.Image{}
+			indexMap := map[string]v1.ImageIndex{}
 			srcList, path := args[:len(args)-1], args[len(args)-1]
 			for _, src := range srcList {
-				img, err := crane.Pull(src, *options...)
+				o := crane.GetOptions(*options...)
+				ref, err := name.ParseReference(src, o.Name...)
 				if err != nil {
-					return fmt.Errorf("pulling %s: %w", src, err)
+					return fmt.Errorf("parsing reference %q: %w", src, err)
+				}
+
+				rmt, err := remote.Get(ref, o.Remote...)
+				if err != nil {
+					return err
+				}
+
+				// If we're writing an index to a layout and --platform hasn't been set,
+				// pull the entire index, not just a child image.
+				if format == "oci" && rmt.MediaType.IsIndex() && o.Platform == nil {
+					idx, err := rmt.ImageIndex()
+					if err != nil {
+						return err
+					}
+					indexMap[src] = idx
+					continue
+				}
+
+				img, err := rmt.Image()
+				if err != nil {
+					return err
 				}
 				if cachePath != "" {
 					img = cache.Image(img, cache.NewFilesystemCache(cachePath))
 				}
-
 				imageMap[src] = img
 			}
 
@@ -59,6 +84,20 @@ func NewCmdPull(options *[]crane.Option) *cobra.Command {
 				if err := crane.MultiSaveOCI(imageMap, path); err != nil {
 					return fmt.Errorf("saving oci image layout %s: %w", path, err)
 				}
+
+				// crane.MultiSaveOCI doesn't support index, so just append these at the end.
+				p, err := layout.FromPath(path)
+				if err != nil {
+					return err
+				}
+				for ref, idx := range indexMap {
+					anns := map[string]string{
+						"dev.ggcr.image.name": ref,
+					}
+					if err := p.AppendIndex(idx, layout.WithAnnotations(anns)); err != nil {
+						return err
+					}
+				}
 			default:
 				return fmt.Errorf("unexpected --format: %q (valid values are: tarball, legacy, and oci)", format)
 			}
diff --git a/pkg/crane/pull.go b/pkg/crane/pull.go
index 7e6e5b7b6..b19ac7c21 100644
--- a/pkg/crane/pull.go
+++ b/pkg/crane/pull.go
@@ -133,8 +133,11 @@ func MultiSaveOCI(imgMap map[string]v1.Image, path string) error {
 			return err
 		}
 	}
-	for _, img := range imgMap {
-		if err = p.AppendImage(img); err != nil {
+	for ref, img := range imgMap {
+		anns := map[string]string{
+			"dev.ggcr.image.name": ref,
+		}
+		if err = p.AppendImage(img, layout.WithAnnotations(anns)); err != nil {
 			return err
 		}
 	}