From 68adee36a920500942007eb4235457dbf61da7aa Mon Sep 17 00:00:00 2001
From: Forrest Marshall <forrest@gravitational.com>
Date: Wed, 18 Nov 2020 13:23:34 -0800
Subject: [PATCH] fix tsh login with trusted clusters

---
 lib/client/api.go    | 18 ++++++++++++++++++
 lib/client/client.go | 13 +++++++++++++
 tool/tsh/tsh.go      | 26 +++++++++++++++++++-------
 3 files changed, 50 insertions(+), 7 deletions(-)

diff --git a/lib/client/api.go b/lib/client/api.go
index 152de17d3ba..fe8286b1714 100644
--- a/lib/client/api.go
+++ b/lib/client/api.go
@@ -956,6 +956,24 @@ func (tc *TeleportClient) NewWatcher(ctx context.Context, watch services.Watch)
 	return proxyClient.NewWatcher(ctx, watch)
 }
 
+// WithRootClusterClient provides a functional interface for making calls against the root cluster's
+// auth server.
+func (tc *TeleportClient) WithRootClusterClient(ctx context.Context, do func(clt auth.ClientI) error) error {
+	proxyClient, err := tc.ConnectToProxy(ctx)
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	defer proxyClient.Close()
+
+	clt, err := proxyClient.ConnectToRootCluster(ctx, false)
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	defer clt.Close()
+
+	return trace.Wrap(do(clt))
+}
+
 // SSH connects to a node and, if 'command' is specified, executes the command on it,
 // otherwise runs interactive shell
 //
diff --git a/lib/client/client.go b/lib/client/client.go
index daa507160cd..676346bf5cb 100644
--- a/lib/client/client.go
+++ b/lib/client/client.go
@@ -359,6 +359,19 @@ func (proxy *ProxyClient) ConnectToCurrentCluster(ctx context.Context, quiet boo
 	return proxy.ConnectToCluster(ctx, cluster.Name, quiet)
 }
 
+// ConnectToRootCluster connects to the auth server of the root cluster
+// cluster via proxy. It returns connected and authenticated auth server client
+//
+// if 'quiet' is set to true, no errors will be printed to stdout, otherwise
+// any connection errors are visible to a user.
+func (proxy *ProxyClient) ConnectToRootCluster(ctx context.Context, quiet bool) (auth.ClientI, error) {
+	clusterName, err := proxy.RootClusterName()
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	return proxy.ConnectToCluster(ctx, clusterName, quiet)
+}
+
 // ConnectToCluster connects to the auth server of the given cluster via proxy.
 // It returns connected and authenticated auth server client
 //
diff --git a/tool/tsh/tsh.go b/tool/tsh/tsh.go
index 4090e7cee2d..1fd2c9bebb6 100644
--- a/tool/tsh/tsh.go
+++ b/tool/tsh/tsh.go
@@ -616,14 +616,26 @@ func onLogin(cf *CLIConf) {
 			tc.Logout()
 			utils.FatalError(err)
 		}
-		for _, roleName := range roleNames {
-			role, err := tc.GetRole(cf.Context, roleName)
-			if err != nil {
-				tc.Logout()
-				utils.FatalError(err)
+		// load all roles from root cluster and collect relevant options.
+		// the normal one-off TeleportClient methods don't re-use the auth server
+		// connection, so we use WithRootClusterClient to speed things up.
+		err = tc.WithRootClusterClient(cf.Context, func(clt auth.ClientI) error {
+			for _, roleName := range roleNames {
+				role, err := clt.GetRole(roleName)
+				if err != nil {
+					return trace.Wrap(err)
+				}
+				reason = reason || role.GetOptions().RequestAccess.RequireReason()
+				auto = auto || role.GetOptions().RequestAccess.ShouldAutoRequest()
+				if prompt == "" {
+					prompt = role.GetOptions().RequestPrompt
+				}
 			}
-			reason = reason || role.GetOptions().RequestAccess.RequireReason()
-			auto = auto || role.GetOptions().RequestAccess.ShouldAutoRequest()
+			return nil
+		})
+		if err != nil {
+			tc.Logout()
+			utils.FatalError(err)
 		}
 		if reason && cf.RequestReason == "" {
 			tc.Logout()