osbase/oci: implement support for OCI index

Previously, only OCI images were supported, now we can also handle
indexes. The new Ref type is either an Image or Index.

Change-Id: I1b282ed6078d53e9a69e7613f601fdbbe64e192b
Reviewed-on: https://review.monogon.dev/c/monogon/+/4475
Tested-by: Jenkins CI
Reviewed-by: Tim Windelschmidt <tim@monogon.tech>
diff --git a/cloud/agent/BUILD.bazel b/cloud/agent/BUILD.bazel
index a25ea38..765d63c 100644
--- a/cloud/agent/BUILD.bazel
+++ b/cloud/agent/BUILD.bazel
@@ -22,6 +22,7 @@
         "//osbase/efivarfs",
         "//osbase/net/proto",
         "//osbase/nvme",
+        "//osbase/oci",
         "//osbase/oci/osimage",
         "//osbase/oci/registry",
         "//osbase/pki",
diff --git a/cloud/agent/e2e/main_test.go b/cloud/agent/e2e/main_test.go
index 53b719b..9dfd129 100644
--- a/cloud/agent/e2e/main_test.go
+++ b/cloud/agent/e2e/main_test.go
@@ -115,7 +115,7 @@
 				Host:       registryAddr.String(),
 				Repository: "testos",
 				Tag:        "latest",
-				Digest:     image.ManifestDigest,
+				Digest:     image.Digest(),
 			},
 			NodeParameters: &mpb.NodeParameters{},
 			RootDevice:     "vda",
@@ -179,7 +179,7 @@
 	grpcListenAddr := grpcLis.Addr().(*net.TCPAddr)
 
 	registryServer := registry.NewServer()
-	registryServer.AddImage("testos", "latest", image)
+	registryServer.AddRef("testos", "latest", image)
 
 	registryLis, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
diff --git a/cloud/agent/install.go b/cloud/agent/install.go
index 9ada49b..9144769 100644
--- a/cloud/agent/install.go
+++ b/cloud/agent/install.go
@@ -21,6 +21,7 @@
 	metropolisInstall "source.monogon.dev/metropolis/installer/install"
 	"source.monogon.dev/osbase/blockdev"
 	"source.monogon.dev/osbase/efivarfs"
+	"source.monogon.dev/osbase/oci"
 	"source.monogon.dev/osbase/oci/osimage"
 	"source.monogon.dev/osbase/oci/registry"
 	"source.monogon.dev/osbase/structfs"
@@ -75,7 +76,7 @@
 		Repository: req.OsImage.Repository,
 	}
 
-	image, err := client.Read(ctx, req.OsImage.Tag, req.OsImage.Digest)
+	image, err := oci.AsImage(client.Read(ctx, req.OsImage.Tag, req.OsImage.Digest))
 	if err != nil {
 		return fmt.Errorf("failed to fetch OS image: %w", err)
 	}
diff --git a/metropolis/cli/metroctl/cmd_install_usb.go b/metropolis/cli/metroctl/cmd_install_usb.go
index 4453ed7..5af158a 100644
--- a/metropolis/cli/metroctl/cmd_install_usb.go
+++ b/metropolis/cli/metroctl/cmd_install_usb.go
@@ -38,7 +38,7 @@
 		if err != nil {
 			return err
 		}
-		image, err := oci.ReadLayout(imagePathResolved)
+		image, err := oci.AsImage(oci.ReadLayout(imagePathResolved))
 		if err != nil {
 			return fmt.Errorf("failed to read OS image: %w", err)
 		}
diff --git a/metropolis/cli/takeover/install.go b/metropolis/cli/takeover/install.go
index 97e8565..9755c55 100644
--- a/metropolis/cli/takeover/install.go
+++ b/metropolis/cli/takeover/install.go
@@ -37,7 +37,7 @@
 		return err
 	}
 
-	image, err := oci.ReadLayout("/osimage")
+	image, err := oci.AsImage(oci.ReadLayout("/osimage"))
 	if err != nil {
 		return fmt.Errorf("failed to read OS image: %w", err)
 	}
diff --git a/metropolis/cli/takeover/takeover.go b/metropolis/cli/takeover/takeover.go
index 805a06d..5ff26bd 100644
--- a/metropolis/cli/takeover/takeover.go
+++ b/metropolis/cli/takeover/takeover.go
@@ -94,7 +94,7 @@
 		return nil, err
 	}
 
-	image, err := oci.ReadLayout(filepath.Join(filepath.Dir(currPath), "osimage"))
+	image, err := oci.AsImage(oci.ReadLayout(filepath.Join(filepath.Dir(currPath), "osimage")))
 	if err != nil {
 		return nil, fmt.Errorf("failed to read OS image: %w", err)
 	}
diff --git a/metropolis/installer/main.go b/metropolis/installer/main.go
index 735a8cd..ecb4d01 100644
--- a/metropolis/installer/main.go
+++ b/metropolis/installer/main.go
@@ -161,7 +161,7 @@
 		return fmt.Errorf("failed to open node parameters from ESP: %w", err)
 	}
 
-	ociImage, err := oci.ReadLayout("/installer/metropolis-installer/osimage")
+	ociImage, err := oci.AsImage(oci.ReadLayout("/installer/metropolis-installer/osimage"))
 	if err != nil {
 		return fmt.Errorf("failed to read OS image from ESP: %w", err)
 	}
diff --git a/metropolis/installer/test/run_test.go b/metropolis/installer/test/run_test.go
index daff8cd..898ad44 100644
--- a/metropolis/installer/test/run_test.go
+++ b/metropolis/installer/test/run_test.go
@@ -156,7 +156,7 @@
 		log.Fatal(err)
 	}
 
