Fixed regressions in "tsh login <clusterName>".

Fixed regressions in how the cluster name was set in the profile when
doing "tsh login <clusterName>". The following examples illustrate
expected behavior.

Suppose root cluster is example.com and the leaf cluster is
leaf.example.com.

* "tsh login" will update profile to example.com.
* "tsh login example.com" will update profile to example.com.
* "tsh login leaf.example.com" will update profile to
  leaf.example.com.
* "tsh login invalid.example.com" will return an error.

Note: The profile is only updated when a identity flag is NOT provided.
This means the following command will NOT update the profile:

tsh --proxy=example.com login --out=certs.pem leaf.example.com

lib/client/api.go

@@ -1730,8 +1730,9 @@ func (tc *TeleportClient) Login(ctx context.Context, activateKey bool) (*Key, er
 	if len(response.HostSigners) <= 0 {
 		return nil, trace.BadParameter("bad response from the server: expected at least one certificate, got 0")
 	}
+
+	// Add the cluster name into the key from the host certificate.
 	key.ClusterName = response.HostSigners[0].ClusterName
-	tc.SiteName = response.HostSigners[0].ClusterName
 
 	if activateKey && tc.localAgent != nil {
 		// save the list of CAs client trusts to ~/.tsh/known_hosts
@@ -1752,15 +1753,58 @@ func (tc *TeleportClient) Login(ctx context.Context, activateKey bool) (*Key, er
 			return nil, trace.Wrap(err)
 		}
 
-		// Connect to the Auth Server of the main cluster
-		// and fetch the known hosts for this cluster.
+		// Connect to the Auth Server of the main cluster and fetch the known hosts
+		// for this cluster.
 		if err := tc.UpdateTrustedCA(ctx, key.ClusterName); err != nil {
 			return nil, trace.Wrap(err)
 		}
+
+		// Update the cluster name (which will be saved in the profile) with the
+		// name of the cluster the caller requested to connect to.
+		tc.SiteName, err = updateClusterName(ctx, tc, tc.SiteName, response.HostSigners)
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
 	}
 	return key, nil
 }
 
+// certGetter is used in updateClusterName to control the response of
+// GetTrustedCA in testing.
+type certGetter interface {
+	// GetTrustedCA returns a list of trusted clusters.
+	GetTrustedCA(context.Context, string) ([]services.CertAuthority, error)
+}
+
+// updateClusterName returns the name of the cluster the user is connected to.
+func updateClusterName(ctx context.Context, certGetter certGetter, clusterName string, certificates []auth.TrustedCerts) (string, error) {
+	// Extract the name of the cluster the caller actually connected to.
+	if len(certificates) == 0 {
+		return "", trace.BadParameter("missing host certificates")
+	}
+	certificateClusterName := certificates[0].ClusterName
+
+	// The caller did not specify a cluster name, for example "tsh login", or
+	// requested the same name that is on the host certificate. In this case
+	// return the cluster name on the host certificate returned.
+	if clusterName == "" || clusterName == certificateClusterName {
+		return certificateClusterName, nil
+	}
+
+	// If the caller requested login to a leaf cluster, make sure the cluster
+	// exists.
+	leafClusters, err := certGetter.GetTrustedCA(ctx, certificateClusterName)
+	if err != nil {
+		return "", trace.Wrap(err)
+	}
+	for _, leafCluster := range leafClusters {
+		if leafCluster.GetClusterName() == clusterName {
+			return clusterName, nil
+		}
+	}
+	return "", trace.BadParameter(`unknown cluster: %q, run "tsh clusters" for a list of clusters`, clusterName)
+}
+
 // GetTrustedCA returns a list of host certificate authorities
 // trusted by the cluster client is authenticated with.
 func (tc *TeleportClient) GetTrustedCA(ctx context.Context, clusterName string) ([]services.CertAuthority, error) {

lib/client/api_test.go

@@ -17,8 +17,11 @@ limitations under the License.
 package client
 
 import (
+	"context"
 	"testing"
 
+	"github.com/gravitational/teleport/lib/auth"
+	"github.com/gravitational/teleport/lib/services"
 	"github.com/gravitational/teleport/lib/utils"
 
 	"gopkg.in/check.v1"
@@ -234,3 +237,92 @@ func (s *APITestSuite) TestDynamicPortsParsing(c *check.C) {
 		c.Assert(specs, check.DeepEquals, tt.output)
 	}
 }
+
+// TestLoginCluster makes sure the cluster name is correctly returned. This is
+// to make sure "tsh login <clusterName>" correctly updates the profile.
+func (s *APITestSuite) TestLoginCluster(c *check.C) {
+	tests := []struct {
+		inClusterName  string
+		inCertGetter   *testCertGetter
+		inCertificates []auth.TrustedCerts
+		outClusterName string
+		outError       bool
+	}{
+		// "tsh login", root cluster: example.com, leaf clusters: none.
+		{
+			inClusterName: "",
+			inCertGetter:  &testCertGetter{},
+			inCertificates: []auth.TrustedCerts{
+				auth.TrustedCerts{
+					ClusterName: "example.com",
+				},
+			},
+			outClusterName: "example.com",
+			outError:       false,
+		},
+		// "tsh login example.com", root cluster: example.com, leafClusters: none.
+		{
+			inClusterName: "example.com",
+			inCertGetter:  &testCertGetter{},
+			inCertificates: []auth.TrustedCerts{
+				auth.TrustedCerts{
+					ClusterName: "example.com",
+				},
+			},
+			outClusterName: "example.com",
+			outError:       false,
+		},
+		// "tsh login leaf.example.com", root cluster: example.com, leafClusters: [leaf.example.com].
+		{
+			inClusterName: "leaf.example.com",
+			inCertGetter: &testCertGetter{
+				clusterNames: []string{"leaf.example.com"},
+			},
+			inCertificates: []auth.TrustedCerts{
+				auth.TrustedCerts{
+					ClusterName: "example.com",
+				},
+			},
+			outClusterName: "leaf.example.com",
+			outError:       false,
+		},
+		// "tsh login invalid.example.com", root cluster: example.com, leafClusters: [leaf.example.com].
+		{
+			inClusterName: "invalid.example.com",
+			inCertGetter: &testCertGetter{
+				clusterNames: []string{"leaf.example.com"},
+			},
+			inCertificates: []auth.TrustedCerts{
+				auth.TrustedCerts{
+					ClusterName: "example.com",
+				},
+			},
+			outClusterName: "",
+			outError:       true,
+		},
+	}
+
+	for _, tt := range tests {
+		clusterName, err := updateClusterName(context.Background(), tt.inCertGetter, tt.inClusterName, tt.inCertificates)
+		c.Assert(clusterName, check.Equals, tt.outClusterName)
+		c.Assert(err != nil, check.Equals, tt.outError)
+	}
+}
+
+// testCertGetter implies the certGetter interface allowing tests to simulate
+// response from auth server.
+type testCertGetter struct {
+	clusterNames []string
+}
+
+// GetTrustedCA returns a list of trusted clusters.
+func (t *testCertGetter) GetTrustedCA(ctx context.Context, clusterName string) ([]services.CertAuthority, error) {
+	var cas []services.CertAuthority
+
+	for _, clusterName := range t.clusterNames {
+		// Only the cluster name is checked in tests, pass in nil for the keys.
+		cas = append(cas, services.NewCertAuthority(services.HostCA, clusterName, nil, nil, nil))
+	}
+
+	return cas, nil
+}