From 4be7a1edd46feca4ec91b35bd323c8834dc575ce Mon Sep 17 00:00:00 2001
From: Giuseppe Scrivano <gscrivan@redhat.com>
Date: Mon, 5 Jun 2023 21:31:25 +0200
Subject: [PATCH] storage: generate different imageID with partial layers

when an image using partial layers is pulled, generate a different
imageID that combines the digest of the TOC for the partial layers
since the partial pull didn't validate the diffID.

Signed-off-by: Giuseppe Scrivano <gscrivan@redhat.com>
---
 manifest/oci.go         | 38 +++++++++++++++++++++++++++++++++++++-
 storage/storage_dest.go | 10 +++++++---
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/manifest/oci.go b/manifest/oci.go
index a85641c36a..6d5acb45d8 100644
--- a/manifest/oci.go
+++ b/manifest/oci.go
@@ -9,6 +9,7 @@ import (
 	compressiontypes "github.com/containers/image/v5/pkg/compression/types"
 	"github.com/containers/image/v5/types"
 	ociencspec "github.com/containers/ocicrypt/spec"
+	chunkedToc "github.com/containers/storage/pkg/chunked/toc"
 	"github.com/opencontainers/go-digest"
 	"github.com/opencontainers/image-spec/specs-go"
 	imgspecv1 "github.com/opencontainers/image-spec/specs-go/v1"
@@ -235,7 +236,7 @@ func (m *OCI1) Inspect(configGetter func(types.BlobInfo) ([]byte, error)) (*type
 }
 
 // ImageID computes an ID which can uniquely identify this image by its contents.
-func (m *OCI1) ImageID([]digest.Digest) (string, error) {
+func (m *OCI1) ImageID(diffIDs []digest.Digest) (string, error) {
 	// The way m.Config.Digest “uniquely identifies” an image is
 	// by containing RootFS.DiffIDs, which identify the layers of the image.
 	// For non-image artifacts, the we can’t expect the config to change
@@ -259,9 +260,44 @@ func (m *OCI1) ImageID([]digest.Digest) (string, error) {
 	if err := m.Config.Digest.Validate(); err != nil {
 		return "", err
 	}
+
+	// If there is any layer that is using partial content, we calculate the image ID
+	// in a different way since the diffID cannot be validated as for regular pulled images.
+	for _, layer := range m.Layers {
+		toc, err := chunkedToc.GetTOCDigest(layer.Annotations)
+		if err != nil {
+			return "", fmt.Errorf("error looking up annotation for layer %q: %w", layer.Digest, err)
+		}
+		if toc != nil {
+			return m.calculateImageIDForPartialImage(diffIDs)
+		}
+	}
+
 	return m.Config.Digest.Hex(), nil
 }
 
+func (m *OCI1) calculateImageIDForPartialImage(diffIDs []digest.Digest) (string, error) {
+	newID := digest.Canonical.Digester()
+	for i, layer := range m.Layers {
+		diffID := diffIDs[i]
+		_, err := newID.Hash().Write([]byte(diffID.Hex()))
+		if err != nil {
+			return "", fmt.Errorf("error writing diffID %q: %w", diffID, err)
+		}
+		toc, err := chunkedToc.GetTOCDigest(layer.Annotations)
+		if err != nil {
+			return "", fmt.Errorf("error looking up annotation for layer %q: %w", layer.Digest, err)
+		}
+		if toc != nil {
+			_, err = newID.Hash().Write([]byte(toc.Hex()))
+			if err != nil {
+				return "", fmt.Errorf("error writing TOC %q: %w", toc, err)
+			}
+		}
+	}
+	return newID.Digest().Hex(), nil
+}
+
 // CanChangeLayerCompression returns true if we can compress/decompress layers with mimeType in the current image
 // (and the code can handle that).
 // NOTE: Even if this returns true, the relevant format might not accept all compression algorithms; the set of accepted
diff --git a/storage/storage_dest.go b/storage/storage_dest.go
index 6eaaba018e..6e1e83ecdb 100644
--- a/storage/storage_dest.go
+++ b/storage/storage_dest.go
@@ -464,9 +464,13 @@ func (s *storageImageDestination) computeID(m manifest.Manifest) string {
 			}
 			diffIDs = append([]digest.Digest{diffID}, diffIDs...)
 		}
-	case *manifest.Schema2, *manifest.OCI1:
-		// We know the ID calculation for these formats doesn't actually use the diffIDs,
-		// so we don't need to populate the diffID list.
+	case *manifest.Schema2:
+		// We know the ID calculation doesn't actually use the diffIDs, so we don't need to populate
+		// the diffID list.
+	case *manifest.OCI1:
+		for _, l := range m.Layers {
+			diffIDs = append(diffIDs, l.Digest)
+		}
 	default:
 		return ""
 	}