update tsh proxy kube cluster selection ux

tool/tsh/common/kube.go

@@ -1298,26 +1298,41 @@ func (c *kubeLoginCommand) selectorsOrWildcard() string {
 
 // checkClusterSelection checks the kube clusters selected by user input.
 func (c *kubeLoginCommand) checkClusterSelection(cf *CLIConf, tc *client.TeleportClient, clusters types.KubeClusters) error {
-	clusters = filterKubeClusters(c.kubeCluster, clusters)
-	switch {
-	// If the user is trying to login to a specific cluster, check that it
-	// exists and that a cluster matched the name/prefix unambiguously.
-	case c.kubeCluster != "" && len(clusters) == 1:
+	clusters = matchClustersByName(c.kubeCluster, clusters)
+	err := checkClusterSelection(cf, clusters, c.kubeCluster)
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	if c.kubeCluster != "" && len(clusters) == 1 {
 		// Populate settings using the selected kube cluster, which contains
 		// the full cluster name.
 		c.kubeCluster = clusters[0].GetName()
 		// Set CLIConf.KubernetesCluster so that the kube cluster's context
 		// is automatically selected.
 		cf.KubernetesCluster = c.kubeCluster
 		tc.KubernetesCluster = c.kubeCluster
+	}
+	return nil
+}
+
+func checkClusterSelection(cf *CLIConf, clusters types.KubeClusters, name string) error {
+	switch {
+	// If the user is trying to login to a specific cluster, check that it
+	// exists and that a cluster matched the name/prefix unambiguously.
+	case name != "" && len(clusters) == 1:
 		return nil
 	// allow multiple selection without a name.
-	case c.kubeCluster == "" && len(clusters) > 0:
+	case name == "" && len(clusters) > 0:
 		return nil
 	}
 
 	// anything else is an error.
-	selectors := c.getSelectors()
+	selectors := resourceSelectors{
+		kind:   "kubernetes cluster",
+		name:   name,
+		labels: cf.Labels,
+		query:  cf.PredicateExpression,
+	}
 	if len(clusters) == 0 {
 		if selectors.IsEmpty() {
 			return trace.NotFound("no kubernetes clusters found, check 'tsh kube ls' for a list of known clusters")
@@ -1337,7 +1352,7 @@ func (c *kubeLoginCommand) getSelectors() resourceSelectors {
 	}
 }
 
-func filterKubeClusters(nameOrPrefix string, clusters types.KubeClusters) types.KubeClusters {
+func matchClustersByName(nameOrPrefix string, clusters types.KubeClusters) types.KubeClusters {
 	if nameOrPrefix == "" {
 		return clusters
 	}

tool/tsh/common/kube_proxy.go

@@ -35,6 +35,7 @@ import (
 
 	"github.com/gravitational/teleport/api/client/proto"
 	apidefaults "github.com/gravitational/teleport/api/defaults"
+	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/api/utils/keys"
 	"github.com/gravitational/teleport/lib/asciitable"
 	"github.com/gravitational/teleport/lib/auth/native"
@@ -53,6 +54,9 @@ type proxyKubeCommand struct {
 	namespace         string
 	port              string
 	format            string
+
+	labels              string
+	predicateExpression string
 }
 
 func newProxyKubeCommand(parent *kingpin.CmdClause) *proxyKubeCommand {
@@ -68,10 +72,15 @@ func newProxyKubeCommand(parent *kingpin.CmdClause) *proxyKubeCommand {
 	c.Flag("kube-namespace", "Configure the default Kubernetes namespace.").Short('n').StringVar(&c.namespace)
 	c.Flag("port", "Specifies the source port used by the proxy listener").Short('p').StringVar(&c.port)
 	c.Flag("format", envVarFormatFlagDescription()).Short('f').Default(envVarDefaultFormat()).EnumVar(&c.format, envVarFormats...)
+	c.Flag("labels", labelHelp).StringVar(&c.labels)
+	c.Flag("query", queryHelp).StringVar(&c.predicateExpression)
 	return c
 }
 
 func (c *proxyKubeCommand) run(cf *CLIConf) error {
+	cf.Labels = c.labels
+	cf.PredicateExpression = c.predicateExpression
+	cf.SiteName = c.siteName
 	tc, err := makeClient(cf)
 	if err != nil {
 		return trace.Wrap(err)
@@ -106,22 +115,36 @@ func (c *proxyKubeCommand) run(cf *CLIConf) error {
 }
 
 func (c *proxyKubeCommand) prepare(cf *CLIConf, tc *client.TeleportClient) (*clientcmdapi.Config, kubeconfig.LocalProxyClusters, error) {
-	defaultConfig, err := kubeconfig.Load("")
+	defaultConfig, err := kubeconfig.Load(getKubeConfigPath(cf, ""))
 	if err != nil {
 		return nil, nil, trace.Wrap(err)
 	}
 
 	// Use kube clusters from arg.
-	if len(c.kubeClusters) > 0 {
-		if c.siteName == "" {
-			c.siteName = tc.SiteName
+	if len(c.kubeClusters) > 0 || cf.Labels != "" || cf.PredicateExpression != "" {
+		_, kubeClusters, err := fetchKubeClusters(cf.Context, tc)
+		if err != nil {
+			return nil, nil, trace.Wrap(err)
+		}
+		switch len(c.kubeClusters) {
+		case 0:
+			// if no names are given, check just the labels/predicate selection.
+			if err := checkClusterSelection(cf, kubeClusters, ""); err != nil {
+				return nil, nil, trace.Wrap(err)
+			}
+		default:
+			// otherwise, check that each name matches exactly one kube cluster.
+			matchMap := matchClustersByNames(kubeClusters, c.kubeClusters...)
+			if err := checkMultipleClusterSelections(cf, matchMap); err != nil {
+				return nil, nil, trace.Wrap(err)
+			}
+			kubeClusters = combineMatchedClusters(matchMap)
 		}
-
 		var clusters kubeconfig.LocalProxyClusters
-		for _, kubeCluster := range c.kubeClusters {
+		for _, kc := range kubeClusters {
 			clusters = append(clusters, kubeconfig.LocalProxyCluster{
-				TeleportCluster:   c.siteName,
-				KubeCluster:       kubeCluster,
+				TeleportCluster:   tc.SiteName,
+				KubeCluster:       kc.GetName(),
 				Impersonate:       c.impersonateUser,
 				ImpersonateGroups: c.impersonateGroups,
 				Namespace:         c.namespace,
@@ -497,6 +520,39 @@ func issueKubeCert(ctx context.Context, tc *client.TeleportClient, proxy *client
 	return cert, nil
 }
 
+// checkMultipleClusterSelections takes a map of name selectors to matched
+// clusters and checks that each matching is valid.
+func checkMultipleClusterSelections(cf *CLIConf, matchMap map[string]types.KubeClusters) error {
+	for name, clusters := range matchMap {
+		err := checkClusterSelection(cf, clusters, name)
+		if err != nil {
+			return trace.Wrap(err)
+		}
+	}
+	return nil
+}
+
+// combineMatchedClusters combineMatchedClusters takes a map from name selector
+// to matched clusters and combines all the matched clusters into a deduplicated
+// slice.
+func combineMatchedClusters(matchMap map[string]types.KubeClusters) types.KubeClusters {
+	var out types.KubeClusters
+	for _, clusters := range matchMap {
+		out = append(out, clusters...)
+	}
+	return types.DeduplicateKubeClusters(out)
+}
+
+// matchClustersByNames maps each name to the clusters it matches by exact name
+// or by prefix.
+func matchClustersByNames(clusters types.KubeClusters, names ...string) map[string]types.KubeClusters {
+	matchesForNames := make(map[string]types.KubeClusters)
+	for _, name := range names {
+		matchesForNames[name] = matchClustersByName(name, clusters)
+	}
+	return matchesForNames
+}
+
 // proxyKubeTemplate is the message that gets printed to a user when a kube proxy is started.
 var proxyKubeTemplate = template.Must(template.New("").
 	Funcs(template.FuncMap{

tool/tsh/common/kube_proxy_test.go

@@ -36,7 +36,10 @@ import (
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/api/utils/keypaths"
 	"github.com/gravitational/teleport/lib/kube/kubeconfig"
+	"github.com/gravitational/teleport/lib/service/servicecfg"
 	"github.com/gravitational/teleport/lib/srv/alpnproxy/common"
+	"github.com/gravitational/teleport/lib/utils"
+	"github.com/gravitational/teleport/tool/teleport/testenv"
 )
 
 func (p *kubeTestPack) testProxyKube(t *testing.T) {
@@ -86,6 +89,190 @@ func (p *kubeTestPack) testProxyKube(t *testing.T) {
 	})
 }
 
+func TestProxyKubeComplexSelectors(t *testing.T) {
+	testenv.WithInsecureDevMode(t, true)
+	testenv.WithResyncInterval(t, 0)
+	kubeFoo := "foo"
+	kubeFooBar := "foo-bar"
+	kubeBaz := "baz"
+	kubeFooLeaf := "foo"
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	s := newTestSuite(t,
+		withRootConfigFunc(func(cfg *servicecfg.Config) {
+			cfg.Auth.NetworkingConfig.SetProxyListenerMode(types.ProxyListenerMode_Multiplex)
+			cfg.SSH.Enabled = false
+			cfg.Kube.Enabled = true
+			cfg.Kube.ListenAddr = utils.MustParseAddr(localListenerAddr())
+			cfg.Kube.KubeconfigPath = newKubeConfigFile(t, kubeFoo, kubeFooBar, kubeBaz)
+			cfg.Kube.StaticLabels = map[string]string{"env": "root"}
+		}),
+		withLeafCluster(),
+		withLeafConfigFunc(
+			func(cfg *servicecfg.Config) {
+				cfg.Auth.NetworkingConfig.SetProxyListenerMode(types.ProxyListenerMode_Multiplex)
+				cfg.SSH.Enabled = false
+				cfg.Kube.Enabled = true
+				cfg.Kube.ListenAddr = utils.MustParseAddr(localListenerAddr())
+				cfg.Kube.KubeconfigPath = newKubeConfigFile(t, kubeFooLeaf)
+				cfg.Kube.StaticLabels = map[string]string{"env": "leaf"}
+			},
+		),
+		withValidationFunc(func(s *suite) bool {
+			rootClusters, err := s.root.GetAuthServer().GetKubernetesServers(ctx)
+			require.NoError(t, err)
+			leafClusters, err := s.leaf.GetAuthServer().GetKubernetesServers(ctx)
+			require.NoError(t, err)
+			return len(rootClusters) == 3 && len(leafClusters) == 1
+		}),
+	)
+	rootClusterName := s.root.Config.Auth.ClusterName.GetClusterName()
+	leafClusterName := s.leaf.Config.Auth.ClusterName.GetClusterName()
+
+	tests := []struct {
+		desc              string
+		makeValidateCmdFn func(*testing.T) func(*exec.Cmd) error
+		args              []string
+		wantErr           string
+	}{
+		{
+			desc: "with full name",
+			makeValidateCmdFn: func(t *testing.T) func(*exec.Cmd) error {
+				return func(cmd *exec.Cmd) error {
+					config := kubeConfigFromCmdEnv(t, cmd)
+					checkKubeLocalProxyConfig(t, s, config, rootClusterName, kubeFoo)
+					return nil
+				}
+			},
+			args: []string{kubeFoo, "--insecure"},
+		},
+		{
+			desc: "with prefix name",
+			makeValidateCmdFn: func(t *testing.T) func(*exec.Cmd) error {
+				return func(cmd *exec.Cmd) error {
+					config := kubeConfigFromCmdEnv(t, cmd)
+					checkKubeLocalProxyConfig(t, s, config, rootClusterName, kubeFooBar)
+					return nil
+				}
+			},
+			args: []string{"foo-b", "--insecure"},
+		},
+		{
+			desc: "with labels",
+			makeValidateCmdFn: func(t *testing.T) func(*exec.Cmd) error {
+				return func(cmd *exec.Cmd) error {
+					config := kubeConfigFromCmdEnv(t, cmd)
+					checkKubeLocalProxyConfig(t, s, config, rootClusterName, kubeFoo)
+					checkKubeLocalProxyConfig(t, s, config, rootClusterName, kubeFooBar)
+					checkKubeLocalProxyConfig(t, s, config, rootClusterName, kubeBaz)
+					return nil
+				}
+			},
+			args: []string{"--labels", "env=root", "--insecure"},
+		},
+		{
+			desc: "with query",
+			makeValidateCmdFn: func(t *testing.T) func(*exec.Cmd) error {
+				return func(cmd *exec.Cmd) error {
+					config := kubeConfigFromCmdEnv(t, cmd)
+					checkKubeLocalProxyConfig(t, s, config, rootClusterName, kubeFoo)
+					return nil
+				}
+			},
+			args: []string{"--query", `labels["env"]=="root"`, "--insecure"},
+		},
+		{
+			desc: "with labels, query, and prefix",
+			makeValidateCmdFn: func(t *testing.T) func(*exec.Cmd) error {
+				return func(cmd *exec.Cmd) error {
+					config := kubeConfigFromCmdEnv(t, cmd)
+					checkKubeLocalProxyConfig(t, s, config, rootClusterName, kubeFoo)
+					return nil
+				}
+			},
+			args: []string{
+				"--labels", "env=root",
+				"--query", `name == "foo"`,
+				"f", // prefix of "foo".
+				"--insecure",
+			},
+		},
+		{
+			desc: "in leaf cluster with prefix name",
+			makeValidateCmdFn: func(t *testing.T) func(*exec.Cmd) error {
+				return func(cmd *exec.Cmd) error {
+					config := kubeConfigFromCmdEnv(t, cmd)
+					checkKubeLocalProxyConfig(t, s, config, leafClusterName, kubeFooLeaf)
+					return nil
+				}
+			},
+			args: []string{
+				"--cluster", leafClusterName,
+				"--insecure",
+				"f", // prefix of "foo" kube cluster in leaf teleport cluster.
+			},
+		},
+		{
+			desc: "ambiguous name prefix is an error",
+			args: []string{
+				"f", // prefix of foo, foo-bar in root cluster.
+				"--insecure",
+			},
+			wantErr: `kubernetes cluster "f" matches multiple`,
+		},
+		{
+			desc: "zero name matches is an error",
+			args: []string{
+				"xxx",
+				"--insecure",
+			},
+			wantErr: `kubernetes cluster "xxx" not found`,
+		},
+		{
+			desc: "zero label matches is an error",
+			args: []string{
+				"--labels", "env=nonexistent",
+				"--insecure",
+			},
+			wantErr: `kubernetes cluster with labels "env=nonexistent" not found`,
+		},
+		{
+			desc: "zero query matches is an error",
+			args: []string{
+				"--query", `labels["env"]=="nonexistent"`,
+				"--insecure",
+			},
+			wantErr: `kubernetes cluster with query (labels["env"]=="nonexistent") not found`,
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.desc, func(t *testing.T) {
+			t.Parallel()
+			// login for each parallel test to avoid races when multiple tsh
+			// clients work in the same profile dir.
+			tshHome, _ := mustLogin(t, s)
+			// Set kubeconfig to a non-exist file to avoid loading other things.
+			kubeConfigPath := path.Join(tshHome, "kube-config")
+			var cmdRunner func(*exec.Cmd) error
+			if test.makeValidateCmdFn != nil {
+				cmdRunner = test.makeValidateCmdFn(t)
+			}
+			err := Run(ctx, append([]string{"proxy", "kube", "--port", ports.Pop()}, test.args...),
+				setCmdRunner(cmdRunner),
+				setHomePath(tshHome),
+				setKubeConfigPath(kubeConfigPath),
+			)
+			if test.wantErr != "" {
+				require.ErrorContains(t, err, test.wantErr)
+				return
+			}
+			require.NoError(t, err)
+		})
+	}
+}
+
 func kubeConfigFromCmdEnv(t *testing.T, cmd *exec.Cmd) *clientcmdapi.Config {
 	t.Helper()
 
@@ -118,6 +305,9 @@ func sendRequestToKubeLocalProxy(t *testing.T, config *clientcmdapi.Config, tele
 
 	contextName := kubeconfig.ContextName(teleportCluster, kubeCluster)
 
+	require.NotNil(t, config)
+	require.NotNil(t, config.Clusters)
+	require.Contains(t, config.Clusters, contextName)
 	proxyURL, err := url.Parse(config.Clusters[contextName].ProxyURL)
 	require.NoError(t, err)