-	image, err := oci.ReadLayout(xImagePath)
+	image, err := oci.AsImage(oci.ReadLayout(xImagePath))
 	if err != nil {
 		log.Fatal(err)
 	}
diff --git a/metropolis/node/core/update/BUILD.bazel b/metropolis/node/core/update/BUILD.bazel
index 73f6773..ed98d21 100644
--- a/metropolis/node/core/update/BUILD.bazel
+++ b/metropolis/node/core/update/BUILD.bazel
@@ -18,6 +18,7 @@
         "//osbase/efivarfs",
         "//osbase/gpt",
         "//osbase/kexec",
+        "//osbase/oci",
         "//osbase/oci/osimage",
         "//osbase/oci/registry",
         "@com_github_cenkalti_backoff_v4//:backoff",
diff --git a/metropolis/node/core/update/e2e/e2e_test.go b/metropolis/node/core/update/e2e/e2e_test.go
index 890a06f..041de3f 100644
--- a/metropolis/node/core/update/e2e/e2e_test.go
+++ b/metropolis/node/core/update/e2e/e2e_test.go
@@ -141,7 +141,7 @@
 		Port: 80,
 	}
 
-	imageX, err := oci.ReadLayout(xImageXPath)
+	imageX, err := oci.AsImage(oci.ReadLayout(xImageXPath))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -160,8 +160,8 @@
 	}
 
 	registryServer := registry.NewServer()
-	registryServer.AddImage("testos", "y", imageY)
-	registryServer.AddImage("testos", "z", imageZ)
+	registryServer.AddRef("testos", "y", imageY)
+	registryServer.AddRef("testos", "z", imageZ)
 	registryLis, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
 		t.Fatal(err)
@@ -225,8 +225,8 @@
 		"-device", "virtio-rng-pci",
 		"-serial", "stdio",
 		"-no-reboot",
-		"-fw_cfg", "name=opt/testos_y_digest,string=" + imageY.ManifestDigest,
-		"-fw_cfg", "name=opt/testos_z_digest,string=" + imageZ.ManifestDigest,
+		"-fw_cfg", "name=opt/testos_y_digest,string=" + imageY.Digest(),
+		"-fw_cfg", "name=opt/testos_z_digest,string=" + imageZ.Digest(),
 	}
 	return qemuArgs
 }
diff --git a/metropolis/node/core/update/update.go b/metropolis/node/core/update/update.go
index 66e7501..047e9f8 100644
--- a/metropolis/node/core/update/update.go
+++ b/metropolis/node/core/update/update.go
@@ -32,6 +32,7 @@
 	"source.monogon.dev/osbase/efivarfs"
 	"source.monogon.dev/osbase/gpt"
 	"source.monogon.dev/osbase/kexec"
+	"source.monogon.dev/osbase/oci"
 	"source.monogon.dev/osbase/oci/osimage"
 	"source.monogon.dev/osbase/oci/registry"
 
@@ -289,7 +290,7 @@
 		Repository: imageRef.Repository,
 	}
 
-	image, err := client.Read(downloadCtx, imageRef.Tag, imageRef.Digest)
+	image, err := oci.AsImage(client.Read(downloadCtx, imageRef.Tag, imageRef.Digest))
 	if err != nil {
 		return fmt.Errorf("failed to fetch OS image: %w", err)
 	}
diff --git a/metropolis/test/launch/cluster.go b/metropolis/test/launch/cluster.go
index 7f64427..c77a101 100644
--- a/metropolis/test/launch/cluster.go
+++ b/metropolis/test/launch/cluster.go
@@ -166,7 +166,7 @@
 	}
 
 	// Initialize the node's storage.
-	ociImage, err := oci.ReadLayout(xNodeImagePath)
+	ociImage, err := oci.AsImage(oci.ReadLayout(xNodeImagePath))
 	if err != nil {
 		return nil, fmt.Errorf("failed to read OS image: %w", err)
 	}
diff --git a/metropolis/test/localregistry/localregistry.go b/metropolis/test/localregistry/localregistry.go
index 694476e..7440aae 100644
--- a/metropolis/test/localregistry/localregistry.go
+++ b/metropolis/test/localregistry/localregistry.go
@@ -32,7 +32,7 @@
 		if err != nil {
 			return nil, fmt.Errorf("failed to read image from %q: %w", i.Path, err)
 		}
-		err = s.AddImage(i.Repository, i.Tag, image)
+		err = s.AddRef(i.Repository, i.Tag, image)
 		if err != nil {
 			return nil, fmt.Errorf("failed to add image: %w", err)
 		}
diff --git a/osbase/oci/layout.go b/osbase/oci/layout.go
index 128c4d1..edea5b0 100644
--- a/osbase/oci/layout.go
+++ b/osbase/oci/layout.go
@@ -8,7 +8,6 @@
 	"fmt"
 	"io"
 	"os"
-	"path"
 	"path/filepath"
 
 	"github.com/opencontainers/go-digest"
@@ -18,8 +17,8 @@
 	"source.monogon.dev/osbase/structfs"
 )
 
