fix: implement thread-safe token caching for self-managed oidc (#677)

Introduce a synchronized TokenCache to prevent concurrent OIDC login requests from colliding and overloading the Pinniped Supervisor or upstream LDAP server.

Author: steven-zou

internal/authctx/client.go

@@ -5,6 +5,9 @@
 package authctx
 
 import (
+	"sync"
+	"time"
+
 	"github.com/pkg/errors"
 
 	"github.com/vmware/terraform-provider-tanzu-mission-control/internal/client"
@@ -43,6 +46,51 @@ type TanzuContext struct {
 	VMWCloudEndPoint string // selfmanaged odic issuer is stored here
 	TMCConnection    *client.TanzuMissionControl
 	TLSConfig        *proxy.TLSConfig
+
+	// Cache for self-managed OIDC tokens
+	smTokenCache *TokenCache
+}
+
+// TokenCache manages the lifecycle of the self-managed OIDC tokens.
+type TokenCache struct {
+	mu            sync.Mutex
+	cachedHeaders map[string]string
+	expiry        time.Time
+}
+
+// NewTokenCache creates a new TokenCache instance.
+func NewTokenCache() *TokenCache {
+	return &TokenCache{}
+}
+
+// GetToken returns the cached headers if they are still valid (with a 1-minute buffer),
+// or fetches new headers using the provided fetch function.
+func (tc *TokenCache) GetToken(fetch func() (map[string]string, time.Time, error)) (map[string]string, error) {
+	tc.mu.Lock()
+	defer tc.mu.Unlock()
+
+	if tc.cachedHeaders != nil && time.Now().Add(1*time.Minute).Before(tc.expiry) {
+		return tc.cachedHeaders, nil
+	}
+
+	headers, expiry, err := fetch()
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to refresh token")
+	}
+
+	tc.cachedHeaders = headers
+	tc.expiry = expiry
+
+	return headers, nil
+}
+
+// Update explicitly updates the cached headers and expiry.
+func (tc *TokenCache) Update(headers map[string]string, expiry time.Time) {
+	tc.mu.Lock()
+	defer tc.mu.Unlock()
+
+	tc.cachedHeaders = headers
+	tc.expiry = expiry
 }
 
 func (cfg *TanzuContext) Setup() (err error) {
@@ -66,6 +114,10 @@ func (cfg *TanzuContext) SetupWithDefaultTransportForTesting() (err error) {
 }
 
 func setup(cfg *TanzuContext) (err error) {
+	if cfg.IsSelfManaged() {
+		cfg.smTokenCache = NewTokenCache()
+	}
+
 	fetchAuthHeaders := getUserAuthCtxHeaders(cfg)
 
 	md, err := fetchAuthHeaders()
@@ -93,18 +145,28 @@ func setup(cfg *TanzuContext) (err error) {
 }
 
 func getUserAuthCtxHeaders(config *TanzuContext) func() (map[string]string, error) {
-	issuerURL := config.VMWCloudEndPoint
-	token := config.Token
-	proxyConfig := config.TLSConfig
-
 	if config.IsSelfManaged() {
-		username := config.SMUsername
-
 		return func() (map[string]string, error) {
-			return getSMUserAuthCtx(issuerURL, username, token, proxyConfig)
+			// For compatibility considerations.
+			if config.smTokenCache == nil {
+				headers, _, err := getSMUserAuthCtx(config.VMWCloudEndPoint, config.SMUsername, config.Token, config.TLSConfig)
+				if err != nil {
+					return nil, errors.Wrap(err, "failed to get self-managed user auth context headers")
+				}
+
+				return headers, nil
+			}
+
+			return config.smTokenCache.GetToken(func() (map[string]string, time.Time, error) {
+				return getSMUserAuthCtx(config.VMWCloudEndPoint, config.SMUsername, config.Token, config.TLSConfig)
+			})
 		}
 	}
 
+	issuerURL := config.VMWCloudEndPoint
+	token := config.Token
+	proxyConfig := config.TLSConfig
+
 	return func() (map[string]string, error) {
 		return getSaaSUserAuthCtx(issuerURL, token, proxyConfig)
 	}

internal/authctx/selfmanaged.go

@@ -45,41 +45,41 @@ type smSession struct {
 }
 
 // todo: proxy support is not added for the self-managed flow. Add it when there is a requirement.
-func getSMUserAuthCtx(pinnipedURL, uName, password string, config *proxy.TLSConfig) (metadata map[string]string, err error) {
+func getSMUserAuthCtx(pinnipedURL, uName, password string, config *proxy.TLSConfig) (metadata map[string]string, expiry time.Time, err error) {
 	if pinnipedURL == "" || uName == "" || password == "" {
-		return nil, errors.New("Invalid auth configuration for self_managed")
+		return nil, time.Time{}, errors.New("Invalid auth configuration for self_managed")
 	}
 
 	tlsConfig, err := proxy.GetConnectorTLSConfig(config)
 	if err != nil {
-		return nil, err
+		return nil, time.Time{}, err
 	}
 
 	session, err := initSession(pinnipedURL, uName, password, tlsConfig)
 	if err != nil {
-		return nil, err
+		return nil, time.Time{}, err
 	}
 
 	expectedRedirectURL, err := url.Parse(session.sharedOauthConfig.RedirectURL)
 	if err != nil {
-		return nil, errors.Wrapf(err, "failed to parse expected redirect URL %s", session.sharedOauthConfig.RedirectURL)
+		return nil, time.Time{}, errors.Wrapf(err, "failed to parse expected redirect URL %s", session.sharedOauthConfig.RedirectURL)
 	}
 
 	actualRedirectURL, err := session.initiateAuthorizeRequestUnamePwd()
 	if err != nil {
-		return nil, errors.Wrapf(err, "failed to initiate authorize request with issuer %s", session.issuerURL)
+		return nil, time.Time{}, errors.Wrapf(err, "failed to initiate authorize request with issuer %s", session.issuerURL)
 	}
 
 	// Check that the redirect was to the expected location.
 	if actualRedirectURL.Scheme != expectedRedirectURL.Scheme ||
 		actualRedirectURL.Host != expectedRedirectURL.Host || actualRedirectURL.Path != expectedRedirectURL.Path {
-		return nil, fmt.Errorf("error getting authorization: redirected to the wrong location: %s",
+		return nil, time.Time{}, fmt.Errorf("error getting authorization: redirected to the wrong location: %s",
 			actualRedirectURL.String())
 	}
 
 	// validate the state param to detect and prevent CSRF attacks.
 	if err := session.stateVal.Validate(actualRedirectURL.Query().Get("state")); err != nil {
-		return nil, errors.Wrap(err, "failed to validate state")
+		return nil, time.Time{}, errors.Wrap(err, "failed to validate state")
 	}
 
 	// Get the auth code or return the error from the server.
@@ -90,10 +90,10 @@ func getSMUserAuthCtx(pinnipedURL, uName, password string, config *proxy.TLSConf
 
 		optionalErrorDescription := actualRedirectURL.Query().Get("error_description")
 		if optionalErrorDescription == "" {
-			return nil, fmt.Errorf("login failed with code %q", requiredErrorCode)
+			return nil, time.Time{}, fmt.Errorf("login failed with code %q", requiredErrorCode)
 		}
 
-		return nil, fmt.Errorf("login failed with code %q: %s", requiredErrorCode, optionalErrorDescription)
+		return nil, time.Time{}, fmt.Errorf("login failed with code %q: %s", requiredErrorCode, optionalErrorDescription)
 	}
 
 	customClient := &http.Client{
@@ -111,7 +111,7 @@ func getSMUserAuthCtx(pinnipedURL, uName, password string, config *proxy.TLSConf
 
 	token, err := session.sharedOauthConfig.Exchange(tokenCtx, authCode, session.pkceCodePair.Verifier())
 	if err != nil {
-		return nil, errors.Wrapf(err, "failed to exchange auth code for oauth tokens")
+		return nil, time.Time{}, errors.Wrapf(err, "failed to exchange auth code for oauth tokens")
 	}
 
 	extraFields := map[string]interface{}{extraIDToken: token.Extra(extraIDToken).(string)}
@@ -123,7 +123,7 @@ func getSMUserAuthCtx(pinnipedURL, uName, password string, config *proxy.TLSConf
 
 	token = token.WithExtra(extraFields)
 
-	return getSMHeaders(token), nil
+	return getSMHeaders(token), token.Expiry, nil
 }
 
 // todo: if slowness is experienced, then we can avoid re-initialising same values again.
@@ -251,7 +251,15 @@ func (s *smSession) getAuthCodeURL() string {
 }
 
 func refreshSMUserAuthCtx(config *TanzuContext) {
-	md, _ := getSMUserAuthCtx(config.VMWCloudEndPoint, config.SMUsername, config.Token, config.TLSConfig)
+	md, expiry, err := getSMUserAuthCtx(config.VMWCloudEndPoint, config.SMUsername, config.Token, config.TLSConfig)
+	if err != nil {
+		return
+	}
+
+	if config.smTokenCache != nil {
+		config.smTokenCache.Update(md, expiry)
+	}
+
 	for key, value := range md {
 		config.TMCConnection.Headers.Set(key, value)
 	}