m/node/core/update: add support for OCI index

The updater now also accepts OCI indexes, and will select the image with
the matching architecture from the index.

This is tested by making on of the update steps a multiarch index.
The testos_multiarch_* targets are tagged "manual" such that test //...
does not build it for all variants.

Change-Id: I2682beb1adf61de0e86c53371c63c4fd9afecf08
Reviewed-on: https://review.monogon.dev/c/monogon/+/4478
Tested-by: Jenkins CI
Reviewed-by: Tim Windelschmidt <tim@monogon.tech>
diff --git a/metropolis/node/core/update/BUILD.bazel b/metropolis/node/core/update/BUILD.bazel
index ed98d21..ff2faea 100644
--- a/metropolis/node/core/update/BUILD.bazel
+++ b/metropolis/node/core/update/BUILD.bazel
@@ -22,6 +22,7 @@
         "//osbase/oci/osimage",
         "//osbase/oci/registry",
         "@com_github_cenkalti_backoff_v4//:backoff",
+        "@com_github_opencontainers_image_spec//specs-go/v1:specs-go",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//status",
         "@org_golang_google_protobuf//proto",
diff --git a/metropolis/node/core/update/e2e/BUILD.bazel b/metropolis/node/core/update/e2e/BUILD.bazel
index 0bc7652..51fc851 100644
--- a/metropolis/node/core/update/e2e/BUILD.bazel
+++ b/metropolis/node/core/update/e2e/BUILD.bazel
@@ -12,13 +12,13 @@
         "//metropolis/node/abloader",
         # For the two update tests
         "//metropolis/node/core/update/e2e/testos:testos_image_y",
-        "//metropolis/node/core/update/e2e/testos:testos_image_z",
+        "//metropolis/node/core/update/e2e/testos:testos_multiarch_z",
         "//build/toolchain/toolchain-bundle:qemu-kvm",
     ],
     x_defs = {
         "xImageXPath": "$(rlocationpath //metropolis/node/core/update/e2e/testos:testos_image_x )",
         "xImageYPath": "$(rlocationpath //metropolis/node/core/update/e2e/testos:testos_image_y )",
-        "xImageZPath": "$(rlocationpath //metropolis/node/core/update/e2e/testos:testos_image_z )",
+        "xImageZPath": "$(rlocationpath //metropolis/node/core/update/e2e/testos:testos_multiarch_z )",
         "xOvmfVarsPath": "$(rlocationpath //third_party/edk2:VARS.fd )",
         "xOvmfCodePath": "$(rlocationpath //third_party/edk2:CODE.fd )",
         "xAbloaderPath": "$(rlocationpath //metropolis/node/abloader )",
diff --git a/metropolis/node/core/update/e2e/testos/testos.bzl b/metropolis/node/core/update/e2e/testos/testos.bzl
index ab7d2f0..2057c95 100644
--- a/metropolis/node/core/update/e2e/testos/testos.bzl
+++ b/metropolis/node/core/update/e2e/testos/testos.bzl
@@ -2,6 +2,7 @@
 load("//osbase/build/genproductinfo:test.bzl", "test_product_info")
 load("//osbase/build/mkerofs:def.bzl", "erofs_image")
 load("//osbase/build/mkoci:def.bzl", "oci_os_image")
+load("//osbase/build/mkoci/index:def.bzl", "oci_index")
 load("//osbase/build/mkpayload:def.bzl", "efi_unified_kernel_image")
 load("//osbase/build/mkverity:def.bzl", "verity_image")
 
@@ -49,6 +50,17 @@
         visibility = ["//metropolis/node/core/update/e2e:__pkg__"],
     )
 
+    oci_index(
+        name = "testos_multiarch_" + variant,
+        src = ":testos_image_" + variant,
+        platforms = [
+            "//build/platforms:linux_x86_64",
+            "//build/platforms:linux_aarch64",
+        ],
+        tags = ["manual"],
+        visibility = ["//metropolis/node/core/update/e2e:__pkg__"],
+    )
+
     go_binary(
         name = "testos_" + variant,
         embed = [":testos_lib"],
diff --git a/metropolis/node/core/update/update.go b/metropolis/node/core/update/update.go
index 047e9f8..35f2240 100644
--- a/metropolis/node/core/update/update.go
+++ b/metropolis/node/core/update/update.go
@@ -15,11 +15,13 @@
 	"os"
 	"path/filepath"
 	"regexp"
+	"runtime"
 	"strconv"
 	"strings"
 	"time"
 
 	"github.com/cenkalti/backoff/v4"
+	ocispecv1 "github.com/opencontainers/image-spec/specs-go/v1"
 	"golang.org/x/sys/unix"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -260,6 +262,30 @@
 	return nil
 }
 
+func selectArchitecture(ref oci.Ref, architecture string) (*oci.Image, error) {
+	switch ref := ref.(type) {
+	case *oci.Image:
+		return ref, nil
+	case *oci.Index:
+		var found *ocispecv1.Descriptor
+		for i := range ref.Manifest.Manifests {
+			descriptor := &ref.Manifest.Manifests[i]
+			if descriptor.Platform != nil && descriptor.Platform.Architecture == architecture {
+				if found != nil {
+					return nil, fmt.Errorf("invalid index, found multiple matching entries")
+				}
+				found = descriptor
+			}
+		}
+		if found == nil {
+			return nil, fmt.Errorf("no matching entry found in index for architecture %s", architecture)
+		}
+		return oci.AsImage(ref.Ref(found))
+	default:
+		return nil, fmt.Errorf("unknown manifest media type %q", ref.MediaType())
+	}
+}
+
 // InstallImage fetches the given image, installs it into the currently inactive
 // slot and sets that slot to boot next. If it doesn't return an error, a reboot
 // boots into the new slot.
@@ -290,7 +316,11 @@
 		Repository: imageRef.Repository,
 	}
 
-	image, err := oci.AsImage(client.Read(downloadCtx, imageRef.Tag, imageRef.Digest))
+	ref, err := client.Read(downloadCtx, imageRef.Tag, imageRef.Digest)
+	if err != nil {
+		return fmt.Errorf("failed to fetch OS image: %w", err)
+	}
+	image, err := selectArchitecture(ref, runtime.GOARCH)
 	if err != nil {
 		return fmt.Errorf("failed to fetch OS image: %w", err)
 	}