Fix the TokenStore getting stuck in a read lock (#3035)

Author: adreed-msft

common/oauthTokenManager.go

@@ -63,10 +63,17 @@ const DefaultActiveDirectoryEndpoint = "https://login.microsoftonline.com"
 
 const TokenCache = "AzCopyTokenCache"
 
+type CredCacheImplementation interface {
+	HasCachedToken() (bool, error)
+	LoadToken() (*OAuthTokenInfo, error)
+	SaveToken(OAuthTokenInfo) error
+	RemoveCachedToken() error
+}
+
 // UserOAuthTokenManager for token management.
 type UserOAuthTokenManager struct {
 	oauthClient *http.Client
-	credCache   *CredCache
+	credCache   CredCacheImplementation
 
 	// Stash the credential info as we delete the environment variable after reading it, and we need to get it multiple times.
 	stashedInfo *OAuthTokenInfo
@@ -479,7 +486,7 @@ func (credInfo *OAuthTokenInfo) Refresh(ctx context.Context) (*Token, error) {
 }
 
 // Single instance token store credential cache shared by entire azcopy process.
-var tokenStoreCredCache = NewCredCacheInternalIntegration(CredCacheOptions{
+var tokenStoreCredCache CredCacheImplementation = NewCredCacheInternalIntegration(CredCacheOptions{
 	KeyName:     "azcopy/aadtoken/" + strconv.Itoa(os.Getpid()),
 	ServiceName: "azcopy",
 	AccountName: "aadtoken/" + strconv.Itoa(os.Getpid()),
@@ -510,8 +517,9 @@ func getAuthorityURL(activeDirectoryEndpoint string) (*url.URL, error) {
 const minimumTokenValidDuration = time.Minute * 5
 
 type TokenStoreCredential struct {
-	token *azcore.AccessToken
-	lock  sync.RWMutex
+	token     *azcore.AccessToken
+	lock      sync.RWMutex
+	credCache CredCacheImplementation
 }
 
 // globalTokenStoreCredential is created to make sure that all
@@ -531,22 +539,26 @@ var globalTokenStoreCredential *TokenStoreCredential
 var globalTsc sync.Once
 
 func (tsc *TokenStoreCredential) GetToken(_ context.Context, _ policy.TokenRequestOptions) (azcore.AccessToken, error) {
-	// if the token we've has not expired, return the same.
+	// if the token we have has not expired, return the same.
 	tsc.lock.RLock()
-	if time.Until(tsc.token.ExpiresOn) > minimumTokenValidDuration {
+	if rem := time.Until(tsc.token.ExpiresOn); rem > minimumTokenValidDuration {
+		tsc.lock.RUnlock() // return path, so we must release the read lock here as well.
 		return *tsc.token, nil
 	}
 	tsc.lock.RUnlock()
 
 	tsc.lock.Lock()
 	defer tsc.lock.Unlock()
-	hasToken, err := tokenStoreCredCache.HasCachedToken()
+
+	hasToken, err := tsc.credCache.HasCachedToken()
 	if err != nil || !hasToken {
+		AzcopyCurrentJobLogger.Log(LogDebug, fmt.Sprintf("no token found %v", err))
 		return azcore.AccessToken{}, fmt.Errorf("no cached token found in Token Store Mode(SE), %w", err)
 	}
 
-	tokenInfo, err := tokenStoreCredCache.LoadToken()
+	tokenInfo, err := tsc.credCache.LoadToken()
 	if err != nil {
+		AzcopyCurrentJobLogger.Log(LogDebug, fmt.Sprintf("get token failed %s", err.Error()))
 		return azcore.AccessToken{}, fmt.Errorf("get cached token failed in Token Store Mode(SE), %w", err)
 	}
 
@@ -568,6 +580,7 @@ func GetTokenStoreCredential(accessToken string, expiresOn time.Time) azcore.Tok
 				Token:     accessToken,
 				ExpiresOn: expiresOn,
 			},
+			credCache: tokenStoreCredCache,
 		}
 	})
 	return globalTokenStoreCredential

common/oauthTokenManager_test.go

@@ -22,11 +22,17 @@ package common
 
 import (
 	"context"
+	"encoding/json"
+	"errors"
 	"fmt"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	"github.com/stretchr/testify/assert"
 	"os"
 	"reflect"
 	"strconv"
 	"testing"
+	"time"
 )
 
 const tokenInfoJson = `{
@@ -183,3 +189,72 @@ func TestUserOAuthTokenManager_GetTokenInfo(t *testing.T) {
 		})
 	}
 }
+
+type TestCredCache struct {
+	stashed *OAuthTokenInfo
+}
+
+func (d *TestCredCache) HasCachedToken() (bool, error) {
+	return d.stashed != nil, nil
+}
+
+func (d *TestCredCache) LoadToken() (*OAuthTokenInfo, error) {
+	if d.stashed == nil {
+		return nil, errors.New("no cached token found")
+	}
+
+	return d.stashed, nil
+}
+
+func (d *TestCredCache) SaveToken(info OAuthTokenInfo) error {
+	d.stashed = &info
+	return nil
+}
+
+func (d *TestCredCache) RemoveCachedToken() error {
+	if d.stashed == nil {
+		return errors.New("no cached token found")
+	}
+
+	d.stashed = nil
+	return nil
+}
+
+func TestTokenStoreCredentialHang(t *testing.T) {
+	tok := &azcore.AccessToken{
+		Token:     "asdf",
+		ExpiresOn: time.Now().Add(minimumTokenValidDuration * 2), // we want to hit that if statement at the start and get it into the read lock
+	}
+
+	tsc := &TokenStoreCredential{
+		token: tok,
+		credCache: &TestCredCache{
+			stashed: &OAuthTokenInfo{
+				Token: Token{
+					AccessToken: "foobar",
+					ExpiresOn:   json.Number(fmt.Sprint(tok.ExpiresOn.Unix())),
+				},
+			},
+		},
+	}
+
+	// Prior to this PR, we'd get locked into a read state doing this, because the if statement didn't contain a way out of the read lock.
+	outTok, err := tsc.GetToken(context.Background(), policy.TokenRequestOptions{})
+	assert.NoError(t, err)
+	assert.Equal(t, tok.Token, outTok.Token)
+
+	// we shouldn't get blocked here, otherwise we have problems.
+	assert.Equal(t, true, tsc.lock.TryLock())
+	tsc.lock.Unlock()
+
+	tok.ExpiresOn = time.Now() // now it should refresh
+
+	// We shouldn't get caught here at all
+	outTok, err = tsc.GetToken(context.Background(), policy.TokenRequestOptions{})
+	assert.NoError(t, err)
+	assert.Equal(t, "foobar", outTok.Token) // ensure it refreshed
+
+	// we shouldn't get blocked here, otherwise we have problems.
+	assert.Equal(t, true, tsc.lock.TryLock())
+	tsc.lock.Unlock()
+}