-// ReadLayout reads an image from an OS path to an OCI layout directory.
-func ReadLayout(path string) (*Image, error) {
+// ReadLayoutIndex reads the index from an OS path to an OCI layout directory.
+func ReadLayoutIndex(path string) (*Index, error) {
 	// Read the oci-layout marker file.
 	layoutBytes, err := os.ReadFile(filepath.Join(path, "oci-layout"))
 	if err != nil {
@@ -35,41 +34,33 @@
 	}
 
 	// Read the index.
-	imageIndexBytes, err := os.ReadFile(filepath.Join(path, "index.json"))
+	indexBytes, err := os.ReadFile(filepath.Join(path, "index.json"))
 	if err != nil {
 		return nil, err
 	}
-	imageIndex := ocispecv1.Index{}
-	err = json.Unmarshal(imageIndexBytes, &imageIndex)
+	blobs := &layoutBlobs{path: path}
+	ref, err := NewRef(indexBytes, ocispecv1.MediaTypeImageIndex, "", blobs)
 	if err != nil {
-		return nil, fmt.Errorf("failed to parse index.json: %w", err)
+		return nil, err
 	}
-	if imageIndex.MediaType != ocispecv1.MediaTypeImageIndex {
-		return nil, fmt.Errorf("unknown index.json mediaType %q", imageIndex.MediaType)
+	return ref.(*Index), nil
+}
+
+// ReadLayout reads a manifest from an OS path to an OCI layout directory.
+// It expects the index to point to exactly one manifest, which is common.
+func ReadLayout(path string) (Ref, error) {
+	index, err := ReadLayoutIndex(path)
+	if err != nil {
+		return nil, err
 	}
-	if len(imageIndex.Manifests) == 0 {
+
+	if len(index.Manifest.Manifests) == 0 {
 		return nil, fmt.Errorf("index.json contains no manifests")
 	}
-	if len(imageIndex.Manifests) != 1 {
+	if len(index.Manifest.Manifests) != 1 {
 		return nil, fmt.Errorf("index.json files containing multiple manifests are not supported")
 	}
-	manifestDescriptor := &imageIndex.Manifests[0]
-	if manifestDescriptor.MediaType != ocispecv1.MediaTypeImageManifest {
-		return nil, fmt.Errorf("unexpected manifest media type %q", manifestDescriptor.MediaType)
-	}
-
-	// Read the image manifest.
-	imageManifestPath, err := layoutBlobPath(path, manifestDescriptor)
-	if err != nil {
-		return nil, err
-	}
-	imageManifestBytes, err := os.ReadFile(imageManifestPath)
-	if err != nil {
-		return nil, fmt.Errorf("failed to read image manifest: %w", err)
-	}
-
-	blobs := &layoutBlobs{path: path}
-	return NewImage(imageManifestBytes, string(manifestDescriptor.Digest), blobs)
+	return index.Ref(&index.Manifest.Manifests[0])
 }
 
 type layoutBlobs struct {
@@ -77,27 +68,43 @@
 }
 
 func (r *layoutBlobs) Blob(descriptor *ocispecv1.Descriptor) (io.ReadCloser, error) {
-	blobPath, err := layoutBlobPath(r.path, descriptor)
+	algorithm, encoded, err := ParseDigest(string(descriptor.Digest))
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse digest in descriptor: %w", err)
+	}
+	return os.Open(filepath.Join(r.path, "blobs", algorithm, encoded))
+}
+
+func (r *layoutBlobs) Manifest(descriptor *ocispecv1.Descriptor) ([]byte, error) {
+	blob, err := r.Blob(descriptor)
 	if err != nil {
 		return nil, err
 	}
-	return os.Open(blobPath)
-}
-
-func layoutBlobPath(layoutPath string, descriptor *ocispecv1.Descriptor) (string, error) {
-	algorithm, encoded, err := ParseDigest(string(descriptor.Digest))
+	defer blob.Close()
+	manifestBytes := make([]byte, descriptor.Size)
+	_, err = io.ReadFull(blob, manifestBytes)
 	if err != nil {
-		return "", fmt.Errorf("failed to parse digest in image manifest: %w", err)
+		return nil, fmt.Errorf("failed to read manifest: %w", err)
 	}
-	return filepath.Join(layoutPath, "blobs", algorithm, encoded), nil
+	return manifestBytes, nil
 }
 
-// CreateLayout builds an OCI layout from an Image.
-func CreateLayout(image *Image) (structfs.Tree, error) {
+func (r *layoutBlobs) Blobs(_ *ocispecv1.Descriptor) (Blobs, error) {
+	return r, nil
+}
+
+// CreateLayout builds an OCI layout from a Ref.
+func CreateLayout(ref Ref) (structfs.Tree, error) {
 	// Build the index.
-	artifactType := image.Manifest.Config.MediaType
-	if artifactType == ocispecv1.MediaTypeImageConfig {
-		artifactType = ""
+	artifactType := ""
+	if image, ok := ref.(*Image); ok {
+		// According to the OCI spec, the artifactType is the config descriptor
+		// mediaType, and is only set when the descriptor references the image
+		// manifest of an artifact.
+		artifactType = image.Manifest.Config.MediaType
+		if artifactType == ocispecv1.MediaTypeImageConfig {
+			artifactType = ""
+		}
 	}
 	imageIndex := ocispecv1.Index{
 		Versioned: ocispec.Versioned{SchemaVersion: 2},
@@ -105,8 +112,8 @@
 		Manifests: []ocispecv1.Descriptor{{
 			MediaType:    ocispecv1.MediaTypeImageManifest,
 			ArtifactType: artifactType,
-			Digest:       digest.Digest(image.ManifestDigest),
-			Size:         int64(len(image.RawManifest)),
+			Digest:       digest.Digest(ref.Digest()),
+			Size:         int64(len(ref.RawManifest())),
 		}},
 	}
 	imageIndexBytes, err := json.MarshalIndent(imageIndex, "", "\t")
@@ -120,33 +127,50 @@
 		structfs.File("index.json", structfs.Bytes(imageIndexBytes)),
 	}
 
-	algorithm, encoded, err := ParseDigest(image.ManifestDigest)
-	if err != nil {
-		return nil, fmt.Errorf("failed to parse manifest digest: %w", err)
+	hasBlob := make(map[string]bool)
+	blobDirs := make(map[string]*structfs.Node)
+	addBlob := func(digest string, blob structfs.Blob) error {
+		if hasBlob[digest] {
+			// If multiple blobs have the same digest, we only need the first one.
+			return nil
+		}
+		hasBlob[digest] = true
+		algorithm, encoded, err := ParseDigest(digest)
+		if err != nil {
+			return fmt.Errorf("failed to parse manifest digest: %w", err)
+		}
+		blobDir, ok := blobDirs[algorithm]
+		if !ok {
+			blobDir = structfs.Dir(algorithm, nil)
+			err = root.Place("blobs", blobDir)
+			if err != nil {
+				return err
+			}
+			blobDirs[algorithm] = blobDir
+		}
+		// root.PlaceFile is not used here because then running time would be
+		// quadratic in the number of blobs.
+		blobDir.Children = append(blobDir.Children, structfs.File(encoded, blob))
+		return nil
 	}
-	imageManifestPath := path.Join("blobs", algorithm, encoded)
-	err = root.PlaceFile(imageManifestPath, structfs.Bytes(image.RawManifest))
+	err = WalkRefs(string(imageIndex.Manifests[0].Digest), ref, func(digest string, ref Ref) error {
+		err := addBlob(digest, structfs.Bytes(ref.RawManifest()))
+		if err != nil {
+			return err
+		}
+		if image, ok := ref.(*Image); ok {
+			for descriptor := range image.Descriptors() {
+				err := addBlob(string(descriptor.Digest), image.StructfsBlob(descriptor))
+				if err != nil {
+					return err
+				}
+			}
+		}
+		return nil
+	})
 	if err != nil {
 		return nil, err
 	}
 
-	hasBlob := map[string]bool{}
-	for descriptor := range image.Descriptors() {
-		algorithm, encoded, err := ParseDigest(string(descriptor.Digest))
-		if err != nil {
-			return nil, fmt.Errorf("failed to parse digest in image manifest: %w", err)
-		}
-		blobPath := path.Join("blobs", algorithm, encoded)
-		if hasBlob[blobPath] {
-			// If multiple blobs have the same hash, we only need the first one.
-			continue
-		}
-		hasBlob[blobPath] = true
-		err = root.PlaceFile(blobPath, image.StructfsBlob(descriptor))
-		if err != nil {
-			return nil, err
-		}
-	}
-
 	return root, nil
 }
diff --git a/osbase/oci/oci.go b/osbase/oci/oci.go
index a62b527..a652ca4 100644
--- a/osbase/oci/oci.go
+++ b/osbase/oci/oci.go
@@ -17,55 +17,210 @@
 	"source.monogon.dev/osbase/structfs"
 )
 
+// Index represents an OCI image index.
+type Index struct {
+	// Manifest contains the parsed index manifest.
+	Manifest    *ocispecv1.Index
+	rawManifest []byte
+	digest      string
+	blobs       Blobs
+}
+
 // Image represents an OCI image.
 type Image struct {
 	// Manifest contains the parsed image manifest.
-	Manifest *ocispecv1.Manifest
-	// RawManifest contains the bytes of the image manifest.
-	RawManifest []byte
-	// ManifestDigest contains the computed digest of RawManifest.
-	ManifestDigest string
-
-	blobs Blobs
+	Manifest    *ocispecv1.Manifest
+	rawManifest []byte
+	digest      string
+	blobs       Blobs
 }
 
-// Blobs is the interface which image sources implement to retrieve the content
-// of blobs.
+// Ref is either an [*Index] or [*Image].
+type Ref interface {
+	// RawManifest returns the bytes of the manifest.
+	// The returned value is shared and must not be modified.
+	RawManifest() []byte
+	// Digest returns the computed digest of RawManifest, in the default digest
+	// algorithm. Only sha256 is supported currently.
+	Digest() string
+	// MediaType returns the media type of the manifest.
+	MediaType() string
+	// isRef is an unexported marker to disallow implementations of the interface
+	// outside this package.
+	isRef()
+}
+
+func (i *Index) RawManifest() []byte { return i.rawManifest }
+func (i *Index) Digest() string      { return i.digest }
+func (i *Index) MediaType() string   { return ocispecv1.MediaTypeImageIndex }
+func (i *Index) isRef()              {}
+
+func (i *Image) RawManifest() []byte { return i.rawManifest }
+func (i *Image) Digest() string      { return i.digest }
+func (i *Image) MediaType() string   { return ocispecv1.MediaTypeImageManifest }
+func (i *Image) isRef()              {}
+
+// Blobs is the interface which image sources implement
+// to retrieve the content of blobs and manifests.
 type Blobs interface {
 	// Blob returns the contents of a blob from its descriptor.
 	// It does not verify the contents against the digest.
+	//
+	// This is only called on images.
 	Blob(*ocispecv1.Descriptor) (io.ReadCloser, error)
+	// Manifest returns the contents of a manifest from its descriptor.
+	// It does not verify the contents against the digest.
+	//
+	// This is only called on indexes.
+	Manifest(*ocispecv1.Descriptor) ([]byte, error)
+	// Blobs returns the [Blobs] for the manifest from its descriptor.
+	// Most implementations simply return the receiver itself, but this
+	// allows combining Refs from different sources into an Index.
+	//
+	// This is only called on indexes.
+	Blobs(*ocispecv1.Descriptor) (Blobs, error)
 }
 
-// NewImage verifies the manifest against the expected digest if not empty,
-// then parses it and returns an [Image].
-func NewImage(rawManifest []byte, expectedDigest string, blobs Blobs) (*Image, error) {
+// NewRef verifies the manifest against the expected digest if not empty,
+// then parses it according to mediaType and returns a [Ref].
+func NewRef(rawManifest []byte, mediaType string, expectedDigest string, blobs Blobs) (Ref, error) {
 	digest := fmt.Sprintf("sha256:%x", sha256.Sum256(rawManifest))
 	if expectedDigest != "" && expectedDigest != digest {
+		if _, _, err := ParseDigest(expectedDigest); err != nil {
+			return nil, err
+		}
 		return nil, fmt.Errorf("failed verification of manifest: expected digest %q, computed %q", expectedDigest, digest)
 	}
 
-	manifest := &ocispecv1.Manifest{}
-	err := json.Unmarshal(rawManifest, &manifest)
-	if err != nil {
-		return nil, fmt.Errorf("failed to parse image manifest: %w", err)
-	}
-	if manifest.MediaType != ocispecv1.MediaTypeImageManifest {
-		return nil, fmt.Errorf("unexpected manifest media type %q", manifest.MediaType)
-	}
-	image := &Image{
-		Manifest:       manifest,
-		RawManifest:    rawManifest,
-		ManifestDigest: digest,
-		blobs:          blobs,
-	}
-	for descriptor := range image.Descriptors() {
-		if descriptor.Size < 0 {
-			return nil, fmt.Errorf("invalid manifest: contains descriptor with negative size")
+	switch mediaType {
+	case ocispecv1.MediaTypeImageManifest:
+		manifest := &ocispecv1.Manifest{}
+		if err := json.Unmarshal(rawManifest, manifest); err != nil {
+			return nil, fmt.Errorf("failed to parse image manifest: %w", err)
 		}
+		if manifest.MediaType != ocispecv1.MediaTypeImageManifest {
+			return nil, fmt.Errorf("unexpected manifest media type %q, expected %q", manifest.MediaType, ocispecv1.MediaTypeImageManifest)
+		}
+		image := &Image{
+			Manifest:    manifest,
+			rawManifest: rawManifest,
+			digest:      digest,
+			blobs:       blobs,
+		}
+		for descriptor := range image.Descriptors() {
+			// We validate this here such that StructfsBlob does not need an error return.
+			if descriptor.Size < 0 {
+				return nil, fmt.Errorf("invalid manifest: contains descriptor with negative size")
+			}
+		}
+		return image, nil
+	case ocispecv1.MediaTypeImageIndex:
+		manifest := &ocispecv1.Index{}
+		if err := json.Unmarshal(rawManifest, manifest); err != nil {
+			return nil, fmt.Errorf("failed to parse index manifest: %w", err)
+		}
+		if manifest.MediaType != ocispecv1.MediaTypeImageIndex {
+			return nil, fmt.Errorf("unexpected manifest media type %q, expected %q", manifest.MediaType, ocispecv1.MediaTypeImageIndex)
+		}
+		index := &Index{
+			Manifest:    manifest,
+			rawManifest: rawManifest,
+			digest:      digest,
+			blobs:       blobs,
+		}
+		return index, nil
+	default:
+		return nil, fmt.Errorf("unknown manifest media type %q", mediaType)
+	}
+}
+
+// AsImage can be conveniently wrapped around a call which returns a [Ref] or
+// error, when only [*Image] can be handled.
+func AsImage(ref Ref, err error) (*Image, error) {
+	if err != nil {
+		return nil, err
+	}
+	image, ok := ref.(*Image)
+	if !ok {
+		return nil, fmt.Errorf("unexpected manifest media type %q, only image is supported", ref.MediaType())
+	}
+	return image, nil
+}
+
+// WalkRefs iterates over all Refs reachable from ref in DFS post-order.
+// Each digest is only visited once, even if reachable multiple times.
+//
+// For each Ref, we also pass the digest by which it is referenced. This may be
+// different from ref.Digest() if we ever support multiple digest algorithms.
+func WalkRefs(digest string, ref Ref, fn func(digest string, ref Ref) error) error {
+	visited := make(map[string]bool)
+	return walkRefs(digest, ref, fn, visited)
+}
+
+func walkRefs(digest string, ref Ref, fn func(digest string, ref Ref) error, visited map[string]bool) error {
+	if visited[digest] {
+		return nil
+	}
+	visited[digest] = true
+	switch ref := ref.(type) {
+	case *Image:
+	case *Index:
+		for i := range ref.Manifest.Manifests {
+			descriptor := &ref.Manifest.Manifests[i]
+			childRef, err := ref.Ref(descriptor)
+			if err != nil {
+				return err
+			}
+			err = walkRefs(string(descriptor.Digest), childRef, fn, visited)
+			if err != nil {
+				return err
+			}
+		}
+	default:
+		return fmt.Errorf("unknown manifest media type %q", ref.MediaType())
+	}
+	return fn(digest, ref)
+}
+
+// Ref reads a manifest from its descriptor and wraps it in a [Ref].
+// The manifest is verified against the digest.
+func (i *Index) Ref(descriptor *ocispecv1.Descriptor) (Ref, error) {
+	if descriptor.Size < 0 {
+		return nil, fmt.Errorf("invalid descriptor size %d", descriptor.Size)
+	}
+	if descriptor.Size > 50*1024*1024 {
+		return nil, fmt.Errorf("refusing to read manifest of size %d into memory", descriptor.Size)
+	}
+	switch descriptor.MediaType {
+	case ocispecv1.MediaTypeImageManifest:
+	case ocispecv1.MediaTypeImageIndex:
+	default:
+		return nil, fmt.Errorf("unknown manifest media type %q", descriptor.MediaType)
+	}
+	if descriptor.Digest == "" { // NewRef treats empty digest as unknown.
+		return nil, fmt.Errorf("invalid digest")
 	}
 
-	return image, nil
+	var rawManifest []byte
+	if int64(len(descriptor.Data)) == descriptor.Size {
+		rawManifest = descriptor.Data
+	} else if len(descriptor.Data) != 0 {
+		return nil, fmt.Errorf("descriptor has embedded data of wrong length")
+	} else {
+		var err error
+		rawManifest, err = i.blobs.Manifest(descriptor)
+		if err != nil {
+			return nil, err
+		}
+	}
+	if int64(len(rawManifest)) != descriptor.Size {
+		return nil, fmt.Errorf("manifest has wrong length, expected %d, got %d bytes", descriptor.Size, len(rawManifest))
+	}
+	blobs, err := i.blobs.Blobs(descriptor)
+	if err != nil {
+		return nil, err
+	}
+	return NewRef(rawManifest, descriptor.MediaType, string(descriptor.Digest), blobs)
 }
 
 // Descriptors returns an iterator over all descriptors in the image (config and
@@ -86,6 +241,9 @@
 // Blob returns the contents of a blob from its descriptor.
 // It does not verify the contents against the digest.
 func (i *Image) Blob(descriptor *ocispecv1.Descriptor) (io.ReadCloser, error) {
+	if descriptor.Size < 0 {
+		return nil, fmt.Errorf("invalid descriptor size %d", descriptor.Size)
+	}
 	if int64(len(descriptor.Data)) == descriptor.Size {
 		return structfs.Bytes(descriptor.Data).Open()
 	} else if len(descriptor.Data) != 0 {
@@ -97,9 +255,6 @@
 // ReadBlobVerified reads a blob into a byte slice and verifies it against the
 // digest.
 func (i *Image) ReadBlobVerified(descriptor *ocispecv1.Descriptor) ([]byte, error) {
-	if descriptor.Size < 0 {
-		return nil, fmt.Errorf("invalid descriptor size %d", descriptor.Size)
-	}
 	if descriptor.Size > 50*1024*1024 {
 		return nil, fmt.Errorf("refusing to read blob of size %d into memory", descriptor.Size)
 	}
diff --git a/osbase/oci/oci_test.go b/osbase/oci/oci_test.go
index 93976ed..573771b 100644
--- a/osbase/oci/oci_test.go
+++ b/osbase/oci/oci_test.go
@@ -33,7 +33,7 @@
 }`
 	// Pass nil for blobs, which means reading can only work if it uses the
 	// embedded content.
-	image, err := NewImage([]byte(manifest), "", nil)
+	image, err := AsImage(NewRef([]byte(manifest), "application/vnd.oci.image.manifest.v1+json", "", nil))
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/osbase/oci/osimage/osimage_test.go b/osbase/oci/osimage/osimage_test.go
index 6ee8afa..26a094a 100644
--- a/osbase/oci/osimage/osimage_test.go
+++ b/osbase/oci/osimage/osimage_test.go
@@ -67,7 +67,7 @@
 	}
 	for _, tC := range testCases {
 		t.Run(tC.desc, func(t *testing.T) {
-			image, err := oci.ReadLayout(tC.path)
+			image, err := oci.AsImage(oci.ReadLayout(tC.path))
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -103,11 +103,11 @@
 
 func TestVerification(t *testing.T) {
 	server := registry.NewServer()
-	srcImage, err := oci.ReadLayout(xImageUncompressedPath)
+	srcImage, err := oci.AsImage(oci.ReadLayout(xImageUncompressedPath))
 	if err != nil {
 		t.Fatal(err)
 	}
-	server.AddImage("test/repo", "test-tag", srcImage)
+	server.AddRef("test/repo", "test-tag", srcImage)
 	corrupter := &corruptingServer{handler: server}
 
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
@@ -135,14 +135,14 @@
 	if err != nil {
 		t.Errorf("Expected reading manifest to succeed when digest not given: %v", err)
 	}
-	_, err = client.Read(context.Background(), "test-tag", srcImage.ManifestDigest)
+	_, err = client.Read(context.Background(), "test-tag", srcImage.Digest())
 	if !strings.Contains(fmt.Sprintf("%v", err), "failed verification") {
 		t.Errorf("Expected failed verification, got %v", err)
 	}
 
 	// Test config verification
 	corrupter.affectedPath = fmt.Sprintf("/v2/test/repo/blobs/%s", srcImage.Manifest.Config.Digest)
-	image, err := client.Read(context.Background(), "test-tag", srcImage.ManifestDigest)
+	image, err := oci.AsImage(client.Read(context.Background(), "test-tag", srcImage.Digest()))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -153,7 +153,7 @@
 
 	// Test payload verification
 	corrupter.affectedPath = fmt.Sprintf("/v2/test/repo/blobs/%s", srcImage.Manifest.Layers[0].Digest)
-	image, err = client.Read(context.Background(), "test-tag", srcImage.ManifestDigest)
+	image, err = oci.AsImage(client.Read(context.Background(), "test-tag", srcImage.Digest()))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -214,7 +214,7 @@
 }
 
 func TestTruncation(t *testing.T) {
-	srcImage, err := oci.ReadLayout(xImageUncompressedPath)
+	srcImage, err := oci.AsImage(oci.ReadLayout(xImageUncompressedPath))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -222,7 +222,7 @@
 		image:  srcImage,
 		length: srcImage.Manifest.Config.Size,
 	}
-	truncatedImage, err := oci.NewImage(srcImage.RawManifest, "", blobs)
+	truncatedImage, err := oci.AsImage(oci.NewRef(srcImage.RawManifest(), ocispecv1.MediaTypeImageManifest, "", blobs))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -264,6 +264,14 @@
 	return reader, nil
 }
 
+func (b *truncatedBlobs) Manifest(_ *ocispecv1.Descriptor) ([]byte, error) {
+	return nil, fmt.Errorf("not implemented")
+}
+
+func (b *truncatedBlobs) Blobs(_ *ocispecv1.Descriptor) (oci.Blobs, error) {
+	return b, nil
+}
+
 type readCloser struct {
 	io.Reader
 	io.Closer
diff --git a/osbase/oci/registry/client.go b/osbase/oci/registry/client.go
index 4e60b7b..4f2b6d2 100644
--- a/osbase/oci/registry/client.go
+++ b/osbase/oci/registry/client.go
@@ -43,6 +43,12 @@
 	DigestRegexp     = regexp.MustCompile(`^` + digestExpr + `$`)
 )
 
+// unknownManifest can be used to parse the media type from a manifest of
+// unknown type.
+type unknownManifest struct {
+	MediaType string `json:"mediaType,omitempty"`
+}
+
 // Client is an OCI registry client.
 type Client struct {
 	// Transport will be used to make requests. For example, this allows
@@ -69,10 +75,10 @@
 	bearerToken string
 }
 
-// Read fetches an image manifest from the registry and returns an [oci.Image].
+// Read fetches a manifest from the registry and returns an [oci.Ref].
 //
-// The context is used for the manifest request and all blob requests made
-// through the Image.
+// The context is used for the manifest request and for all blob and manifest
+// requests made through the Ref.
 //
 // At least one of tag and digest must be set. If only tag is set, then you are
 // trusting the registry to return the right content. Otherwise, the digest is
@@ -80,7 +86,7 @@
 // used in the request, and the digest is used to verify the response. The
 // advantage of fetching by tag is that it allows a pull through cache to
 // display tags to a user inspecting the cache contents.
-func (c *Client) Read(ctx context.Context, tag, digest string) (*oci.Image, error) {
+func (c *Client) Read(ctx context.Context, tag, digest string) (oci.Ref, error) {
 	if !RepositoryRegexp.MatchString(c.Repository) {
 		return nil, fmt.Errorf("invalid repository %q", c.Repository)
 	}
@@ -102,13 +108,14 @@
 	}
 
 	manifestPath := fmt.Sprintf("/v2/%s/manifests/%s", c.Repository, reference)
-	var imageManifestBytes []byte
+	var manifestBytes []byte
+	var manifestMediaType string
 	err := c.retry(ctx, func() error {
 		req, err := c.newGet(manifestPath)
 		if err != nil {
 			return err
 		}
-		req.Header.Set("Accept", ocispecv1.MediaTypeImageManifest)
+		req.Header.Set("Accept", ocispecv1.MediaTypeImageManifest+","+ocispecv1.MediaTypeImageIndex)
 		resp, err := c.doGet(ctx, req)
 		if err != nil {
 			return err
@@ -117,18 +124,34 @@
 			return readClientError(resp, req)
 		}
 		defer resp.Body.Close()
-		imageManifestBytes, err = readFullBody(resp, 50*1024*1024)
+		manifestMediaType = resp.Header.Get("Content-Type")
+		manifestBytes, err = readFullBody(resp, 50*1024*1024)
 		return err
 	})
 	if err != nil {
 		return nil, err
 	}
 
+	// Remove any parameters from the Content-Type header.
+	manifestMediaType, _, _ = strings.Cut(manifestMediaType, ";")
+	switch manifestMediaType {
+	case ocispecv1.MediaTypeImageManifest, ocispecv1.MediaTypeImageIndex:
+		// The Content-Type header is valid, use it.
+	default:
+		// We need to parse the manifest to extract the media type, then parse it
+		// again for that media type.
+		var manifest unknownManifest
+		if err := json.Unmarshal(manifestBytes, &manifest); err != nil {
+			return nil, fmt.Errorf("failed to parse manifest: %w", err)
+		}
+		manifestMediaType = manifest.MediaType
+	}
+
 	blobs := &clientBlobs{
 		ctx:    ctx,
 		client: c,
 	}
-	return oci.NewImage(imageManifestBytes, digest, blobs)
+	return oci.NewRef(manifestBytes, manifestMediaType, digest, blobs)
 }
 
 type clientBlobs struct {
@@ -136,6 +159,41 @@
 	client *Client
 }
 
+func (r *clientBlobs) Manifest(descriptor *ocispecv1.Descriptor) ([]byte, error) {
+	digest := string(descriptor.Digest)
+	if _, _, err := oci.ParseDigest(digest); err != nil {
+		return nil, err
+	}
+
+	manifestPath := fmt.Sprintf("/v2/%s/manifests/%s", r.client.Repository, digest)
+	var manifestBytes []byte
+	err := r.client.retry(r.ctx, func() error {
+		req, err := r.client.newGet(manifestPath)
+		if err != nil {
+			return err
+		}
+		req.Header.Set("Accept", ocispecv1.MediaTypeImageManifest+","+ocispecv1.MediaTypeImageIndex)
+		resp, err := r.client.doGet(r.ctx, req)
+		if err != nil {
+			return err
+		}
+		if resp.StatusCode != http.StatusOK {
+			return readClientError(resp, req)
+		}
+		defer resp.Body.Close()
+		manifestBytes, err = readKnownSizeBody(resp, int(descriptor.Size))
+		return err
+	})
+	if err != nil {
+		return nil, err
+	}
+	return manifestBytes, nil
+}
+
+func (r *clientBlobs) Blobs(_ *ocispecv1.Descriptor) (oci.Blobs, error) {
+	return r, nil
+}
+
 func (r *clientBlobs) Blob(descriptor *ocispecv1.Descriptor) (io.ReadCloser, error) {
 	if !DigestRegexp.MatchString(string(descriptor.Digest)) {
 		return nil, fmt.Errorf("invalid blob digest %q", descriptor.Digest)
@@ -480,3 +538,15 @@
 		return nil, backoff.Permanent(fmt.Errorf("HTTP response of size %d exceeds limit of %d bytes", resp.ContentLength, limit))
 	}
 }
+
+func readKnownSizeBody(resp *http.Response, size int) ([]byte, error) {
+	if resp.ContentLength >= 0 && resp.ContentLength != int64(size) {
+		return nil, backoff.Permanent(fmt.Errorf("HTTP response has size %d, expected %d bytes", resp.ContentLength, size))
+	}
+	content := make([]byte, size)
+	_, err := io.ReadFull(resp.Body, content)
+	if err != nil {
+		return nil, err
+	}
+	return content, nil
+}
diff --git a/osbase/oci/registry/client_test.go b/osbase/oci/registry/client_test.go
index 328d2fa..622a3ab 100644
--- a/osbase/oci/registry/client_test.go
+++ b/osbase/oci/registry/client_test.go
@@ -39,12 +39,12 @@
 }
 
 func TestRetries(t *testing.T) {
-	srcImage, err := oci.ReadLayout(xImagePath)
+	srcImage, err := oci.AsImage(oci.ReadLayout(xImagePath))
 	if err != nil {
 		t.Fatal(err)
 	}
 	server := NewServer()
-	server.AddImage("test/repo", "test-tag", srcImage)
+	server.AddRef("test/repo", "test-tag", srcImage)
 	wrapper := &unreliableServer{
 		handler:   server,
 		blobLimit: srcImage.Manifest.Config.Size / 2,
@@ -71,7 +71,7 @@
 		Repository: "test/repo",
 	}
 
-	image, err := client.Read(context.Background(), "test-tag", srcImage.ManifestDigest)
+	image, err := oci.AsImage(client.Read(context.Background(), "test-tag", srcImage.Digest()))
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/osbase/oci/registry/server.go b/osbase/oci/registry/server.go
index 9c99c40..d7765b1 100644
--- a/osbase/oci/registry/server.go
+++ b/osbase/oci/registry/server.go
@@ -15,8 +15,6 @@
 	"strings"
 	"time"
 
-	ocispecv1 "github.com/opencontainers/image-spec/specs-go/v1"
-
 	"source.monogon.dev/osbase/oci"
 	"source.monogon.dev/osbase/structfs"
 )
@@ -49,16 +47,24 @@
 	}
 }
 
-// AddImage adds an image to the server in the specified repository.
+// AddRef adds a Ref to the server in the specified repository.
 //
 // If the tag is empty, the image can only be fetched by digest.
-func (s *Server) AddImage(repository string, tag string, image *oci.Image) error {
+func (s *Server) AddRef(repository string, tag string, ref oci.Ref) error {
 	if !RepositoryRegexp.MatchString(repository) {
 		return fmt.Errorf("invalid repository %q", repository)
 	}
 	if tag != "" && !TagRegexp.MatchString(tag) {
 		return fmt.Errorf("invalid tag %q", tag)
 	}
+	var refs []oci.Ref
+	err := oci.WalkRefs(ref.Digest(), ref, func(_ string, r oci.Ref) error {
+		refs = append(refs, r)
+		return nil
+	})
+	if err != nil {
+		return err
+	}
 
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -71,17 +77,22 @@
 		}
 		s.repositories[repository] = repo
 	}
-	if _, ok := repo.manifests[image.ManifestDigest]; !ok {
-		for descriptor := range image.Descriptors() {
-			repo.blobs[string(descriptor.Digest)] = image.StructfsBlob(descriptor)
+	for _, ref := range refs {
+		if _, ok := repo.manifests[ref.Digest()]; ok {
+			continue
 		}
-		repo.manifests[image.ManifestDigest] = serverManifest{
-			contentType: ocispecv1.MediaTypeImageManifest,
-			content:     image.RawManifest,
+		if image, ok := ref.(*oci.Image); ok {
+			for descriptor := range image.Descriptors() {
+				repo.blobs[string(descriptor.Digest)] = image.StructfsBlob(descriptor)
+			}
+		}
+		repo.manifests[ref.Digest()] = serverManifest{
+			contentType: ref.MediaType(),
+			content:     ref.RawManifest(),
 		}
 	}
 	if tag != "" {
-		repo.tags[tag] = image.ManifestDigest
+		repo.tags[tag] = ref.Digest()
 	}
 	return nil
 }