[azidentity] Making ChainedTokenCredential re-use the first successful credential

sdk/azidentity/CHANGELOG.md

@@ -6,6 +6,9 @@
 
 ### Breaking Changes
 
+* Instances of `ChainedTokenCredential` will now skip looping through the list of source credentials and re-use the first successful credential on subsequent calls to `GetToken`.
+  * If the `RetryAllSources` option is set to true, it will not assume the first successful credential should be always used.
+
 ### Bugs Fixed
 
 ### Other Changes

sdk/azidentity/chained_token_credential.go

@@ -14,15 +14,20 @@ import (
 
 // ChainedTokenCredentialOptions contains optional parameters for ChainedTokenCredential.
 type ChainedTokenCredentialOptions struct {
-	// placeholder for future options
+	// If true, it will not assume the first successful credential should be always used.
+	RetryAllSources bool
 }
 
 // ChainedTokenCredential is a chain of credentials that enables fallback behavior when a credential can't authenticate.
 type ChainedTokenCredential struct {
-	sources []azcore.TokenCredential
+	sources              []azcore.TokenCredential
+	successfulCredential azcore.TokenCredential
+	retryAllSources      bool
 }
 
 // NewChainedTokenCredential creates a ChainedTokenCredential.
+// By default, this credential will assume that the first successful credential should be the only credential used on future requests.
+// If the `RetryAllSources` option is set to true, it will always try to get a token using all of the available credentials.
 // sources: Credential instances to comprise the chain. GetToken() will invoke them in the given order.
 // options: Optional configuration.
 func NewChainedTokenCredential(sources []azcore.TokenCredential, options *ChainedTokenCredentialOptions) (*ChainedTokenCredential, error) {
@@ -36,32 +41,50 @@ func NewChainedTokenCredential(sources []azcore.TokenCredential, options *Chaine
 	}
 	cp := make([]azcore.TokenCredential, len(sources))
 	copy(cp, sources)
-	return &ChainedTokenCredential{sources: cp}, nil
+	credentialOptions := ChainedTokenCredentialOptions{}
+	if options != nil {
+		credentialOptions = *options
+	}
+	return &ChainedTokenCredential{sources: cp, retryAllSources: credentialOptions.RetryAllSources}, nil
 }
 
 // GetToken calls GetToken on the chained credentials in turn, stopping when one returns a token. This method is called automatically by Azure SDK clients.
 // ctx: Context controlling the request lifetime.
 // opts: Options for the token request, in particular the desired scope of the access token.
 func (c *ChainedTokenCredential) GetToken(ctx context.Context, opts policy.TokenRequestOptions) (token *azcore.AccessToken, err error) {
 	var errList []CredentialUnavailableError
+
+	formatError := func(err error) error {
+		var authFailed AuthenticationFailedError
+		if errors.As(err, &authFailed) {
+			err = fmt.Errorf("Authentication failed:\n%s\n%s"+createChainedErrorMessage(errList), err)
+			authErr := newAuthenticationFailedError(err, authFailed.RawResponse())
+			return authErr
+		}
+		return err
+	}
+
+	if c.successfulCredential != nil && !c.retryAllSources {
+		token, err = c.successfulCredential.GetToken(ctx, opts)
+		if err != nil {
+			return nil, formatError(err)
+		}
+		return token, nil
+	}
 	for _, cred := range c.sources {
 		token, err = cred.GetToken(ctx, opts)
 		var credErr CredentialUnavailableError
 		if errors.As(err, &credErr) {
 			errList = append(errList, credErr)
 		} else if err != nil {
-			var authFailed AuthenticationFailedError
-			if errors.As(err, &authFailed) {
-				err = fmt.Errorf("Authentication failed:\n%s\n%s"+createChainedErrorMessage(errList), err)
-				authErr := newAuthenticationFailedError(err, authFailed.RawResponse())
-				return nil, authErr
-			}
-			return nil, err
+			return nil, formatError(err)
 		} else {
 			logGetTokenSuccess(c, opts)
+			c.successfulCredential = cred
 			return token, nil
 		}
 	}
+
 	// if we reach this point it means that all of the credentials in the chain returned CredentialUnavailableError
 	credErr := newCredentialUnavailableError("Chained Token Credential", createChainedErrorMessage(errList))
 	// skip adding the stack trace here as it was already logged by other calls to GetToken()

sdk/azidentity/chained_token_credential_test.go

@@ -153,6 +153,220 @@ func TestChainedTokenCredential_GetTokenWithUnavailableCredentialInChain(t *test
 	}
 }
 
+func TestChainedTokenCredential_ChecksThatSuccessfulCredentialIsSet(t *testing.T) {
+	err := initEnvironmentVarsForTest()
+	if err != nil {
+		t.Fatalf("Could not set environment variables for testing: %v", err)
+	}
+	srv, close := mock.NewTLSServer()
+	defer close()
+	srv.AppendResponse(mock.WithBody([]byte(accessTokenRespSuccess)))
+	options := ClientSecretCredentialOptions{}
+	options.AuthorityHost = AuthorityHost(srv.URL())
+	options.Transport = srv
+	secCred, err := NewClientSecretCredential(tenantID, clientID, secret, &options)
+	if err != nil {
+		t.Fatalf("Unable to create credential. Received: %v", err)
+	}
+	envCred, err := NewEnvironmentCredential(&EnvironmentCredentialOptions{
+		ClientOptions: azcore.ClientOptions{Transport: srv},
+		AuthorityHost: AuthorityHost(srv.URL()),
+	})
+	if err != nil {
+		t.Fatalf("Failed to create environment credential: %v", err)
+	}
+	cred, err := NewChainedTokenCredential([]azcore.TokenCredential{secCred, envCred}, nil)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	tk, err := cred.GetToken(context.Background(), policy.TokenRequestOptions{Scopes: []string{scope}})
+	if err != nil {
+		t.Fatalf("Received an error when attempting to get a token but expected none")
+	}
+	if tk.Token != tokenValue {
+		t.Fatalf("Received an incorrect access token")
+	}
+	if tk.ExpiresOn.IsZero() {
+		t.Fatalf("Received an incorrect time in the response")
+	}
+	if cred.successfulCredential == nil {
+		t.Fatalf("The successful credential was not assigned")
+	}
+	if cred.successfulCredential != secCred {
+		t.Fatalf("The successful credential should have been the secret credential")
+	}
+}
+
+/**
+ * Helps count the number of times a credential is called.
+ */
+type TestCountPolicy struct{ count int }
+
+/**
+ * Helps count the number of times a credential is called.
+ */
+func (p *TestCountPolicy) Do(req *policy.Request) (*http.Response, error) {
+	p.count += 1
+	return req.Next()
+}
+
+func TestChainedTokenCredential_RepeatedGetTokenWithSuccessfulCredential(t *testing.T) {
+	err := initEnvironmentVarsForTest()
+	if err != nil {
+		t.Fatalf("Could not set environment variables for testing: %v", err)
+	}
+	srv, close := mock.NewTLSServer()
+	defer close()
+
+	secretCountPolicy := &TestCountPolicy{}
+	environmentCountPolicy := &TestCountPolicy{}
+
+	srv.AppendResponse(mock.WithBody([]byte(accessTokenRespSuccess)))
+	srv.AppendResponse(mock.WithBody([]byte(accessTokenRespSuccess)))
+	options := ClientSecretCredentialOptions{}
+	options.AuthorityHost = AuthorityHost(srv.URL())
+	options.PerCallPolicies = []policy.Policy{secretCountPolicy}
+	options.Transport = srv
+	secCred, err := NewClientSecretCredential(tenantID, clientID, secret, &options)
+	if err != nil {
+		t.Fatalf("Unable to create credential. Received: %v", err)
+	}
+	envCred, err := NewEnvironmentCredential(&EnvironmentCredentialOptions{
+		ClientOptions: azcore.ClientOptions{
+			Transport:       srv,
+			PerCallPolicies: []policy.Policy{environmentCountPolicy},
+		},
+		AuthorityHost: AuthorityHost(srv.URL()),
+	})
+	if err != nil {
+		t.Fatalf("Failed to create environment credential: %v", err)
+	}
+	cred, err := NewChainedTokenCredential([]azcore.TokenCredential{secCred, envCred}, nil)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	tk, err := cred.GetToken(context.Background(), policy.TokenRequestOptions{Scopes: []string{scope}})
+	if err != nil {
+		t.Fatalf("Received an error when attempting to get a token but expected none")
+	}
+	if tk.Token != tokenValue {
+		t.Fatalf("Received an incorrect access token")
+	}
+	if tk.ExpiresOn.IsZero() {
+		t.Fatalf("Received an incorrect time in the response")
+	}
+	if cred.successfulCredential == nil {
+		t.Fatalf("The successful credential was not assigned")
+	}
+	if cred.successfulCredential != secCred {
+		t.Fatalf("The successful credential should have been the secret credential")
+	}
+	if secretCountPolicy.count != 1 {
+		t.Fatalf("The secret credential policies should have been triggered once")
+	}
+	if environmentCountPolicy.count != 0 {
+		t.Fatalf("The environment credential policies should not have been triggered")
+	}
+	tk2, err2 := cred.GetToken(context.Background(), policy.TokenRequestOptions{Scopes: []string{scope}})
+	if err2 != nil {
+		t.Fatalf("Received an error when attempting to get a token but expected none. Error: %v", err2)
+	}
+	if tk2.Token != tokenValue {
+		t.Fatalf("Received an incorrect access token")
+	}
+	if tk2.ExpiresOn.IsZero() {
+		t.Fatalf("Received an incorrect time in the response")
+	}
+	if secretCountPolicy.count != 2 {
+		t.Fatalf("The secret credential policies should have been triggered twice")
+	}
+	if environmentCountPolicy.count != 0 {
+		t.Fatalf("The environment credential policies should not have been triggered")
+	}
+}
+
+// A credential that always throws a CredentialUnavailableError
+type UnavailableCredential struct {
+	callCount int
+}
+
+func NewUnavailableCredential() (*UnavailableCredential, error) {
+	return &UnavailableCredential{callCount: 0}, nil
+}
+func (c *UnavailableCredential) GetToken(ctx context.Context, opts policy.TokenRequestOptions) (token *azcore.AccessToken, err error) {
+	c.callCount += 1
+	return nil, newCredentialUnavailableError("UnavailableCredential", "Expected CredentialUnavailableError")
+}
+
+func TestChainedTokenCredential_RepeatedGetTokenWithSuccessfulCredentialWithRetryAllSources(t *testing.T) {
+	err := initEnvironmentVarsForTest()
+	if err != nil {
+		t.Fatalf("Could not set environment variables for testing: %v", err)
+	}
+	srv, close := mock.NewTLSServer()
+	defer close()
+
+	secretCountPolicy := &TestCountPolicy{}
+
+	srv.AppendResponse(mock.WithBody([]byte(accessTokenRespSuccess)))
+	srv.AppendResponse(mock.WithBody([]byte(accessTokenRespSuccess)))
+
+	options := ClientSecretCredentialOptions{}
+	options.AuthorityHost = AuthorityHost(srv.URL())
+	options.PerCallPolicies = []policy.Policy{secretCountPolicy}
+	options.Transport = srv
+
+	unavailableCred, _ := NewUnavailableCredential()
+	secCred, err := NewClientSecretCredential(tenantID, clientID, secret, &options)
+	if err != nil {
+		t.Fatalf("Unable to create credential. Received: %v", err)
+	}
+
+	// Backwards order: envCred first, secCred later, to check that envCred is always called when RetryAllSources is set to true.
+	cred, err := NewChainedTokenCredential([]azcore.TokenCredential{unavailableCred, secCred}, &ChainedTokenCredentialOptions{RetryAllSources: true})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	tk, err := cred.GetToken(context.Background(), policy.TokenRequestOptions{Scopes: []string{scope}})
+	if err != nil {
+		t.Fatalf("Received an error when attempting to get a token but expected none. Error: %v", err)
+	}
+	if tk.Token != tokenValue {
+		t.Fatalf("Received an incorrect access token")
+	}
+	if tk.ExpiresOn.IsZero() {
+		t.Fatalf("Received an incorrect time in the response")
+	}
+	if cred.successfulCredential == nil {
+		t.Fatalf("The successful credential was not assigned")
+	}
+	if cred.successfulCredential != secCred {
+		t.Fatalf("The successful credential should have been the secret credential")
+	}
+	if secretCountPolicy.count != 1 {
+		t.Fatalf("The secret credential policies should have been triggered once")
+	}
+	if unavailableCred.callCount != 1 {
+		t.Fatalf("The environment credential policies should have been triggered once")
+	}
+	tk2, err2 := cred.GetToken(context.Background(), policy.TokenRequestOptions{Scopes: []string{scope}})
+	if err2 != nil {
+		t.Fatalf("Received an error when attempting to get a token but expected none. Error: %v", err2)
+	}
+	if tk2.Token != tokenValue {
+		t.Fatalf("Received an incorrect access token")
+	}
+	if tk2.ExpiresOn.IsZero() {
+		t.Fatalf("Received an incorrect time in the response")
+	}
+	if secretCountPolicy.count != 2 {
+		t.Fatalf("The secret credential policies should have been triggered twice")
+	}
+	if unavailableCred.callCount != 2 {
+		t.Fatalf("The environment credential policies should have been triggered twice")
+	}
+}
+
 func TestBearerPolicy_ChainedTokenCredential(t *testing.T) {
 	err := initEnvironmentVarsForTest()
 	if err != nil {