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

lib/client/api.go

@@ -2550,6 +2550,82 @@ func (tc *TeleportClient) ListApps(ctx context.Context, customFilter *proto.List
 	return types.DeduplicateApps(apps), nil
 }
 
+// GetAppWithLogins returns an application with available logins.
+func (tc *TeleportClient) GetAppWithLogins(ctx context.Context, predicateExpression string) (types.Application, []string, error) {
+	ctx, span := tc.Tracer.Start(
+		ctx,
+		"teleportClient/GetAppWithLogins",
+		oteltrace.WithSpanKind(oteltrace.SpanKindClient),
+		oteltrace.WithAttributes(attribute.String("predicate", predicateExpression)),
+	)
+	defer span.End()
+
+	clt, err := tc.ConnectToCluster(ctx)
+	if err != nil {
+		return nil, nil, trace.Wrap(err)
+	}
+	defer clt.Close()
+
+	apps, err := tc.getAppsWithLogins(ctx, clt.AuthClient, predicateExpression)
+	if err != nil {
+		return nil, nil, trace.Wrap(err)
+	}
+
+	if len(apps) == 0 {
+		return nil, nil, trace.NotFound("application not found")
+	}
+
+	return apps[0].app, apps[0].logins, nil
+}
+
+// appWithLogins represents an application with its available logins.
+type appWithLogins struct {
+	app    types.Application
+	logins []string
+}
+
+// getAppsWithLogins returns a list of apps with their available logins.
+func (tc *TeleportClient) getAppsWithLogins(ctx context.Context, clt client.GetResourcesClient, predicateExpression string) ([]appWithLogins, error) {
+	ctx, span := tc.Tracer.Start(
+		ctx,
+		"teleportClient/getAppsWithLogins",
+		oteltrace.WithSpanKind(oteltrace.SpanKindClient),
+	)
+	defer span.End()
+
+	req := proto.ListResourcesRequest{
+		ResourceType:        types.KindAppServer,
+		SortBy:              types.SortBy{Field: types.ResourceMetadataName},
+		PredicateExpression: predicateExpression,
+		IncludeLogins:       true,
+	}
+
+	var apps []appWithLogins
+	for {
+		res, err := client.GetEnrichedResourcePage(ctx, clt, &req)
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+
+		for _, r := range res.Resources {
+			appServer, ok := r.ResourceWithLabels.(types.AppServer)
+			if !ok {
+				log.Warnf("expected types.AppServer but received unexpected type %T", r.ResourceWithLabels)
+				continue
+			}
+
+			apps = append(apps, appWithLogins{app: appServer.GetApp(), logins: r.Logins})
+		}
+
+		req.StartKey = res.NextKey
+		if req.StartKey == "" {
+			break
+		}
+	}
+
+	return apps, nil
+}
+
 // DeleteAppSession removes the specified application access session.
 func (tc *TeleportClient) DeleteAppSession(ctx context.Context, sessionID string) error {
 	ctx, span := tc.Tracer.Start(

lib/client/api_test.go

@@ -26,6 +26,7 @@ import (
 	"io"
 	"math"
 	"os"
+	"strconv"
 	"testing"
 	"time"
 
@@ -1258,16 +1259,35 @@ func TestIsErrorResolvableWithRelogin(t *testing.T) {
 type fakeResourceClient struct {
 	apiclient.GetResourcesClient
 
-	nodes []*types.ServerV2
+	// resources represents the pages returned by GetResources.
+	resources    [][]*proto.PaginatedResource
+	resourcesErr error
 }
 
 func (f fakeResourceClient) GetResources(ctx context.Context, req *proto.ListResourcesRequest) (*proto.ListResourcesResponse, error) {
-	out := make([]*proto.PaginatedResource, 0, len(f.nodes))
-	for _, n := range f.nodes {
-		out = append(out, &proto.PaginatedResource{Resource: &proto.PaginatedResource_Node{Node: n}})
+	if f.resourcesErr != nil {
+		return nil, f.resourcesErr
 	}
 
-	return &proto.ListResourcesResponse{Resources: out}, nil
+	if len(f.resources) == 0 {
+		return &proto.ListResourcesResponse{}, nil
+	}
+
+	pageIdx, err := strconv.Atoi(req.StartKey)
+	if req.StartKey != "" && (err != nil || pageIdx >= len(f.resources)) {
+		return &proto.ListResourcesResponse{}, nil
+	}
+
+	currPage := f.resources[pageIdx]
+	var nextKey string
+	if pageIdx+1 < len(f.resources) {
+		nextKey = strconv.Itoa(pageIdx + 1)
+	}
+
+	return &proto.ListResourcesResponse{
+		Resources: currPage,
+		NextKey:   nextKey,
+	}, nil
 }
 
 func TestGetTargetNodes(t *testing.T) {
@@ -1299,19 +1319,25 @@ func TestGetTargetNodes(t *testing.T) {
 			name:     "labels",
 			labels:   map[string]string{"foo": "bar"},
 			expected: []targetNode{{hostname: "labels", addr: "abcd:0"}},
-			clt:      fakeResourceClient{nodes: []*types.ServerV2{{Metadata: types.Metadata{Name: "abcd"}, Spec: types.ServerSpecV2{Hostname: "labels"}}}},
+			clt: fakeResourceClient{resources: [][]*proto.PaginatedResource{{
+				paginatedNode(&types.ServerV2{Metadata: types.Metadata{Name: "abcd"}, Spec: types.ServerSpecV2{Hostname: "labels"}}),
+			}}},
 		},
 		{
 			name:     "search",
 			search:   []string{"foo", "bar"},
 			expected: []targetNode{{hostname: "search", addr: "abcd:0"}},
-			clt:      fakeResourceClient{nodes: []*types.ServerV2{{Metadata: types.Metadata{Name: "abcd"}, Spec: types.ServerSpecV2{Hostname: "search"}}}},
+			clt: fakeResourceClient{resources: [][]*proto.PaginatedResource{{
+				paginatedNode(&types.ServerV2{Metadata: types.Metadata{Name: "abcd"}, Spec: types.ServerSpecV2{Hostname: "search"}}),
+			}}},
 		},
 		{
 			name:      "predicate",
 			predicate: `resource.spec.hostname == "test"`,
 			expected:  []targetNode{{hostname: "predicate", addr: "abcd:0"}},
-			clt:       fakeResourceClient{nodes: []*types.ServerV2{{Metadata: types.Metadata{Name: "abcd"}, Spec: types.ServerSpecV2{Hostname: "predicate"}}}},
+			clt: fakeResourceClient{resources: [][]*proto.PaginatedResource{{
+				paginatedNode(&types.ServerV2{Metadata: types.Metadata{Name: "abcd"}, Spec: types.ServerSpecV2{Hostname: "predicate"}}),
+			}}},
 		},
 	}
 
@@ -1334,3 +1360,78 @@ func TestGetTargetNodes(t *testing.T) {
 		})
 	}
 }
+
+func TestGetAppsWithLogin(t *testing.T) {
+	ctx := context.Background()
+
+	for name, tc := range map[string]struct {
+		clt          fakeResourceClient
+		expectedApps []appWithLogins
+		expectedErr  string
+	}{
+		"single app": {
+			clt: fakeResourceClient{resources: [][]*proto.PaginatedResource{
+				{paginatedAppServer("sample", []string{"llama", "bird"})},
+			}},
+			expectedApps: []appWithLogins{
+				{app: appWithNameURI("sample"), logins: []string{"llama", "bird"}},
+			},
+		},
+		"multiple pages": {
+			clt: fakeResourceClient{resources: [][]*proto.PaginatedResource{
+				{paginatedAppServer("first", []string{"llama", "bird"}),
+					paginatedAppServer("second", []string{"bob", "alice"})},
+				{paginatedAppServer("third", []string{}),
+					paginatedAppServer("forth", []string{"bob"})},
+			}},
+			expectedApps: []appWithLogins{
+				{app: appWithNameURI("first"), logins: []string{"llama", "bird"}},
+				{app: appWithNameURI("second"), logins: []string{"bob", "alice"}},
+				{app: appWithNameURI("third"), logins: []string{}},
+				{app: appWithNameURI("forth"), logins: []string{"bob"}},
+			},
+		},
+		"no apps": {
+			clt: fakeResourceClient{resources: [][]*proto.PaginatedResource{}},
+		},
+		"list error": {
+			clt:         fakeResourceClient{resourcesErr: errors.New("failure")},
+			expectedErr: "failure",
+		},
+	} {
+		t.Run(name, func(t *testing.T) {
+			clt := TeleportClient{Config: Config{Tracer: tracing.NoopTracer("")}}
+			// Given we're mocking the list result, the expression won't be
+			// considered.
+			res, err := clt.getAppsWithLogins(ctx, tc.clt, "")
+			if tc.expectedErr != "" {
+				require.ErrorContains(t, err, tc.expectedErr)
+				return
+			}
+
+			require.NoError(t, err)
+			require.EqualValues(t, tc.expectedApps, res)
+		})
+	}
+}
+
+func paginatedNode(node *types.ServerV2) *proto.PaginatedResource {
+	return &proto.PaginatedResource{Resource: &proto.PaginatedResource_Node{Node: node}}
+}
+
+func appWithNameURI(name string) *types.AppV3 {
+	return &types.AppV3{Metadata: types.Metadata{Name: name}, Spec: types.AppSpecV3{URI: name}}
+}
+
+func paginatedAppServer(name string, logins []string) *proto.PaginatedResource {
+	return &proto.PaginatedResource{
+		Logins: logins,
+		Resource: &proto.PaginatedResource_AppServer{
+			AppServer: &types.AppServerV3{
+				Spec: types.AppServerSpecV3{
+					App: appWithNameURI(name),
+				},
+			},
+		},
+	}
+}

tool/teleport/testenv/test_server.go

@@ -253,6 +253,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) {
@@ -291,6 +295,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

@@ -518,11 +518,15 @@ 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, tc, cf.AppName)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
 
+	if len(logins) == 0 && app.IsAWSConsole() {
+		logins = getARNFromRoles(cf, tc, profile, app)
+	}
+
 	appInfo := &appInfo{
 		profile: profile,
 		RouteToApp: proto.RouteToApp{
@@ -537,7 +541,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)
 		}
@@ -592,7 +596,7 @@ func (a *appInfo) GetApp(ctx context.Context, tc *client.TeleportClient) (types.
 		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, tc, a.Name)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -601,23 +605,20 @@ 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
+func getApp(ctx context.Context, tc *client.TeleportClient, name string) (app types.Application, logins []string, err error) {
 	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),
-		})
+		app, logins, err = tc.GetAppWithLogins(ctx, fmt.Sprintf(`name == "%s"`, name))
 		return trace.Wrap(err)
 	})
 	if err != nil {
-		return 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 trace.IsNotFound(err) {
+			return nil, nil, trace.NotFound("app %q not found, use `tsh apps ls` to see registered apps", name)
+		}
+
+		return nil, nil, trace.Wrap(err)
 	}
-	return apps[0], nil
+
+	return
 }
 
 // pickActiveApp returns the app the current profile is logged into.