support custom AWS credential provider

Author: qingyang-hu

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

@@ -52,20 +52,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 +136,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/aws_provider.go

@@ -0,0 +1,58 @@
+// Copyright (C) MongoDB, Inc. 2025-present.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+package credproviders
+
+import (
+	"context"
+
+	"go.mongodb.org/mongo-driver/v2/internal/aws/credentials"
+)
+
+const awsProviderName = "AwsProvider"
+
+// AwsCredentialsProvider is a function that retrieves AWS credentials.
+type AwsCredentialsProvider func(context.Context) (AwsCredentials, error)
+
+// AwsCredentials represents AWS credentials with an expiration callback.
+type AwsCredentials struct {
+	AccessKeyID        string
+	SecretAccessKey    string
+	SessionToken       string
+	ExpirationCallback func() bool
+}
+
+// AwsProvider retrieves credentials from the given AWS credentials provider.
+type AwsProvider struct {
+	credentials *AwsCredentials
+
+	Provider AwsCredentialsProvider
+}
+
+// Retrieve retrieves the keys from the given AWS credentials provider.
+func (a *AwsProvider) Retrieve(ctx context.Context) (credentials.Value, error) {
+	var value credentials.Value
+	if a.credentials == nil {
+		creds, err := a.Provider(ctx)
+		if err != nil {
+			return value, err
+		}
+		a.credentials = &creds
+	}
+	value.AccessKeyID = a.credentials.AccessKeyID
+	value.SecretAccessKey = a.credentials.SecretAccessKey
+	value.SessionToken = a.credentials.SessionToken
+	value.ProviderName = awsProviderName
+	return value, nil
+}
+
+// IsExpired returns true if the credentials have not been retrieved.
+func (a *AwsProvider) IsExpired() bool {
+	if a.credentials == nil || a.credentials.ExpirationCallback == nil {
+		return true
+	}
+	return a.credentials.ExpirationCallback()
+}

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())