From 406984ca8d12277abacf02e0736fb3e5b43a1db9 Mon Sep 17 00:00:00 2001
From: Danielle Lancashire <dani@hashicorp.com>
Date: Wed, 18 Dec 2019 12:20:32 +0100
Subject: [PATCH] csimanager: Fingerprint controller capabilities

---
 client/pluginmanager/csimanager/instance.go   | 13 ++++
 .../pluginmanager/csimanager/instance_test.go | 76 +++++++++++++++++++
 nomad/structs/node.go                         | 18 ++++-
 3 files changed, 106 insertions(+), 1 deletion(-)

diff --git a/client/pluginmanager/csimanager/instance.go b/client/pluginmanager/csimanager/instance.go
index 9de20ce8b..5313bc07c 100644
--- a/client/pluginmanager/csimanager/instance.go
+++ b/client/pluginmanager/csimanager/instance.go
@@ -127,6 +127,13 @@ func (i *instanceManager) runLoop() {
 	}
 }
 
+func applyCapabilitySetToControllerInfo(cs *csi.ControllerCapabilitySet, info *structs.CSIControllerInfo) {
+	info.SupportsReadOnlyAttach = cs.HasPublishReadonly
+	info.SupportsAttachDetach = cs.HasPublishUnpublishVolume
+	info.SupportsListVolumes = cs.HasListVolumes
+	info.SupportsListVolumesAttachedNodes = cs.HasListVolumesPublishedNodes
+}
+
 func (i *instanceManager) buildControllerFingerprint(ctx context.Context, base *structs.CSIInfo) (*structs.CSIInfo, error) {
 	fp := base.Copy()
 
@@ -136,6 +143,12 @@ func (i *instanceManager) buildControllerFingerprint(ctx context.Context, base *
 	}
 	fp.SetHealthy(healthy)
 
+	caps, err := i.client.ControllerGetCapabilities(ctx)
+	if err != nil {
+		return fp, err
+	}
+	applyCapabilitySetToControllerInfo(caps, fp.ControllerInfo)
+
 	return fp, nil
 }
 
diff --git a/client/pluginmanager/csimanager/instance_test.go b/client/pluginmanager/csimanager/instance_test.go
index ca30a321e..ba0e2d572 100644
--- a/client/pluginmanager/csimanager/instance_test.go
+++ b/client/pluginmanager/csimanager/instance_test.go
@@ -157,3 +157,79 @@ func TestBuildBasicFingerprint_Node(t *testing.T) {
 		})
 	}
 }
+
+func TestBuildControllerFingerprint(t *testing.T) {
+	tt := []struct {
+		Name string
+
+		Capabilities          *csi.ControllerCapabilitySet
+		CapabilitiesErr       error
+		CapabilitiesCallCount int64
+
+		ProbeResponse  bool
+		ProbeErr       error
+		ProbeCallCount int64
+
+		ExpectedControllerInfo *structs.CSIControllerInfo
+		ExpectedErr            error
+	}{
+		{
+			Name: "Minimal successful response",
+
+			Capabilities:          &csi.ControllerCapabilitySet{},
+			CapabilitiesCallCount: 1,
+
+			ProbeResponse:  true,
+			ProbeCallCount: 1,
+
+			ExpectedControllerInfo: &structs.CSIControllerInfo{},
+		},
+		{
+			Name: "Successful response with capabilities",
+
+			Capabilities: &csi.ControllerCapabilitySet{
+				HasListVolumes: true,
+			},
+			CapabilitiesCallCount: 1,
+
+			ProbeResponse:  true,
+			ProbeCallCount: 1,
+
+			ExpectedControllerInfo: &structs.CSIControllerInfo{
+				SupportsListVolumes: true,
+			},
+		},
+		{
+			Name: "ControllerGetCapabilities Failed",
+
+			CapabilitiesErr:       errors.New("request failed"),
+			CapabilitiesCallCount: 1,
+
+			ProbeResponse:  true,
+			ProbeCallCount: 1,
+
+			ExpectedControllerInfo: &structs.CSIControllerInfo{},
+			ExpectedErr:            errors.New("request failed"),
+		},
+	}
+
+	for _, test := range tt {
+		t.Run(test.Name, func(t *testing.T) {
+			client, im := setupTestNodeInstanceManager(t)
+
+			client.NextControllerGetCapabilitiesResponse = test.Capabilities
+			client.NextControllerGetCapabilitiesErr = test.CapabilitiesErr
+
+			client.NextPluginProbeResponse = test.ProbeResponse
+			client.NextPluginProbeErr = test.ProbeErr
+
+			info, err := im.buildControllerFingerprint(context.TODO(), &structs.CSIInfo{ControllerInfo: &structs.CSIControllerInfo{}})
+
+			require.Equal(t, test.ExpectedControllerInfo, info.ControllerInfo)
+			require.Equal(t, test.ExpectedErr, err)
+
+			require.Equal(t, test.CapabilitiesCallCount, client.ControllerGetCapabilitiesCallCount)
+			require.Equal(t, test.ProbeCallCount, client.PluginProbeCallCount)
+		})
+	}
+}
diff --git a/nomad/structs/node.go b/nomad/structs/node.go
index 7143c42e2..e5960dbbf 100644
--- a/nomad/structs/node.go
+++ b/nomad/structs/node.go
@@ -100,7 +100,23 @@ func (n *CSINodeInfo) Copy() *CSINodeInfo {
 // CSIControllerInfo is the fingerprinted data from a CSI Plugin that is specific to
 // the Controller API.
 type CSIControllerInfo struct {
-	// Currently empty
+	// SupportsReadOnlyAttach is set to true when the controller returns the
+	// ATTACH_READONLY capability.
+	SupportsReadOnlyAttach bool
+
+	// SupportsPublishVolume is true when the controller implements the methods
+	// required to attach and detach volumes. If this is false Nomad should skip
+	// the controller attachment flow.
+	SupportsAttachDetach bool
+
+	// SupportsListVolums is true when the controller implements the ListVolumes
+	// RPC. NOTE: This does not guaruntee that attached nodes will be returned
+	// unless SupportsListVolumesAttachedNodes is also true.
+	SupportsListVolumes bool
+
+	// SupportsListVolumesAttachedNodes indicates whether the plugin will return
+	// attached nodes data when making ListVolume RPCs
+	SupportsListVolumesAttachedNodes bool
 }
 
 func (c *CSIControllerInfo) Copy() *CSIControllerInfo {