feat!: allow pulling images by index sha

src/pkg/images/common.go

@@ -18,9 +18,11 @@ import (
 	"github.com/defenseunicorns/pkg/helpers/v2"
 	"github.com/google/go-containerregistry/pkg/authn"
 	"github.com/google/go-containerregistry/pkg/crane"
+	"github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/zarf-dev/zarf/src/pkg/logger"
 	"github.com/zarf-dev/zarf/src/pkg/state"
+	"oras.land/oras-go/v2/content"
 	"oras.land/oras-go/v2/registry/remote/auth"
 	"oras.land/oras-go/v2/registry/remote/retry"
 )
@@ -119,15 +121,17 @@ func ShouldUsePlainHTTP(ctx context.Context, registryURL string, client *auth.Cl
 	return true, nil
 }
 
-func isManifest(mediaType string) bool {
+// IsManifest reports whether the media type represents an OCI manifest.
+func IsManifest(mediaType string) bool {
 	switch mediaType {
 	case ocispec.MediaTypeImageManifest, DockerMediaTypeManifest:
 		return true
 	}
 	return false
 }
 
-func isIndex(mediaType string) bool {
+// IsIndex reports whether the media type represents an OCI image index.
+func IsIndex(mediaType string) bool {
 	switch mediaType {
 	case ocispec.MediaTypeImageIndex, DockerMediaTypeManifestList:
 		return true
@@ -214,12 +218,106 @@ func WithPushAuth(ri state.RegistryInfo) crane.Option {
 	return WithBasicAuth(ri.PushUsername, ri.PushPassword)
 }
 
-func getSizeOfImage(manifestDesc ocispec.Descriptor, manifest ocispec.Manifest) int64 {
-	var totalSize int64
-	totalSize += manifestDesc.Size
+// formatPlatform renders an ocispec.Platform as "arch[/variant]". Empty input returns "".
+func formatPlatform(p *ocispec.Platform) string {
+	if p == nil || p.Architecture == "" {
+		return ""
+	}
+	s := p.Architecture
+	if p.Variant != "" {
+		s += "/" + p.Variant
+	}
+	return s
+}
+
+// getSizeOfManifest returns the total byte size of a manifest plus its config and layers.
+func getSizeOfManifest(manifestDesc ocispec.Descriptor, manifestBytes []byte) (int64, error) {
+	var manifest ocispec.Manifest
+	if err := json.Unmarshal(manifestBytes, &manifest); err != nil {
+		return 0, fmt.Errorf("unable to unmarshal manifest %s: %w", manifestDesc.Digest, err)
+	}
+	totalSize := manifestDesc.Size
 	for _, layer := range manifest.Layers {
 		totalSize += layer.Size
 	}
 	totalSize += manifest.Config.Size
-	return totalSize
+	return totalSize, nil
+}
+
+// inspectIndex walks an OCI image index (recursing into nested indexes) and returns the total
+// byte size of every uniquely-referenced blob and one "arch[/variant]" string per leaf manifest.
+func inspectIndex(ctx context.Context, fetcher content.Fetcher, indexDesc ocispec.Descriptor, indexBytes []byte) (int64, []string, error) {
+	return walkIndex(ctx, fetcher, indexDesc, indexBytes, map[digest.Digest]struct{}{})
+}
+
+func walkIndex(ctx context.Context, fetcher content.Fetcher, indexDesc ocispec.Descriptor, indexBytes []byte, seen map[digest.Digest]struct{}) (int64, []string, error) {
+	if _, ok := seen[indexDesc.Digest]; ok {
+		return 0, nil, nil
+	}
+	seen[indexDesc.Digest] = struct{}{}
+	var idx ocispec.Index
+	if err := json.Unmarshal(indexBytes, &idx); err != nil {
+		return 0, nil, fmt.Errorf("unable to unmarshal index: %w", err)
+	}
+	totalSize := indexDesc.Size
+	var platforms []string
+	for _, child := range idx.Manifests {
+		if _, ok := seen[child.Digest]; ok {
+			continue
+		}
+		switch {
+		case IsIndex(child.MediaType):
+			b, err := content.FetchAll(ctx, fetcher, child)
+			if err != nil {
+				return 0, nil, fmt.Errorf("failed to fetch nested index %s: %w", child.Digest, err)
+			}
+			size, childPlatforms, err := walkIndex(ctx, fetcher, child, b, seen)
+			if err != nil {
+				return 0, nil, err
+			}
+			totalSize += size
+			platforms = append(platforms, childPlatforms...)
+		case IsManifest(child.MediaType):
+			seen[child.Digest] = struct{}{}
+			b, err := content.FetchAll(ctx, fetcher, child)
+			if err != nil {
+				return 0, nil, fmt.Errorf("failed to fetch child manifest %s: %w", child.Digest, err)
+			}
+			size, err := uniqueManifestSize(child, b, seen)
+			if err != nil {
+				return 0, nil, err
+			}
+			totalSize += size
+			if s := formatPlatform(child.Platform); s != "" {
+				platforms = append(platforms, s)
+			}
+		default:
+			seen[child.Digest] = struct{}{}
+			totalSize += child.Size
+		}
+	}
+	return totalSize, platforms, nil
+}
+
+// uniqueManifestSize returns manifestDesc.Size plus the size of any config/layer blobs whose
+// digests have not yet been seen. The seen set is mutated to include every newly-counted digest.
+// The manifest's own digest must already be in seen.
+func uniqueManifestSize(manifestDesc ocispec.Descriptor, manifestBytes []byte, seen map[digest.Digest]struct{}) (int64, error) {
+	var manifest ocispec.Manifest
+	if err := json.Unmarshal(manifestBytes, &manifest); err != nil {
+		return 0, fmt.Errorf("unable to unmarshal manifest %s: %w", manifestDesc.Digest, err)
+	}
+	total := manifestDesc.Size
+	if _, ok := seen[manifest.Config.Digest]; !ok {
+		seen[manifest.Config.Digest] = struct{}{}
+		total += manifest.Config.Size
+	}
+	for _, layer := range manifest.Layers {
+		if _, ok := seen[layer.Digest]; ok {
+			continue
+		}
+		seen[layer.Digest] = struct{}{}
+		total += layer.Size
+	}
+	return total, nil
 }

src/pkg/images/common_test.go

@@ -0,0 +1,123 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: 2021-Present The Zarf Authors
+
+package images
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"github.com/stretchr/testify/require"
+	"github.com/zarf-dev/zarf/src/test/testutil"
+	"oras.land/oras-go/v2/content"
+	"oras.land/oras-go/v2/registry/remote"
+)
+
+func fetchAll(ctx context.Context, t *testing.T, fetcher content.Fetcher, desc ocispec.Descriptor) []byte {
+	t.Helper()
+	b, err := content.FetchAll(ctx, fetcher, desc)
+	require.NoError(t, err)
+	return b
+}
+
+// pushImageWithLayer pushes a single-platform image referencing the given layer, tags the manifest
+// with linux/arch, and returns its descriptor plus the naive (non-deduplicated) total of every
+// blob it references.
+func pushImageWithLayer(ctx context.Context, t *testing.T, repo *remote.Repository, arch string, layer ocispec.Descriptor) (ocispec.Descriptor, int64) {
+	t.Helper()
+	desc := testutil.PushSinglePlatformImageWithLayer(ctx, t, repo, arch, layer)
+	desc.Platform = &ocispec.Platform{OS: "linux", Architecture: arch}
+	config := manifestConfig(ctx, t, repo, desc)
+	return desc, desc.Size + config.Size + layer.Size
+}
+
+// manifestConfig fetches a manifest and returns its config descriptor.
+func manifestConfig(ctx context.Context, t *testing.T, fetcher content.Fetcher, manDesc ocispec.Descriptor) ocispec.Descriptor {
+	t.Helper()
+	b := fetchAll(ctx, t, fetcher, manDesc)
+	var m ocispec.Manifest
+	require.NoError(t, json.Unmarshal(b, &m))
+	return m.Config
+}
+
+func TestInspectIndex(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.TestContext(t)
+	upstream := testutil.SetupInMemoryRegistryDynamic(ctx, t)
+
+	t.Run("flat oci index sums every blob and reports each leaf platform", func(t *testing.T) {
+		t.Parallel()
+		repo := testutil.NewRepo(t, upstream+"/inspect/flat")
+		layerA := testutil.PushBlob(ctx, t, repo, ocispec.MediaTypeImageLayer, testutil.RandomBytes(t, 128))
+		layerB := testutil.PushBlob(ctx, t, repo, ocispec.MediaTypeImageLayer, testutil.RandomBytes(t, 256))
+		manA, sizeA := pushImageWithLayer(ctx, t, repo, "amd64", layerA)
+		manB, sizeB := pushImageWithLayer(ctx, t, repo, "arm64", layerB)
+		idx := testutil.PushIndex(ctx, t, repo, []ocispec.Descriptor{manA, manB})
+
+		size, platforms, err := inspectIndex(ctx, repo, idx, fetchAll(ctx, t, repo, idx))
+		require.NoError(t, err)
+		require.Equal(t, idx.Size+sizeA+sizeB, size)
+		require.ElementsMatch(t, []string{"amd64", "arm64"}, platforms)
+	})
+
+	t.Run("nested oci index recurses into inner index", func(t *testing.T) {
+		t.Parallel()
+		repo := testutil.NewRepo(t, upstream+"/inspect/nested")
+		layerA := testutil.PushBlob(ctx, t, repo, ocispec.MediaTypeImageLayer, testutil.RandomBytes(t, 128))
+		layerB := testutil.PushBlob(ctx, t, repo, ocispec.MediaTypeImageLayer, testutil.RandomBytes(t, 256))
+		manA, sizeA := pushImageWithLayer(ctx, t, repo, "amd64", layerA)
+		manB, sizeB := pushImageWithLayer(ctx, t, repo, "arm64", layerB)
+		innerIdx := testutil.PushIndex(ctx, t, repo, []ocispec.Descriptor{manA, manB})
+		outerIdx := testutil.PushIndex(ctx, t, repo, []ocispec.Descriptor{innerIdx})
+
+		size, platforms, err := inspectIndex(ctx, repo, outerIdx, fetchAll(ctx, t, repo, outerIdx))
+		require.NoError(t, err)
+		require.Equal(t, outerIdx.Size+innerIdx.Size+sizeA+sizeB, size)
+		require.ElementsMatch(t, []string{"amd64", "arm64"}, platforms)
+	})
+
+	t.Run("docker manifest list is treated as an index", func(t *testing.T) {
+		t.Parallel()
+		repo := testutil.NewRepo(t, upstream+"/inspect/docker")
+		layerA := testutil.PushBlob(ctx, t, repo, ocispec.MediaTypeImageLayer, testutil.RandomBytes(t, 128))
+		manA, sizeA := pushImageWithLayer(ctx, t, repo, "amd64", layerA)
+		list := pushDockerManifestList(ctx, t, repo, []ocispec.Descriptor{manA})
+
+		size, platforms, err := inspectIndex(ctx, repo, list, fetchAll(ctx, t, repo, list))
+		require.NoError(t, err)
+		require.Equal(t, list.Size+sizeA, size)
+		require.ElementsMatch(t, []string{"amd64"}, platforms)
+	})
+
+	t.Run("dedups a layer referenced by sibling manifests", func(t *testing.T) {
+		t.Parallel()
+		repo := testutil.NewRepo(t, upstream+"/inspect/dedup-layer")
+		shared := testutil.PushBlob(ctx, t, repo, ocispec.MediaTypeImageLayer, testutil.RandomBytes(t, 512))
+		manA, _ := pushImageWithLayer(ctx, t, repo, "amd64", shared)
+		manB, _ := pushImageWithLayer(ctx, t, repo, "arm64", shared)
+		idx := testutil.PushIndex(ctx, t, repo, []ocispec.Descriptor{manA, manB})
+
+		size, _, err := inspectIndex(ctx, repo, idx, fetchAll(ctx, t, repo, idx))
+		require.NoError(t, err)
+
+		configA := manifestConfig(ctx, t, repo, manA)
+		configB := manifestConfig(ctx, t, repo, manB)
+		expected := idx.Size + manA.Size + configA.Size + shared.Size + manB.Size + configB.Size
+		require.Equal(t, expected, size)
+	})
+
+	t.Run("dedups a manifest listed twice in an index", func(t *testing.T) {
+		t.Parallel()
+		repo := testutil.NewRepo(t, upstream+"/inspect/dup-manifest")
+		layer := testutil.PushBlob(ctx, t, repo, ocispec.MediaTypeImageLayer, testutil.RandomBytes(t, 256))
+		man, manSize := pushImageWithLayer(ctx, t, repo, "amd64", layer)
+		idx := testutil.PushIndex(ctx, t, repo, []ocispec.Descriptor{man, man})
+
+		size, platforms, err := inspectIndex(ctx, repo, idx, fetchAll(ctx, t, repo, idx))
+		require.NoError(t, err)
+		require.Equal(t, idx.Size+manSize, size)
+		require.ElementsMatch(t, []string{"amd64"}, platforms)
+	})
+}