Skip to content

GODRIVER-3567 Update the Retrieve interface. #2481

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
qingyang-hu wants to merge 1 commit into
base: master
Choose a base branch
from
+213 −517
Open

GODRIVER-3567 Update the Retrieve interface. #2481

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 10 additions & 27 deletions internal/aws/credentials/chain_provider.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,25 +11,20 @@
package credentials

import (
"context"
"errors"

"go.mongodb.org/mongo-driver/v2/internal/aws/awserr"
)

// A ChainProvider will search for a provider which returns credentials
// and cache that provider until Retrieve is called again.
//
// The ChainProvider provides a way of chaining multiple providers together
// which will pick the first available using priority order of the Providers
// in the list.
//
// If none of the Providers retrieve valid credentials Value, ChainProvider's
// Retrieve() will return the error ErrNoValidProvidersFoundInChain.
//
// If a Provider is found which returns valid credentials Value ChainProvider
// will cache that Provider for all calls to IsExpired(), until Retrieve is
// called again.
// Retrieve() will return an error.
type ChainProvider struct {
Providers []Provider
curr Provider
}

// NewChainCredentials returns a pointer to a new Credentials object
Expand All @@ -42,31 +37,19 @@ func NewChainCredentials(providers []Provider) *Credentials {

// Retrieve returns the credentials value or error if no provider returned
// without error.
//
// 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) {
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
if !creds.Expired() && creds.HasKeys() {
return creds, nil
}
err = errors.New("credentials are invalid")
}
errs = append(errs, err)
}
c.curr = nil

err := awserr.NewBatchError("NoCredentialProviders", "no valid providers in chain", errs)
return Value{}, err
}

// IsExpired will returned the expired state of the currently cached provider
// if there is one. If there is no current provider, true will be returned.
func (c *ChainProvider) IsExpired() bool {
if c.curr != nil {
return c.curr.IsExpired()
}

return true
}
99 changes: 6 additions & 93 deletions internal/aws/credentials/chain_provider_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,34 +11,19 @@
package credentials

import (
"context"
"reflect"
"testing"

"go.mongodb.org/mongo-driver/v2/internal/aws/awserr"
)

type secondStubProvider struct {
creds Value
expired bool
err error
}

func (s *secondStubProvider) Retrieve() (Value, error) {
s.expired = false
s.creds.ProviderName = "secondStubProvider"
return s.creds, s.err
}

func (s *secondStubProvider) IsExpired() bool {
return s.expired
}

