GODRIVER-3454 Support custom AWS credential provider.

.evergreen/config.yml

@@ -517,6 +517,20 @@ functions:
         binary: "bash"
         include_expansions_in_env: [KMS_TLS_TESTCASE]
         args: [*task-runner, evg-test-kms]
+  run-custom-aws-credentials-test:
+    - command: subprocess.exec
+      params:
+        binary: "bash"
+        env:
+          GO_BUILD_TAGS: cse
+        include_expansions_in_env: [AUTH, SSL, MONGODB_URI, TOPOLOGY, MONGO_GO_DRIVER_COMPRESSOR]
+        args: [*task-runner, setup-test]
+    - command: subprocess.exec
+      type: test
+      retry_on_failure: true
+      params:
+        binary: "bash"
+        args: [*task-runner, evg-test-custom-aws-credentials]
   run-kmip-tests:
     - command: subprocess.exec
       params:
@@ -1438,6 +1452,16 @@ tasks:
           TOPOLOGY: "server"
           AUTH: "noauth"
           SSL: "nossl"
+  - name: "test-custom-aws-credentials"
+    tags: ["kms-test"]
+    commands:
+      - func: bootstrap-mongo-orchestration
+        vars:
+          TOPOLOGY: "server"
+          AUTH: "noauth"
+          SSL: "nossl"
+      - func: start-cse-servers
+      - func: run-custom-aws-credentials-test
   - name: "test-kms-kmip"
     tags: ["kms-kmip"]
     commands:

Taskfile.yml

@@ -96,6 +96,8 @@ tasks:
     - go test -exec "env PKG_CONFIG_PATH=${PKG_CONFIG_PATH} LD_LIBRARY_PATH=${LD_LIBRARY_PATH} DYLD_LIBRARY_PATH=${MACOS_LIBRARY_PATH}" ${BUILD_TAGS} -v -timeout {{.TEST_TIMEOUT}}s ./internal/integration -run TestClientSideEncryptionProse/kms_tls_options_test >> test.suite
   evg-test-kms:
     - go test -exec "env PKG_CONFIG_PATH=${PKG_CONFIG_PATH} LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" ${BUILD_TAGS} -v -timeout {{.TEST_TIMEOUT}}s ./internal/integration -run TestClientSideEncryptionProse/kms_tls_tests >> test.suite
+  evg-test-custom-aws-credentials:
+    - go test -exec "env PKG_CONFIG_PATH=${PKG_CONFIG_PATH} LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" ${BUILD_TAGS} -v -timeout {{.TEST_TIMEOUT}}s ./internal/integration -run TestCustomAwsCredentialsProse >> test.suite
   evg-test-retry-kms-requests:
     - go test -exec "env PKG_CONFIG_PATH=${PKG_CONFIG_PATH} LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" ${BUILD_TAGS} -v -timeout {{.TEST_TIMEOUT}}s ./internal/integration -run TestClientSideEncryptionProse/kms_retry_tests >> test.suite
   evg-test-load-balancers:

go.work

@@ -9,4 +9,5 @@ use (
 	./internal/cmd/faas/awslambda/mongodb
 	./internal/test/compilecheck
 	./internal/test/goleak
+	./modules/mongoaws
 )

internal/aws/credentials/chain_provider.go

@@ -11,6 +11,8 @@
 package credentials
 
 import (
+	"context"
+
 	"go.mongodb.org/mongo-driver/v2/internal/aws/awserr"
 )
 
