From 6355186c3e7a82608a438a0c54da9ad4f8cf134e Mon Sep 17 00:00:00 2001
From: Cheng Pan <chengpan@amazon.com>
Date: Thu, 8 Nov 2018 17:28:37 +0000
Subject: [PATCH] Add In-tree to CSI migration logic for Mount and Unmount

---
 pkg/controller/.import-restrictions           |  2 +
 .../operationexecutor/operation_generator.go  | 41 ++++++++++++++++---
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/pkg/controller/.import-restrictions b/pkg/controller/.import-restrictions
index 9fa385d43a..6000753734 100644
--- a/pkg/controller/.import-restrictions
+++ b/pkg/controller/.import-restrictions
@@ -265,6 +265,7 @@
         "k8s.io/csi-api/pkg/client/clientset/versioned",
         "k8s.io/csi-api/pkg/client/listers/csi/v1alpha1",
         "k8s.io/csi-api/pkg/client/informers/externalversions",
+        "k8s.io/csi-api/pkg/client/informers/externalversions/csi/v1alpha1",
         "k8s.io/heapster/metrics/api/v1/types",
         "k8s.io/kube-controller-manager/config/v1alpha1",
         "k8s.io/metrics/pkg/apis/custom_metrics/v1beta2",
@@ -291,6 +292,7 @@
         "golang.org/x/time/rate",
         "golang.org/x/sys/unix",
         "golang.org/x/oauth2",
+        "golang.org/x/net/context",
         "google.golang.org/api/compute/v1",
         "google.golang.org/api/googleapi",
         "google.golang.org/api/compute/v0.alpha",
diff --git a/pkg/volume/util/operationexecutor/operation_generator.go b/pkg/volume/util/operationexecutor/operation_generator.go
index f49c0dcb8b..7a9fffab4a 100644
--- a/pkg/volume/util/operationexecutor/operation_generator.go
+++ b/pkg/volume/util/operationexecutor/operation_generator.go
@@ -514,6 +514,17 @@ func (og *operationGenerator) GenerateMountVolumeFunc(
 	}
 
 	mountVolumeFunc := func() (error, error) {
+		originalSpec := volumeToMount.VolumeSpec
+		// Get mounter plugin
+		if useCSIPlugin(og.volumePluginMgr, volumeToMount.VolumeSpec) {
+			csiSpec, err := translateSpec(volumeToMount.VolumeSpec)
+			if err != nil {
+				return volumeToMount.GenerateError("MountVolume.TranslateSpec failed", err)
+			}
+			volumeToMount.VolumeSpec = csiSpec
+		}
+
+		volumePlugin, err := og.volumePluginMgr.FindPluginBySpec(volumeToMount.VolumeSpec)
 		if err != nil || volumePlugin == nil {
 			return volumeToMount.GenerateError("MountVolume.FindPluginBySpec failed", err)
 		}
@@ -643,7 +654,7 @@ func (og *operationGenerator) GenerateMountVolumeFunc(
 			nil,
 			volumeToMount.OuterVolumeSpecName,
 			volumeToMount.VolumeGidValue,
-			volumeToMount.VolumeSpec)
+			originalSpec)
 		if markVolMountedErr != nil {
 			// On failure, return error. Caller will log and retry.
 			return volumeToMount.GenerateError("MountVolume.MarkVolumeAsMounted failed", markVolMountedErr)
@@ -719,13 +730,19 @@ func (og *operationGenerator) GenerateUnmountVolumeFunc(
 	volumeToUnmount MountedVolume,
 	actualStateOfWorld ActualStateOfWorldMounterUpdater,
 	podsDir string) (volumetypes.GeneratedOperations, error) {
+
+	var pluginName string
+	if useCSIPlugin(og.volumePluginMgr, volumeToUnmount.VolumeSpec) {
+		pluginName = csi.CSIPluginName
+	} else {
+		pluginName = volumeToUnmount.PluginName
+	}
+
 	// Get mountable plugin
-	volumePlugin, err :=
-		og.volumePluginMgr.FindPluginByName(volumeToUnmount.PluginName)
+	volumePlugin, err := og.volumePluginMgr.FindPluginByName(pluginName)
 	if err != nil || volumePlugin == nil {
 		return volumetypes.GeneratedOperations{}, volumeToUnmount.GenerateErrorDetailed("UnmountVolume.FindPluginByName failed", err)
 	}
-
 	volumeUnmounter, newUnmounterErr := volumePlugin.NewUnmounter(
 		volumeToUnmount.InnerVolumeSpecName, volumeToUnmount.PodUID)
 	if newUnmounterErr != nil {
@@ -780,12 +797,26 @@ func (og *operationGenerator) GenerateUnmountDeviceFunc(
 	deviceToDetach AttachedVolume,
 	actualStateOfWorld ActualStateOfWorldMounterUpdater,
 	mounter mount.Interface) (volumetypes.GeneratedOperations, error) {
+
+	var pluginName string
+	if useCSIPlugin(og.volumePluginMgr, deviceToDetach.VolumeSpec) {
+		pluginName = csi.CSIPluginName
+		csiSpec, err := translateSpec(deviceToDetach.VolumeSpec)
+		if err != nil {
+			return volumetypes.GeneratedOperations{}, deviceToDetach.GenerateErrorDetailed("UnmountDevice.TranslateSpec failed", err)
+		}
+		deviceToDetach.VolumeSpec = csiSpec
+	} else {
+		pluginName = deviceToDetach.PluginName
+	}
+
 	// Get DeviceMounter plugin
 	deviceMountableVolumePlugin, err :=
-		og.volumePluginMgr.FindDeviceMountablePluginByName(deviceToDetach.PluginName)
+		og.volumePluginMgr.FindDeviceMountablePluginByName(pluginName)
 	if err != nil || deviceMountableVolumePlugin == nil {
 		return volumetypes.GeneratedOperations{}, deviceToDetach.GenerateErrorDetailed("UnmountDevice.FindDeviceMountablePluginByName failed", err)
 	}
+
 	volumeDeviceUmounter, err := deviceMountableVolumePlugin.NewDeviceUnmounter()
 	if err != nil {
 		return volumetypes.GeneratedOperations{}, deviceToDetach.GenerateErrorDetailed("UnmountDevice.NewDeviceUmounter failed", err)