From 20f79300cf35fe453a4fcbd0ea79d2c90e6f5892 Mon Sep 17 00:00:00 2001
From: Derek McGowan <derek@mcg.dev>
Date: Tue, 29 Aug 2023 17:16:19 -0700
Subject: [PATCH] push: inherit distribution sources from parent

When a blob does not exist locally, rather than erroring on info
lookup, inherit the parent distribution sources. Push is able
to succeed even if the blob does not exist locally when a cross
repository mount is done. This is a common operation pushing a
multi-platform image to the same registry but different namespace.

Signed-off-by: Derek McGowan <derek@mcg.dev>
---
 remotes/handlers.go | 32 +++++++++++++++++++++++++-------
 1 file changed, 25 insertions(+), 7 deletions(-)

diff --git a/remotes/handlers.go b/remotes/handlers.go
index 75707703b..874c83415 100644
--- a/remotes/handlers.go
+++ b/remotes/handlers.go
@@ -204,8 +204,9 @@ func push(ctx context.Context, provider content.Provider, pusher Pusher, desc oc
 // Base handlers can be provided which will be called before any push specific
 // handlers.
 //
-// If the passed in content.Provider is also a content.Manager then this will
-// also annotate the distribution sources in the manager.
+// If the passed in content.Provider is also a content.InfoProvider (such as
+// content.Manager) then this will also annotate the distribution sources using
+// labels prefixed with "containerd.io/distribution.source".
 func PushContent(ctx context.Context, pusher Pusher, desc ocispec.Descriptor, store content.Provider, limiter *semaphore.Weighted, platform platforms.MatchComparer, wrapper func(h images.Handler) images.Handler) error {
 
 	var m sync.Mutex
@@ -234,7 +235,7 @@ func PushContent(ctx context.Context, pusher Pusher, desc ocispec.Descriptor, st
 	platformFilterhandler := images.FilterPlatforms(images.ChildrenHandler(store), platform)
 
 	var handler images.Handler
-	if m, ok := store.(content.Manager); ok {
+	if m, ok := store.(content.InfoProvider); ok {
 		annotateHandler := annotateDistributionSourceHandler(platformFilterhandler, m)
 		handler = images.Handlers(annotateHandler, filterHandler, pushHandler)
 	} else {
@@ -340,14 +341,15 @@ func FilterManifestByPlatformHandler(f images.HandlerFunc, m platforms.Matcher)
 
 // annotateDistributionSourceHandler add distribution source label into
 // annotation of config or blob descriptor.
-func annotateDistributionSourceHandler(f images.HandlerFunc, manager content.Manager) images.HandlerFunc {
+func annotateDistributionSourceHandler(f images.HandlerFunc, provider content.InfoProvider) images.HandlerFunc {
 	return func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
 		children, err := f(ctx, desc)
 		if err != nil {
 			return nil, err
 		}
 
-		// only add distribution source for the config or blob data descriptor
+		// Distribution source is only used for config or blob but may be inherited from
+		// a manifest or manifest list
 		switch desc.MediaType {
 		case images.MediaTypeDockerSchema2Manifest, ocispec.MediaTypeImageManifest,
 			images.MediaTypeDockerSchema2ManifestList, ocispec.MediaTypeImageIndex:
@@ -355,12 +357,28 @@ func annotateDistributionSourceHandler(f images.HandlerFunc, manager content.Man
 			return children, nil
 		}
 
+		// parentInfo can be used to inherit info for non-existent blobs
+		var parentInfo *content.Info
+
 		for i := range children {
 			child := children[i]
 
-			info, err := manager.Info(ctx, child.Digest)
+			info, err := provider.Info(ctx, child.Digest)
 			if err != nil {
-				return nil, err
+				if !errdefs.IsNotFound(err) {
+					return nil, err
+				}
+				if parentInfo == nil {
+					pi, err := provider.Info(ctx, desc.Digest)
+					if err != nil {
+						return nil, err
+					}
+					parentInfo = &pi
+				}
+				// Blob may not exist locally, annotate with parent labels for cross repo
+				// mount or fetch. Parent sources may apply to all children since most
+				// registries enforce that children exist before the manifests.
+				info = *parentInfo
 			}
 
 			for k, v := range info.Labels {