diff --git a/go.mod b/go.mod index 1a03a4e..1138d74 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,7 @@ require ( github.com/coder/websocket v1.8.15 github.com/stretchr/testify v1.11.1 golang.org/x/oauth2 v0.36.0 + golang.org/x/sync v0.22.0 google.golang.org/grpc v1.81.1 google.golang.org/protobuf v1.36.11 ) diff --git a/go.sum b/go.sum index dc14bb3..4450854 100644 --- a/go.sum +++ b/go.sum @@ -34,6 +34,8 @@ golang.org/x/net v0.51.0 h1:94R/GTO7mt3/4wIKpcR5gkGmRLOuE/2hNGeWq/GBIFo= golang.org/x/net v0.51.0/go.mod h1:aamm+2QF5ogm02fjy5Bb7CQ0WMt1/WVM7FtyaTLlA9Y= golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= +golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= +golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk= diff --git a/openshell/v1/auth_refresh.go b/openshell/v1/auth_refresh.go index a23a9eb..4b25764 100644 --- a/openshell/v1/auth_refresh.go +++ b/openshell/v1/auth_refresh.go @@ -10,11 +10,16 @@ import ( "time" "golang.org/x/oauth2" + "golang.org/x/sync/singleflight" "github.com/rhuss/openshell-sdk-go/openshell/v1/types" ) -const defaultLeeway = 10 * time.Second +const ( + defaultLeeway = 10 * time.Second + initialBackoff = 1 * time.Second + maxBackoff = 30 * time.Second +) var errNilTokenSource = errors.New("openshell: TokenSource must not be nil") @@ -52,11 +57,14 @@ func WithLogger(l types.Logger) RefreshOption { } type refreshableAuth struct { - source oauth2.TokenSource - mu sync.RWMutex - tok *oauth2.Token - leeway time.Duration - logger types.Logger + source oauth2.TokenSource + mu sync.Mutex + group singleflight.Group + tok *oauth2.Token + leeway time.Duration + logger types.Logger + nextRetry time.Time + backoff time.Duration } func (r *refreshableAuth) isTokenValid() bool { @@ -70,25 +78,46 @@ func (r *refreshableAuth) isTokenValid() bool { } func (r *refreshableAuth) GetRequestMetadata(_ context.Context, _ ...string) (map[string]string, error) { - // Fast path: RLock, return cached token if valid. - r.mu.RLock() + r.mu.Lock() if r.isTokenValid() { tok := r.tok.AccessToken - r.mu.RUnlock() + r.mu.Unlock() return map[string]string{"authorization": "Bearer " + tok}, nil } - r.mu.RUnlock() - // Slow path: Lock, re-check, fetch if still stale. + if !r.nextRetry.IsZero() && time.Now().Before(r.nextRetry) { + if r.tok != nil { + tok := r.tok.AccessToken + r.mu.Unlock() + return map[string]string{"authorization": "Bearer " + tok}, nil + } + r.mu.Unlock() + return nil, errors.New("openshell: token refresh failed and backoff is active") + } + r.mu.Unlock() + + val, err, _ := r.group.Do("refresh", func() (any, error) { + return r.source.Token() + }) + r.mu.Lock() defer r.mu.Unlock() - if r.isTokenValid() { - return map[string]string{"authorization": "Bearer " + r.tok.AccessToken}, nil - } - - newTok, err := r.source.Token() if err != nil { + if r.nextRetry.IsZero() || !time.Now().Before(r.nextRetry) { + bo := r.backoff + if bo == 0 { + bo = initialBackoff + } else { + bo *= 2 + if bo > maxBackoff { + bo = maxBackoff + } + } + r.backoff = bo + r.nextRetry = time.Now().Add(bo) + } + if r.tok != nil { if r.logger != nil { r.logger.Error(err, "token refresh failed, using cached token") @@ -98,7 +127,13 @@ func (r *refreshableAuth) GetRequestMetadata(_ context.Context, _ ...string) (ma return nil, err } - r.tok = newTok + tok, ok := val.(*oauth2.Token) + if !ok || tok == nil { + return nil, errors.New("openshell: token source returned nil token without error") + } + r.tok = tok + r.backoff = 0 + r.nextRetry = time.Time{} return map[string]string{"authorization": "Bearer " + r.tok.AccessToken}, nil } @@ -108,7 +143,8 @@ func (r *refreshableAuth) RequireTransportSecurity() bool { // RefreshableToken returns an AuthProvider that caches tokens from src // and refreshes them before expiry. Concurrent callers share a single -// refresh call (coalesced via RWMutex double-checked locking). +// in-flight refresh via singleflight. Failed refreshes trigger exponential +// backoff (1s, 2s, 4s, ..., 30s cap) to avoid amplifying token endpoint outages. func RefreshableToken(src oauth2.TokenSource, opts ...RefreshOption) (AuthProvider, error) { if src == nil { return nil, errNilTokenSource diff --git a/openshell/v1/auth_refresh_test.go b/openshell/v1/auth_refresh_test.go index 451e8e6..828ba20 100644 --- a/openshell/v1/auth_refresh_test.go +++ b/openshell/v1/auth_refresh_test.go @@ -121,22 +121,24 @@ func TestGetRequestMetadata_ConcurrentSingleFlight(t *testing.T) { src := &mockTokenSource{ tokenFunc: func() (*oauth2.Token, error) { fetchCount.Add(1) - time.Sleep(10 * time.Millisecond) // simulate slow token fetch + time.Sleep(100 * time.Millisecond) return &oauth2.Token{AccessToken: "shared-token", Expiry: time.Now().Add(time.Hour)}, nil }, } provider, err := RefreshableToken(src) require.NoError(t, err) - const goroutines = 1000 + const goroutines = 20 var wg sync.WaitGroup wg.Add(goroutines) + ready := make(chan struct{}) results := make([]string, goroutines) errs := make([]error, goroutines) for i := range goroutines { go func(idx int) { defer wg.Done() + <-ready md, e := provider.GetRequestMetadata(context.Background()) errs[idx] = e if md != nil { @@ -144,6 +146,7 @@ func TestGetRequestMetadata_ConcurrentSingleFlight(t *testing.T) { } }(i) } + close(ready) wg.Wait() for i := range goroutines { @@ -318,6 +321,132 @@ func TestGetRequestMetadata_ZeroExpiryNeverRefreshes(t *testing.T) { assert.Equal(t, 1, src.calls(), "zero-expiry token should never be refreshed") } +// --- Backoff tests --- + +func TestGetRequestMetadata_BackoffSkipsRefreshDuringWindow(t *testing.T) { + var callNum atomic.Int32 + src := &mockTokenSource{ + tokenFunc: func() (*oauth2.Token, error) { + n := callNum.Add(1) + if n == 1 { + return &oauth2.Token{AccessToken: "stale", Expiry: time.Now().Add(-time.Minute)}, nil + } + return nil, fmt.Errorf("idp down") + }, + } + provider, err := RefreshableToken(src, WithLeeway(0)) + require.NoError(t, err) + + _, err = provider.GetRequestMetadata(context.Background()) + require.NoError(t, err) + _, err = provider.GetRequestMetadata(context.Background()) + require.NoError(t, err) + beforeCount := src.calls() + require.Equal(t, 2, beforeCount) + + ra := provider.(*refreshableAuth) + ra.mu.Lock() + ra.nextRetry = time.Now().Add(time.Minute) + ra.mu.Unlock() + + md, err := provider.GetRequestMetadata(context.Background()) + require.NoError(t, err) + assert.Equal(t, "Bearer stale", md["authorization"]) + assert.Equal(t, beforeCount, src.calls(), "should not call Token() during backoff window") +} + +func TestGetRequestMetadata_BackoffResetsOnSuccess(t *testing.T) { + var callNum atomic.Int32 + src := &mockTokenSource{ + tokenFunc: func() (*oauth2.Token, error) { + n := callNum.Add(1) + switch n { + case 1: + return &oauth2.Token{AccessToken: "initial", Expiry: time.Now().Add(-time.Second)}, nil + case 2: + return nil, fmt.Errorf("fail once") + default: + return &oauth2.Token{AccessToken: "recovered", Expiry: time.Now().Add(time.Hour)}, nil + } + }, + } + provider, err := RefreshableToken(src, WithLeeway(0)) + require.NoError(t, err) + + _, _ = provider.GetRequestMetadata(context.Background()) + _, _ = provider.GetRequestMetadata(context.Background()) + + ra := provider.(*refreshableAuth) + ra.mu.Lock() + ra.nextRetry = time.Time{} + ra.mu.Unlock() + + md, err := provider.GetRequestMetadata(context.Background()) + require.NoError(t, err) + assert.Equal(t, "Bearer recovered", md["authorization"]) + + ra.mu.Lock() + assert.Equal(t, time.Duration(0), ra.backoff, "backoff should reset after success") + assert.True(t, ra.nextRetry.IsZero(), "nextRetry should be zero after success") + ra.mu.Unlock() +} + +func TestGetRequestMetadata_BackoffCapsAt30s(t *testing.T) { + src := &mockTokenSource{ + tokenFunc: func() (*oauth2.Token, error) { + return nil, fmt.Errorf("always fail") + }, + } + provider, err := RefreshableToken(src, WithLeeway(0)) + require.NoError(t, err) + + ra := provider.(*refreshableAuth) + + for range 10 { + ra.mu.Lock() + ra.nextRetry = time.Time{} + ra.mu.Unlock() + + _, _ = provider.GetRequestMetadata(context.Background()) + } + + ra.mu.Lock() + assert.Equal(t, maxBackoff, ra.backoff, "backoff should cap at maxBackoff") + ra.mu.Unlock() +} + +func TestGetRequestMetadata_ConcurrentFailureBackoffNotOverIncremented(t *testing.T) { + src := &mockTokenSource{ + tokenFunc: func() (*oauth2.Token, error) { + time.Sleep(10 * time.Millisecond) + return nil, fmt.Errorf("idp down") + }, + } + provider, err := RefreshableToken(src, WithLeeway(0)) + require.NoError(t, err) + + const goroutines = 20 + var wg sync.WaitGroup + wg.Add(goroutines) + ready := make(chan struct{}) + for range goroutines { + go func() { + defer wg.Done() + <-ready + _, _ = provider.GetRequestMetadata(context.Background()) + }() + } + close(ready) + wg.Wait() + + ra := provider.(*refreshableAuth) + ra.mu.Lock() + assert.Equal(t, initialBackoff, ra.backoff, + "a single coalesced failure should set backoff to initialBackoff, not escalate") + ra.mu.Unlock() + assert.Equal(t, 1, src.calls(), "singleflight should coalesce to 1 call") +} + // --- benchmarks --- func BenchmarkGetRequestMetadata_CachedToken(b *testing.B) {