1 package bearer 2 3 import ( 4 "context" 5 "fmt" 6 "sync/atomic" 7 "time" 8 9 smithycontext "github.com/aws/smithy-go/context" 10 "github.com/aws/smithy-go/internal/sync/singleflight" 11 ) 12 13 // package variable that can be override in unit tests. 14 var timeNow = time.Now 15 16 // TokenCacheOptions provides a set of optional configuration options for the 17 // TokenCache TokenProvider. 18 type TokenCacheOptions struct { 19 // The duration before the token will expire when the credentials will be 20 // refreshed. If DisableAsyncRefresh is true, the RetrieveBearerToken calls 21 // will be blocking. 22 // 23 // Asynchronous refreshes are deduplicated, and only one will be in-flight 24 // at a time. If the token expires while an asynchronous refresh is in 25 // flight, the next call to RetrieveBearerToken will block on that refresh 26 // to return. 27 RefreshBeforeExpires time.Duration 28 29 // The timeout the underlying TokenProvider's RetrieveBearerToken call must 30 // return within, or will be canceled. Defaults to 0, no timeout. 31 // 32 // If 0 timeout, its possible for the underlying tokenProvider's 33 // RetrieveBearerToken call to block forever. Preventing subsequent 34 // TokenCache attempts to refresh the token. 35 // 36 // If this timeout is reached all pending deduplicated calls to 37 // TokenCache RetrieveBearerToken will fail with an error. 38 RetrieveBearerTokenTimeout time.Duration 39 40 // The minimum duration between asynchronous refresh attempts. If the next 41 // asynchronous recent refresh attempt was within the minimum delay 42 // duration, the call to retrieve will return the current cached token, if 43 // not expired. 44 // 45 // The asynchronous retrieve is deduplicated across multiple calls when 46 // RetrieveBearerToken is called. The asynchronous retrieve is not a 47 // periodic task. It is only performed when the token has not yet expired, 48 // and the current item is within the RefreshBeforeExpires window, and the 49 // TokenCache's RetrieveBearerToken method is called. 50 // 51 // If 0, (default) there will be no minimum delay between asynchronous 52 // refresh attempts. 53 // 54 // If DisableAsyncRefresh is true, this option is ignored. 55 AsyncRefreshMinimumDelay time.Duration 56 57 // Sets if the TokenCache will attempt to refresh the token in the 58 // background asynchronously instead of blocking for credentials to be 59 // refreshed. If disabled token refresh will be blocking. 60 // 61 // The first call to RetrieveBearerToken will always be blocking, because 62 // there is no cached token. 63 DisableAsyncRefresh bool 64 } 65 66 // TokenCache provides an utility to cache Bearer Authentication tokens from a 67 // wrapped TokenProvider. The TokenCache can be has options to configure the 68 // cache's early and asynchronous refresh of the token. 69 type TokenCache struct { 70 options TokenCacheOptions 71 provider TokenProvider 72 73 cachedToken atomic.Value 74 lastRefreshAttemptTime atomic.Value 75 sfGroup singleflight.Group 76 } 77 78 // NewTokenCache returns a initialized TokenCache that implements the 79 // TokenProvider interface. Wrapping the provider passed in. Also taking a set 80 // of optional functional option parameters to configure the token cache. 81 func NewTokenCache(provider TokenProvider, optFns ...func(*TokenCacheOptions)) *TokenCache { 82 var options TokenCacheOptions 83 for _, fn := range optFns { 84 fn(&options) 85 } 86 87 return &TokenCache{ 88 options: options, 89 provider: provider, 90 } 91 } 92 93 // RetrieveBearerToken returns the token if it could be obtained, or error if a 94 // valid token could not be retrieved. 95 // 96 // The passed in Context's cancel/deadline/timeout will impacting only this 97 // individual retrieve call and not any other already queued up calls. This 98 // means underlying provider's RetrieveBearerToken calls could block for ever, 99 // and not be canceled with the Context. Set RetrieveBearerTokenTimeout to 100 // provide a timeout, preventing the underlying TokenProvider blocking forever. 101 // 102 // By default, if the passed in Context is canceled, all of its values will be 103 // considered expired. The wrapped TokenProvider will not be able to lookup the 104 // values from the Context once it is expired. This is done to protect against 105 // expired values no longer being valid. To disable this behavior, use 106 // smithy-go's context.WithPreserveExpiredValues to add a value to the Context 107 // before calling RetrieveBearerToken to enable support for expired values. 108 // 109 // Without RetrieveBearerTokenTimeout there is the potential for a underlying 110 // Provider's RetrieveBearerToken call to sit forever. Blocking in subsequent 111 // attempts at refreshing the token. 112 func (p *TokenCache) RetrieveBearerToken(ctx context.Context) (Token, error) { 113 cachedToken, ok := p.getCachedToken() 114 if !ok || cachedToken.Expired(timeNow()) { 115 return p.refreshBearerToken(ctx) 116 } 117 118 // Check if the token should be refreshed before it expires. 119 refreshToken := cachedToken.Expired(timeNow().Add(p.options.RefreshBeforeExpires)) 120 if !refreshToken { 121 return cachedToken, nil 122 } 123 124 if p.options.DisableAsyncRefresh { 125 return p.refreshBearerToken(ctx) 126 } 127 128 p.tryAsyncRefresh(ctx) 129 130 return cachedToken, nil 131 } 132 133 // tryAsyncRefresh attempts to asynchronously refresh the token returning the 134 // already cached token. If it AsyncRefreshMinimumDelay option is not zero, and 135 // the duration since the last refresh is less than that value, nothing will be 136 // done. 137 func (p *TokenCache) tryAsyncRefresh(ctx context.Context) { 138 if p.options.AsyncRefreshMinimumDelay != 0 { 139 var lastRefreshAttempt time.Time 140 if v := p.lastRefreshAttemptTime.Load(); v != nil { 141 lastRefreshAttempt = v.(time.Time) 142 } 143 144 if timeNow().Before(lastRefreshAttempt.Add(p.options.AsyncRefreshMinimumDelay)) { 145 return 146 } 147 } 148 149 // Ignore the returned channel so this won't be blocking, and limit the 150 // number of additional goroutines created. 151 p.sfGroup.DoChan("async-refresh", func() (interface{}, error) { 152 res, err := p.refreshBearerToken(ctx) 153 if p.options.AsyncRefreshMinimumDelay != 0 { 154 var refreshAttempt time.Time 155 if err != nil { 156 refreshAttempt = timeNow() 157 } 158 p.lastRefreshAttemptTime.Store(refreshAttempt) 159 } 160 161 return res, err 162 }) 163 } 164 165 func (p *TokenCache) refreshBearerToken(ctx context.Context) (Token, error) { 166 resCh := p.sfGroup.DoChan("refresh-token", func() (interface{}, error) { 167 ctx := smithycontext.WithSuppressCancel(ctx) 168 if v := p.options.RetrieveBearerTokenTimeout; v != 0 { 169 var cancel func() 170 ctx, cancel = context.WithTimeout(ctx, v) 171 defer cancel() 172 } 173 return p.singleRetrieve(ctx) 174 }) 175 176 select { 177 case res := <-resCh: 178 return res.Val.(Token), res.Err 179 case <-ctx.Done(): 180 return Token{}, fmt.Errorf("retrieve bearer token canceled, %w", ctx.Err()) 181 } 182 } 183 184 func (p *TokenCache) singleRetrieve(ctx context.Context) (interface{}, error) { 185 token, err := p.provider.RetrieveBearerToken(ctx) 186 if err != nil { 187 return Token{}, fmt.Errorf("failed to retrieve bearer token, %w", err) 188 } 189 190 p.cachedToken.Store(&token) 191 return token, nil 192 } 193 194 // getCachedToken returns the currently cached token and true if found. Returns 195 // false if no token is cached. 196 func (p *TokenCache) getCachedToken() (Token, bool) { 197 v := p.cachedToken.Load() 198 if v == nil { 199 return Token{}, false 200 } 201 202 t := v.(*Token) 203 if t == nil || t.Value == "" { 204 return Token{}, false 205 } 206 207 return *t, true 208 } 209