func TestChainProviderWithNames(t *testing.T) {
func TestChainProviderRetrieve(t *testing.T) {
p := &ChainProvider{
Providers: []Provider{
&stubProvider{err: awserr.New("FirstError", "first provider error", nil)},
&stubProvider{err: awserr.New("SecondError", "second provider error", nil)},
&secondStubProvider{
&stubProvider{
creds: Value{
AccessKeyID: "AKIF",
SecretAccessKey: "NOSECRET",
Expand All @@ -55,13 +40,10 @@ 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)
}
if e, a := "secondStubProvider", creds.ProviderName; e != a {
t.Errorf("Expect provider name to match, %v got, %v", e, a)
}

// Also check credentials
if e, a := "AKIF", creds.AccessKeyID; e != a {
Expand All @@ -75,78 +57,12 @@ func TestChainProviderWithNames(t *testing.T) {
}
}

func TestChainProviderGet(t *testing.T) {
p := &ChainProvider{
Providers: []Provider{
&stubProvider{err: awserr.New("FirstError", "first provider error", nil)},
&stubProvider{err: awserr.New("SecondError", "second provider error", nil)},
&stubProvider{
creds: Value{
AccessKeyID: "AKID",
SecretAccessKey: "SECRET",
SessionToken: "",
},
},
},
}

creds, err := p.Retrieve()
if err != nil {
t.Errorf("Expect no error, got %v", err)
}
if e, a := "AKID", creds.AccessKeyID; e != a {
t.Errorf("Expect access key ID to match, %v got %v", e, a)
}
if e, a := "SECRET", creds.SecretAccessKey; e != a {
t.Errorf("Expect secret access key to match, %v got %v", e, a)
}
if v := creds.SessionToken; len(v) != 0 {
t.Errorf("Expect session token to be empty, %v", v)
}
}

func TestChainProviderIsExpired(t *testing.T) {
stubProvider := &stubProvider{expired: true}
p := &ChainProvider{
Providers: []Provider{
stubProvider,
},
}

if !p.IsExpired() {
t.Errorf("Expect expired to be true before any Retrieve")
}
_, err := p.Retrieve()
if err != nil {
t.Errorf("Expect no error, got %v", err)
}
if p.IsExpired() {
t.Errorf("Expect not expired after retrieve")
}

stubProvider.expired = true
if !p.IsExpired() {
t.Errorf("Expect return of expired provider")
}

_, err = p.Retrieve()
if err != nil {
t.Errorf("Expect no error, got %v", err)
}
if p.IsExpired() {
t.Errorf("Expect not expired after retrieve")
}
}

func TestChainProviderWithNoProvider(t *testing.T) {
p := &ChainProvider{
Providers: []Provider{},
}

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)
}
Expand All @@ -164,10 +80,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) {
Expand Down
124 changes: 51 additions & 73 deletions internal/aws/credentials/credentials.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ package credentials

import (
"context"
"sync"
"sync/atomic"
"time"

"go.mongodb.org/mongo-driver/v2/internal/aws/awserr"
Expand All @@ -33,8 +33,30 @@ type Value struct {
// AWS Session Token
SessionToken string

// Provider used to get credentials
ProviderName 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
}

// Expired returns if the credentials have expired.
func (v Value) Expired() bool {
if v.CanExpire {
// Calling Round(0) on the current time will truncate the monotonic
// reading only. Ensures credential expiry time is always based on
// reported wall-clock time.
return !v.Expires.After(time.Now().Round(0))
}

return false
}

// HasKeys returns if the credentials Value has both AccessKeyID and
Expand All @@ -52,18 +74,7 @@ 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)

// 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)
Retrieve(context.Context) (Value, error)
}

// A Credentials provides concurrency safe retrieval of AWS credentials Value.
Expand All @@ -79,13 +90,12 @@ type ProviderWithContext interface {
//
// The first Credentials.Get() will always call Provider.Retrieve() to get the
// first instance of the credentials Value. All calls to Get() after that
// will return the cached credentials Value until IsExpired() returns true.
// will return the cached credentials Value until Expired() returns true.
type Credentials struct {
sf singleflight.Group

m sync.RWMutex
creds Value
provider Provider

creds atomic.Value
sf singleflight.Group
}

// NewCredentials returns a pointer to a new Credentials with the provider set.
Expand All @@ -96,34 +106,18 @@ func NewCredentials(provider Provider) *Credentials {
return c
}

// GetWithContext returns the credentials value, or error if the credentials
// Get returns the credentials value, or error if the credentials
// Value failed to be retrieved. Will return early if the passed in context is
// canceled.
//
// Will return the cached credentials Value if it has not expired. If the
// credentials Value has expired the Provider's Retrieve() will be called
// to refresh the credentials.
//
// If Credentials.Expire() was called the credentials Value will be force
// expired, and the next call to Get() will cause them to be refreshed.
func (c *Credentials) GetWithContext(ctx context.Context) (Value, error) {
// Check if credentials are cached, and not expired.
select {
case curCreds, ok := <-c.asyncIsExpired():
// ok will only be true, of the credentials were not expired. ok will
// be false and have no value if the credentials are expired.
if ok {
return curCreds, nil
}
case <-ctx.Done():
return Value{}, awserr.New("RequestCanceled",
"request context canceled", ctx.Err())
}

func (c *Credentials) Get(ctx context.Context) (Value, error) {
// Cannot pass context down to the actual retrieve, because the first
// context would cancel the whole group when there is not direct
// association of items in the group.
resCh := c.sf.DoChan("", func() (interface{}, error) {
resCh := c.sf.DoChan("", func() (any, error) {
return c.singleRetrieve(&suppressedContext{ctx})
})
select {
Expand All @@ -135,49 +129,33 @@ func (c *Credentials) GetWithContext(ctx context.Context) (Value, error) {
}
}

func (c *Credentials) singleRetrieve(ctx context.Context) (interface{}, error) {
c.m.Lock()
defer c.m.Unlock()

if curCreds := c.creds; !c.isExpiredLocked(curCreds) {
return curCreds, nil
func (c *Credentials) singleRetrieve(ctx context.Context) (any, error) {
if currCreds, ok := c.getCreds(); ok && !currCreds.Expired() {
return currCreds, 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()
}
newCreds, err := c.provider.Retrieve(ctx)
if err == nil {
c.creds = creds
c.creds.Store(&newCreds)
}

return creds, err
return newCreds, err
}

// asyncIsExpired returns a channel of credentials Value. If the channel is
// closed the credentials are expired and credentials value are not empty.
func (c *Credentials) asyncIsExpired() <-chan Value {
ch := make(chan Value, 1)
go func() {
c.m.RLock()
defer c.m.RUnlock()

if curCreds := c.creds; !c.isExpiredLocked(curCreds) {
ch <- curCreds
}

close(ch)
}()
// getCreds returns the currently stored credentials and true. Returning false
// if no credentials were stored.
func (c *Credentials) getCreds() (Value, bool) {
v := c.creds.Load()
if v == nil {
return Value{}, false
}

return ch
}
val := v.(*Value)
if val == nil || !val.HasKeys() {
return Value{}, false
}
Comment on lines +153 to +156

// isExpiredLocked helper method wrapping the definition of expired credentials.
func (c *Credentials) isExpiredLocked(creds interface{}) bool {
return creds == nil || creds.(Value) == Value{} || c.provider.IsExpired()
return *val, true
}

type suppressedContext struct {
Expand Down
Loading
Loading