Update AWS roles ARNs displayed on tsh app login for AWS console apps

tool/teleport/testenv/test_server.go

@@ -258,6 +258,10 @@ func waitForServices(t *testing.T, auth *service.TeleportProcess, cfg *servicecf
 	if cfg.Auth.Enabled && cfg.Databases.Enabled {
 		waitForDatabases(t, auth, cfg.Databases.Databases)
 	}
+
+	if cfg.Auth.Enabled && cfg.Apps.Enabled {
+		waitForApps(t, auth, cfg.Apps.Apps)
+	}
 }
 
 func waitForEvents(t *testing.T, svc service.Supervisor, events ...string) {
@@ -296,6 +300,34 @@ func waitForDatabases(t *testing.T, auth *service.TeleportProcess, dbs []service
 	}
 }
 
+func waitForApps(t *testing.T, auth *service.TeleportProcess, apps []servicecfg.App) {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	for {
+		select {
+		case <-time.After(500 * time.Millisecond):
+			all, err := auth.GetAuthServer().GetApplicationServers(ctx, apidefaults.Namespace)
+			require.NoError(t, err)
+
+			var registered int
+			for _, app := range apps {
+				for _, a := range all {
+					if a.GetName() == app.Name {
+						registered++
+						break
+					}
+				}
+			}
+
+			if registered == len(apps) {
+				return
+			}
+		case <-ctx.Done():
+			t.Fatal("Apps not registered after 10s")
+		}
+	}
+}
+
 type TestServersOpts struct {
 	Bootstrap   []types.Resource
 	ConfigFuncs []func(cfg *servicecfg.Config)

tool/tsh/common/app.go

@@ -32,6 +32,7 @@ import (
 	"github.com/gravitational/trace"
 
 	"github.com/gravitational/teleport"
+	apiclient "github.com/gravitational/teleport/api/client"
 	"github.com/gravitational/teleport/api/client/proto"
 	"github.com/gravitational/teleport/api/mfa"
 	"github.com/gravitational/teleport/api/types"
@@ -49,12 +50,36 @@ func onAppLogin(cf *CLIConf) error {
 		return trace.Wrap(err)
 	}
 
-	appInfo, err := getAppInfo(cf, tc, nil /*matchRouteToApp*/)
-	if err != nil {
+	var (
+		clusterClient *client.ClusterClient
+		appInfo       *appInfo
+		app           types.Application
+	)
+	if err := client.RetryWithRelogin(cf.Context, tc, func() error {
+		var err error
+		profile, err := tc.ProfileStatus()
+		if err != nil {
+			return trace.Wrap(err)
+		}
+
+		clusterClient, err = tc.ConnectToCluster(cf.Context)
+		if err != nil {
+			return trace.Wrap(err)
+		}
+
+		appInfo, err = getAppInfo(cf, clusterClient.AuthClient, profile, tc.SiteName, nil /*matchRouteToApp*/)
+		if err != nil {
+			return trace.Wrap(err)
+		}
+
+		app, err = appInfo.GetApp(cf.Context, clusterClient.AuthClient)
+		return trace.Wrap(err)
+	}); err != nil {
 		return trace.Wrap(err)
 	}
+	defer clusterClient.Close()
 
-	app, err := appInfo.GetApp(cf.Context, tc)
+	rootClient, err := clusterClient.ConnectToRootCluster(cf.Context)
 	if err != nil {
 		return trace.Wrap(err)
 	}
@@ -65,15 +90,6 @@ func onAppLogin(cf *CLIConf) error {
 		AccessRequests: appInfo.profile.ActiveRequests.AccessRequests,
 	}
 
-	clusterClient, err := tc.ConnectToCluster(cf.Context)
-	if err != nil {
-		return trace.Wrap(err)
-	}
-	rootClient, err := clusterClient.ConnectToRootCluster(cf.Context)
-	if err != nil {
-		return trace.Wrap(err)
-	}
-
 	key, err := appLogin(cf.Context, tc, clusterClient, rootClient, appCertParams)
 	if err != nil {
 		return trace.Wrap(err)
@@ -186,7 +202,6 @@ func printAppCommand(cf *CLIConf, tc *client.TeleportClient, app types.Applicati
 		if err != nil {
 			return trace.Wrap(err)
 		}
-
 		curlCmd, err := formatAppConfig(tc, profile, routeToApp, appFormatCURL)
 		if err != nil {
 			return trace.Wrap(err)
@@ -499,19 +514,10 @@ func serializeAppConfig(configInfo *appConfigInfo, format string) (string, error
 }
 
 // getAppInfo fetches app information using the user's tsh profile,
-// command line args, and the ListApps endpoint if necessary. If
+// command line args, and the list resources endpoint if necessary. If
 // provided, the matcher will be used to filter active apps in the
-// tsh profile. getAppInfo will also perform re-login if necessary.
-func getAppInfo(cf *CLIConf, tc *client.TeleportClient, matchRouteToApp func(tlsca.RouteToApp) bool) (*appInfo, error) {
-	var profile *client.ProfileStatus
-	if err := client.RetryWithRelogin(cf.Context, tc, func() error {
-		var err error
-		profile, err = tc.ProfileStatus()
-		return trace.Wrap(err)
-	}); err != nil {
-		return nil, trace.Wrap(err)
-	}
-
+// tsh profile.
+func getAppInfo(cf *CLIConf, clt authclient.ClientI, profile *client.ProfileStatus, siteName string, matchRouteToApp func(tlsca.RouteToApp) bool) (*appInfo, error) {
 	activeRoutes := profile.Apps
 	if matchRouteToApp != nil {
 		var filteredRoutes []tlsca.RouteToApp
@@ -534,17 +540,21 @@ func getAppInfo(cf *CLIConf, tc *client.TeleportClient, matchRouteToApp func(tls
 	}
 
 	// If we didn't find an active profile for the app, get info from server.
-	app, err := getApp(cf.Context, tc, cf.AppName)
+	app, logins, err := getApp(cf.Context, clt, cf.AppName)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
 
+	if len(logins) == 0 && app.IsAWSConsole() {
+		logins = getARNFromRoles(cf, clt, profile, siteName, app)
+	}
+
 	appInfo := &appInfo{
 		profile: profile,
 		RouteToApp: proto.RouteToApp{
 			Name:        app.GetName(),
 			PublicAddr:  app.GetPublicAddr(),
-			ClusterName: tc.SiteName,
+			ClusterName: siteName,
 			URI:         app.GetURI(),
 		},
 		app: app,
@@ -553,7 +563,7 @@ func getAppInfo(cf *CLIConf, tc *client.TeleportClient, matchRouteToApp func(tls
 	// If this is a cloud app, set additional applicable fields from CLI flags or roles.
 	switch {
 	case app.IsAWSConsole():
-		awsRoleARN, err := getARNFromFlags(cf, profile, app)
+		awsRoleARN, err := getARNFromFlags(cf, app, logins)
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
@@ -601,14 +611,14 @@ func (a *appInfo) appLocalCAPath(cluster string) string {
 
 // GetApp returns the cached app or fetches it using the app route and
 // caches the result.
-func (a *appInfo) GetApp(ctx context.Context, tc *client.TeleportClient) (types.Application, error) {
+func (a *appInfo) GetApp(ctx context.Context, clt apiclient.GetResourcesClient) (types.Application, error) {
 	a.appMu.Lock()
 	defer a.appMu.Unlock()
 	if a.app != nil {
 		return a.app.Copy(), nil
 	}
 	// holding mutex across the api call to avoid multiple redundant api calls.
-	app, err := getApp(ctx, tc, a.Name)
+	app, _, err := getApp(ctx, clt, a.Name)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -617,23 +627,30 @@ func (a *appInfo) GetApp(ctx context.Context, tc *client.TeleportClient) (types.
 }
 
 // getApp returns the registered application with the specified name.
-func getApp(ctx context.Context, tc *client.TeleportClient, name string) (app types.Application, err error) {
-	var apps []types.Application
-	err = client.RetryWithRelogin(ctx, tc, func() error {
-		apps, err = tc.ListApps(ctx, &proto.ListResourcesRequest{
-			Namespace:           tc.Namespace,
-			ResourceType:        types.KindAppServer,
-			PredicateExpression: fmt.Sprintf(`name == "%s"`, name),
-		})
-		return trace.Wrap(err)
+func getApp(ctx context.Context, clt apiclient.GetResourcesClient, name string) (app types.Application, logins []string, err error) {
+	// When listing a single app we only need to grab one page.
+	res, err := apiclient.GetEnrichedResourcePage(ctx, clt, &proto.ListResourcesRequest{
+		ResourceType:        types.KindAppServer,
+		SortBy:              types.SortBy{Field: types.ResourceMetadataName},
+		PredicateExpression: fmt.Sprintf(`name == "%s"`, name),
+		Limit:               1,
+		IncludeLogins:       true,
 	})
 	if err != nil {
-		return nil, trace.Wrap(err)
+		return nil, nil, trace.Wrap(err)
 	}
-	if len(apps) == 0 {
-		return nil, trace.NotFound("app %q not found, use `tsh apps ls` to see registered apps", name)
+
+	if len(res.Resources) == 0 {
+		return nil, nil, trace.NotFound("app %q not found, use `tsh apps ls` to see registered apps", name)
 	}
-	return apps[0], nil
+
+	appServer, ok := res.Resources[0].ResourceWithLabels.(types.AppServer)
+	if !ok {
+		log.Warnf("expected types.AppServer but received unexpected type %T", res.Resources[0].ResourceWithLabels)
+		return nil, nil, trace.NotFound("app %q not found, use `tsh apps ls` to see registered apps", name)
+	}
+
+	return appServer.GetApp(), res.Resources[0].Logins, nil
 }
 
 // pickActiveApp returns the app the current profile is logged into.

tool/tsh/common/app_aws.go

@@ -21,6 +21,7 @@ package common
 import (
 	"context"
 	"fmt"
+	"io"
 	"os"
 	"os/exec"
 	"strings"
@@ -34,6 +35,7 @@ import (
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/lib/asciitable"
 	"github.com/gravitational/teleport/lib/client"
+	"github.com/gravitational/teleport/lib/services"
 	"github.com/gravitational/teleport/lib/srv/alpnproxy"
 	"github.com/gravitational/teleport/lib/tlsca"
 	awsutils "github.com/gravitational/teleport/lib/utils/aws"
@@ -284,7 +286,7 @@ func (a *awsApp) RunCommand(cmd *exec.Cmd) error {
 	return nil
 }
 
-func printAWSRoles(roles awsutils.Roles) {
+func printAWSRoles(w io.Writer, roles awsutils.Roles) {
 	if len(roles) == 0 {
 		return
 	}
@@ -296,22 +298,22 @@ func printAWSRoles(roles awsutils.Roles) {
 		t.AddRow([]string{role.Display, role.ARN})
 	}
 
-	fmt.Println("Available AWS roles:")
-	fmt.Println(t.AsBuffer().String())
+	fmt.Fprintln(w, "Available AWS roles:")
+	fmt.Fprintln(w, t.AsBuffer().String())
 }
 
-func getARNFromFlags(cf *CLIConf, profile *client.ProfileStatus, app types.Application) (string, error) {
+func getARNFromFlags(cf *CLIConf, app types.Application, logins []string) (string, error) {
 	// Filter AWS roles by AWS account ID. If AWS account ID is empty, all
 	// roles are returned.
-	roles := awsutils.FilterAWSRoles(profile.AWSRolesARNs, app.GetAWSAccountID())
+	roles := awsutils.FilterAWSRoles(logins, app.GetAWSAccountID())
 
 	if cf.AWSRole == "" {
 		if len(roles) == 1 {
 			log.Infof("AWS Role %v is selected by default as it is the only role configured for this AWS app.", roles[0].Display)
 			return roles[0].ARN, nil
 		}
 
-		printAWSRoles(roles)
+		printAWSRoles(cf.Stdout(), roles)
 		return "", trace.BadParameter("--aws-role flag is required")
 	}
 
@@ -321,7 +323,7 @@ func getARNFromFlags(cf *CLIConf, profile *client.ProfileStatus, app types.Appli
 			return role.ARN, nil
 		}
 
-		printAWSRoles(roles)
+		printAWSRoles(cf.Stdout(), roles)
 		return "", trace.NotFound("failed to find the %q role ARN", cf.AWSRole)
 	}
 
@@ -331,15 +333,38 @@ func getARNFromFlags(cf *CLIConf, profile *client.ProfileStatus, app types.Appli
 	case 1:
 		return rolesMatched[0].ARN, nil
 	case 0:
-		printAWSRoles(roles)
+		printAWSRoles(cf.Stdout(), roles)
 		return "", trace.NotFound("failed to find the %q role name", cf.AWSRole)
 	default:
 		// Print roles matched the provided role name.
-		printAWSRoles(rolesMatched)
+		printAWSRoles(cf.Stdout(), rolesMatched)
 		return "", trace.BadParameter("provided role name %q is ambiguous, please specify full role ARN", cf.AWSRole)
 	}
 }
 
+// getARNFromRoles fetches the available AWS ARNs logins for given app.
+// If any step of fetching the roles ARNs fail, fallback into returning the
+// profile ARNs.
+//
+// TODO(gabrielcorado): DELETE IN V18.0.0
+// This is here for backward compatibility in case the auth server
+// does not support enriched resources yet.
+func getARNFromRoles(cf *CLIConf, roleGetter services.CurrentUserRoleGetter, profile *client.ProfileStatus, siteName string, app types.Application) []string {
+	accessChecker, err := services.NewAccessCheckerForRemoteCluster(cf.Context, profile.AccessInfo(), siteName, roleGetter)
+	if err != nil {
+		log.WithError(err).Debugf("Failed to fetch user roles.")
+		return profile.AWSRolesARNs
+	}
+
+	logins, err := accessChecker.GetAllowedLoginsForResource(app)
+	if err != nil {
+		log.WithError(err).Debugf("Failed to fetch app logins.")
+		return profile.AWSRolesARNs
+	}
+
+	return logins
+}
+
 func matchAWSApp(app tlsca.RouteToApp) bool {
 	return app.AWSRoleARN != ""
 }
@@ -350,8 +375,23 @@ func pickAWSApp(cf *CLIConf) (*awsApp, error) {
 		return nil, trace.Wrap(err)
 	}
 
-	appInfo, err := getAppInfo(cf, tc, matchAWSApp)
-	if err != nil {
+	var appInfo *appInfo
+	if err := client.RetryWithRelogin(cf.Context, tc, func() error {
+		var err error
+		profile, err := tc.ProfileStatus()
+		if err != nil {
+			return trace.Wrap(err)
+		}
+
+		clusterClient, err := tc.ConnectToCluster(cf.Context)
+		if err != nil {
+			return trace.Wrap(err)
+		}
+		defer clusterClient.Close()
+
+		appInfo, err = getAppInfo(cf, clusterClient.AuthClient, profile, tc.SiteName, matchAWSApp)
+		return trace.Wrap(err)
+	}); err != nil {
 		return nil, trace.Wrap(err)
 	}