From 113ca0c23ec74de208206cd0ef112ed73adab0a9 Mon Sep 17 00:00:00 2001 From: Paul Querna Date: Sat, 5 Sep 2026 02:47:21 +0000 Subject: [PATCH] Allow embedding applications to inject AWS credential clients Co-authored-by: c1-squire-dev[bot] --- pkg/connector/aws_client_factory.go | 3 +- pkg/connector/connector.go | 69 ++++++++++++-- pkg/connector/connector_options_test.go | 116 ++++++++++++++++++++++++ 3 files changed, 178 insertions(+), 10 deletions(-) create mode 100644 pkg/connector/connector_options_test.go diff --git a/pkg/connector/aws_client_factory.go b/pkg/connector/aws_client_factory.go index f389bf61..0b45adb7 100644 --- a/pkg/connector/aws_client_factory.go +++ b/pkg/connector/aws_client_factory.go @@ -7,7 +7,6 @@ import ( "sync" awsSdk "github.com/aws/aws-sdk-go-v2/aws" - awsConfig "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/service/iam" awsOrgs "github.com/aws/aws-sdk-go-v2/service/organizations" "github.com/aws/aws-sdk-go-v2/service/sts" @@ -60,7 +59,7 @@ func (f *AWSClientFactory) getConfig(ctx context.Context, accountId string) (aws opts := GetAwsConfigOptionsForAssumeRole(output, f.baseClient, f.config) - baseConfig, err := awsConfig.LoadDefaultConfig(ctx, opts...) + baseConfig, err := f.aws.loadAWSConfig(ctx, opts...) if err != nil { return awsSdk.Config{}, err } diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index 1a5e750d..fd0b10cc 100644 --- a/pkg/connector/connector.go +++ b/pkg/connector/connector.go @@ -30,6 +30,41 @@ import ( v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" ) +// AWSConfigLoader resolves the AWS SDK configuration used by the connector. +type AWSConfigLoader func(context.Context, ...func(*awsConfig.LoadOptions) error) (awsSdk.Config, error) + +// STSClient is the STS operation set the connector needs for credential +// exchanges and caller-identity validation. +type STSClient interface { + stscreds.AssumeRoleAPIClient + GetCallerIdentity(context.Context, *sts.GetCallerIdentityInput, ...func(*sts.Options)) (*sts.GetCallerIdentityOutput, error) +} + +// STSClientFactory constructs an STS client for role assumptions. +type STSClientFactory func(awsSdk.Config) STSClient + +type options struct { + loadAWSConfig AWSConfigLoader + newSTSClient STSClientFactory +} + +// Option configures the AWS connector. +type Option func(*options) + +// WithAWSConfigLoader replaces the default AWS SDK configuration loader. +func WithAWSConfigLoader(loader AWSConfigLoader) Option { + return func(o *options) { + o.loadAWSConfig = loader + } +} + +// WithSTSClientFactory replaces the STS client constructor used for role assumptions. +func WithSTSClientFactory(factory STSClientFactory) Option { + return func(o *options) { + o.newSTSClient = factory + } +} + type Config struct { UseAssumeRole bool GlobalBindingExternalID string @@ -69,6 +104,8 @@ type AWS struct { _onceCallingConfig map[string]*sync.Once _callingConfig map[string]awsSdk.Config _callingConfigError map[string]error + loadAWSConfig AWSConfigLoader + newSTSClient STSClientFactory _identityInstancesCacheMtx sync.Mutex _identityInstancesCacheErr error @@ -117,12 +154,12 @@ func (o *AWS) getSSOSCIMClient(ctx context.Context) (*awsIdentityCenterSCIMClien }, nil } -func (o *AWS) getSTSClient(ctx context.Context) (*sts.Client, error) { +func (o *AWS) getSTSClient(ctx context.Context) (STSClient, error) { callingConfig, err := o.getCallingConfig(ctx, o.globalRegion) if err != nil { return nil, err } - return sts.NewFromConfig(callingConfig), nil + return o.newSTSClient(callingConfig), nil } func (o *AWS) getCallingConfig(ctx context.Context, region string) (awsSdk.Config, error) { @@ -137,7 +174,7 @@ func (o *AWS) getCallingConfig(ctx context.Context, region string) (awsSdk.Confi l := ctxzap.Extract(ctx) // ok, if we are an instance, we do the assumeRole twice, first time from our Instance role, INTO the binding account // and from there, into the customer account. - stsSvc := sts.NewFromConfig(o.baseConfig) + stsSvc := o.newSTSClient(o.baseConfig) bindingCreds := awsSdk.NewCredentialsCache(stscreds.NewAssumeRoleProvider(stsSvc, o.globalRoleARN, func(aro *stscreds.AssumeRoleOptions) { if o.globalBindingExternalID != "" { aro.ExternalID = awsSdk.String(o.globalBindingExternalID) @@ -159,7 +196,7 @@ func (o *AWS) getCallingConfig(ctx context.Context, region string) (awsSdk.Confi stsConfig := o.baseConfig.Copy() stsConfig.Credentials = bindingCreds - callingSTSService := sts.NewFromConfig(stsConfig) + callingSTSService := o.newSTSClient(stsConfig) callingConfig := awsSdk.Config{ HTTPClient: o.baseClient, @@ -181,15 +218,29 @@ func (o *AWS) getCallingConfig(ctx context.Context, region string) (awsSdk.Confi return o._callingConfig[region], o._callingConfigError[region] } -func New(ctx context.Context, config Config) (*AWS, error) { +func New(ctx context.Context, config Config, optFns ...Option) (*AWS, error) { + opts := options{ + loadAWSConfig: awsConfig.LoadDefaultConfig, + newSTSClient: func(cfg awsSdk.Config) STSClient { + return sts.NewFromConfig(cfg) + }, + } + for _, fn := range optFns { + if fn != nil { + fn(&opts) + } + } + if opts.loadAWSConfig == nil || opts.newSTSClient == nil { + return nil, fmt.Errorf("aws connector: AWS config loader and STS client factory are required") + } + httpClient, err := uhttp.NewClient(ctx, uhttp.WithLogger(true, ctxzap.Extract(ctx))) if err != nil { return nil, err } - opts := GetAwsConfigOptions(httpClient, config) - - baseConfig, err := awsConfig.LoadDefaultConfig(ctx, opts...) + awsOpts := GetAwsConfigOptions(httpClient, config) + baseConfig, err := opts.loadAWSConfig(ctx, awsOpts...) if err != nil { return nil, fmt.Errorf("aws connector: config load failure: %w", err) } @@ -215,6 +266,8 @@ func New(ctx context.Context, config Config) (*AWS, error) { _callingConfig: map[string]awsSdk.Config{}, _callingConfigError: map[string]error{}, syncSecrets: config.SyncSecrets, + loadAWSConfig: opts.loadAWSConfig, + newSTSClient: opts.newSTSClient, } rv.awsClientFactory = NewAWSClientFactory(config, rv, httpClient) diff --git a/pkg/connector/connector_options_test.go b/pkg/connector/connector_options_test.go new file mode 100644 index 00000000..f5135de6 --- /dev/null +++ b/pkg/connector/connector_options_test.go @@ -0,0 +1,116 @@ +package connector + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + awsSdk "github.com/aws/aws-sdk-go-v2/aws" + awsConfig "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/service/sts" + ststypes "github.com/aws/aws-sdk-go-v2/service/sts/types" + "github.com/stretchr/testify/require" +) + +type injectedSTSClient struct { + calls int +} + +func (c *injectedSTSClient) AssumeRole(context.Context, *sts.AssumeRoleInput, ...func(*sts.Options)) (*sts.AssumeRoleOutput, error) { + c.calls++ + expiresAt := time.Now().Add(time.Hour) + return &sts.AssumeRoleOutput{Credentials: &ststypes.Credentials{ + AccessKeyId: awsSdk.String("access-key"), + SecretAccessKey: awsSdk.String("secret-key"), + SessionToken: awsSdk.String("session-token"), + Expiration: &expiresAt, + }}, nil +} + +func (c *injectedSTSClient) GetCallerIdentity(context.Context, *sts.GetCallerIdentityInput, ...func(*sts.Options)) (*sts.GetCallerIdentityOutput, error) { + return &sts.GetCallerIdentityOutput{}, nil +} + +func TestNewUsesInjectedAWSConfigLoader(t *testing.T) { + t.Parallel() + + loaderErr := errors.New("injected config loader") + called := false + _, err := New(t.Context(), Config{}, WithAWSConfigLoader(func(context.Context, ...func(*awsConfig.LoadOptions) error) (awsSdk.Config, error) { + called = true + return awsSdk.Config{}, loaderErr + })) + + require.ErrorIs(t, err, loaderErr) + require.True(t, called) +} + +func TestInjectedSTSClientFactoryDrivesAssumeRole(t *testing.T) { + t.Parallel() + + stsClient := &injectedSTSClient{} + factoryCalls := 0 + connector := &AWS{ + useAssumeRole: true, + globalRoleARN: "arn:aws:iam::123456789012:role/binding", + roleARN: "arn:aws:iam::123456789012:role/customer", + baseConfig: awsSdk.Config{Region: "us-west-2"}, + newSTSClient: func(awsSdk.Config) STSClient { + factoryCalls++ + return stsClient + }, + _onceCallingConfig: map[string]*sync.Once{}, + _callingConfig: map[string]awsSdk.Config{}, + _callingConfigError: map[string]error{}, + } + + _, err := connector.getCallingConfig(t.Context(), "us-west-2") + + require.NoError(t, err) + require.Equal(t, 2, factoryCalls) + require.Equal(t, 2, stsClient.calls) +} + +func TestInjectedHooksDriveChildAccountAssumeRole(t *testing.T) { + t.Parallel() + + stsClient := &injectedSTSClient{} + loaderCalls := 0 + factoryCalls := 0 + connector := &AWS{ + useAssumeRole: true, + globalRoleARN: "arn:aws:iam::123456789012:role/binding", + roleARN: "arn:aws:iam::123456789012:role/customer", + baseConfig: awsSdk.Config{Region: "us-west-2"}, + loadAWSConfig: func(context.Context, ...func(*awsConfig.LoadOptions) error) (awsSdk.Config, error) { + loaderCalls++ + return awsSdk.Config{Region: "us-west-2"}, nil + }, + newSTSClient: func(awsSdk.Config) STSClient { + factoryCalls++ + return stsClient + }, + _onceCallingConfig: map[string]*sync.Once{}, + _callingConfig: map[string]awsSdk.Config{}, + _callingConfigError: map[string]error{}, + } + + _, err := NewAWSClientFactory(Config{IamAssumeRoleName: "OrganizationAccountAccessRole"}, connector, nil).getConfig(t.Context(), "210987654321") + + require.NoError(t, err) + require.Equal(t, 1, loaderCalls) + require.Equal(t, 3, factoryCalls) + require.Equal(t, 3, stsClient.calls) +} + +func TestNewRejectsNilAWSHooks(t *testing.T) { + t.Parallel() + + _, err := New(t.Context(), Config{}, WithAWSConfigLoader(nil)) + require.EqualError(t, err, "aws connector: AWS config loader and STS client factory are required") + + _, err = New(t.Context(), Config{}, WithSTSClientFactory(nil)) + require.EqualError(t, err, "aws connector: AWS config loader and STS client factory are required") +}