From c1b7bcf395d1334333e254ead5b6cd7dfa84cbb9 Mon Sep 17 00:00:00 2001
From: ktock <ktokunaga.mail@gmail.com>
Date: Mon, 23 Mar 2020 10:46:34 +0900
Subject: [PATCH] Enable to pass additional handler on pull for stargz-based
 remote snapshots

Throughout container lifecycle, pulling image is one of the time-consuming
steps. Recently, containerd community started to tackle this issue with
stargz-based remote snapshots, as a non-core
subproject(https://github.com/containerd/stargz-snapshotter).

This snapshotter is implemented as a standard proxy plugin but it requires the
client to pass some additional information (image ref and layer digest) for each
pull operation to query layer contents on the registry. Stargz snapshotter
project provides an image handler to do this and stargz snapshot users need to
pass this handler to containerd client.

This commit enables to use stargz-based remote snapshots through CRI by passing
the handler to containerd client on pull operation.

Signed-off-by: Kohei Tokunaga <ktokunaga.mail@gmail.com>
---
 pkg/config/config.go     |  5 +++++
 pkg/server/image_pull.go | 43 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/pkg/config/config.go b/pkg/config/config.go
index 1115ee6fc..2a0dbeac8 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -73,6 +73,11 @@ type ContainerdConfig struct {
 	// NoPivot disables pivot-root (linux only), required when running a container in a RamDisk with runc
 	// This only works for runtime type "io.containerd.runtime.v1.linux".
 	NoPivot bool `toml:"no_pivot" json:"noPivot"`
+
+	// DisableSnapshotAnnotations disables to pass additional annotations (image
+	// related information) to snapshotters. These annotations are required by
+	// stargz snapshotter (https://github.com/containerd/stargz-snapshotter).
+	DisableSnapshotAnnotations bool `toml:"disable_snapshot_annotations" json:"disableSnapshotAnnotations"`
 }
 
 // CniConfig contains toml config related to cni
diff --git a/pkg/server/image_pull.go b/pkg/server/image_pull.go
index d2457a267..e776c050e 100644
--- a/pkg/server/image_pull.go
+++ b/pkg/server/image_pull.go
@@ -120,6 +120,10 @@ func (c *criService) PullImage(ctx context.Context, r *runtime.PullImageRequest)
 	}
 
 	pullOpts = append(pullOpts, c.encryptedImagesPullOpts()...)
+	if !c.config.ContainerdConfig.DisableSnapshotAnnotations {
+		pullOpts = append(pullOpts,
+			containerd.WithImageHandlerWrapper(appendInfoHandlerWrapper(ref)))
+	}
 
 	image, err := c.client.Pull(ctx, ref, pullOpts...)
 	if err != nil {
@@ -422,3 +426,42 @@ func (c *criService) encryptedImagesPullOpts() []containerd.RemoteOpt {
 	}
 	return nil
 }
+
+const (
+	// targetRefLabel is a label which contains image reference and will be passed
+	// to snapshotters.
+	targetRefLabel = "containerd.io/snapshot/cri.image-ref"
+	// targetDigestLabel is a label which contains layer digest and will be passed
+	// to snapshotters.
+	targetDigestLabel = "containerd.io/snapshot/cri.layer-digest"
+)
+
+// appendInfoHandlerWrapper makes a handler which appends some basic information
+// of images to each layer descriptor as annotations during unpack. These
+// annotations will be passed to snapshotters as labels. These labels will be
+// used mainly by stargz-based snapshotters for querying image contents from the
+// registry.
+func appendInfoHandlerWrapper(ref string) func(f containerdimages.Handler) containerdimages.Handler {
+	return func(f containerdimages.Handler) containerdimages.Handler {
+		return containerdimages.HandlerFunc(func(ctx context.Context, desc imagespec.Descriptor) ([]imagespec.Descriptor, error) {
+			children, err := f.Handle(ctx, desc)
+			if err != nil {
+				return nil, err
+			}
+			switch desc.MediaType {
+			case imagespec.MediaTypeImageManifest, containerdimages.MediaTypeDockerSchema2Manifest:
+				for i := range children {
+					c := &children[i]
+					if containerdimages.IsLayerType(c.MediaType) {
+						if c.Annotations == nil {
+							c.Annotations = make(map[string]string)
+						}
+						c.Annotations[targetRefLabel] = ref
+						c.Annotations[targetDigestLabel] = c.Digest.String()
+					}
+				}
+			}
+			return children, nil
+		})
+	}
+}