Disables kubeconfig update on tsh login when --kube-cluster is empty (#19694)

This PR disables the update of kubeconfig when `tsh login` does not have the `--kube-cluster` flag.

The new behavior will update the kubeconfig only if the user provided a non-empty value in  `--kube-cluster`.

This PR also introduces a mechanism to restore the current context on `tsh logout`. On `tsh login --kube-cluster=<cluster>` or `tsh kube login <cluster>`, Teleport updates the selected context - if it's not from Teleport - to add an extension whose key is: `selected_context`. Once `tsh logout` is called and the active context was a Teleport cluster, we restore the current context to the cluster that has the `selected_context` key in extensions.

Fixes #9718

lib/kube/kubeconfig/kubeconfig.go

@@ -24,6 +24,7 @@ import (
 
 	"github.com/gravitational/trace"
 	"github.com/sirupsen/logrus"
+	"golang.org/x/exp/maps"
 	"k8s.io/client-go/tools/clientcmd"
 	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
 
@@ -165,6 +166,7 @@ func Update(path string, v Values, storeAllCAs bool) error {
 			if _, ok := config.Contexts[contextName]; !ok {
 				return trace.BadParameter("can't switch kubeconfig context to cluster %q, run 'tsh kube ls' to see available clusters", v.SelectCluster)
 			}
+			setSelectedExtension(config.Contexts, config.CurrentContext, v.TeleportClusterName)
 			config.CurrentContext = contextName
 		}
 	} else {
@@ -202,6 +204,7 @@ func Update(path string, v Values, storeAllCAs bool) error {
 				ClientKeyData:         rsaKeyPEM,
 			}
 			setContext(config.Contexts, contextName, clusterName, contextName, v.Namespace)
+			setSelectedExtension(config.Contexts, config.CurrentContext, clusterName)
 			config.CurrentContext = contextName
 		} else if !trace.IsBadParameter(err) {
 			return trace.Wrap(err)
@@ -236,7 +239,7 @@ func setContext(contexts map[string]*clientcmdapi.Context, name, cluster, auth s
 //
 // If `path` is empty, Remove will try to guess it based on the environment or
 // known defaults.
-func Remove(path, name string) error {
+func Remove(path, clusterName string) error {
 	// Load existing kubeconfig from disk.
 	config, err := Load(path)
 	if err != nil {
@@ -244,17 +247,22 @@ func Remove(path, name string) error {
 	}
 
 	// Remove Teleport related AuthInfos, Clusters, and Contexts from kubeconfig.
-	delete(config.AuthInfos, name)
-	delete(config.Clusters, name)
-	delete(config.Contexts, name)
-
+	maps.DeleteFunc(
+		config.Contexts,
+		func(key string, val *clientcmdapi.Context) bool {
+			if !strings.HasPrefix(key, clusterName) {
+				return false
+			}
+			delete(config.AuthInfos, val.AuthInfo)
+			delete(config.Clusters, val.Cluster)
+			return true
+		},
+	)
+	prevSelectedCluster := searchForSelectedCluster(config.Contexts)
 	// Take an element from the list of contexts and make it the current
 	// context, unless current context points to something else.
-	if config.CurrentContext == name && len(config.Contexts) > 0 {
-		for name := range config.Contexts {
-			config.CurrentContext = name
-			break
-		}
+	if strings.HasPrefix(config.CurrentContext, clusterName) {
+		config.CurrentContext = prevSelectedCluster
 	}
 
 	// Update kubeconfig on disk.
@@ -359,9 +367,53 @@ func SelectContext(teleportCluster, kubeCluster string) error {
 	if _, ok := kc.Contexts[kubeContext]; !ok {
 		return trace.NotFound("kubeconfig context %q not found", kubeContext)
 	}
+	setSelectedExtension(kc.Contexts, kc.CurrentContext, teleportCluster)
 	kc.CurrentContext = kubeContext
 	if err := Save("", *kc); err != nil {
 		return trace.Wrap(err)
 	}
 	return nil
 }
+
+const selectedExtension = "teleport-prev-selec-ctx"
+
+// setSelectedExtension sets an extension to indentify that the current non-teleport
+// context was selected before introducing Teleport contexts in kubeconfig.
+// If the currentContext is not from Teleport, this function adds the following
+// extensions:
+//   - extension: null
+//     name: teleport-prev-selec-ctx
+//
+// Only one context is allowed to have the selected extension. If other context has it,
+// this function deletes it and introduces it in the desired context.
+func setSelectedExtension(contexts map[string]*clientcmdapi.Context, prevCluster string, teleportCluster string) {
+	selected, ok := contexts[prevCluster]
+	if !ok || strings.HasPrefix(prevCluster, teleportCluster) || len(prevCluster) == 0 {
+		return
+	}
+	for _, v := range contexts {
+		delete(v.Extensions, selectedExtension)
+	}
+
+	selected.Extensions[selectedExtension] = nil
+}
+
+// searchForSelectedCluster looks for contexts that were previously selected
+// in order to restore the the CurrentContext value.
+// If no such key is found or multiple keys exist, it returns an empty selected
+// cluster.
+func searchForSelectedCluster(contexts map[string]*clientcmdapi.Context) string {
+	count := 0
+	selected := ""
+	for k, v := range contexts {
+		if _, ok := v.Extensions[selectedExtension]; ok {
+			delete(v.Extensions, selectedExtension)
+			count++
+			selected = k
+		}
+	}
+	if count != 1 {
+		return ""
+	}
+	return selected
+}

lib/kube/kubeconfig/kubeconfig_test.go

@@ -182,6 +182,9 @@ func TestUpdate(t *testing.T) {
 	require.NoError(t, err)
 
 	wantConfig := initialConfig.DeepCopy()
+	wantConfig.Contexts[wantConfig.CurrentContext].Extensions = map[string]runtime.Object{
+		selectedExtension: nil,
+	}
 	wantConfig.Clusters[clusterName] = &clientcmdapi.Cluster{
 		Server:                   clusterAddr,
 		CertificateAuthorityData: caCertPEM,
@@ -425,6 +428,7 @@ func TestRemove(t *testing.T) {
 		clusterAddr = "https://1.2.3.6:3080"
 	)
 	kubeconfigPath, initialConfig := setup(t)
+
 	creds, _, err := genUserKey("localhost")
 	require.NoError(t, err)
 

tool/tsh/tsh.go

@@ -1487,7 +1487,7 @@ func onLogin(cf *CLIConf) error {
 			if err != nil {
 				return trace.Wrap(err)
 			}
-			if err := updateKubeConfig(cf, tc, ""); err != nil {
+			if err := updateKubeConfigOnLogin(cf, tc, ""); err != nil {
 				return trace.Wrap(err)
 			}
 			env := getTshEnv()
@@ -1510,7 +1510,7 @@ func onLogin(cf *CLIConf) error {
 
 			// Try updating kube config. If it fails, then we may have
 			// switched to an inactive profile. Continue to normal login.
-			if err := updateKubeConfig(cf, tc, ""); err == nil {
+			if err := updateKubeConfigOnLogin(cf, tc, ""); err == nil {
 				return trace.Wrap(onStatus(cf))
 			}
 
@@ -1533,7 +1533,7 @@ func onLogin(cf *CLIConf) error {
 			if err := tc.SaveProfile(true); err != nil {
 				return trace.Wrap(err)
 			}
-			if err := updateKubeConfig(cf, tc, ""); err != nil {
+			if err := updateKubeConfigOnLogin(cf, tc, ""); err != nil {
 				return trace.Wrap(err)
 			}
 
@@ -1549,7 +1549,7 @@ func onLogin(cf *CLIConf) error {
 			if err := executeAccessRequest(cf, tc); err != nil {
 				return trace.Wrap(err)
 			}
-			if err := updateKubeConfig(cf, tc, ""); err != nil {
+			if err := updateKubeConfigOnLogin(cf, tc, ""); err != nil {
 				return trace.Wrap(err)
 			}
 			return trace.Wrap(onStatus(cf))
@@ -1632,7 +1632,7 @@ func onLogin(cf *CLIConf) error {
 
 	// If the proxy is advertising that it supports Kubernetes, update kubeconfig.
 	if tc.KubeProxyAddr != "" {
-		if err := updateKubeConfig(cf, tc, ""); err != nil {
+		if err := updateKubeConfigOnLogin(cf, tc, ""); err != nil {
 			return trace.Wrap(err)
 		}
 	}
@@ -3920,7 +3920,7 @@ func reissueWithRequests(cf *CLIConf, tc *client.TeleportClient, newRequests []s
 	if err := tc.SaveProfile(true); err != nil {
 		return trace.Wrap(err)
 	}
-	if err := updateKubeConfig(cf, tc, ""); err != nil {
+	if err := updateKubeConfigOnLogin(cf, tc, ""); err != nil {
 		return trace.Wrap(err)
 	}
 	return nil
@@ -4260,3 +4260,14 @@ func forEachProfile(cf *CLIConf, fn func(tc *client.TeleportClient, profile *cli
 
 	return trace.NewAggregate(errors...)
 }
+
+// updateKubeConfigOnLogin checks if the `--kube-cluster` flag was provided to
+// tsh login call and updates the default kubeconfig with its value,
+// otherwise does nothing.
+func updateKubeConfigOnLogin(cf *CLIConf, tc *client.TeleportClient, path string) error {
+	if len(cf.KubernetesCluster) == 0 {
+		return nil
+	}
+	err := updateKubeConfig(cf, tc, "")
+	return trace.Wrap(err)
+}