From 91919fce12907e336a894e3dcb587995c4869c13 Mon Sep 17 00:00:00 2001
From: Tim Gross <tgross@hashicorp.com>
Date: Fri, 4 Mar 2022 15:04:21 -0500
Subject: [PATCH] csi: fix prefix queries for plugin list RPC

The `CSIPlugin.List` RPC was intended to accept a prefix to filter the
list of plugins being listed. This was being accidentally being done
in the state store instead, which contributed to incorrect filtering
behavior for plugins in the `volume snapshot list` command.

Move the prefix matching into the RPC so that it calls the
prefix-matching method in the state store if we're looking for a
prefix.

Update the `plugin status command` to accept a prefix for the plugin
ID argument so that it matches the expected behavior of other commands.
---
 command/plugin_status_csi.go | 23 +++++++++++++++++++++++
 nomad/csi_endpoint.go        | 18 ++++++++++++++----
 nomad/csi_endpoint_test.go   |  3 ++-
 nomad/state/state_store.go   |  2 +-
 4 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/command/plugin_status_csi.go b/command/plugin_status_csi.go
index b3a6ac36b39f..cdf8e03c5ef4 100644
--- a/command/plugin_status_csi.go
+++ b/command/plugin_status_csi.go
@@ -37,6 +37,29 @@ func (c *PluginStatusCommand) csiStatus(client *api.Client, id string) int {
 		return 0
 	}
 
+	// filter by plugin if a plugin ID was passed
+	plugs, _, err := client.CSIPlugins().List(&api.QueryOptions{Prefix: id})
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Error querying CSI plugins: %s", err))
+		return 1
+	}
+	if len(plugs) == 0 {
+		c.Ui.Error(fmt.Sprintf("No plugins(s) with prefix or ID %q found", id))
+		return 1
+	}
+	if len(plugs) > 1 {
+		if id != plugs[0].ID {
+			out, err := c.csiFormatPlugins(plugs)
+			if err != nil {
+				c.Ui.Error(fmt.Sprintf("Error formatting: %s", err))
+				return 1
+			}
+			c.Ui.Error(fmt.Sprintf("Prefix matched multiple plugins\n\n%s", out))
+			return 1
+		}
+	}
+	id = plugs[0].ID
+
 	// Lookup matched a single plugin
 	plug, _, err := client.CSIPlugins().Info(id, nil)
 	if err != nil {
diff --git a/nomad/csi_endpoint.go b/nomad/csi_endpoint.go
index 66083452fa82..ed0f95131e8f 100644
--- a/nomad/csi_endpoint.go
+++ b/nomad/csi_endpoint.go
@@ -1294,10 +1294,20 @@ func (v *CSIPlugin) List(args *structs.CSIPluginListRequest, reply *structs.CSIP
 		queryOpts: &args.QueryOptions,
 		queryMeta: &reply.QueryMeta,
 		run: func(ws memdb.WatchSet, state *state.StateStore) error {
-			// Query all plugins
-			iter, err := state.CSIPlugins(ws)
-			if err != nil {
-				return err
+
+			var iter memdb.ResultIterator
+			var err error
+			if args.Prefix != "" {
+				iter, err = state.CSIPluginsByIDPrefix(ws, args.Prefix)
+				if err != nil {
+					return err
+				}
+			} else {
+				// Query all plugins
+				iter, err = state.CSIPlugins(ws)
+				if err != nil {
+					return err
+				}
 			}
 
 			// Collect results
diff --git a/nomad/csi_endpoint_test.go b/nomad/csi_endpoint_test.go
index 4b8275064984..4c846d549fa7 100644
--- a/nomad/csi_endpoint_test.go
+++ b/nomad/csi_endpoint_test.go
@@ -1100,6 +1100,7 @@ func TestCSIVolumeEndpoint_ListExternal(t *testing.T) {
 	// List external volumes; note that none of these exist in the state store
 
 	req := &structs.CSIVolumeExternalListRequest{
+		PluginID: "minnie",
 		QueryOptions: structs.QueryOptions{
 			Region:    "global",
 			Namespace: structs.DefaultNamespace,
@@ -1371,8 +1372,8 @@ func TestCSIVolumeEndpoint_ListSnapshots(t *testing.T) {
 	require.NoError(t, state.UpsertNode(structs.MsgTypeTestSetup, index, node))
 
 	// List snapshots
-
 	req := &structs.CSISnapshotListRequest{
+		PluginID: "minnie",
 		Secrets: structs.CSISecrets{
 			"secret-key-1": "secret-val-1",
 		},
diff --git a/nomad/state/state_store.go b/nomad/state/state_store.go
index 4cc1902aa224..8ebbc6bc2fac 100644
--- a/nomad/state/state_store.go
+++ b/nomad/state/state_store.go
@@ -2696,7 +2696,7 @@ func (s *StateStore) CSIPluginByID(ws memdb.WatchSet, id string) (*structs.CSIPl
 // CSIPluginByIDTxn returns a named CSIPlugin
 func (s *StateStore) CSIPluginByIDTxn(txn Txn, ws memdb.WatchSet, id string) (*structs.CSIPlugin, error) {
 
-	watchCh, obj, err := txn.FirstWatch("csi_plugins", "id_prefix", id)
+	watchCh, obj, err := txn.FirstWatch("csi_plugins", "id", id)
 	if err != nil {
 		return nil, fmt.Errorf("csi_plugin lookup failed: %s %v", id, err)
 	}