From 17da6df50d200a4202ee652a22674df17373886a Mon Sep 17 00:00:00 2001 From: girish-cheedala Date: Tue, 10 Mar 2026 22:49:48 +0530 Subject: [PATCH 1/4] refactor(aws): enhance createAssumedRoleSession to accept account ID as parameter and improve error handling in Verify method - Updated createAssumedRoleSession to take accountId as an argument, allowing for more flexible role assumption. - Modified Verify method to iterate over AccountIds, capturing errors for each account and returning the last encountered error if all attempts fail. --- pkg/providers/aws/aws.go | 32 +++++++++++++++++--------------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/pkg/providers/aws/aws.go b/pkg/providers/aws/aws.go index 4ebaf51..97c2fa5 100644 --- a/pkg/providers/aws/aws.go +++ b/pkg/providers/aws/aws.go @@ -249,7 +249,7 @@ func New(block schema.OptionBlock) (*Provider, error) { if err != nil && options.AssumeRoleName != "" && len(options.AccountIds) > 0 { // Base user doesn't have DescribeRegions permission, try with assumed role - tempSession, err := createAssumedRoleSession(options, sess, config) + tempSession, err := createAssumedRoleSession(options, sess, config, options.AccountIds[0]) if err != nil { return nil, errors.Wrap(err, "could not create assumed role session") } @@ -270,12 +270,9 @@ func New(block schema.OptionBlock) (*Provider, error) { return provider, nil } -func createAssumedRoleSession(options *ProviderOptions, sess *session.Session, config *aws.Config) (*session.Session, error) { - if len(options.AccountIds) == 0 { - return nil, errors.New("no account IDs provided for assume role") - } +func createAssumedRoleSession(options *ProviderOptions, sess *session.Session, config *aws.Config, accountId string) (*session.Session, error) { stsClient := sts.New(sess) - roleArn := fmt.Sprintf("arn:aws:iam::%s:role/%s", options.AccountIds[0], options.AssumeRoleName) + roleArn := fmt.Sprintf("arn:aws:iam::%s:role/%s", accountId, options.AssumeRoleName) roleInput := &sts.AssumeRoleInput{ RoleArn: aws.String(roleArn), @@ -524,16 +521,21 @@ func (p *Provider) Verify(ctx context.Context) error { } if p.options.AssumeRoleName != "" && len(p.options.AccountIds) > 0 { - tempSession, err := createAssumedRoleSession(p.options, p.session, p.session.Config) - if err != nil { - return err - } - p.initServices(tempSession) - err = p.verify() - if err != nil { - return err + var lastErr error + for _, accountId := range p.options.AccountIds { + tempSession, err := createAssumedRoleSession(p.options, p.session, p.session.Config, accountId) + if err != nil { + lastErr = err + continue + } + p.initServices(tempSession) + if err := p.verify(); err != nil { + lastErr = err + continue + } + return nil } - return nil + return errors.Wrap(lastErr, "failed to verify credentials across all accounts") } return err } From 031f0652e8737ae58637f032c308ba4acbf6936b Mon Sep 17 00:00:00 2001 From: girish-cheedala Date: Wed, 11 Mar 2026 12:22:49 +0530 Subject: [PATCH 2/4] refactor(aws): enhance createAssumedRoleSession to accept account ID as parameter and improve error handling in Verify method - Updated createAssumedRoleSession to take accountId as an argument, allowing for more flexible role assumption. - Modified Verify method to iterate over AccountIds, capturing errors for each account and returning the last encountered error if all attempts fail. --- pkg/providers/aws/aws.go | 43 +++++++++++++++++++++++++++++----------- 1 file changed, 31 insertions(+), 12 deletions(-) diff --git a/pkg/providers/aws/aws.go b/pkg/providers/aws/aws.go index 97c2fa5..1c2e9e8 100644 --- a/pkg/providers/aws/aws.go +++ b/pkg/providers/aws/aws.go @@ -521,21 +521,40 @@ func (p *Provider) Verify(ctx context.Context) error { } if p.options.AssumeRoleName != "" && len(p.options.AccountIds) > 0 { - var lastErr error + var mu sync.Mutex + var failedAccounts []string + var wg sync.WaitGroup + for _, accountId := range p.options.AccountIds { - tempSession, err := createAssumedRoleSession(p.options, p.session, p.session.Config, accountId) - if err != nil { - lastErr = err - continue - } - p.initServices(tempSession) - if err := p.verify(); err != nil { - lastErr = err - continue + wg.Add(1) + go func(id string) { + defer wg.Done() + tempSession, err := createAssumedRoleSession(p.options, p.session, p.session.Config, id) + if err != nil { + mu.Lock() + failedAccounts = append(failedAccounts, id) + mu.Unlock() + return + } + tempProvider := &Provider{options: p.options, session: tempSession} + tempProvider.initServices(tempSession) + if err := tempProvider.verify(); err != nil { + mu.Lock() + failedAccounts = append(failedAccounts, id) + mu.Unlock() + } + }(accountId) + } + wg.Wait() + + if len(failedAccounts) > 0 { + msg := fmt.Sprintf("failed to assume role %s in accounts: %s", p.options.AssumeRoleName, strings.Join(failedAccounts, ", ")) + if p.options.OrgDiscoveryRoleArn != "" { + msg += ". Add these to exclude_account_ids if they should not be part of discovery" } - return nil + return errors.New(msg) } - return errors.Wrap(lastErr, "failed to verify credentials across all accounts") + return nil } return err } From c8d044f1652a46e6084424d5d4ecb17fa02ac705 Mon Sep 17 00:00:00 2001 From: girish-cheedala Date: Wed, 11 Mar 2026 12:42:02 +0530 Subject: [PATCH 3/4] refactor(aws): improve region retrieval logic in New function by iterating over AccountIds - Enhanced the New function to attempt region retrieval using multiple AccountIds. - Implemented error handling to continue attempts with subsequent accounts if the initial one fails. - Ensured that a comprehensive error message is returned if all attempts to retrieve regions fail. --- pkg/providers/aws/aws.go | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/pkg/providers/aws/aws.go b/pkg/providers/aws/aws.go index 1c2e9e8..d231ffb 100644 --- a/pkg/providers/aws/aws.go +++ b/pkg/providers/aws/aws.go @@ -249,16 +249,21 @@ func New(block schema.OptionBlock) (*Provider, error) { if err != nil && options.AssumeRoleName != "" && len(options.AccountIds) > 0 { // Base user doesn't have DescribeRegions permission, try with assumed role - tempSession, err := createAssumedRoleSession(options, sess, config, options.AccountIds[0]) - if err != nil { - return nil, errors.Wrap(err, "could not create assumed role session") + var regionErr error + for _, accountId := range options.AccountIds { + tempSession, err := createAssumedRoleSession(options, sess, config, accountId) + if err != nil { + regionErr = err + continue + } + tempRC := ec2.New(tempSession) + regions, regionErr = tempRC.DescribeRegions(&ec2.DescribeRegionsInput{}) + if regionErr == nil { + break + } } - - // Use assumed role session for DescribeRegions - tempRC := ec2.New(tempSession) - regions, err = tempRC.DescribeRegions(&ec2.DescribeRegionsInput{}) - if err != nil { - return nil, errors.Wrap(err, "could not get list of regions even with assumed role") + if regionErr != nil { + return nil, errors.Wrap(regionErr, "could not get list of regions with any account") } } else if err != nil { return nil, errors.Wrap(err, "could not get list of regions") From 1d7e6049ef863be0b9fa3b6d2d846ac6668e81d8 Mon Sep 17 00:00:00 2001 From: girish-cheedala Date: Wed, 11 Mar 2026 18:19:52 +0530 Subject: [PATCH 4/4] limit go routines to 200 to prevent rate limit --- pkg/providers/aws/aws.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pkg/providers/aws/aws.go b/pkg/providers/aws/aws.go index d231ffb..5c63ef0 100644 --- a/pkg/providers/aws/aws.go +++ b/pkg/providers/aws/aws.go @@ -529,11 +529,14 @@ func (p *Provider) Verify(ctx context.Context) error { var mu sync.Mutex var failedAccounts []string var wg sync.WaitGroup + sem := make(chan struct{}, 200) for _, accountId := range p.options.AccountIds { wg.Add(1) + sem <- struct{}{} go func(id string) { defer wg.Done() + defer func() { <-sem }() tempSession, err := createAssumedRoleSession(p.options, p.session, p.session.Config, id) if err != nil { mu.Lock()