From a05b0ac923cb18949296c81575bebafddc93c258 Mon Sep 17 00:00:00 2001
From: Avi Deitcher <avi@deitcher.net>
Date: Mon, 21 Jun 2021 23:33:34 +0300
Subject: [PATCH] check for arch when pulling to cache, push by descriptor

Signed-off-by: Avi Deitcher <avi@deitcher.net>
---
 src/cmd/linuxkit/cache/pull.go  | 18 ++++++++++++++++++
 src/cmd/linuxkit/cache/push.go  |  8 ++------
 src/cmd/linuxkit/cache/write.go |  7 ++-----
 3 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/cmd/linuxkit/cache/pull.go b/src/cmd/linuxkit/cache/pull.go
index ef3618367..28853ccf3 100644
--- a/src/cmd/linuxkit/cache/pull.go
+++ b/src/cmd/linuxkit/cache/pull.go
@@ -2,6 +2,7 @@ package cache
 
 import (
 	"errors"
+	"fmt"
 
 	"github.com/containerd/containerd/reference"
 	"github.com/google/go-containerregistry/pkg/v1"
@@ -28,6 +29,9 @@ func (p *Provider) ValidateImage(ref *reference.Spec, architecture string) (lkts
 			if desc, err = partial.Descriptor(img); err != nil {
 				return ImageSource{}, errors.New("image could not create valid descriptor")
 			}
+			if desc.Platform == nil || desc.Platform.Architecture != architecture || desc.Platform.OS != "linux" {
+				return ImageSource{}, fmt.Errorf("image was not for requested architecture: linux/%s", architecture)
+			}
 		} else {
 			ii, err := root.ImageIndex()
 			if err == nil {
@@ -36,6 +40,20 @@ func (p *Provider) ValidateImage(ref *reference.Spec, architecture string) (lkts
 					return ImageSource{}, errors.New("index could not create valid descriptor")
 				}
 			}
+			// check that the index has a manifest for our arch
+			im, err := imageIndex.IndexManifest()
+			if err != nil {
+				return ImageSource{}, fmt.Errorf("could not get index manifest: %v", err)
+			}
+			for _, m := range im.Manifests {
+				if m.Platform != nil && m.Platform.Architecture == architecture && m.Platform.OS == "linux" {
+					return p.NewSource(
+						ref,
+						architecture,
+						desc,
+					), nil
+				}
+			}
 		}
 	}
 	// three possibilities now:
diff --git a/src/cmd/linuxkit/cache/push.go b/src/cmd/linuxkit/cache/push.go
index cf9c44598..fc49bc414 100644
--- a/src/cmd/linuxkit/cache/push.go
+++ b/src/cmd/linuxkit/cache/push.go
@@ -58,13 +58,9 @@ func (p *Provider) Push(name string) error {
 			if err != nil {
 				return fmt.Errorf("could not create a valid arch-specific tag %s: %v", archTag, err)
 			}
-			image, err := p.FindRoot(archTag)
+			img, err := p.cache.Image(m.Digest)
 			if err != nil {
-				return fmt.Errorf("could not find arch-specific image in cache %s: %v", archTag, err)
-			}
-			img, err := image.Image()
-			if err != nil {
-				return fmt.Errorf("found arch-specific image in cache %s, but could not resolve to actual image: %v", archTag, err)
+				return fmt.Errorf("could not find arch-specific image in cache %s: %v", m.Digest, err)
 			}
 			log.Debugf("pushing image %s", tag)
 			if err := remote.Tag(tag, img, options...); err != nil {
diff --git a/src/cmd/linuxkit/cache/write.go b/src/cmd/linuxkit/cache/write.go
index f90534e31..4d485d2cb 100644
--- a/src/cmd/linuxkit/cache/write.go
+++ b/src/cmd/linuxkit/cache/write.go
@@ -82,11 +82,8 @@ func (p *Provider) ImagePull(ref *reference.Spec, trustedRef, architecture strin
 	if err != nil {
 		return ImageSource{}, fmt.Errorf("unable to save image to cache: %v", err)
 	}
-	return p.NewSource(
-		ref,
-		architecture,
-		&desc.Descriptor,
-	), nil
+	// ensure it includes our architecture
+	return p.ValidateImage(ref, architecture)
 }
 
 // ImageLoad takes an OCI format image tar stream and writes it locally. It should be