Use Proxy gRPC API during login

lib/client/api.go

@@ -2863,6 +2863,7 @@ func (tc *TeleportClient) ConnectToCluster(ctx context.Context) (*ClusterClient,
 		AuthClient:  authClient,
 		Tracer:      tc.Tracer,
 		cluster:     cluster,
+		root:        root,
 	}, nil
 }
 
@@ -4079,7 +4080,7 @@ func (tc *TeleportClient) ssoLogin(ctx context.Context, priv *keys.PrivateKey, c
 
 // ConnectToRootCluster activates the provided key and connects to the
 // root cluster with its credentials.
-func (tc *TeleportClient) ConnectToRootCluster(ctx context.Context, key *Key) (*ProxyClient, auth.ClientI, error) {
+func (tc *TeleportClient) ConnectToRootCluster(ctx context.Context, key *Key) (*ClusterClient, auth.ClientI, error) {
 	ctx, span := tc.Tracer.Start(
 		ctx,
 		"teleportClient/ConnectToRootCluster",
@@ -4091,21 +4092,21 @@ func (tc *TeleportClient) ConnectToRootCluster(ctx context.Context, key *Key) (*
 		return nil, nil, trace.Wrap(err)
 	}
 
-	proxyClient, err := tc.ConnectToProxy(ctx)
+	clusterClient, err := tc.ConnectToCluster(ctx)
 	if err != nil {
 		return nil, nil, trace.Wrap(err)
 	}
 
-	rootAuthClient, err := proxyClient.ConnectToRootCluster(ctx)
+	rootAuthClient, err := clusterClient.ConnectToRootCluster(ctx)
 	if err != nil {
-		return nil, nil, trace.NewAggregate(err, proxyClient.Close())
+		return nil, nil, trace.NewAggregate(err, clusterClient.Close())
 	}
 
 	if err := tc.UpdateTrustedCA(ctx, rootAuthClient); err != nil {
-		return nil, nil, trace.NewAggregate(err, rootAuthClient.Close(), proxyClient.Close())
+		return nil, nil, trace.NewAggregate(err, rootAuthClient.Close(), clusterClient.Close())
 	}
 
-	return proxyClient, rootAuthClient, nil
+	return clusterClient, rootAuthClient, nil
 }
 
 // activateKey saves the target session cert into the local

lib/client/cluster_client.go

@@ -38,6 +38,7 @@ type ClusterClient struct {
 	AuthClient  auth.ClientI
 	Tracer      oteltrace.Tracer
 	cluster     string
+	root        string
 }
 
 // ClusterName returns the name of the cluster that the client
@@ -46,6 +47,34 @@ func (c *ClusterClient) ClusterName() string {
 	return c.cluster
 }
 
+// CurrentCluster returns an authenticated auth server client for the local cluster.
+// The returned auth server client does not need to be closed, it will be closed
+// when the ClusterClient is closed.
+func (c *ClusterClient) CurrentCluster() auth.ClientI {
+	// The auth.ClientI is wrapped in an sharedAuthClient to prevent callers from
+	// being able to close the client. The auth.ClientI is only to be closed
+	// when the ClusterClient is closed.
+	return sharedAuthClient{ClientI: c.AuthClient}
+}
+
+// ConnectToRootCluster connects to the auth server of the root cluster
+// via proxy. It returns connected and authenticated auth server client.
+func (c *ClusterClient) ConnectToRootCluster(ctx context.Context) (auth.ClientI, error) {
+	root, err := c.ConnectToCluster(ctx, c.root)
+	return root, trace.Wrap(err)
+}
+
+// ConnectToCluster connects to the auth server of the given cluster via proxy. It returns connected and authenticated auth server client
+func (c *ClusterClient) ConnectToCluster(ctx context.Context, clusterName string) (auth.ClientI, error) {
+	if c.cluster == clusterName {
+		return c.CurrentCluster(), nil
+	}
+
+	clientConfig := c.ProxyClient.ClientConfig(ctx, clusterName)
+	authClient, err := auth.NewClient(clientConfig)
+	return authClient, trace.Wrap(err)
+}
+
 // Close terminates the connections to Auth and Proxy.
 func (c *ClusterClient) Close() error {
 	// close auth client first since it is tunneled through the proxy client
@@ -97,6 +126,7 @@ func (c *ClusterClient) SessionSSHConfig(ctx context.Context, user string, targe
 			AuthClient:  authClient,
 			Tracer:      c.Tracer,
 			cluster:     rootClusterName,
+			root:        rootClusterName,
 		}
 		// only close the new auth client and not the copied cluster client.
 		defer authClient.Close()

tool/tsh/common/tsh.go

@@ -1903,13 +1903,13 @@ func onLogin(cf *CLIConf) error {
 	// "authoritative" source.
 	cf.Username = tc.Username
 
-	proxyClient, rootAuthClient, err := tc.ConnectToRootCluster(cf.Context, key)
+	clusterClient, rootAuthClient, err := tc.ConnectToRootCluster(cf.Context, key)
 	if err != nil {
 		return trace.Wrap(err)
 	}
 	defer func() {
 		rootAuthClient.Close()
-		proxyClient.Close()
+		clusterClient.Close()
 	}()
 
 	// TODO(fspmarshall): Refactor access request & cert reissue logic to allow
@@ -2017,7 +2017,7 @@ func onLogin(cf *CLIConf) error {
 	// don't use the alert API very heavily. If we start to make more use of it, we
 	// could probably add a separate `tsh alerts ls` command, and truncate the list
 	// with a message like "run 'tsh alerts ls' to view N additional alerts".
-	if err := common.ShowClusterAlerts(cf.Context, proxyClient.CurrentCluster(), os.Stderr, map[string]string{
+	if err := common.ShowClusterAlerts(cf.Context, clusterClient.CurrentCluster(), os.Stderr, map[string]string{
 		types.AlertOnLogin: "yes",
 	}, types.AlertSeverity_LOW); err != nil {
 		log.WithError(err).Warn("Failed to display cluster alerts.")