Support multi-account teleport discovery bootstrap (#56693)

This change updates teleport discovery bootstrap to be able to
bootstrap policies/ssm documents in multiple accounts at once if
the required assume_role_arns are set in the discovery service's matcher config.

lib/configurators/aws/aws.go

@@ -288,6 +288,9 @@ type awsConfigurator struct {
 	// actions list of the configurator actions, those are populated on the
 	// `build` function.
 	actions []configurators.ConfiguratorAction
+	// targetAccounts is a list of AWS account IDs that will be affected by
+	// configuration.
+	targetAccounts []string
 }
 
 type ConfiguratorConfig struct {
@@ -295,17 +298,73 @@ type ConfiguratorConfig struct {
 	Flags configurators.BootstrapFlags
 	// ServiceConfig Teleport database service config.
 	ServiceConfig *servicecfg.Config
-	// Policies instance of the `Policies` that the actions use.
-	Policies awslib.Policies
-	// Identity is the current AWS credentials chain identity.
-	Identity awslib.Identity
+	// awsConfigs is a cache of AWS configs.
+	awsConfigs *awsconfig.Cache
+	// identity is a cached AWS identity.
+	identity awslib.Identity
+	// getPolicies gets the AWS policy client for the specified assume role ARN
+	// and external ID. assumeRoleARN and externalID may be empty.
+	// Overridden in tests.
+	getPolicies func(ctx context.Context, assumeRoleARN, externalID string) (awslib.Policies, error)
+	// getIAMClient gets the AWS IAM client for the specified assume role ARN
+	// and external ID. assumeRoleARN and externalID may be empty.
+	// Overridden in tests.
+	getIAMClient func(ctx context.Context, assumeRoleARN, externalID string) (iamClient, error)
+	// getSSMClient gets the AWS SSM client for the specified assume role ARN,
+	// external ID, and region. assumeRoleARN and externalID may be empty.
+	// Overridden in tests.
+	getSSMClient func(ctx context.Context, region, assumeRoleARN, externalID string) (ssmClient, error)
+}
 
-	// awsCfg is the configuration used for AWS service clients.
-	awsCfg *aws.Config
-	// iamClient is an AWS IAM client.
-	iamClient iamClient
-	// ssmClients is a mapping of region -> AWS SSM client
-	ssmClients map[string]ssmClient
+// getAWSConfig gets the cached AWS config for the specified assume role ARN
+// and external ID. assumeRoleARN and externalID may be empty.
+func (c *ConfiguratorConfig) getAWSConfig(ctx context.Context, assumeRoleARN, externalID string) (aws.Config, error) {
+	if c.Flags.Manual {
+		return aws.Config{}, trace.BadParameter("GetAWSConfig not allowed in manual mode")
+	}
+	cfg, err := c.awsConfigs.GetConfig(
+		ctx,
+		"", /* get region from env > profile > fallback func */
+		awsconfig.WithFallbackRegionResolver(func(ctx context.Context) (string, error) {
+			return getFallbackRegion(ctx, os.Stdout, nil), nil
+		}),
+		awsconfig.WithAmbientCredentials(),
+		awsconfig.WithAssumeRole(assumeRoleARN, externalID),
+	)
+	if err != nil {
+		return aws.Config{}, trace.Wrap(err)
+	}
+	return cfg, nil
+}
+
+// getIdentity gets the cached AWS identity for the specified assume role ARN
+// and external ID. assumeRoleARN and externalID may be empty.
+func (c *ConfiguratorConfig) getIdentity(ctx context.Context, assumeRoleARN, externalID string) (awslib.Identity, error) {
+	// Assumed roles can be determined from the ARN.
+	if assumeRoleARN != "" {
+		identity, err := awslib.IdentityFromArn(assumeRoleARN)
+		return identity, trace.Wrap(err)
+	}
+	// Return a placeholder in manual mode.
+	if c.Flags.Manual {
+		identity, err := awslib.IdentityFromArn(buildIAMARN(targetIdentityARNSectionPlaceholder, targetIdentityARNSectionPlaceholder, "user", defaultAttachUser))
+		return identity, trace.Wrap(err)
+	}
+	// Check cache.
+	if c.identity != nil {
+		return c.identity, nil
+	}
+	// Fetch identity.
+	awsCfg, err := c.getAWSConfig(ctx, assumeRoleARN, externalID)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	identity, err := awslib.GetIdentityWithClient(ctx, getSTSClient(awsCfg))
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	c.identity = identity
+	return identity, nil
 }
 
 type iamClient interface {
@@ -359,94 +418,114 @@ To avoid seeing this warning, please provide a region in your AWS config or thro
 
 // CheckAndSetDefaults checks and set configuration default values.
 func (c *ConfiguratorConfig) CheckAndSetDefaults() error {
-	ctx := context.Background()
 	if c.ServiceConfig == nil {
 		return trace.BadParameter("config file is required")
 	}
-
-	// When running the command in manual mode, we want to have zero dependency
-	// with AWS configurations (like awscli or environment variables), so that
-	// the user can run this command and generate the instructions without any
-	// pre-requisite.
-	if !c.Flags.Manual {
-		var err error
-
-		if c.awsCfg == nil {
-			cfg, err := awsconfig.GetConfig(
-				ctx,
-				"", /* get region from env > profile > fallback func */
-				awsconfig.WithFallbackRegionResolver(func(ctx context.Context) (string, error) {
-					return getFallbackRegion(ctx, os.Stdout, nil), nil
-				}),
-				awsconfig.WithAmbientCredentials(),
-				awsconfig.WithAssumeRole(c.Flags.AssumeRoleARN, c.Flags.ExternalID),
-			)
-			if err != nil {
-				return trace.Wrap(err)
-			}
-			c.awsCfg = &cfg
+	if c.awsConfigs == nil {
+		cache, err := awsconfig.NewCache()
+		if err != nil {
+			return trace.Wrap(err)
 		}
-
-		if c.iamClient == nil {
-			c.iamClient = iamutils.NewFromConfig(*c.awsCfg, func(o *iam.Options) {
+		c.awsConfigs = cache
+	}
+	if c.getPolicies == nil {
+		c.getPolicies = func(ctx context.Context, assumeRoleARN, externalID string) (awslib.Policies, error) {
+			awsCfg, err := c.getAWSConfig(ctx, assumeRoleARN, externalID)
+			if err != nil {
+				return nil, trace.Wrap(err)
+			}
+			identity, err := c.getIdentity(ctx, assumeRoleARN, externalID)
+			if err != nil {
+				return nil, trace.Wrap(err)
+			}
+			iamClient := iamutils.NewFromConfig(awsCfg, func(o *iam.Options) {
 				o.TracerProvider = smithyoteltracing.Adapt(otel.GetTracerProvider())
 			})
+			partition := identity.GetPartition()
+			accountID := identity.GetAccountID()
+			return awslib.NewPolicies(partition, accountID, iamClient), nil
 		}
-		if c.Identity == nil {
-			c.Identity, err = awslib.GetIdentityWithClient(ctx, getSTSClient(*c.awsCfg))
+	}
+	if c.getIAMClient == nil {
+		c.getIAMClient = func(ctx context.Context, assumeRoleARN, externalID string) (iamClient, error) {
+			awsCfg, err := c.getAWSConfig(ctx, assumeRoleARN, externalID)
 			if err != nil {
-				return trace.Wrap(err)
+				return nil, trace.Wrap(err)
 			}
+			return iamutils.NewFromConfig(awsCfg, func(o *iam.Options) {
+				o.TracerProvider = smithyoteltracing.Adapt(otel.GetTracerProvider())
+			}), nil
 		}
-
-		if c.ssmClients == nil {
-			c.ssmClients = make(map[string]ssmClient)
-			for _, matcher := range c.ServiceConfig.Discovery.AWSMatchers {
-				if !slices.Contains(matcher.Types, types.AWSMatcherEC2) {
-					continue
-				}
-				for _, region := range matcher.Regions {
-					if _, ok := c.ssmClients[region]; ok {
-						continue
-					}
-					withRegion := func(o *ssm.Options) {
-						o.Region = region
-					}
-					c.ssmClients[region] = ssm.NewFromConfig(*c.awsCfg, withRegion, func(o *ssm.Options) {
-						o.TracerProvider = smithyoteltracing.Adapt(otel.GetTracerProvider())
-					})
-				}
+	}
+	if c.getSSMClient == nil {
+		c.getSSMClient = func(ctx context.Context, region, assumeRoleARN, externalID string) (ssmClient, error) {
+			awsCfg, err := c.getAWSConfig(ctx, assumeRoleARN, externalID)
+			if err != nil {
+				return nil, trace.Wrap(err)
 			}
-
-		}
-
-		if c.Policies == nil {
-			partition := c.Identity.GetPartition()
-			accountID := c.Identity.GetAccountID()
-			iamClient := iamutils.NewFromConfig(*c.awsCfg, func(o *iam.Options) {
+			return ssm.NewFromConfig(awsCfg, func(o *ssm.Options) {
+				o.Region = region
 				o.TracerProvider = smithyoteltracing.Adapt(otel.GetTracerProvider())
-			})
-			c.Policies = awslib.NewPolicies(partition, accountID, iamClient)
+			}), nil
 		}
 	}
-
 	return nil
 }
 
+// getDistinctAssumedRoles gets a list of the distinct roles that can be assumed
+// from the AWS matchers. If there are no AWS matchers, one AssumeRole
+// is returned for the current identity.
+func (c *ConfiguratorConfig) getDistinctAssumedRoles() []types.AssumeRole {
+	defaultAssumeRole := types.AssumeRole{
+		RoleARN:    c.Flags.AssumeRoleARN,
+		ExternalID: c.Flags.ExternalID,
+	}
+	matchers := awsMatchersFromConfig(c.Flags, c.ServiceConfig)
+	if len(matchers) == 0 {
+		return []types.AssumeRole{defaultAssumeRole}
+	}
+
+	assumedRoles := make([]types.AssumeRole, 0, len(matchers))
+	for _, matcher := range matchers {
+		if ar := matcher.AssumeRole; ar == nil {
+			assumedRoles = append(assumedRoles, defaultAssumeRole)
+		} else {
+			assumedRoles = append(assumedRoles, *ar)
+		}
+	}
+	return apiutils.DeduplicateAny(assumedRoles, func(ar1, ar2 types.AssumeRole) bool {
+		return ar1.RoleARN == ar2.RoleARN && ar1.ExternalID == ar2.ExternalID
+	})
+}
+
 // NewAWSConfigurator creates an instance of awsConfigurator and builds its
 // actions.
-func NewAWSConfigurator(config ConfiguratorConfig) (configurators.Configurator, error) {
+func NewAWSConfigurator(ctx context.Context, config ConfiguratorConfig) (configurators.Configurator, error) {
 	err := config.CheckAndSetDefaults()
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
 
-	actions, err := buildActions(config)
+	actions, err := buildActions(ctx, config)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
 
-	return &awsConfigurator{config, actions}, nil
+	assumedRoles := config.getDistinctAssumedRoles()
+	targetAccounts := make([]string, 0, len(assumedRoles))
+	for _, ar := range assumedRoles {
+		identity, err := config.getIdentity(ctx, ar.RoleARN, ar.ExternalID)
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+		targetAccounts = append(targetAccounts, identity.GetAccountID())
+	}
+
+	return &awsConfigurator{
+		config:         config,
+		actions:        actions,
+		targetAccounts: apiutils.Deduplicate(targetAccounts),
+	}, nil
 }
 
 // IsEmpty checks if the configurator has no actions.
@@ -461,7 +540,7 @@ func (a *awsConfigurator) Name() string {
 
 // Description returns a brief description of the configurator.
 func (a *awsConfigurator) Description() string {
-	return "Configure AWS for " + a.config.Flags.Service.Name()
+	return fmt.Sprintf("Configure AWS for %s for accounts: [%s]", a.config.Flags.Service.Name(), strings.Join(a.targetAccounts, ", "))
 }
 
 // Actions list of configurator actions.
@@ -478,11 +557,13 @@ type awsPolicyCreator struct {
 	policy *awslib.Policy
 	// formattedPolicy human-readable representation of the policy document.
 	formattedPolicy string
+	// accountID is the account that the policy will be created in.
+	accountID string
 }
 
 // Description returns what the action will perform.
 func (a *awsPolicyCreator) Description() string {
-	return fmt.Sprintf("Create IAM Policy %q", a.policy.Name)
+	return fmt.Sprintf("[%s] Create IAM Policy %q", a.accountID, a.policy.Name)
 }
 
 // Details returns the policy document that will be created.
@@ -515,7 +596,7 @@ type awsPoliciesAttacher struct {
 
 // Description returns what the action will perform.
 func (a *awsPoliciesAttacher) Description() string {
-	return fmt.Sprintf("Attach IAM policies to %q", a.target.GetName())
+	return fmt.Sprintf("[%s] Attach IAM policies to %q", a.target.GetAccountID(), a.target.GetName())
 }
 
 // Details attacher doesn't have any extra detail, this function returns an
@@ -543,8 +624,8 @@ func (a *awsPoliciesAttacher) Execute(ctx context.Context, actionCtx *configurat
 	return nil
 }
 
-func buildDiscoveryActions(config ConfiguratorConfig, targetCfg targetConfig) ([]configurators.ConfiguratorAction, error) {
-	actions, err := buildCommonActions(config, targetCfg)
+func buildDiscoveryActions(ctx context.Context, config ConfiguratorConfig, targetCfg targetConfig) ([]configurators.ConfiguratorAction, error) {
+	actions, err := buildCommonActions(ctx, config, targetCfg)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -554,11 +635,16 @@ func buildDiscoveryActions(config ConfiguratorConfig, targetCfg targetConfig) ([
 		return nil, err
 	}
 
-	actions = append(actions, buildSSMDocumentCreators(config.ssmClients, targetCfg, proxyAddr)...)
+	ssmActions, err := buildSSMDocumentCreators(ctx, config, targetCfg, proxyAddr)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	actions = append(actions, ssmActions...)
 	return actions, nil
 }
 
-func buildCommonActions(config ConfiguratorConfig, targetCfg targetConfig) ([]configurators.ConfiguratorAction, error) {
+func buildCommonActions(ctx context.Context, config ConfiguratorConfig, targetCfg targetConfig) ([]configurators.ConfiguratorAction, error) {
 	// Generate policies.
 	policy, err := buildPolicyDocument(config.Flags, targetCfg)
 	if err != nil {
@@ -577,81 +663,132 @@ func buildCommonActions(config ConfiguratorConfig, targetCfg targetConfig) ([]co
 		return nil, trace.Wrap(err)
 	}
 	var actions []configurators.ConfiguratorAction
+	var policies awslib.Policies
+	if !config.Flags.Manual {
+		policies, err = config.getPolicies(ctx, targetCfg.assumeRole.RoleARN, targetCfg.assumeRole.ExternalID)
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+	}
 
 	// Create IAM Policy.
 	actions = append(actions, &awsPolicyCreator{
-		policies:        config.Policies,
+		policies:        policies,
 		policy:          policy,
 		formattedPolicy: formattedPolicy,
+		accountID:       targetCfg.identity.GetAccountID(),
 	})
 
 	// Attach the policy to the target.
-	actions = append(actions, &awsPoliciesAttacher{policies: config.Policies, target: targetCfg.identity})
+	actions = append(actions, &awsPoliciesAttacher{policies: policies, target: targetCfg.identity})
 	return actions, nil
 }
 
 // buildActions generates the policy documents and configurator actions.
-func buildActions(config ConfiguratorConfig) ([]configurators.ConfiguratorAction, error) {
-	// Identity is going to be empty (`nil`) when running the command on
-	// `Manual` mode, place a wildcard to keep the generated policies valid.
-	accountID := targetIdentityARNSectionPlaceholder
-	partitionID := targetIdentityARNSectionPlaceholder
-	if config.Identity != nil {
-		accountID = config.Identity.GetAccountID()
-		partitionID = config.Identity.GetPartition()
+func buildActions(ctx context.Context, config ConfiguratorConfig) ([]configurators.ConfiguratorAction, error) {
+	var allActions []configurators.ConfiguratorAction
+	for _, assumeRole := range config.getDistinctAssumedRoles() {
+		target, err := policiesTarget(ctx, config, assumeRole)
+		if err != nil {
+			var unreachableErr unreachablePolicyTargetError
+			if errors.As(err, &unreachableErr) {
+				fmt.Printf("⚠️ Skipping matchers with identity %q: %s\n", unreachableErr.from.GetName(), unreachableErr.Error())
+				continue
+			}
+			return nil, trace.Wrap(err)
+		}
+		targetCfg, err := getTargetConfig(config.Flags, config.ServiceConfig, target, assumeRole)
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+		var actions []configurators.ConfiguratorAction
+		if config.Flags.Service.IsDiscovery() {
+			actions, err = buildDiscoveryActions(ctx, config, targetCfg)
+		} else {
+			actions, err = buildCommonActions(ctx, config, targetCfg)
+		}
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+		allActions = append(allActions, actions...)
 	}
 
-	// Define the target and target type.
-	target, err := policiesTarget(config.Flags, accountID, partitionID, config.Identity, config.iamClient)
-	if err != nil {
-		return nil, trace.Wrap(err)
-	}
+	return allActions, nil
+}
 
-	targetCfg, err := getTargetConfig(config.Flags, config.ServiceConfig, target)
-	if err != nil {
-		return nil, trace.Wrap(err)
-	}
+// unreachablePolicyTargetError indicates that a target identity could not be
+// accessed from another identity (typically due to them being in different
+// accounts).
+type unreachablePolicyTargetError struct {
+	target awslib.Identity
+	from   awslib.Identity
+}
 
-	if config.Flags.Service.IsDiscovery() {
-		return buildDiscoveryActions(config, targetCfg)
-	}
-	return buildCommonActions(config, targetCfg)
+func (e unreachablePolicyTargetError) Error() string {
+	return fmt.Sprintf(
+		"%q is unreachable from %q",
+		e.target, e.from,
+	)
 }
 
 // policiesTarget defines which target and its type the policies will be
 // attached to.
-func policiesTarget(flags configurators.BootstrapFlags, accountID string, partitionID string, identity awslib.Identity, iamClient iamClient) (awslib.Identity, error) {
-	if flags.AttachToUser != "" {
-		userArn := flags.AttachToUser
-		if !arn.IsARN(flags.AttachToUser) {
-			userArn = buildIAMARN(partitionID, accountID, "user", flags.AttachToUser)
+func policiesTarget(ctx context.Context, config ConfiguratorConfig, targetAssumeRole types.AssumeRole) (awslib.Identity, error) {
+	baseIdentity, err := config.getIdentity(ctx, targetAssumeRole.RoleARN, targetAssumeRole.ExternalID)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	defaultPartitionID := baseIdentity.GetPartition()
+	defaultAccountID := baseIdentity.GetAccountID()
+
+	// Attach to user if provided.
+	attachToUser := config.Flags.AttachToUser
+	if attachToUser != "" {
+		userArn := attachToUser
+		if !arn.IsARN(attachToUser) {
+			userArn = buildIAMARN(defaultPartitionID, defaultAccountID, "user", attachToUser)
 		}
-
-		return awslib.IdentityFromArn(userArn)
-	}
-
-	if flags.AttachToRole != "" {
-		roleArn := flags.AttachToRole
-		if !arn.IsARN(flags.AttachToRole) {
-			roleArn = buildIAMARN(partitionID, accountID, "role", flags.AttachToRole)
+		userIdentity, err := awslib.IdentityFromArn(userArn)
+		if err != nil {
+			return nil, trace.Wrap(err)
 		}
-
-		return awslib.IdentityFromArn(roleArn)
+		if defaultAccountID != userIdentity.GetAccountID() {
+			return nil, unreachablePolicyTargetError{target: userIdentity, from: baseIdentity}
+		}
+		return userIdentity, nil
 	}
 
-	if identity == nil {
-		return awslib.IdentityFromArn(buildIAMARN(partitionID, accountID, "user", defaultAttachUser))
+	// Attach to role if provided.
+	attachToRole := config.Flags.AttachToRole
+	if attachToRole != "" {
+		roleArn := attachToRole
+		if !arn.IsARN(attachToRole) {
+			roleArn = buildIAMARN(defaultPartitionID, defaultAccountID, "role", attachToRole)
+		}
+		roleIdentity, err := awslib.IdentityFromArn(roleArn)
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+		if defaultAccountID != roleIdentity.GetAccountID() {
+			return nil, unreachablePolicyTargetError{target: roleIdentity, from: baseIdentity}
+		}
+		return roleIdentity, nil
 	}
 
-	if identity.GetType() == awslib.ResourceTypeAssumedRole {
-		roleIdentity, err := getRoleARNForAssumedRole(iamClient, identity)
+	// Attach to current identity.
+	if baseIdentity.GetType() == awslib.ResourceTypeAssumedRole {
+		baseIAMClient, err := config.getIAMClient(ctx, targetAssumeRole.RoleARN, targetAssumeRole.ExternalID)
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+		roleIdentity, err := getRoleARNForAssumedRole(baseIAMClient, baseIdentity)
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
 		return roleIdentity, nil
 	}
 
-	return identity, nil
+	return baseIdentity, nil
 }
 
 // buildIAMARN constructs an AWS IAM ARN string from the given partition,
@@ -686,6 +823,9 @@ func failedToResolveAssumeRoleARN(roleIdentity string) string {
 // This is necessary since the assumed-role ARN does not include the role path,
 // so we cannot reliably reconstruct the role ARN from the assumed-role ARN.
 func getRoleARNForAssumedRole(iamClient iamClient, identity awslib.Identity) (awslib.Identity, error) {
+	if iamClient == nil {
+		return nil, trace.BadParameter("missing iamClient")
+	}
 	failedToResolveAssumeRoleARN := failedToResolveAssumeRoleARN(identity.GetName())
 
 	out, err := iamClient.GetRole(context.Background(), &iam.GetRoleInput{
@@ -810,22 +950,31 @@ func getProxyAddrFromConfig(cfg *servicecfg.Config, flags configurators.Bootstra
 	return "", trace.NotFound("proxy address not found, please provide --proxy, or set either teleport.proxy_server or proxy_service.public_addr in the teleport config")
 }
 
-func buildSSMDocumentCreators(ssm map[string]ssmClient, targetCfg targetConfig, proxyAddr string) []configurators.ConfiguratorAction {
+func buildSSMDocumentCreators(ctx context.Context, config ConfiguratorConfig, targetCfg targetConfig, proxyAddr string) ([]configurators.ConfiguratorAction, error) {
 	var creators []configurators.ConfiguratorAction
 	for _, matcher := range targetCfg.awsMatchers {
 		if !slices.Contains(matcher.Types, types.AWSMatcherEC2) {
 			continue
 		}
 		for _, region := range matcher.Regions {
+			var ssmClient ssmClient
+			if !config.Flags.Manual {
+				var err error
+				ssmClient, err = config.getSSMClient(ctx, region, targetCfg.assumeRole.RoleARN, targetCfg.assumeRole.ExternalID)
+				if err != nil {
+					return nil, trace.Wrap(err)
+				}
+			}
 			ssmCreator := awsSSMDocumentCreator{
-				ssm:      ssm[region],
-				Name:     matcher.SSM.DocumentName,
-				Contents: awslib.EC2DiscoverySSMDocument(proxyAddr),
+				ssm:       ssmClient,
+				Name:      matcher.SSM.DocumentName,
+				Contents:  awslib.EC2DiscoverySSMDocument(proxyAddr),
+				accountID: targetCfg.identity.GetAccountID(),
 			}
 			creators = append(creators, &ssmCreator)
 		}
 	}
-	return creators
+	return creators, nil
 }
 
 func isEC2AutoDiscoveryEnabled(flags configurators.BootstrapFlags, matchers []types.AWSMatcher) bool {
@@ -1104,14 +1253,15 @@ func buildARN(target awslib.Identity, service, resource string) string {
 }
 
 type awsSSMDocumentCreator struct {
-	Contents string
-	ssm      ssmClient
-	Name     string
+	Contents  string
+	ssm       ssmClient
+	Name      string
+	accountID string
 }
 
 // Description returns what the action will perform.
 func (a *awsSSMDocumentCreator) Description() string {
-	return fmt.Sprintf("Create SSM Document %q", a.Name)
+	return fmt.Sprintf("[%s] Create SSM Document %q", a.accountID, a.Name)
 }
 
 // Details returns the policy document that will be created.
@@ -1121,6 +1271,9 @@ func (a *awsSSMDocumentCreator) Details() string {
 
 // Execute upserts the policy and store its ARN in the action context.
 func (a *awsSSMDocumentCreator) Execute(ctx context.Context, actionCtx *configurators.ConfiguratorActionContext) error {
+	if a.ssm == nil {
+		return trace.BadParameter("ssm client not initialized")
+	}
 	_, err := a.ssm.CreateDocument(ctx, &ssm.CreateDocumentInput{
 		Content:        aws.String(a.Contents),
 		Name:           aws.String(a.Name),
@@ -1146,6 +1299,8 @@ func (a *awsSSMDocumentCreator) Execute(ctx context.Context, actionCtx *configur
 type targetConfig struct {
 	// identity is the target identity.
 	identity awslib.Identity
+	// assumeRole is the role that should be assumed while configuring the target, if any.
+	assumeRole types.AssumeRole
 	// awsMatchers are the AWS matchers associated with the target identity.
 	awsMatchers []types.AWSMatcher
 	// databases are the databases associated with the target identity.
@@ -1157,7 +1312,7 @@ type targetConfig struct {
 
 // getTargetConfig gets the resources that are relevant to the target identity
 // from cli flags and file configuration.
-func getTargetConfig(flags configurators.BootstrapFlags, cfg *servicecfg.Config, target awslib.Identity) (targetConfig, error) {
+func getTargetConfig(flags configurators.BootstrapFlags, cfg *servicecfg.Config, target awslib.Identity, assumeRole types.AssumeRole) (targetConfig, error) {
 	forcedRoles, err := parseForcedAWSRoles(flags, target)
 	if err != nil {
 		return targetConfig{}, trace.Wrap(err)
@@ -1173,6 +1328,7 @@ func getTargetConfig(flags configurators.BootstrapFlags, cfg *servicecfg.Config,
 	}
 	return targetConfig{
 		identity:        target,
+		assumeRole:      assumeRole,
 		awsMatchers:     matchersForTarget(awsMatchers, target, targetIsAssumeRole),
 		databases:       databasesForTarget(databases, target, targetIsAssumeRole),
 		assumesAWSRoles: targetAssumesRoles,

lib/configurators/aws/aws_test.go

@@ -34,6 +34,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/gravitational/trace"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/gravitational/teleport/api/types"
@@ -52,7 +53,67 @@ var sortStringsTrans = cmp.Transformer("SortStrings", func(in []string) []string
 	return out
 })
 
+func TestGetIdentity(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name           string
+		config         ConfiguratorConfig
+		roleARN        string
+		externalID     string
+		assert         assert.ErrorAssertionFunc
+		expectIdentity awslib.Identity
+	}{
+		{
+			name: "identity from assume role",
+			config: ConfiguratorConfig{
+				Flags: configurators.BootstrapFlags{
+					Manual: true,
+				},
+				identity: identityFromArn(t, "arn:aws:iam::123456789012:role/not-this-one"),
+			},
+			roleARN:        "arn:aws:iam::123456789012:role/example-role",
+			externalID:     "foobar",
+			assert:         assert.NoError,
+			expectIdentity: identityFromArn(t, "arn:aws:iam::123456789012:role/example-role"),
+		},
+		{
+			name: "placeholder identity in manual mode",
+			config: ConfiguratorConfig{
+				Flags: configurators.BootstrapFlags{
+					Manual: true,
+				},
+				identity: identityFromArn(t, "arn:aws:iam::123456789012:role/not-this-one"),
+			},
+			assert:         assert.NoError,
+			expectIdentity: identityFromArn(t, buildIAMARN(targetIdentityARNSectionPlaceholder, targetIdentityARNSectionPlaceholder, "user", defaultAttachUser)),
+		},
+		{
+			name: "cached identity",
+			config: ConfiguratorConfig{
+				identity: identityFromArn(t, "arn:aws:iam::123456789012:role/example-role"),
+			},
+			assert:         assert.NoError,
+			expectIdentity: identityFromArn(t, "arn:aws:iam::123456789012:role/example-role"),
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			identity, err := tc.config.getIdentity(t.Context(), tc.roleARN, tc.externalID)
+			tc.assert(t, err)
+			if tc.expectIdentity == nil {
+				assert.Nil(t, identity)
+				return
+			}
+			if assert.NotNil(t, identity) {
+				assert.Equal(t, tc.expectIdentity.String(), identity.String())
+				assert.Equal(t, tc.expectIdentity.GetType(), identity.GetType())
+			}
+		})
+	}
+}
+
 func TestAWSIAMDocuments(t *testing.T) {
+	t.Parallel()
 	userTarget, err := awslib.IdentityFromArn("arn:aws:iam::123456789012:user/example-user")
 	require.NoError(t, err)
 
@@ -1085,7 +1146,7 @@ func mustGetTargetConfig(t *testing.T, flags configurators.BootstrapFlags, roleT
 	require.NoError(t, fileConfig.CheckAndSetDefaults())
 	serviceCfg := servicecfg.MakeDefaultConfig()
 	require.NoError(t, config.ApplyFileConfig(fileConfig, serviceCfg))
-	targetCfg, err := getTargetConfig(flags, serviceCfg, roleTarget)
+	targetCfg, err := getTargetConfig(flags, serviceCfg, roleTarget, types.AssumeRole{})
 	require.NoError(t, err)
 	return targetCfg
 }
@@ -1099,7 +1160,8 @@ func mustBuildPolicyDocument(t *testing.T, flags configurators.BootstrapFlags, t
 }
 
 func TestAWSPolicyCreator(t *testing.T) {
-	ctx := context.Background()
+	t.Parallel()
+	ctx := t.Context()
 
 	tests := map[string]struct {
 		returnError bool
@@ -1137,7 +1199,8 @@ func TestAWSPolicyCreator(t *testing.T) {
 }
 
 func TestAWSPoliciesAttacher(t *testing.T) {
-	ctx := context.Background()
+	t.Parallel()
+	ctx := t.Context()
 	userTarget, err := awslib.IdentityFromArn("arn:aws:iam::1234567:user/example-user")
 	require.NoError(t, err)
 
@@ -1196,33 +1259,56 @@ func TestAWSPoliciesAttacher(t *testing.T) {
 	}
 }
 
+func makeIAMClientGetter(expectedRoleARN, expectedExternalID string, clt iamClient) func(ctx context.Context, assumeRoleARN, externalID string) (iamClient, error) {
+	return func(ctx context.Context, assumeRoleARN, externalID string) (iamClient, error) {
+		if assumeRoleARN != expectedRoleARN || externalID != expectedExternalID {
+			return nil, trace.NotFound("no IAM client for assume role %q with external ID %q", expectedRoleARN, expectedExternalID)
+		}
+		return clt, nil
+	}
+}
+
+func makePoliciesGetter(expectedRoleARN, expectedExternalID string, policies awslib.Policies) func(ctx context.Context, assumeRoleARN, externalID string) (awslib.Policies, error) {
+	return func(ctx context.Context, assumeRoleARN, externalID string) (awslib.Policies, error) {
+		if assumeRoleARN != expectedRoleARN || externalID != expectedExternalID {
+			return nil, trace.NotFound("no policies client for assume role %q with external ID %q", expectedRoleARN, expectedExternalID)
+		}
+		return policies, nil
+	}
+}
+
+func makeSSMClientGetter(expectedRegion, expectedRoleARN, expectedExternalID string, ssm ssmClient) func(ctx context.Context, region, assumeRoleARN, externalID string) (ssmClient, error) {
+	return func(ctx context.Context, region, assumeRoleARN, externalID string) (ssmClient, error) {
+		if region != expectedRegion || assumeRoleARN != expectedRoleARN || externalID != expectedExternalID {
+			return nil, trace.NotFound("no IAM client for assume role %q with external ID %q", expectedRoleARN, expectedExternalID)
+		}
+		return ssm, nil
+	}
+}
+
 func TestAWSPoliciesTarget(t *testing.T) {
-	userIdentity, err := awslib.IdentityFromArn("arn:aws:iam::123456789012:user/example-user")
-	require.NoError(t, err)
-
-	roleIdentity, err := awslib.IdentityFromArn("arn:aws:iam::123456789012:role/example-role")
-	require.NoError(t, err)
-
-	assumedRoleIdentity, err := awslib.IdentityFromArn("arn:aws:sts::123456789012:assumed-role/example-role/i-12345")
-	require.NoError(t, err)
+	t.Parallel()
+	userIdentity := identityFromArn(t, "arn:aws:iam::123456789012:user/example-user")
+	roleIdentity := identityFromArn(t, "arn:aws:iam::123456789012:role/example-role")
+	assumedRoleIdentity := identityFromArn(t, "arn:aws:sts::123456789012:assumed-role/example-role/i-12345")
+	altAssumedRoleIdentity := identityFromArn(t, "arn:aws:sts::123456789012:assumed-role/alternate-role/i-12345")
 
+	defaultIdentity := identityFromArn(t, "arn:aws:iam::123456789012:user/me")
 	tests := map[string]struct {
-		flags             configurators.BootstrapFlags
-		identity          awslib.Identity
-		accountID         string
-		partitionID       string
+		config            ConfiguratorConfig
+		assumeRole        types.AssumeRole
 		targetType        awslib.Identity
 		targetName        string
 		targetAccountID   string
 		targetPartitionID string
 		targetString      string
-		iamClient         iamClient
 		wantErrContains   string
 	}{
 		"UserNameFromFlags": {
-			flags:             configurators.BootstrapFlags{AttachToUser: "example-user"},
-			accountID:         "123456",
-			partitionID:       "aws",
+			config: ConfiguratorConfig{
+				Flags:    configurators.BootstrapFlags{AttachToUser: "example-user"},
+				identity: identityFromArn(t, buildIAMARN("aws", "123456", "user", defaultAttachUser)),
+			},
 			targetType:        awslib.User{},
 			targetName:        "example-user",
 			targetAccountID:   "123456",
@@ -1230,9 +1316,10 @@ func TestAWSPoliciesTarget(t *testing.T) {
 			targetString:      "arn:aws:iam::123456:user/example-user",
 		},
 		"UserNameWithPathFromFlags": {
-			flags:             configurators.BootstrapFlags{AttachToUser: "/some/path/example-user"},
-			accountID:         "123456",
-			partitionID:       "aws",
+			config: ConfiguratorConfig{
+				Flags:    configurators.BootstrapFlags{AttachToUser: "/some/path/example-user"},
+				identity: identityFromArn(t, buildIAMARN("aws", "123456", "user", defaultAttachUser)),
+			},
 			targetType:        awslib.User{},
 			targetName:        "example-user",
 			targetAccountID:   "123456",
@@ -1240,17 +1327,31 @@ func TestAWSPoliciesTarget(t *testing.T) {
 			targetString:      "arn:aws:iam::123456:user/some/path/example-user",
 		},
 		"UserARNFromFlags": {
-			flags:             configurators.BootstrapFlags{AttachToUser: "arn:aws:iam::123456789012:user/example-user"},
+			config: ConfiguratorConfig{
+				Flags:    configurators.BootstrapFlags{AttachToUser: "arn:aws:iam::123456789012:user/example-user"},
+				identity: defaultIdentity,
+			},
 			targetType:        awslib.User{},
 			targetName:        "example-user",
 			targetAccountID:   "123456789012",
 			targetPartitionID: "aws",
 			targetString:      "arn:aws:iam::123456789012:user/example-user",
 		},
+		"UserARNFromFlagsWrongAccount": {
+			config: ConfiguratorConfig{
+				Flags:    configurators.BootstrapFlags{AttachToUser: "arn:aws:iam::987654321098:user/alt-user"},
+				identity: defaultIdentity,
+			},
+			wantErrContains: unreachablePolicyTargetError{
+				target: identityFromArn(t, "arn:aws:iam::987654321098:user/alt-user"),
+				from:   defaultIdentity,
+			}.Error(),
+		},
 		"RoleNameFromFlags": {
-			flags:             configurators.BootstrapFlags{AttachToRole: "example-role"},
-			accountID:         "123456789012",
-			partitionID:       "aws",
+			config: ConfiguratorConfig{
+				Flags:    configurators.BootstrapFlags{AttachToRole: "example-role"},
+				identity: identityFromArn(t, buildIAMARN("aws", "123456789012", "user", defaultAttachUser)),
+			},
 			targetType:        awslib.Role{},
 			targetName:        "example-role",
 			targetAccountID:   "123456789012",
@@ -1258,9 +1359,10 @@ func TestAWSPoliciesTarget(t *testing.T) {
 			targetString:      "arn:aws:iam::123456789012:role/example-role",
 		},
 		"RoleNameWithPathFromFlags": {
-			flags:             configurators.BootstrapFlags{AttachToRole: "/some/path/example-role"},
-			accountID:         "123456789012",
-			partitionID:       "aws",
+			config: ConfiguratorConfig{
+				Flags:    configurators.BootstrapFlags{AttachToRole: "/some/path/example-role"},
+				identity: identityFromArn(t, buildIAMARN("aws", "123456789012", "user", defaultAttachUser)),
+			},
 			targetType:        awslib.Role{},
 			targetName:        "example-role",
 			targetAccountID:   "123456789012",
@@ -1268,16 +1370,31 @@ func TestAWSPoliciesTarget(t *testing.T) {
 			targetString:      "arn:aws:iam::123456789012:role/some/path/example-role",
 		},
 		"RoleARNFromFlags": {
-			flags:             configurators.BootstrapFlags{AttachToRole: "arn:aws:iam::123456789012:role/example-role"},
+			config: ConfiguratorConfig{
+				Flags:    configurators.BootstrapFlags{AttachToRole: "arn:aws:iam::123456789012:role/example-role"},
+				identity: defaultIdentity,
+			},
 			targetType:        awslib.Role{},
 			targetName:        "example-role",
 			targetAccountID:   "123456789012",
 			targetPartitionID: "aws",
 			targetString:      "arn:aws:iam::123456789012:role/example-role",
 		},
+		"RoleARNFromFlagsWrongAccount": {
+			config: ConfiguratorConfig{
+				Flags:    configurators.BootstrapFlags{AttachToRole: "arn:aws:iam::987654321098:role/alt-role"},
+				identity: defaultIdentity,
+			},
+			wantErrContains: unreachablePolicyTargetError{
+				target: identityFromArn(t, "arn:aws:iam::987654321098:role/alt-role"),
+				from:   defaultIdentity,
+			}.Error(),
+		},
 		"UserFromIdentity": {
-			flags:             configurators.BootstrapFlags{},
-			identity:          userIdentity,
+			config: ConfiguratorConfig{
+
+				identity: userIdentity,
+			},
 			targetType:        awslib.User{},
 			targetName:        userIdentity.GetName(),
 			targetAccountID:   userIdentity.GetAccountID(),
@@ -1285,8 +1402,9 @@ func TestAWSPoliciesTarget(t *testing.T) {
 			targetString:      "arn:aws:iam::123456789012:user/example-user",
 		},
 		"RoleFromIdentity": {
-			flags:             configurators.BootstrapFlags{},
-			identity:          roleIdentity,
+			config: ConfiguratorConfig{
+				identity: roleIdentity,
+			},
 			targetType:        awslib.Role{},
 			targetName:        roleIdentity.GetName(),
 			targetAccountID:   roleIdentity.GetAccountID(),
@@ -1294,9 +1412,6 @@ func TestAWSPoliciesTarget(t *testing.T) {
 			targetString:      "arn:aws:iam::123456789012:role/example-role",
 		},
 		"DefaultTarget": {
-			flags:             configurators.BootstrapFlags{},
-			accountID:         "*",
-			partitionID:       "*",
 			targetType:        awslib.User{},
 			targetName:        defaultAttachUser,
 			targetAccountID:   "*",
@@ -1304,51 +1419,157 @@ func TestAWSPoliciesTarget(t *testing.T) {
 			targetString:      "arn:*:iam::*:user/username",
 		},
 		"AssumedRoleIdentity": {
-			flags:             configurators.BootstrapFlags{},
-			identity:          assumedRoleIdentity,
+			config: ConfiguratorConfig{
+				identity: assumedRoleIdentity,
+				getIAMClient: makeIAMClientGetter("", "", &iamMock{
+					partition: "aws",
+					account:   "123456789012",
+				}),
+			},
 			targetType:        awslib.Role{},
 			targetName:        assumedRoleIdentity.GetName(),
 			targetAccountID:   assumedRoleIdentity.GetAccountID(),
 			targetPartitionID: assumedRoleIdentity.GetPartition(),
 			targetString:      "arn:aws:iam::123456789012:role/example-role",
-			iamClient:         &iamMock{partition: "aws", account: "123456789012"},
 		},
 		"AssumedRoleIdentityForRoleWithPath": {
-			flags:             configurators.BootstrapFlags{},
-			identity:          assumedRoleIdentity,
+			config: ConfiguratorConfig{
+				identity: assumedRoleIdentity,
+				getIAMClient: makeIAMClientGetter("", "", &iamMock{
+					partition: "aws",
+					account:   "123456789012",
+					addPath:   "/some/path/",
+				}),
+			},
 			targetType:        awslib.Role{},
 			targetName:        assumedRoleIdentity.GetName(),
 			targetAccountID:   assumedRoleIdentity.GetAccountID(),
 			targetPartitionID: assumedRoleIdentity.GetPartition(),
 			targetString:      "arn:aws:iam::123456789012:role/some/path/example-role",
-			iamClient:         &iamMock{partition: "aws", account: "123456789012", addPath: "/some/path/"},
 		},
 		"AssumedRoleIdentityWithoutIAMPermissions": {
-			flags:             configurators.BootstrapFlags{},
-			identity:          assumedRoleIdentity,
-			targetType:        awslib.Role{},
-			targetName:        assumedRoleIdentity.GetName(),
-			targetAccountID:   assumedRoleIdentity.GetAccountID(),
-			targetPartitionID: assumedRoleIdentity.GetPartition(),
-			targetString:      "arn:aws:iam::123456789012:role/example-role",
-			iamClient:         &iamMock{unauthorized: true},
-			wantErrContains:   "policies cannot be attached to an assumed-role",
+			config: ConfiguratorConfig{
+				identity:     assumedRoleIdentity,
+				getIAMClient: makeIAMClientGetter("", "", &iamMock{unauthorized: true}),
+			},
+			wantErrContains: "policies cannot be attached to an assumed-role",
 		},
 		"AssumedRoleIdentityWithRoleFromFlags": {
-			flags:             configurators.BootstrapFlags{AttachToRole: "arn:aws:iam::123456789012:role/some/path/example-role"},
-			identity:          assumedRoleIdentity,
+			config: ConfiguratorConfig{
+				Flags:        configurators.BootstrapFlags{AttachToRole: "arn:aws:iam::123456789012:role/some/path/example-role"},
+				identity:     assumedRoleIdentity,
+				getIAMClient: makeIAMClientGetter("", "", &iamMock{unauthorized: true}),
+			},
 			targetType:        awslib.Role{},
 			targetName:        "example-role",
 			targetAccountID:   "123456789012",
 			targetPartitionID: "aws",
 			targetString:      "arn:aws:iam::123456789012:role/some/path/example-role",
-			iamClient:         &iamMock{unauthorized: true},
+		},
+		"MatcherAssumeRole": {
+			config: ConfiguratorConfig{
+				getIAMClient: makeIAMClientGetter(assumedRoleIdentity.String(), "", &iamMock{
+					account:   assumedRoleIdentity.GetAccountID(),
+					partition: assumedRoleIdentity.GetPartition(),
+					addPath:   "/some/path/",
+				}),
+				identity: defaultIdentity,
+			},
+			assumeRole:        types.AssumeRole{RoleARN: assumedRoleIdentity.String()},
+			targetType:        awslib.Role{},
+			targetName:        roleIdentity.GetName(),
+			targetAccountID:   roleIdentity.GetAccountID(),
+			targetPartitionID: roleIdentity.GetPartition(),
+			targetString:      "arn:aws:iam::123456789012:role/some/path/example-role",
+		},
+		"MatcherAssumeRoleWithUserFlag": {
+			config: ConfiguratorConfig{
+				Flags:    configurators.BootstrapFlags{AttachToUser: userIdentity.String()},
+				identity: defaultIdentity,
+			},
+			assumeRole:        types.AssumeRole{RoleARN: assumedRoleIdentity.String()},
+			targetType:        awslib.User{},
+			targetName:        userIdentity.GetName(),
+			targetAccountID:   userIdentity.GetAccountID(),
+			targetPartitionID: userIdentity.GetPartition(),
+			targetString:      userIdentity.String(),
+		},
+		"MatcherAssumeRoleWithUserFlagPath": {
+			config: ConfiguratorConfig{
+				Flags:    configurators.BootstrapFlags{AttachToUser: "/some/path/example-user"},
+				identity: defaultIdentity,
+			},
+			assumeRole:        types.AssumeRole{RoleARN: assumedRoleIdentity.String()},
+			targetType:        awslib.User{},
+			targetName:        userIdentity.GetName(),
+			targetAccountID:   userIdentity.GetAccountID(),
+			targetPartitionID: userIdentity.GetPartition(),
+			targetString:      "arn:aws:iam::123456789012:user/some/path/example-user",
+		},
+		"MatcherAssumeRoleWithUserFlagWrongAccount": {
+			config: ConfiguratorConfig{
+				Flags:    configurators.BootstrapFlags{AttachToUser: "arn:aws:iam::5678567856782:user/example-user"},
+				identity: defaultIdentity,
+			},
+			assumeRole: types.AssumeRole{RoleARN: assumedRoleIdentity.String()},
+			wantErrContains: unreachablePolicyTargetError{
+				target: identityFromArn(t, "arn:aws:iam::5678567856782:user/example-user"),
+				from:   assumedRoleIdentity,
+			}.Error(),
+		},
+		"MatcherAssumeRoleWithRoleFlag": {
+			config: ConfiguratorConfig{
+				Flags:    configurators.BootstrapFlags{AttachToRole: roleIdentity.String()},
+				identity: defaultIdentity,
+			},
+			assumeRole:        types.AssumeRole{RoleARN: altAssumedRoleIdentity.String()},
+			targetType:        awslib.Role{},
+			targetName:        roleIdentity.GetName(),
+			targetAccountID:   roleIdentity.GetAccountID(),
+			targetPartitionID: roleIdentity.GetPartition(),
+			targetString:      roleIdentity.String(),
+		},
+		"MatcherAssumeRoleWithRoleFlagPath": {
+			config: ConfiguratorConfig{
+				Flags:    configurators.BootstrapFlags{AttachToRole: "/some/path/example-role"},
+				identity: defaultIdentity,
+			},
+			assumeRole:        types.AssumeRole{RoleARN: altAssumedRoleIdentity.String()},
+			targetType:        awslib.Role{},
+			targetName:        roleIdentity.GetName(),
+			targetAccountID:   roleIdentity.GetAccountID(),
+			targetPartitionID: roleIdentity.GetPartition(),
+			targetString:      "arn:aws:iam::123456789012:role/some/path/example-role",
+		},
+		"MatcherAssumeRoleWithRoleFlagWrongAccount": {
+			config: ConfiguratorConfig{
+				Flags:    configurators.BootstrapFlags{AttachToRole: "arn:aws:iam::567856785678:role/example-role"},
+				identity: defaultIdentity,
+			},
+			assumeRole: types.AssumeRole{RoleARN: altAssumedRoleIdentity.String()},
+			wantErrContains: unreachablePolicyTargetError{
+				target: identityFromArn(t, "arn:aws:iam::567856785678:role/example-role"),
+				from:   altAssumedRoleIdentity,
+			}.Error(),
+		},
+		"MatcherAssumeRoleWithoutIAMPermission": {
+			config: ConfiguratorConfig{
+				getIAMClient: makeIAMClientGetter(assumedRoleIdentity.String(), "", &iamMock{unauthorized: true}),
+				identity:     defaultIdentity,
+			},
+			assumeRole:      types.AssumeRole{RoleARN: assumedRoleIdentity.String()},
+			wantErrContains: failedToResolveAssumeRoleARN(assumedRoleIdentity.GetName()),
 		},
 	}
 
 	for name, test := range tests {
 		t.Run(name, func(t *testing.T) {
-			target, err := policiesTarget(test.flags, test.accountID, test.partitionID, test.identity, test.iamClient)
+			test.config.ServiceConfig = &servicecfg.Config{}
+			require.NoError(t, test.config.CheckAndSetDefaults())
+			if test.config.identity == nil {
+				test.config.identity = identityFromArn(t, buildIAMARN(targetIdentityARNSectionPlaceholder, targetIdentityARNSectionPlaceholder, "user", defaultAttachUser))
+			}
+			target, err := policiesTarget(t.Context(), test.config, test.assumeRole)
 			if test.wantErrContains != "" {
 				require.ErrorContains(t, err, test.wantErrContains)
 				return
@@ -1371,13 +1592,17 @@ func identityFromArn(t *testing.T, arn string) awslib.Identity {
 }
 
 func TestAWSDocumentConfigurator(t *testing.T) {
+	t.Parallel()
 	var err error
-	ctx := context.Background()
+	ctx := t.Context()
 	fileConfig := &config.FileConfig{
 		Proxy: config.Proxy{
 			PublicAddr: []string{"proxy.example.org:443"},
 		},
 		Discovery: config.Discovery{
+			Service: config.Service{
+				EnabledFlag: "yes",
+			},
 			AWSMatchers: []config.AWSMatcher{
 				{
 					Types:   []string{"ec2"},
@@ -1392,29 +1617,26 @@ func TestAWSDocumentConfigurator(t *testing.T) {
 	require.NoError(t, config.ApplyFileConfig(fileConfig, serviceConfig))
 
 	config := ConfiguratorConfig{
-		awsCfg:    &aws.Config{},
-		iamClient: &iamMock{},
-		Identity:  identityFromArn(t, "arn:aws:iam::1234567:role/example-role"),
-		ssmClients: map[string]ssmClient{
-			"eu-central-1": &ssmMock{
-				t: t,
-				expectedInput: &ssm.CreateDocumentInput{
-					Content:        aws.String(awslib.EC2DiscoverySSMDocument("https://proxy.example.org:443")),
-					DocumentType:   ssmtypes.DocumentTypeCommand,
-					DocumentFormat: ssmtypes.DocumentFormatYaml,
-					Name:           aws.String("document"),
-				}},
-		},
+		getIAMClient: makeIAMClientGetter("", "", &iamMock{}),
+		identity:     identityFromArn(t, "arn:aws:iam::1234567:role/example-role"),
+		getSSMClient: makeSSMClientGetter("eu-central-1", "", "", &ssmMock{
+			t: t,
+			expectedInput: &ssm.CreateDocumentInput{
+				Content:        aws.String(awslib.EC2DiscoverySSMDocument("https://proxy.example.org:443")),
+				DocumentType:   ssmtypes.DocumentTypeCommand,
+				DocumentFormat: ssmtypes.DocumentFormatYaml,
+				Name:           aws.String("document"),
+			}}),
 		ServiceConfig: serviceConfig,
 		Flags: configurators.BootstrapFlags{
 			Service:             configurators.DiscoveryService,
 			ForceEC2Permissions: true,
 		},
-		Policies: &policiesMock{
+		getPolicies: makePoliciesGetter("", "", &policiesMock{
 			upsertArn: "policies-arn",
-		},
+		}),
 	}
-	configurator, err := NewAWSConfigurator(config)
+	configurator, err := NewAWSConfigurator(t.Context(), config)
 	require.NoError(t, err)
 	require.False(t, configurator.IsEmpty())
 
@@ -1429,25 +1651,25 @@ func TestAWSDocumentConfigurator(t *testing.T) {
 
 // TestAWSConfigurator tests all actions together.
 func TestAWSConfigurator(t *testing.T) {
+	t.Parallel()
 	var err error
-	ctx := context.Background()
+	ctx := t.Context()
 
 	config := ConfiguratorConfig{
-		awsCfg:        &aws.Config{},
-		iamClient:     &iamMock{},
-		Identity:      identityFromArn(t, "arn:aws:iam::1234567:role/example-role"),
-		ssmClients:    map[string]ssmClient{"eu-central-1": &ssmMock{}},
+		getIAMClient:  makeIAMClientGetter("", "", &iamMock{}),
+		identity:      identityFromArn(t, "arn:aws:iam::1234567:role/example-role"),
+		getSSMClient:  makeSSMClientGetter("eu-central-1", "", "", &ssmMock{}),
 		ServiceConfig: &servicecfg.Config{},
 		Flags: configurators.BootstrapFlags{
 			AttachToUser:        "some-user",
 			ForceRDSPermissions: true,
 		},
-		Policies: &policiesMock{
+		getPolicies: makePoliciesGetter("", "", &policiesMock{
 			upsertArn: "policies-arn",
-		},
+		}),
 	}
 
-	configurator, err := NewAWSConfigurator(config)
+	configurator, err := NewAWSConfigurator(t.Context(), config)
 	require.NoError(t, err)
 	require.False(t, configurator.IsEmpty())
 
@@ -1462,7 +1684,7 @@ func TestAWSConfigurator(t *testing.T) {
 	config.Flags.ForceEC2Permissions = true
 	config.Flags.Proxy = "proxy.xyz"
 
-	configurator, err = NewAWSConfigurator(config)
+	configurator, err = NewAWSConfigurator(t.Context(), config)
 	require.NoError(t, err)
 	require.False(t, configurator.IsEmpty())
 
@@ -1751,7 +1973,7 @@ func TestExtractTargetConfig(t *testing.T) {
 	}
 	for name, tt := range tests {
 		t.Run(name, func(t *testing.T) {
-			got, err := getTargetConfig(tt.flags, tt.cfg, tt.target)
+			got, err := getTargetConfig(tt.flags, tt.cfg, tt.target, types.AssumeRole{})
 			if tt.wantErr != "" {
 				require.ErrorContains(t, err, tt.wantErr)
 				return
@@ -1888,7 +2110,7 @@ func (m *ssmMock) CreateDocument(ctx context.Context, input *ssm.CreateDocumentI
 		cmp.Diff(m.expectedInput.Content, input.Content),
 		"Document content diff (-want +got)")
 	require.Empty(m.t,
-		cmp.Diff(m.expectedInput, input, cmpopts.IgnoreFields(ssm.CreateDocumentInput{}, "Content")),
+		cmp.Diff(m.expectedInput, input, cmpopts.IgnoreUnexported(ssm.CreateDocumentInput{})),
 		"Document diff (-want +got)")
 
 	return nil, nil
@@ -1925,6 +2147,7 @@ func (m mockLocalRegionGetter) GetRegion(context.Context) (string, error) {
 }
 
 func Test_getFallbackRegion(t *testing.T) {
+	t.Parallel()
 	tests := []struct {
 		name              string
 		localRegionGetter localRegionGetter
@@ -1948,8 +2171,83 @@ func Test_getFallbackRegion(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			region := getFallbackRegion(context.Background(), io.Discard, test.localRegionGetter)
+			region := getFallbackRegion(t.Context(), io.Discard, test.localRegionGetter)
 			require.Equal(t, test.wantRegion, region)
 		})
 	}
 }
+
+func TestGetDistinctAssumedRoles(t *testing.T) {
+	t.Parallel()
+	defaultAssumeRole := types.AssumeRole{
+		RoleARN:    "arn:aws:iam::123456789012:role/example-role",
+		ExternalID: "foobar",
+	}
+	tests := []struct {
+		name                string
+		matchers            []types.AWSMatcher
+		expectedAssumeRoles []types.AssumeRole
+	}{
+		{
+			name:                "empty",
+			expectedAssumeRoles: []types.AssumeRole{defaultAssumeRole},
+		},
+		{
+			name: "multiple",
+			matchers: []types.AWSMatcher{
+				{},
+				{AssumeRole: &types.AssumeRole{RoleARN: "12345678", ExternalID: "foo"}},
+				{AssumeRole: &types.AssumeRole{RoleARN: "87654321"}},
+			},
+			expectedAssumeRoles: []types.AssumeRole{
+				defaultAssumeRole,
+				{RoleARN: "12345678", ExternalID: "foo"},
+				{RoleARN: "87654321"},
+			},
+		},
+		{
+			name: "filter out duplicates",
+			matchers: []types.AWSMatcher{
+				{},
+				{AssumeRole: &types.AssumeRole{RoleARN: "12345678"}},
+				{AssumeRole: &types.AssumeRole{RoleARN: "87654321", ExternalID: "foo"}},
+				{AssumeRole: &types.AssumeRole{RoleARN: "12345678"}},
+				{AssumeRole: &types.AssumeRole{RoleARN: "87654321", ExternalID: "foo"}},
+				{},
+			},
+			expectedAssumeRoles: []types.AssumeRole{
+				defaultAssumeRole,
+				{RoleARN: "12345678"},
+				{RoleARN: "87654321", ExternalID: "foo"},
+			},
+		},
+		{
+			name: "preserve duplicate arn when external id differs",
+			matchers: []types.AWSMatcher{
+				{AssumeRole: &types.AssumeRole{RoleARN: "12345678", ExternalID: "foo"}},
+				{AssumeRole: &types.AssumeRole{RoleARN: "12345678", ExternalID: "bar"}},
+			},
+			expectedAssumeRoles: []types.AssumeRole{
+				{RoleARN: "12345678", ExternalID: "foo"},
+				{RoleARN: "12345678", ExternalID: "bar"},
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			config := ConfiguratorConfig{
+				Flags: configurators.BootstrapFlags{
+					Service:       configurators.DiscoveryService,
+					AssumeRoleARN: defaultAssumeRole.RoleARN,
+					ExternalID:    defaultAssumeRole.ExternalID,
+				},
+				ServiceConfig: &servicecfg.Config{
+					Discovery: servicecfg.DiscoveryConfig{
+						AWSMatchers: tc.matchers,
+					},
+				},
+			}
+			require.ElementsMatch(t, tc.expectedAssumeRoles, config.getDistinctAssumedRoles())
+		})
+	}
+}

lib/configurators/configuratorbuilder/builder.go

@@ -19,6 +19,8 @@
 package configuratorbuilder
 
 import (
+	"context"
+
 	"github.com/gravitational/trace"
 
 	"github.com/gravitational/teleport/lib/config"
@@ -29,7 +31,7 @@ import (
 
 // BuildConfigurators reads the configuration and returns a list of
 // configurators. Configurators that are "empty" are not returned.
-func BuildConfigurators(flags configurators.BootstrapFlags) ([]configurators.Configurator, error) {
+func BuildConfigurators(ctx context.Context, flags configurators.BootstrapFlags) ([]configurators.Configurator, error) {
 	fileConfig, err := config.ReadFromFile(flags.ConfigPath)
 	if err != nil {
 		return nil, trace.Wrap(err)
@@ -39,7 +41,7 @@ func BuildConfigurators(flags configurators.BootstrapFlags) ([]configurators.Con
 		return nil, trace.Wrap(err)
 	}
 
-	awsConfigurator, err := aws.NewAWSConfigurator(aws.ConfiguratorConfig{
+	awsConfigurator, err := aws.NewAWSConfigurator(ctx, aws.ConfiguratorConfig{
 		Flags:         flags,
 		ServiceConfig: serviceCfg,
 	})

tool/teleport/common/configurator.go

@@ -138,11 +138,10 @@ func makeDatabaseServiceBootstrapFlagsWithDiscoveryServiceConfig(flags configure
 
 // onConfigureDiscoveryBootstrap subcommand that bootstraps configuration for
 // discovery  agents.
-func onConfigureDiscoveryBootstrap(flags configureDiscoveryBootstrapFlags) error {
+func onConfigureDiscoveryBootstrap(ctx context.Context, flags configureDiscoveryBootstrapFlags) error {
 	fmt.Printf("Reading configuration at %q...\n", flags.config.ConfigPath)
 
-	ctx := context.TODO()
-	configurators, err := configuratorbuilder.BuildConfigurators(flags.config)
+	configurators, err := configuratorbuilder.BuildConfigurators(ctx, flags.config)
 	if err != nil {
 		return trace.Wrap(err)
 	}
@@ -152,7 +151,7 @@ func onConfigureDiscoveryBootstrap(flags configureDiscoveryBootstrapFlags) error
 	// service config.
 	if flags.config.Service.IsDiscovery() && flags.databaseServiceRole != "" {
 		config := makeDatabaseServiceBootstrapFlagsWithDiscoveryServiceConfig(flags)
-		dbConfigurators, err := configuratorbuilder.BuildConfigurators(config)
+		dbConfigurators, err := configuratorbuilder.BuildConfigurators(ctx, config)
 		if err != nil {
 			return trace.Wrap(err)
 		}
@@ -245,7 +244,7 @@ type configureDatabaseAWSPrintFlags struct {
 
 // buildAWSConfigurator builds the database configurator used on AWS-specific
 // commands.
-func buildAWSConfigurator(manual bool, flags configureDatabaseAWSFlags) (configurators.Configurator, error) {
+func buildAWSConfigurator(ctx context.Context, manual bool, flags configureDatabaseAWSFlags) (configurators.Configurator, error) {
 	err := flags.CheckAndSetDefaults()
 	if err != nil {
 		return nil, trace.Wrap(err)
@@ -284,7 +283,7 @@ func buildAWSConfigurator(manual bool, flags configureDatabaseAWSFlags) (configu
 		}
 	}
 
-	configurator, err := awsconfigurators.NewAWSConfigurator(awsconfigurators.ConfiguratorConfig{
+	configurator, err := awsconfigurators.NewAWSConfigurator(ctx, awsconfigurators.ConfiguratorConfig{
 		Flags:         configuratorFlags,
 		ServiceConfig: &servicecfg.Config{},
 	})
@@ -297,8 +296,8 @@ func buildAWSConfigurator(manual bool, flags configureDatabaseAWSFlags) (configu
 
 // onConfigureDatabasesAWSPrint is a subcommand used to print AWS IAM access
 // Teleport requires to run databases discovery on AWS.
-func onConfigureDatabasesAWSPrint(flags configureDatabaseAWSPrintFlags) error {
-	configurator, err := buildAWSConfigurator(true, flags.configureDatabaseAWSFlags)
+func onConfigureDatabasesAWSPrint(ctx context.Context, flags configureDatabaseAWSPrintFlags) error {
+	configurator, err := buildAWSConfigurator(ctx, true, flags.configureDatabaseAWSFlags)
 	if err != nil {
 		return trace.Wrap(err)
 	}
@@ -335,9 +334,8 @@ type configureDatabaseAWSCreateFlags struct {
 
 // onConfigureDatabasesAWSCreates is a subcommand used to create AWS IAM access
 // for Teleport to run databases discovery on AWS.
-func onConfigureDatabasesAWSCreate(flags configureDatabaseAWSCreateFlags) error {
-	ctx := context.TODO()
-	configurator, err := buildAWSConfigurator(false, flags.configureDatabaseAWSFlags)
+func onConfigureDatabasesAWSCreate(ctx context.Context, flags configureDatabaseAWSCreateFlags) error {
+	configurator, err := buildAWSConfigurator(ctx, false, flags.configureDatabaseAWSFlags)
 	if err != nil {
 		return trace.Wrap(err)
 	}

tool/teleport/common/teleport.go

@@ -737,19 +737,19 @@ Examples:
 	case dbConfigureCreate.FullCommand():
 		err = onDumpDatabaseConfig(dbConfigCreateFlags)
 	case dbConfigureAWSPrintIAM.FullCommand():
-		err = onConfigureDatabasesAWSPrint(configureDatabaseAWSPrintFlags)
+		err = onConfigureDatabasesAWSPrint(ctx, configureDatabaseAWSPrintFlags)
 	case dbConfigureAWSCreateIAM.FullCommand():
-		err = onConfigureDatabasesAWSCreate(configureDatabaseAWSCreateFlags)
+		err = onConfigureDatabasesAWSCreate(ctx, configureDatabaseAWSCreateFlags)
 	case dbConfigureBootstrap.FullCommand():
 		configureDiscoveryBootstrapFlags.config.Service = configurators.DatabaseService
-		err = onConfigureDiscoveryBootstrap(configureDiscoveryBootstrapFlags)
+		err = onConfigureDiscoveryBootstrap(ctx, configureDiscoveryBootstrapFlags)
 	case systemdInstall.FullCommand():
 		err = onDumpSystemdUnitFile(systemdInstallFlags)
 	case installAutoDiscoverNode.FullCommand():
 		err = onInstallAutoDiscoverNode(installAutoDiscoverNodeFlags)
 	case discoveryBootstrapCmd.FullCommand():
 		configureDiscoveryBootstrapFlags.config.Service = configurators.DiscoveryService
-		err = onConfigureDiscoveryBootstrap(configureDiscoveryBootstrapFlags)
+		err = onConfigureDiscoveryBootstrap(ctx, configureDiscoveryBootstrapFlags)
 	case joinOpenSSH.FullCommand():
 		err = onJoinOpenSSH(ccf, conf)
 	case integrationConfDeployServiceCmd.FullCommand():