fix region picker for AWS Organizations

elastic · Mar 25, 2024 · 5b582f0 · 5b582f0
1 parent 47082fb
commit 5b582f0
Showing 1 changed file with 31 additions and 9 deletions.
diff --git a/internal/flavors/benchmark/aws_org.go b/internal/flavors/benchmark/aws_org.go
@@ -71,23 +71,30 @@ func (a *AWSOrg) initialize(ctx context.Context, log *logp.Logger, cfg *config.C
 		return nil, nil, nil, err
 	}
 
-	// TODO: make this mock-able
-	awsConfig, err := aws.InitializeAWSConfig(cfg.CloudConfig.Aws.Cred)
-	if err != nil {
-		return nil, nil, nil, fmt.Errorf("failed to initialize AWS credentials: %w", err)
-	}
+	var (
+		awsConfig   *awssdk.Config
+		awsIdentity *cloud.Identity
+		err         error
+	)
 
-	a.IAMProvider = iam.NewIAMProvider(log, awsConfig, nil)
+	awsConfig, awsIdentity, err = a.getIdentity(ctx, cfg)
+	if err != nil && cfg.CloudConfig.Aws.Cred.DefaultRegion == "" {
+		log.Warn("failed to initialize identity; retrying to check AWS Gov Cloud regions")
+		cfg.CloudConfig.Aws.Cred.DefaultRegion = awslib.DefaultGovRegion
+		awsConfig, awsIdentity, err = a.getIdentity(ctx, cfg)
+	}
 
-	awsIdentity, err := a.IdentityProvider.GetIdentity(ctx, awsConfig)
 	if err != nil {
-		return nil, nil, nil, fmt.Errorf("failed to get AWS identity: %w", err)
+		return nil, nil, nil, fmt.Errorf("failed to get AWS Identity: %w", err)
 	}
+	log.Info("successfully retrieved AWS Identity")
+
+	a.IAMProvider = iam.NewIAMProvider(log, *awsConfig, nil)
 
 	cache := make(map[string]registry.FetchersMap)
 	reg := registry.NewRegistry(log, registry.WithUpdater(
 		func() (registry.FetchersMap, error) {
-			accounts, err := a.getAwsAccounts(ctx, log, awsConfig, awsIdentity)
+			accounts, err := a.getAwsAccounts(ctx, log, *awsConfig, awsIdentity)
 			if err != nil {
 				return nil, fmt.Errorf("failed to get AWS accounts: %w", err)
 			}
@@ -211,6 +218,21 @@ func (a *AWSOrg) pickManagementAccountRole(ctx context.Context, log *logp.Logger
 	return config, nil
 }
 
+func (a *AWSOrg) getIdentity(ctx context.Context, cfg *config.Config) (*awssdk.Config, *cloud.Identity, error) {
+	// TODO: make this mock-able
+	awsConfig, err := aws.InitializeAWSConfig(cfg.CloudConfig.Aws.Cred)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to initialize AWS credentials: %w", err)
+	}
+
+	awsIdentity, err := a.IdentityProvider.GetIdentity(ctx, awsConfig)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to get AWS identity: %w", err)
+	}
+
+	return &awsConfig, awsIdentity, nil
+}
+
 func (a *AWSOrg) checkDependencies() error {
 	if a.IAMProvider == nil {
 		return errors.New("aws iam provider is uninitialized")