token_source.go

package oauthctx

import (
	"context"

	"golang.org/x/oauth2"
)

// TokenSource is the context-aware analogue of oauth2.TokenSource.
// Implementations must propagate ctx to all network calls so that caller
// deadlines and cancellations are respected.
// Concurrent access is only safe when the implementation is wrapped in
// [ReuseTokenSource].
type TokenSource interface {
	// TokenContext returns a valid access token. ctx is used for deadline and
	// cancellation propagation and must not be stored beyond the call.
	TokenContext(ctx context.Context) (*oauth2.Token, error)
}

// ReuseTokenSource returns a TokenSource which repeatedly returns the
// same token as long as it's valid, starting with t.
// When its cached token is invalid, a new token is obtained from src.
func ReuseTokenSource(t *oauth2.Token, src TokenSource) TokenSource {
	if t == nil {
		if rt, ok := src.(*reuseTokenSource); ok {
			//if t == nil {
			// Just use it directly.
			return rt
			//}
			// not secure, as rt.new is not guarded by single mutex anymore
			//src = rt.new
		}
	}
	ts := &reuseTokenSource{
		new: src,
		mu:  make(chan struct{}, 1),
		t:   t,
	}
	ts.mu <- struct{}{}
	return ts
}

type reuseTokenSource struct {
	new TokenSource   // called when t is expired.
	mu  chan struct{} // guards t. use chain to select on mu and context
	t   *oauth2.Token
}

func (s *reuseTokenSource) TokenContext(ctx context.Context) (*oauth2.Token, error) {
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	default:
	}

	select {
	case <-s.mu:
		defer func() { s.mu <- struct{}{} }()
	case <-ctx.Done():
		return nil, ctx.Err()
	}

	if s.t.Valid() && !isExpired(ctx) {
		return s.t, nil
	}

	t, err := s.new.TokenContext(ctx)
	if err != nil {
		return nil, err
	}
	s.t = t
	return t, nil
}

type expiredTokenKey struct{}

// WithExpiredToken returns a context that causes [ReuseTokenSource] to bypass
// its cached token and fetch a fresh one on the next call to TokenContext.
// Use this when you know the current token has been revoked externally and
// you need an immediate refresh.
func WithExpiredToken(ctx context.Context) context.Context {
	return context.WithValue(ctx, expiredTokenKey{}, expiredTokenKey{})
}

func isExpired(ctx context.Context) bool {
	return ctx.Value(expiredTokenKey{}) != nil
}