Remove context.TODO from production code (#2707)

internal/flavors/assetinventory/strategy_aws.go

@@ -54,7 +54,7 @@ func (s *strategy) initAwsFetchers(ctx context.Context) ([]inventory.AssetFetche
 
 	// Early exit if we're scanning the entire account.
 	if s.cfg.CloudConfig.Aws.AccountType == config.SingleAccount {
-		return awsfetcher.New(s.logger, awsIdentity, *awsConfig), nil
+		return awsfetcher.New(ctx, s.logger, awsIdentity, *awsConfig), nil
 	}
 
 	// Assume audit roles per selected account and generate fetchers for them
@@ -81,7 +81,7 @@ func (s *strategy) initAwsFetchers(ctx context.Context) ([]inventory.AssetFetche
 			s.logger.Infof("Skipping identity on purpose %+v", identity)
 			continue
 		}
-		accountFetchers := awsfetcher.New(s.logger, &identity, assumedRoleConfig)
+		accountFetchers := awsfetcher.New(ctx, s.logger, &identity, assumedRoleConfig)
 		fetchers = append(fetchers, accountFetchers...)
 	}
 

internal/flavors/benchmark/aws.go

@@ -79,7 +79,7 @@ func (a *AWS) initialize(ctx context.Context, log *logp.Logger, cfg *config.Conf
 
 	return registry.NewRegistry(
 		log,
-		registry.WithFetchersMap(preset.NewCisAwsFetchers(log, *awsConfig, ch, awsIdentity)),
+		registry.WithFetchersMap(preset.NewCisAwsFetchers(ctx, log, *awsConfig, ch, awsIdentity)),
 	), cloud.NewDataProvider(cloud.WithAccount(*awsIdentity)), nil, nil
 }
 

internal/flavors/benchmark/aws_org.go

@@ -88,7 +88,7 @@ func (a *AWSOrg) initialize(ctx context.Context, log *logp.Logger, cfg *config.C
 	}
 	log.Info("successfully retrieved AWS Identity")
 
-	a.IAMProvider = iam.NewIAMProvider(log, *awsConfig, nil)
+	a.IAMProvider = iam.NewIAMProvider(ctx, log, *awsConfig, nil)
 
 	cache := make(map[string]registry.FetchersMap)
 	reg := registry.NewRegistry(log, registry.WithUpdater(

internal/flavors/benchmark/eks.go

@@ -97,7 +97,7 @@ func (k *EKS) initialize(ctx context.Context, log *logp.Logger, cfg *config.Conf
 
 	return registry.NewRegistry(
 		log,
-		registry.WithFetchersMap(preset.NewCisEksFetchers(log, awsConfig, ch, k.leaderElector, kubeClient, awsIdentity)),
+		registry.WithFetchersMap(preset.NewCisEksFetchers(ctx, log, awsConfig, ch, k.leaderElector, kubeClient, awsIdentity)),
 	), dp, idp, nil
 }
 

internal/flavors/vulnerability.go

@@ -127,7 +127,7 @@ func (bt *vulnerability) Run(*beat.Beat) error {
 }
 
 func (bt *vulnerability) runIteration() error {
-	worker, err := vuln.NewVulnerabilityWorker(bt.log, bt.config, bt.bdp, bt.cdp)
+	worker, err := vuln.NewVulnerabilityWorker(bt.ctx, bt.log, bt.config, bt.bdp, bt.cdp)
 	if err != nil {
 		bt.log.Warn("vulnerability.runIteration worker creation failed")
 		bt.cancel()

internal/inventory/awsfetcher/awsfetchers.go

@@ -18,6 +18,8 @@
 package awsfetcher
 
 import (
+	"context"
+
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/elastic/elastic-agent-libs/logp"
 
@@ -34,15 +36,15 @@ import (
 	"github.com/elastic/cloudbeat/internal/resources/providers/awslib/sns"
 )
 
-func New(logger *logp.Logger, identity *cloud.Identity, cfg aws.Config) []inventory.AssetFetcher {
-	ec2Provider := ec2.NewEC2Provider(logger, identity.Account, cfg, &awslib.MultiRegionClientFactory[ec2.Client]{})
-	elbProvider := elb.NewElbProvider(logger, identity.Account, cfg, &awslib.MultiRegionClientFactory[elb.Client]{})
-	elbv2Provider := elbv2.NewElbV2Provider(logger, cfg, &awslib.MultiRegionClientFactory[elbv2.Client]{})
-	iamProvider := iam.NewIAMProvider(logger, cfg, &awslib.MultiRegionClientFactory[iam.AccessAnalyzerClient]{})
-	lambdaProvider := lambda.NewLambdaProvider(logger, cfg, &awslib.MultiRegionClientFactory[lambda.Client]{})
-	rdsProvider := rds.NewProvider(logger, cfg, &awslib.MultiRegionClientFactory[rds.Client]{}, ec2Provider)
-	s3Provider := s3.NewProvider(logger, cfg, &awslib.MultiRegionClientFactory[s3.Client]{}, identity.Account)
-	snsProvider := sns.NewSNSProvider(logger, cfg, &awslib.MultiRegionClientFactory[sns.Client]{})
+func New(ctx context.Context, logger *logp.Logger, identity *cloud.Identity, cfg aws.Config) []inventory.AssetFetcher {
+	ec2Provider := ec2.NewEC2Provider(ctx, logger, identity.Account, cfg, &awslib.MultiRegionClientFactory[ec2.Client]{})
+	elbProvider := elb.NewElbProvider(ctx, logger, identity.Account, cfg, &awslib.MultiRegionClientFactory[elb.Client]{})
+	elbv2Provider := elbv2.NewElbV2Provider(ctx, logger, cfg, &awslib.MultiRegionClientFactory[elbv2.Client]{})
+	iamProvider := iam.NewIAMProvider(ctx, logger, cfg, &awslib.MultiRegionClientFactory[iam.AccessAnalyzerClient]{})
+	lambdaProvider := lambda.NewLambdaProvider(ctx, logger, cfg, &awslib.MultiRegionClientFactory[lambda.Client]{})
+	rdsProvider := rds.NewProvider(ctx, logger, cfg, &awslib.MultiRegionClientFactory[rds.Client]{}, ec2Provider)
+	s3Provider := s3.NewProvider(ctx, logger, cfg, &awslib.MultiRegionClientFactory[s3.Client]{}, identity.Account)
+	snsProvider := sns.NewSNSProvider(ctx, logger, cfg, &awslib.MultiRegionClientFactory[sns.Client]{})
 
 	return []inventory.AssetFetcher{
 		newEc2InstancesFetcher(logger, identity, ec2Provider),

internal/resources/fetching/preset/aws_org_preset.go

@@ -65,7 +65,7 @@ func NewCisAwsOrganizationFetchers(ctx context.Context, log *logp.Logger, rootCh
 }
 
 // awsFactory is the same function type as NewCisAwsFetchers, and it's used to mock the function in tests
-type awsFactory func(*logp.Logger, aws.Config, chan fetching.ResourceInfo, *cloud.Identity) registry.FetchersMap
+type awsFactory func(context.Context, *logp.Logger, aws.Config, chan fetching.ResourceInfo, *cloud.Identity) registry.FetchersMap
 
 func newCisAwsOrganizationFetchers(
 	ctx context.Context,
@@ -117,6 +117,7 @@ func newCisAwsOrganizationFetchers(
 		}(account.Identity)
 
 		f := factory(
+			ctx,
 			log.Named("aws").WithOptions(zap.Fields(zap.String("cloud.account.id", account.Identity.Account))),
 			account.Config,
 			ch,

internal/resources/fetching/preset/aws_org_preset_test.go

@@ -70,7 +70,7 @@ func subtest(t *testing.T, drain bool) { //revive:disable-line:flag-parameter
 	ctx, cancel := context.WithCancel(context.Background())
 
 	factory := mockFactory(nAccounts,
-		func(_ *logp.Logger, _ aws.Config, ch chan fetching.ResourceInfo, _ *cloud.Identity) registry.FetchersMap {
+		func(_ context.Context, _ *logp.Logger, _ aws.Config, ch chan fetching.ResourceInfo, _ *cloud.Identity) registry.FetchersMap {
 			if drain {
 				// create some resources if we are testing for that
 				go func() {
@@ -152,7 +152,7 @@ func TestNewCisAwsOrganizationFetchers_LeakContextDone(t *testing.T) {
 		}},
 		nil,
 		mockFactory(1,
-			func(_ *logp.Logger, _ aws.Config, ch chan fetching.ResourceInfo, _ *cloud.Identity) registry.FetchersMap {
+			func(_ context.Context, _ *logp.Logger, _ aws.Config, ch chan fetching.ResourceInfo, _ *cloud.Identity) registry.FetchersMap {
 				ch <- fetching.ResourceInfo{
 					Resource:      mockResource(),
 					CycleMetadata: cycle.Metadata{Sequence: 1},
@@ -181,7 +181,7 @@ func TestNewCisAwsOrganizationFetchers_CloseChannel(t *testing.T) {
 		}},
 		nil,
 		mockFactory(1,
-			func(_ *logp.Logger, _ aws.Config, ch chan fetching.ResourceInfo, _ *cloud.Identity) registry.FetchersMap {
+			func(_ context.Context, _ *logp.Logger, _ aws.Config, ch chan fetching.ResourceInfo, _ *cloud.Identity) registry.FetchersMap {
 				defer close(ch)
 				return registry.FetchersMap{"fetcher": registry.RegisteredFetcher{}}
 			},
@@ -214,7 +214,7 @@ func TestNewCisAwsOrganizationFetchers_Cache(t *testing.T) {
 		},
 		cache,
 		mockFactory(1,
-			func(_ *logp.Logger, _ aws.Config, _ chan fetching.ResourceInfo, identity *cloud.Identity) registry.FetchersMap {
+			func(_ context.Context, _ *logp.Logger, _ aws.Config, _ chan fetching.ResourceInfo, identity *cloud.Identity) registry.FetchersMap {
 				assert.Equal(t, "2", identity.Account)
 				return registry.FetchersMap{"fetcher": registry.RegisteredFetcher{}}
 			},
@@ -241,6 +241,6 @@ func mockResource() *fetching.MockResource {
 
 func mockFactory(times int, f awsFactory) awsFactory {
 	factory := mockAwsFactory{}
-	factory.EXPECT().Execute(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(f).Times(times)
+	factory.EXPECT().Execute(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(f).Times(times)
 	return factory.Execute
 }

internal/resources/fetching/preset/aws_preset.go

@@ -18,6 +18,8 @@
 package preset
 
 import (
+	"context"
+
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/elastic/elastic-agent-libs/logp"
 
@@ -41,24 +43,25 @@ import (
 	"github.com/elastic/cloudbeat/internal/resources/providers/awslib/sns"
 )
 
-func NewCisAwsFetchers(log *logp.Logger, cfg aws.Config, ch chan fetching.ResourceInfo, identity *cloud.Identity) registry.FetchersMap {
+func NewCisAwsFetchers(ctx context.Context, log *logp.Logger, cfg aws.Config, ch chan fetching.ResourceInfo, identity *cloud.Identity) registry.FetchersMap {
 	log.Infof("Initializing AWS fetchers for account: '%s'", identity.Account)
 
 	m := make(registry.FetchersMap)
-	iamProvider := iam.NewIAMProvider(log, cfg, &awslib.MultiRegionClientFactory[iam.AccessAnalyzerClient]{})
+	iamProvider := iam.NewIAMProvider(ctx, log, cfg, &awslib.MultiRegionClientFactory[iam.AccessAnalyzerClient]{})
 	iamFetcher := fetchers.NewIAMFetcher(log, iamProvider, ch, identity)
 	m[fetching.IAMType] = registry.RegisteredFetcher{Fetcher: iamFetcher}
 
-	kmsProvider := kms.NewKMSProvider(log, cfg, &awslib.MultiRegionClientFactory[kms.Client]{})
+	kmsProvider := kms.NewKMSProvider(ctx, log, cfg, &awslib.MultiRegionClientFactory[kms.Client]{})
 	kmsFetcher := fetchers.NewKMSFetcher(log, kmsProvider, ch)
 	m[fetching.KmsType] = registry.RegisteredFetcher{Fetcher: kmsFetcher}
 
-	loggingProvider := logging.NewProvider(log, cfg, &awslib.MultiRegionClientFactory[cloudtrail.Client]{}, &awslib.MultiRegionClientFactory[s3.Client]{}, identity.Account)
-	configserviceProvider := configservice.NewProvider(log, cfg, &awslib.MultiRegionClientFactory[configservice.Client]{}, identity.Account)
+	loggingProvider := logging.NewProvider(ctx, log, cfg, &awslib.MultiRegionClientFactory[cloudtrail.Client]{}, &awslib.MultiRegionClientFactory[s3.Client]{}, identity.Account)
+	configserviceProvider := configservice.NewProvider(ctx, log, cfg, &awslib.MultiRegionClientFactory[configservice.Client]{}, identity.Account)
 	loggingFetcher := fetchers.NewLoggingFetcher(log, loggingProvider, configserviceProvider, ch, identity)
 	m[fetching.TrailType] = registry.RegisteredFetcher{Fetcher: loggingFetcher}
 
 	monitoringProvider := monitoring.NewProvider(
+		ctx,
 		log,
 		cfg,
 		&awslib.MultiRegionClientFactory[cloudtrail.Client]{},
@@ -67,19 +70,19 @@ func NewCisAwsFetchers(log *logp.Logger, cfg aws.Config, ch chan fetching.Resour
 		&awslib.MultiRegionClientFactory[sns.Client]{},
 	)
 
-	securityHubProvider := securityhub.NewProvider(log, cfg, &awslib.MultiRegionClientFactory[securityhub.Client]{}, identity.Account)
+	securityHubProvider := securityhub.NewProvider(ctx, log, cfg, &awslib.MultiRegionClientFactory[securityhub.Client]{}, identity.Account)
 	monitoringFetcher := fetchers.NewMonitoringFetcher(log, monitoringProvider, securityHubProvider, ch, identity)
 	m[fetching.AwsMonitoringType] = registry.RegisteredFetcher{Fetcher: monitoringFetcher}
 
-	ec2Provider := ec2.NewEC2Provider(log, identity.Account, cfg, &awslib.MultiRegionClientFactory[ec2.Client]{})
+	ec2Provider := ec2.NewEC2Provider(ctx, log, identity.Account, cfg, &awslib.MultiRegionClientFactory[ec2.Client]{})
 	networkFetcher := fetchers.NewNetworkFetcher(log, ec2Provider, ch)
 	m[fetching.EC2NetworkingType] = registry.RegisteredFetcher{Fetcher: networkFetcher}
 
-	rdsProvider := rds.NewProvider(log, cfg, &awslib.MultiRegionClientFactory[rds.Client]{}, ec2Provider)
+	rdsProvider := rds.NewProvider(ctx, log, cfg, &awslib.MultiRegionClientFactory[rds.Client]{}, ec2Provider)
 	rdsFetcher := fetchers.NewRdsFetcher(log, rdsProvider, ch)
 	m[fetching.RdsType] = registry.RegisteredFetcher{Fetcher: rdsFetcher}
 
-	s3Provider := s3.NewProvider(log, cfg, &awslib.MultiRegionClientFactory[s3.Client]{}, identity.Account)
+	s3Provider := s3.NewProvider(ctx, log, cfg, &awslib.MultiRegionClientFactory[s3.Client]{}, identity.Account)
 	s3Fetcher := fetchers.NewS3Fetcher(log, s3Provider, ch)
 	m[fetching.S3Type] = registry.RegisteredFetcher{Fetcher: s3Fetcher}
 

internal/resources/fetching/preset/eks_preset.go

@@ -18,6 +18,7 @@
 package preset
 
 import (
+	"context"
 	"fmt"
 	"regexp"
 
@@ -45,16 +46,17 @@ var (
 	eksRequiredProcesses = k8sfetchers.ProcessesConfigMap{"kubelet": {ConfigFileArguments: []string{"config"}}}
 	eksFsPatterns        = []string{
 		"/hostfs/etc/kubernetes/kubelet/kubelet-config.json",
-		"/hostfs/var/lib/kubelet/kubeconfig"}
+		"/hostfs/var/lib/kubelet/kubeconfig",
+	}
 )
 
-func NewCisEksFetchers(log *logp.Logger, awsConfig aws.Config, ch chan fetching.ResourceInfo, le uniqueness.Manager, k8sClient k8s.Interface, identity *cloud.Identity) registry.FetchersMap {
+func NewCisEksFetchers(ctx context.Context, log *logp.Logger, awsConfig aws.Config, ch chan fetching.ResourceInfo, le uniqueness.Manager, k8sClient k8s.Interface, identity *cloud.Identity) registry.FetchersMap {
 	log.Infof("Initializing EKS fetchers")
 	m := make(registry.FetchersMap)
 
 	if identity != nil {
 		log.Info("Initialize aws-related fetchers")
-		ecrPrivateProvider := ecr.NewEcrProvider(log, awsConfig, &awslib.MultiRegionClientFactory[ecr.Client]{})
+		ecrPrivateProvider := ecr.NewEcrProvider(ctx, log, awsConfig, &awslib.MultiRegionClientFactory[ecr.Client]{})
 		privateRepoRegex := fmt.Sprintf(awsfetchers.PrivateRepoRegexTemplate, identity.Account)
 
 		ecrPodDescriber := awsfetchers.PodDescriber{
@@ -65,7 +67,7 @@ func NewCisEksFetchers(log *logp.Logger, awsConfig aws.Config, ch chan fetching.
 		ecrFetcher := awsfetchers.NewEcrFetcher(log, ch, k8sClient, ecrPodDescriber)
 		m[fetching.EcrType] = registry.RegisteredFetcher{Fetcher: ecrFetcher, Condition: []fetching.Condition{condition.NewIsLeader(le)}}
 
-		elbProvider := elb.NewElbProvider(log, identity.Account, awsConfig, &awslib.MultiRegionClientFactory[elb.Client]{})
+		elbProvider := elb.NewElbProvider(ctx, log, identity.Account, awsConfig, &awslib.MultiRegionClientFactory[elb.Client]{})
 		loadBalancerRegex := fmt.Sprintf(elbRegexTemplate, awsConfig.Region)
 		elbFetcher := awsfetchers.NewElbFetcher(log, ch, k8sClient, elbProvider, identity, loadBalancerRegex)
 		m[fetching.ElbType] = registry.RegisteredFetcher{Fetcher: elbFetcher, Condition: []fetching.Condition{condition.NewIsLeader(le)}}

internal/resources/providers/aws_cis/logging/logging.go

@@ -39,6 +39,7 @@ type Provider struct {
 }
 
 func NewProvider(
+	ctx context.Context,
 	log *logp.Logger,
 	cfg aws.Config,
 	multiRegionTrailFactory awslib.CrossRegionFactory[cloudtrail.Client],
@@ -47,7 +48,7 @@ func NewProvider(
 ) *Provider {
 	return &Provider{
 		log:           log,
-		s3Provider:    s3.NewProvider(log, cfg, multiRegionS3Factory, accountId),
-		trailProvider: cloudtrail.NewProvider(log, cfg, multiRegionTrailFactory),
+		s3Provider:    s3.NewProvider(ctx, log, cfg, multiRegionS3Factory, accountId),
+		trailProvider: cloudtrail.NewProvider(ctx, log, cfg, multiRegionTrailFactory),
 	}
 }

internal/resources/providers/aws_cis/monitoring/monitoring.go

@@ -64,12 +64,12 @@ type (
 	}
 )
 
-func NewProvider(log *logp.Logger, awsConfig aws.Config, trailCrossRegionFactory awslib.CrossRegionFactory[cloudtrail.Client], cloudwatchCrossResignFactory awslib.CrossRegionFactory[cloudwatch.Client], cloudwatchlogsCrossRegionFactory awslib.CrossRegionFactory[logs.Client], snsCrossRegionFactory awslib.CrossRegionFactory[sns.Client]) *Provider {
+func NewProvider(ctx context.Context, log *logp.Logger, awsConfig aws.Config, trailCrossRegionFactory awslib.CrossRegionFactory[cloudtrail.Client], cloudwatchCrossResignFactory awslib.CrossRegionFactory[cloudwatch.Client], cloudwatchlogsCrossRegionFactory awslib.CrossRegionFactory[logs.Client], snsCrossRegionFactory awslib.CrossRegionFactory[sns.Client]) *Provider {
 	return &Provider{
-		Cloudtrail:     cloudtrail.NewProvider(log, awsConfig, trailCrossRegionFactory),
-		Cloudwatch:     cloudwatch.NewProvider(log, awsConfig, cloudwatchCrossResignFactory),
-		Cloudwatchlogs: logs.NewCloudwatchLogsProvider(log, awsConfig, cloudwatchlogsCrossRegionFactory),
-		Sns:            sns.NewSNSProvider(log, awsConfig, snsCrossRegionFactory),
+		Cloudtrail:     cloudtrail.NewProvider(ctx, log, awsConfig, trailCrossRegionFactory),
+		Cloudwatch:     cloudwatch.NewProvider(ctx, log, awsConfig, cloudwatchCrossResignFactory),
+		Cloudwatchlogs: logs.NewCloudwatchLogsProvider(ctx, log, awsConfig, cloudwatchlogsCrossRegionFactory),
+		Sns:            sns.NewSNSProvider(ctx, log, awsConfig, snsCrossRegionFactory),
 		Log:            log,
 	}
 }