From 76c8880e45ebd3dd61e9756a0548f65492435b9d Mon Sep 17 00:00:00 2001
From: ChengyuZhu6 <chengyu.zhu@intel.com>
Date: Sun, 28 Apr 2024 15:41:23 +0800
Subject: [PATCH] snapshot: fix error on proxy driver when switching different
 snapshotter

Fixes: #592

Signed-off-by: ChengyuZhu6 <chengyu.zhu@intel.com>
---
 snapshot/process.go  | 10 +++++++
 snapshot/snapshot.go | 68 +++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/snapshot/process.go b/snapshot/process.go
index 9f0da66f2d..22a9c64e76 100644
--- a/snapshot/process.go
+++ b/snapshot/process.go
@@ -57,6 +57,12 @@ func chooseProcessor(ctx context.Context, logger *logrus.Entry,
 		}
 	}
 
+	proxyhandler := func() (bool, []mount.Mount, error) {
+		logger.Infof("Nydus proxy handler")
+		mounts, err := sn.mountProxy(ctx, s)
+		return false, mounts, err
+	}
+
 	// OCI image is also marked with "containerd.io/snapshot.ref" by Containerd
 	target, isRoLayer := labels[label.TargetSnapshotRef]
 
@@ -118,6 +124,10 @@ func chooseProcessor(ctx context.Context, logger *logrus.Entry,
 		// It should not be committed during this Prepare() operation.
 
 		pID, pInfo, _, pErr := snapshot.GetSnapshotInfo(ctx, sn.ms, parent)
+		if checkErr := checkLabelsWithDriver(pInfo.Labels); checkErr != nil {
+			logger.Infof("not found proxy label: %v, err = %v", pInfo.Labels, checkErr)
+			handler = proxyhandler
+		}
 		if pErr == nil && label.IsNydusProxyMode(pInfo.Labels) {
 			logger.Infof("Prepare active snapshot %s in proxy mode", key)
 			handler = remoteHandler(pID, pInfo.Labels)
diff --git a/snapshot/snapshot.go b/snapshot/snapshot.go
index f829d3132c..0d6d6f3bb3 100644
--- a/snapshot/snapshot.go
+++ b/snapshot/snapshot.go
@@ -22,9 +22,9 @@ import (
 	"github.com/containerd/containerd/snapshots"
 	"github.com/containerd/containerd/snapshots/storage"
 	"github.com/containerd/continuity/fs"
-
 	"github.com/containerd/nydus-snapshotter/config"
 	"github.com/containerd/nydus-snapshotter/config/daemonconfig"
+	rafs "github.com/containerd/nydus-snapshotter/pkg/rafs"
 
 	"github.com/containerd/nydus-snapshotter/pkg/cache"
 	"github.com/containerd/nydus-snapshotter/pkg/cgroup"
@@ -432,6 +432,11 @@ func (o *snapshotter) Mounts(ctx context.Context, key string) ([]mount.Mount, er
 		return nil, errors.Wrapf(err, "get snapshot %s", key)
 	}
 
+	if checkErr := checkLabelsWithDriver(info.Labels); checkErr != nil {
+		log.L.Debug("[Mounts] not found proxy label")
+		return o.mountProxy(ctx, *snap)
+	}
+
 	if needRemoteMounts {
 		return o.mountRemote(ctx, info.Labels, *snap, metaSnapshotID, key)
 	}
@@ -838,6 +843,52 @@ func overlayMount(options []string) []mount.Mount {
 	}
 }
 
+func (o *snapshotter) mountProxy(ctx context.Context, s storage.Snapshot) ([]mount.Mount, error) {
+	var overlayOptions []string
+	if s.Kind == snapshots.KindActive {
+		overlayOptions = append(overlayOptions,
+			fmt.Sprintf("workdir=%s", o.workPath(s.ID)),
+			fmt.Sprintf("upperdir=%s", o.upperPath(s.ID)),
+		)
+	}
+
+	log.G(ctx).Debugf("len(s.ParentIDs) = %v", len(s.ParentIDs))
+	parentPaths := make([]string, 0, len(s.ParentIDs)+1)
+	if len(s.ParentIDs) == 0 {
+		parentPaths = append(parentPaths, config.GetSnapshotsRootDir())
+	} else {
+		for _, id := range s.ParentIDs {
+			parentPaths = append(parentPaths, o.upperPath(id))
+		}
+	}
+
+	lowerDirOption := fmt.Sprintf("lowerdir=%s", strings.Join(parentPaths, ":"))
+	overlayOptions = append(overlayOptions, lowerDirOption)
+	log.G(ctx).Infof("proxy mount options %v", overlayOptions)
+	options, err := o.mountWithProxyVolume(rafs.Rafs{
+		FsDriver:    config.GetFsDriver(),
+		ImageID:     "",
+		SnapshotID:  "",
+		SnapshotDir: "",
+		Annotations: make(map[string]string),
+	})
+	if err != nil {
+		return []mount.Mount{}, errors.Wrapf(err, "create kata volume for proxy")
+	}
+	if len(options) > 0 {
+		overlayOptions = append(overlayOptions, options...)
+	}
+	log.G(ctx).Debugf("fuse.nydus-overlayfs mount options %v", overlayOptions)
+	mounts := []mount.Mount{
+		{
+			Type:    "fuse.nydus-overlayfs",
+			Source:  "overlay",
+			Options: overlayOptions,
+		},
+	}
+	return mounts, nil
+}
+
 // `s` is the upmost snapshot and `id` refers to the nydus meta snapshot
 // `s` and `id` can represent a different layer, it's useful when View an image
 func (o *snapshotter) mountRemote(ctx context.Context, labels map[string]string, s storage.Snapshot, id, key string) ([]mount.Mount, error) {
@@ -1011,3 +1062,18 @@ func (o *snapshotter) snapshotRoot() string {
 func (o *snapshotter) snapshotDir(id string) string {
 	return filepath.Join(o.snapshotRoot(), id)
 }
+
+func checkLabelsWithDriver(labels map[string]string) error {
+	isProxyDriver := config.GetFsDriver() == config.FsDriverProxy
+	isProxyLabel := label.IsNydusProxyMode(labels)
+	_, isProxyImage := labels[label.CRIImageRef]
+	log.G(context.Background()).Debugf("isProxyDriver = %t, isProxyLabel = %t, isProxyImage = %t", isProxyDriver, isProxyLabel, isProxyImage)
+	switch {
+	case isProxyDriver && isProxyImage:
+		return nil
+	case isProxyDriver != isProxyLabel:
+		return errors.Errorf("check Labels With Driver failed, driver = %q, labels = %q", config.GetFsDriver(), labels)
+	default:
+		return nil
+	}
+}