@@ -45,10 +47,10 @@ func NewChainCredentials(providers []Provider) *Credentials {
 //
 // If a provider is found it will be cached and any calls to IsExpired()
 // will return the expired state of the cached provider.
-func (c *ChainProvider) Retrieve() (Value, error) {
+func (c *ChainProvider) Retrieve(ctx context.Context) (Value, error) {
 	var errs = make([]error, 0, len(c.Providers))
 	for _, p := range c.Providers {
-		creds, err := p.Retrieve()
+		creds, err := p.Retrieve(ctx)
 		if err == nil {
 			c.curr = p
 			return creds, nil

internal/aws/credentials/chain_provider_test.go

@@ -11,6 +11,7 @@
 package credentials
 
 import (
+	"context"
 	"reflect"
 	"testing"
 
@@ -23,7 +24,7 @@ type secondStubProvider struct {
 	err     error
 }
 
-func (s *secondStubProvider) Retrieve() (Value, error) {
+func (s *secondStubProvider) Retrieve(_ context.Context) (Value, error) {
 	s.expired = false
 	s.creds.ProviderName = "secondStubProvider"
 	return s.creds, s.err
@@ -54,7 +55,7 @@ func TestChainProviderWithNames(t *testing.T) {
 		},
 	}
 
-	creds, err := p.Retrieve()
+	creds, err := p.Retrieve(context.Background())
 	if err != nil {
 		t.Errorf("Expect no error, got %v", err)
 	}
@@ -90,7 +91,7 @@ func TestChainProviderGet(t *testing.T) {
 		},
 	}
 
-	creds, err := p.Retrieve()
+	creds, err := p.Retrieve(context.Background())
 	if err != nil {
 		t.Errorf("Expect no error, got %v", err)
 	}
@@ -113,10 +114,12 @@ func TestChainProviderIsExpired(t *testing.T) {
 		},
 	}
 
+	ctx := context.Background()
+
 	if !p.IsExpired() {
 		t.Errorf("Expect expired to be true before any Retrieve")
 	}
-	_, err := p.Retrieve()
+	_, err := p.Retrieve(ctx)
 	if err != nil {
 		t.Errorf("Expect no error, got %v", err)
 	}
@@ -129,7 +132,7 @@ func TestChainProviderIsExpired(t *testing.T) {
 		t.Errorf("Expect return of expired provider")
 	}
 
-	_, err = p.Retrieve()
+	_, err = p.Retrieve(ctx)
 	if err != nil {
 		t.Errorf("Expect no error, got %v", err)
 	}
@@ -146,7 +149,7 @@ func TestChainProviderWithNoProvider(t *testing.T) {
 	if !p.IsExpired() {
 		t.Errorf("Expect expired with no providers")
 	}
-	_, err := p.Retrieve()
+	_, err := p.Retrieve(context.Background())
 	if err.Error() != "NoCredentialProviders: no valid providers in chain" {
 		t.Errorf("Expect no providers error returned, got %v", err)
 	}
@@ -167,7 +170,7 @@ func TestChainProviderWithNoValidProvider(t *testing.T) {
 	if !p.IsExpired() {
 		t.Errorf("Expect expired with no providers")
 	}
-	_, err := p.Retrieve()
+	_, err := p.Retrieve(context.Background())
 
 	expectErr := awserr.NewBatchError("NoCredentialProviders", "no valid providers in chain", errs)
 	if e, a := expectErr, err; !reflect.DeepEqual(e, a) {

internal/aws/credentials/credentials.go

@@ -33,6 +33,19 @@ type Value struct {
 	// AWS Session Token
 	SessionToken string
 
+	// Source of the credentials
+	Source string
+
+	// States if the credentials can expire or not.
+	CanExpire bool
+
+	// The time the credentials will expire at. Should be ignored if CanExpire
+	// is false.
+	Expires time.Time
+
+	// The ID of the account for the credentials.
+	AccountID string
+
 	// Provider used to get credentials
 	ProviderName string
 }
@@ -52,20 +65,13 @@ func (v Value) HasKeys() bool {
 type Provider interface {
 	// Retrieve returns nil if it successfully retrieved the value.
 	// Error is returned if the value were not obtainable, or empty.
-	Retrieve() (Value, error)
+	Retrieve(context.Context) (Value, error)
 
 	// IsExpired returns if the credentials are no longer valid, and need
 	// to be retrieved.
 	IsExpired() bool
 }
 
-// ProviderWithContext is a Provider that can retrieve credentials with a Context
-type ProviderWithContext interface {
-	Provider
-
-	RetrieveWithContext(context.Context) (Value, error)
-}
-
 // A Credentials provides concurrency safe retrieval of AWS credentials Value.
 //
 // A Credentials is also used to fetch Azure credentials Value.
@@ -143,13 +149,7 @@ func (c *Credentials) singleRetrieve(ctx context.Context) (interface{}, error) {
 		return curCreds, nil
 	}
 
-	var creds Value
-	var err error
-	if p, ok := c.provider.(ProviderWithContext); ok {
-		creds, err = p.RetrieveWithContext(ctx)
-	} else {
-		creds, err = c.provider.Retrieve()
-	}
+	creds, err := c.provider.Retrieve(ctx)
 	if err == nil {
 		c.creds = creds
 	}

internal/aws/credentials/credentials_test.go

@@ -33,7 +33,7 @@ type stubProvider struct {
 	err            error
 }
 
-func (s *stubProvider) Retrieve() (Value, error) {
+func (s *stubProvider) Retrieve(_ context.Context) (Value, error) {
 	s.retrievedCount++
 	s.expired = false
 	s.creds.ProviderName = "stubProvider"
@@ -133,7 +133,7 @@ func (e *MockProvider) IsExpired() bool {
 	return e.expiration.Before(curTime())
 }
 
-func (*MockProvider) Retrieve() (Value, error) {
+func (*MockProvider) Retrieve(_ context.Context) (Value, error) {
 	return Value{}, nil
 }
 
@@ -162,9 +162,9 @@ type stubProviderConcurrent struct {
 	done chan struct{}
 }
 
-func (s *stubProviderConcurrent) Retrieve() (Value, error) {
+func (s *stubProviderConcurrent) Retrieve(ctx context.Context) (Value, error) {
 	<-s.done
-	return s.stubProvider.Retrieve()
+	return s.stubProvider.Retrieve(ctx)
 }
 
 func TestCredentialsGetConcurrent(t *testing.T) {

internal/credproviders/assume_role_provider.go

@@ -57,8 +57,8 @@ func NewAssumeRoleProvider(httpClient *http.Client, expiryWindow time.Duration)
 	}
 }
 
-// RetrieveWithContext retrieves the keys from the AWS service.
-func (a *AssumeRoleProvider) RetrieveWithContext(ctx context.Context) (credentials.Value, error) {
+// Retrieve retrieves the keys from the AWS service.
+func (a *AssumeRoleProvider) Retrieve(ctx context.Context) (credentials.Value, error) {
 	const defaultHTTPTimeout = 10 * time.Second
 
 	v := credentials.Value{ProviderName: assumeRoleProviderName}
@@ -137,11 +137,6 @@ func (a *AssumeRoleProvider) RetrieveWithContext(ctx context.Context) (credentia
 	return v, nil
 }
 
-// Retrieve retrieves the keys from the AWS service.
-func (a *AssumeRoleProvider) Retrieve() (credentials.Value, error) {
-	return a.RetrieveWithContext(context.Background())
-}
-
 // IsExpired returns true if the credentials are expired.
 func (a *AssumeRoleProvider) IsExpired() bool {
 	return a.expiration.Before(time.Now())

internal/credproviders/ec2_provider.go

@@ -146,8 +146,8 @@ func (e *EC2Provider) getCredentials(ctx context.Context, token string, role str
 	return v, ec2Resp.Expiration, nil
 }
 
-// RetrieveWithContext retrieves the keys from the AWS service.
-func (e *EC2Provider) RetrieveWithContext(ctx context.Context) (credentials.Value, error) {
+// Retrieve retrieves the keys from the AWS service.
+func (e *EC2Provider) Retrieve(ctx context.Context) (credentials.Value, error) {
 	v := credentials.Value{ProviderName: ec2ProviderName}
 
 	token, err := e.getToken(ctx)
@@ -172,11 +172,6 @@ func (e *EC2Provider) RetrieveWithContext(ctx context.Context) (credentials.Valu
 	return v, nil
 }
 
-// Retrieve retrieves the keys from the AWS service.
-func (e *EC2Provider) Retrieve() (credentials.Value, error) {
-	return e.RetrieveWithContext(context.Background())
-}
-
 // IsExpired returns true if the credentials are expired.
 func (e *EC2Provider) IsExpired() bool {
 	return e.expiration.Before(time.Now())

internal/credproviders/ecs_provider.go

@@ -49,8 +49,8 @@ func NewECSProvider(httpClient *http.Client, expiryWindow time.Duration) *ECSPro
 	}
 }
 
-// RetrieveWithContext retrieves the keys from the AWS service.
-func (e *ECSProvider) RetrieveWithContext(ctx context.Context) (credentials.Value, error) {
+// Retrieve retrieves the keys from the AWS service.
+func (e *ECSProvider) Retrieve(ctx context.Context) (credentials.Value, error) {
 	const defaultHTTPTimeout = 10 * time.Second
 
 	v := credentials.Value{ProviderName: ecsProviderName}
@@ -101,11 +101,6 @@ func (e *ECSProvider) RetrieveWithContext(ctx context.Context) (credentials.Valu
 	return v, nil
 }
 
-// Retrieve retrieves the keys from the AWS service.
-func (e *ECSProvider) Retrieve() (credentials.Value, error) {
-	return e.RetrieveWithContext(context.Background())
-}
-
 // IsExpired returns true if the credentials are expired.
 func (e *ECSProvider) IsExpired() bool {
 	return e.expiration.Before(time.Now())

internal/credproviders/env_provider.go

@@ -7,6 +7,7 @@
 package credproviders
 
 import (
+	"context"
 	"os"
 
 	"go.mongodb.org/mongo-driver/v2/internal/aws/credentials"
@@ -46,7 +47,7 @@ func NewEnvProvider() *EnvProvider {
 }
 
 // Retrieve retrieves the keys from the environment.
-func (e *EnvProvider) Retrieve() (credentials.Value, error) {
+func (e *EnvProvider) Retrieve(_ context.Context) (credentials.Value, error) {
 	e.retrieved = false
 
 	v := credentials.Value{

internal/credproviders/imds_provider.go

@@ -41,8 +41,8 @@ func NewAzureProvider(httpClient *http.Client, expiryWindow time.Duration) *Azur
 	}
 }
 
-// RetrieveWithContext retrieves the keys from the Azure service.
-func (a *AzureProvider) RetrieveWithContext(ctx context.Context) (credentials.Value, error) {
+// Retrieve retrieves the keys from the Azure service.
+func (a *AzureProvider) Retrieve(ctx context.Context) (credentials.Value, error) {
 	v := credentials.Value{ProviderName: AzureProviderName}
 	req, err := http.NewRequest(http.MethodGet, azureURI, nil)
 	if err != nil {
@@ -92,11 +92,6 @@ func (a *AzureProvider) RetrieveWithContext(ctx context.Context) (credentials.Va
 	return v, err
 }
 
-// Retrieve retrieves the keys from the Azure service.
-func (a *AzureProvider) Retrieve() (credentials.Value, error) {
-	return a.RetrieveWithContext(context.Background())
-}
-
 // IsExpired returns if the credentials have been retrieved.
 func (a *AzureProvider) IsExpired() bool {
 	return a.expiration.Before(time.Now())

internal/credproviders/static_provider.go

@@ -7,6 +7,7 @@
 package credproviders
 
 import (
+	"context"
 	"errors"
 
 	"go.mongodb.org/mongo-driver/v2/internal/aws/credentials"
@@ -42,7 +43,7 @@ func verify(v credentials.Value) error {
 }
 
 // Retrieve returns the credentials or error if the credentials are invalid.
-func (s *StaticProvider) Retrieve() (credentials.Value, error) {
+func (s *StaticProvider) Retrieve(_ context.Context) (credentials.Value, error) {
 	if !s.verified {
 		s.err = verify(s.Value)
 		s.ProviderName = staticProviderName