Refactor: Abstract the way to handle remote snapshotter

pkg/imgutil/imgutil.go

@@ -22,23 +22,19 @@ import (
 	"fmt"
 	"io"
 	"reflect"
-	"strings"
 
 	"github.com/containerd/containerd"
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/images"
 	"github.com/containerd/containerd/platforms"
 	refdocker "github.com/containerd/containerd/reference/docker"
 	"github.com/containerd/containerd/remotes"
-	"github.com/containerd/containerd/snapshots"
 	"github.com/containerd/imgcrypt"
 	"github.com/containerd/imgcrypt/images/encryption"
 	"github.com/containerd/nerdctl/pkg/errutil"
 	"github.com/containerd/nerdctl/pkg/idutil/imagewalker"
 	"github.com/containerd/nerdctl/pkg/imgutil/dockerconfigresolver"
 	"github.com/containerd/nerdctl/pkg/imgutil/pull"
-	nyduslabel "github.com/containerd/nydus-snapshotter/pkg/label"
-	"github.com/containerd/stargz-snapshotter/fs/source"
 	"github.com/docker/docker/errdefs"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 
@@ -77,7 +73,7 @@ func GetExistingImage(ctx context.Context, client *containerd.Client, snapshotte
 				Image:       image,
 				ImageConfig: *imgConfig,
 				Snapshotter: snapshotter,
-				Remote:      isStargz(snapshotter) || isOverlaybd(snapshotter),
+				Remote:      getSnapshotterOpts(snapshotter).isRemote(),
 			}
 			if unpacked, err := image.IsUnpacked(ctx, snapshotter); err == nil && !unpacked {
 				if err := image.Unpack(ctx, snapshotter); err != nil {
@@ -222,7 +218,7 @@ func PullImage(ctx context.Context, client *containerd.Client, stdout, stderr io
 		unpackB = len(ocispecPlatforms) == 1
 	}
 
-	var sgz, overlaybd, nydus bool
+	snOpt := getSnapshotterOpts(snapshotter)
 	if unpackB {
 		logrus.Debugf("The image will be unpacked for platform %q, snapshotter %q.", ocispecPlatforms[0], snapshotter)
 		imgcryptPayload := imgcrypt.Payload{}
@@ -231,35 +227,8 @@ func PullImage(ctx context.Context, client *containerd.Client, stdout, stderr io
 			containerd.WithPullUnpack,
 			containerd.WithUnpackOpts([]containerd.UnpackOpt{imgcryptUnpackOpt}))
 
-		sgz = isStargz(snapshotter)
-		if sgz {
-			// TODO: support "skip-content-verify"
-			config.RemoteOpts = append(
-				config.RemoteOpts,
-				containerd.WithImageHandlerWrapper(source.AppendDefaultLabelsHandlerWrapper(ref, 10*1024*1024)),
-			)
-		}
-		nydus = isNydus(snapshotter)
-		if nydus {
-			config.RemoteOpts = append(
-				config.RemoteOpts,
-				containerd.WithImageHandlerWrapper(nyduslabel.AppendLabelsHandlerWrapper(ref)),
-			)
-		}
-		overlaybd = isOverlaybd(snapshotter)
-		if overlaybd {
-			snlabel := map[string]string{"containerd.io/snapshot/image-ref": ref}
-			logrus.Debugf("append remote opts: %s", snlabel)
-
-			config.RemoteOpts = append(
-				config.RemoteOpts,
-				containerd.WithPullSnapshotter(snapshotter, snapshots.WithLabels(snlabel)),
-			)
-		} else {
-			config.RemoteOpts = append(
-				config.RemoteOpts,
-				containerd.WithPullSnapshotter(snapshotter))
-		}
+		// different remote snapshotters will update pull.Config separately
+		snOpt.apply(config, ref)
 	} else {
 		logrus.Debugf("The image will not be unpacked. Platforms=%v.", ocispecPlatforms)
 	}
@@ -276,30 +245,12 @@ func PullImage(ctx context.Context, client *containerd.Client, stdout, stderr io
 		Image:       containerdImage,
 		ImageConfig: *imgConfig,
 		Snapshotter: snapshotter,
-		Remote:      (sgz || overlaybd || nydus),
+		Remote:      snOpt.isRemote(),
 	}
 	return res, nil
 
 }
 
-func isStargz(sn string) bool {
-	if !strings.Contains(sn, "stargz") {
-		return false
-	}
-	if sn != "stargz" {
-		logrus.Debugf("assuming %q to be a stargz-compatible snapshotter", sn)
-	}
-	return true
-}
-
-func isOverlaybd(sn string) bool {
-	return sn == "overlaybd"
-}
-
-func isNydus(sn string) bool {
-	return sn == "nydus"
-}
-
 func getImageConfig(ctx context.Context, image containerd.Image) (*ocispec.ImageConfig, error) {
 	desc, err := image.Config(ctx)
 	if err != nil {

pkg/imgutil/snapshotter.go

@@ -0,0 +1,134 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package imgutil
+
+import (
+	"strings"
+
+	"github.com/containerd/containerd"
+	"github.com/containerd/containerd/snapshots"
+	"github.com/containerd/nerdctl/pkg/imgutil/pull"
+	nyduslabel "github.com/containerd/nydus-snapshotter/pkg/label"
+	"github.com/containerd/stargz-snapshotter/fs/source"
+	"github.com/sirupsen/logrus"
+)
+
+const (
+	snapshotterNameOverlaybd = "overlaybd"
+	snapshotterNameStargz    = "stargz"
+	snapshotterNameNydus     = "nydus"
+
+	// prefetch size for stargz
+	prefetchSize = 10 * 1024 * 1024
+
+	overlaybdLabelImageRef = "containerd.io/snapshot/image-ref"
+)
+
+// remote snapshotters explicitly handled by nerdctl
+var builtinRemoteSnapshotterOpts = map[string]snapshotterOpts{
+	snapshotterNameOverlaybd: &overlaybdSnapshotterOpts{},
+	snapshotterNameStargz:    &stargzSnapshotterOpts{},
+	snapshotterNameNydus:     &nydusSnapshotterOpts{},
+}
+
+// snapshotterOpts is used to update pull config
+// for different snapshotters
+type snapshotterOpts interface {
+	apply(config *pull.Config, ref string)
+	isRemote() bool
+}
+
+// getSnapshotterOpts get snapshotter opts by fuzzy matching of the snapshotter name
+func getSnapshotterOpts(snapshotter string) snapshotterOpts {
+	for sn, sno := range builtinRemoteSnapshotterOpts {
+		if strings.Contains(snapshotter, sn) {
+			if snapshotter != sn {
+				logrus.Debugf("assuming %s to be a %s-compatible snapshotter", snapshotter, sn)
+			}
+			return sno
+		}
+	}
+
+	return &defaultSnapshotterOpts{snapshotter: snapshotter}
+}
+
+// remoteSnapshotter is used as a default implementation for
+// interface `snapshotterOpts.isRemote()` function
+type remoteSnapshotter struct{}
+
+func (rs *remoteSnapshotter) isRemote() bool {
+	return true
+}
+
+// defaultSnapshotterOpts is for snapshotters that
+// not handled separately
+type defaultSnapshotterOpts struct {
+	snapshotter string
+}
+
+func (dsn *defaultSnapshotterOpts) apply(config *pull.Config, _ref string) {
+	config.RemoteOpts = append(
+		config.RemoteOpts,
+		containerd.WithPullSnapshotter(dsn.snapshotter))
+}
+
+// defaultSnapshotterOpts is not a remote snapshotter
+func (dsn *defaultSnapshotterOpts) isRemote() bool {
+	return false
+}
+
+// stargzSnapshotterOpts for stargz snapshotter
+type stargzSnapshotterOpts struct {
+	remoteSnapshotter
+}
+
+func (ssn *stargzSnapshotterOpts) apply(config *pull.Config, ref string) {
+	// TODO: support "skip-content-verify"
+	config.RemoteOpts = append(
+		config.RemoteOpts,
+		containerd.WithImageHandlerWrapper(source.AppendDefaultLabelsHandlerWrapper(ref, prefetchSize)),
+		containerd.WithPullSnapshotter(snapshotterNameStargz),
+	)
+}
+
+// overlaybdSnapshotterOpts for overlaybd snapshotter
+type overlaybdSnapshotterOpts struct {
+	remoteSnapshotter
+}
+
+func (osn *overlaybdSnapshotterOpts) apply(config *pull.Config, ref string) {
+	snlabel := map[string]string{overlaybdLabelImageRef: ref}
+	logrus.Debugf("append remote opts: %s", snlabel)
+
+	config.RemoteOpts = append(
+		config.RemoteOpts,
+		containerd.WithPullSnapshotter(snapshotterNameOverlaybd, snapshots.WithLabels(snlabel)),
+	)
+}
+
+// nydusSnapshotterOpts for nydus snapshotter
+type nydusSnapshotterOpts struct {
+	remoteSnapshotter
+}
+
+func (nsn *nydusSnapshotterOpts) apply(config *pull.Config, ref string) {
+	config.RemoteOpts = append(
+		config.RemoteOpts,
+		containerd.WithImageHandlerWrapper(nyduslabel.AppendLabelsHandlerWrapper(ref)),
+		containerd.WithPullSnapshotter(snapshotterNameNydus),
+	)
+}