Use underlying error when checking for context errors (Azure#20404)

* Use underlying error when checking for context errors

There is a race between error propagation and closing a context's
deadline channel.  The result is that while the API that takes a context
returns the correct error, context.Err() might return nil.

See golang/go#31863 for more info.

* further refinements

* revert ordering of error check

* add changelog entry

Author: jhendrixMSFT

sdk/azidentity/CHANGELOG.md

@@ -9,6 +9,7 @@
 ### Breaking Changes
 
 ### Bugs Fixed
+* Fixed an issue in `DefaultAzureCredential` that could cause the managed identity endpoint check to fail in rare circumstances.
 
 ### Other Changes
 

sdk/azidentity/chained_token_credential.go

@@ -117,7 +117,7 @@ func (c *ChainedTokenCredential) GetToken(ctx context.Context, opts policy.Token
 			err = newCredentialUnavailableError(c.name, msg)
 		} else {
 			res := getResponseFromError(err)
-			err = newAuthenticationFailedError(c.name, msg, res)
+			err = newAuthenticationFailedError(c.name, msg, res, err)
 		}
 	}
 	return token, err

sdk/azidentity/chained_token_credential_test.go

@@ -113,7 +113,7 @@ func TestChainedTokenCredential_GetTokenSuccess(t *testing.T) {
 
 func TestChainedTokenCredential_GetTokenFail(t *testing.T) {
 	c := NewFakeCredential()
-	c.SetResponse(azcore.AccessToken{}, newAuthenticationFailedError("test", "something went wrong", nil))
+	c.SetResponse(azcore.AccessToken{}, newAuthenticationFailedError("test", "something went wrong", nil, nil))
 	cred, err := NewChainedTokenCredential([]azcore.TokenCredential{c}, nil)
 	if err != nil {
 		t.Fatal(err)
@@ -158,7 +158,7 @@ func TestChainedTokenCredential_MultipleCredentialsGetTokenAuthenticationFailed(
 	c2 := NewFakeCredential()
 	c2.SetResponse(azcore.AccessToken{}, newCredentialUnavailableError("unavailableCredential2", "Unavailable expected error"))
 	c3 := NewFakeCredential()
-	c3.SetResponse(azcore.AccessToken{}, newAuthenticationFailedError("authenticationFailedCredential3", "Authentication failed expected error", nil))
+	c3.SetResponse(azcore.AccessToken{}, newAuthenticationFailedError("authenticationFailedCredential3", "Authentication failed expected error", nil, nil))
 	cred, err := NewChainedTokenCredential([]azcore.TokenCredential{c1, c2, c3}, nil)
 	if err != nil {
 		t.Fatal(err)
@@ -263,7 +263,7 @@ func TestChainedTokenCredential_Race(t *testing.T) {
 	successFake := NewFakeCredential()
 	successFake.SetResponse(azcore.AccessToken{Token: "*", ExpiresOn: time.Now().Add(time.Hour)}, nil)
 	authFailFake := NewFakeCredential()
-	authFailFake.SetResponse(azcore.AccessToken{}, newAuthenticationFailedError("", "", nil))
+	authFailFake.SetResponse(azcore.AccessToken{}, newAuthenticationFailedError("", "", nil, nil))
 	unavailableFake := NewFakeCredential()
 	unavailableFake.SetResponse(azcore.AccessToken{}, newCredentialUnavailableError("", ""))
 	tro := policy.TokenRequestOptions{Scopes: []string{liveTestScope}}

sdk/azidentity/default_azure_credential.go

@@ -182,7 +182,7 @@ func (w *timeoutWrapper) GetToken(ctx context.Context, opts policy.TokenRequestO
 		c, cancel := context.WithTimeout(ctx, w.timeout)
 		defer cancel()
 		tk, err = w.mic.GetToken(c, opts)
-		if ce := c.Err(); errors.Is(ce, context.DeadlineExceeded) {
+		if isAuthFailedDueToContext(err) {
 			err = newCredentialUnavailableError(credNameManagedIdentity, "managed identity timed out")
 		} else {
 			// some managed identity implementation is available, so don't apply the timeout to future calls
@@ -193,3 +193,15 @@ func (w *timeoutWrapper) GetToken(ctx context.Context, opts policy.TokenRequestO
 	}
 	return tk, err
 }
+
+// unwraps nested AuthenticationFailedErrors to get the root error
+func isAuthFailedDueToContext(err error) bool {
+	for {
+		var authFailedErr *AuthenticationFailedError
+		if !errors.As(err, &authFailedErr) {
+			break
+		}
+		err = authFailedErr.err
+	}
+	return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)
+}

sdk/azidentity/default_azure_credential_test.go

@@ -171,7 +171,14 @@ type delayPolicy struct {
 }
 
 func (p *delayPolicy) Do(req *policy.Request) (resp *http.Response, err error) {
-	time.Sleep(p.delay)
+	if p.delay > 0 {
+		select {
+		case <-req.Raw().Context().Done():
+			return nil, req.Raw().Context().Err()
+		case <-time.After(p.delay):
+			// delay has elapsed, continue on
+		}
+	}
 	return req.Next()
 }
 
@@ -180,7 +187,7 @@ func TestDefaultAzureCredential_timeoutWrapper(t *testing.T) {
 	defer close()
 	srv.SetResponse(mock.WithBody(accessTokenRespSuccess))
 
-	timeout := 5 * time.Millisecond
+	timeout := 100 * time.Millisecond
 	dp := delayPolicy{2 * timeout}
 	mic, err := NewManagedIdentityCredential(&ManagedIdentityCredentialOptions{
 		ClientOptions: policy.ClientOptions{
@@ -201,7 +208,7 @@ func TestDefaultAzureCredential_timeoutWrapper(t *testing.T) {
 		// expecting credentialUnavailableError because delay exceeds the wrapper's timeout
 		_, err = chain.GetToken(context.Background(), policy.TokenRequestOptions{Scopes: []string{liveTestScope}})
 		if _, ok := err.(*credentialUnavailableError); !ok {
-			t.Fatalf("expected credentialUnavailableError, got %v", err)
+			t.Fatalf("expected credentialUnavailableError, got %T: %v", err, err)
 		}
 	}
 

sdk/azidentity/errors.go

@@ -39,15 +39,16 @@ type AuthenticationFailedError struct {
 
 	credType string
 	message  string
+	err      error
 }
 
-func newAuthenticationFailedError(credType string, message string, resp *http.Response) error {
-	return &AuthenticationFailedError{credType: credType, message: message, RawResponse: resp}
+func newAuthenticationFailedError(credType string, message string, resp *http.Response, err error) error {
+	return &AuthenticationFailedError{credType: credType, message: message, RawResponse: resp, err: err}
 }
 
 func newAuthenticationFailedErrorFromMSALError(credType string, err error) error {
 	res := getResponseFromError(err)
-	return newAuthenticationFailedError(credType, err.Error(), res)
+	return newAuthenticationFailedError(credType, err.Error(), res, err)
 }
 
 // Error implements the error interface. Note that the message contents are not contractual and can change over time.

sdk/azidentity/errors_test.go

@@ -31,7 +31,7 @@ func TestAuthenticationFailedErrorInterface(t *testing.T) {
 		Body:       io.NopCloser(bytes.NewBufferString(resBodyString)),
 		Request:    req,
 	}
-	err = newAuthenticationFailedError(credNameAzureCLI, "error message", res)
+	err = newAuthenticationFailedError(credNameAzureCLI, "error message", res, nil)
 	if e, ok := err.(*AuthenticationFailedError); ok {
 		if e.RawResponse == nil {
 			t.Fatal("expected a non-nil RawResponse")

sdk/azidentity/managed_identity_client.go

@@ -168,7 +168,7 @@ func (c *managedIdentityClient) authenticate(ctx context.Context, id ManagedIDKi
 
 	resp, err := c.pipeline.Do(msg)
 	if err != nil {
-		return azcore.AccessToken{}, newAuthenticationFailedError(credNameManagedIdentity, err.Error(), nil)
+		return azcore.AccessToken{}, newAuthenticationFailedError(credNameManagedIdentity, err.Error(), nil, err)
 	}
 
 	if runtime.HasStatusCode(resp, http.StatusOK, http.StatusCreated) {
@@ -177,12 +177,12 @@ func (c *managedIdentityClient) authenticate(ctx context.Context, id ManagedIDKi
 
 	if c.msiType == msiTypeIMDS && resp.StatusCode == 400 {
 		if id != nil {
-			return azcore.AccessToken{}, newAuthenticationFailedError(credNameManagedIdentity, "the requested identity isn't assigned to this resource", resp)
+			return azcore.AccessToken{}, newAuthenticationFailedError(credNameManagedIdentity, "the requested identity isn't assigned to this resource", resp, nil)
 		}
 		return azcore.AccessToken{}, newCredentialUnavailableError(credNameManagedIdentity, "no default identity is assigned to this resource")
 	}
 
-	return azcore.AccessToken{}, newAuthenticationFailedError(credNameManagedIdentity, "authentication failed", resp)
+	return azcore.AccessToken{}, newAuthenticationFailedError(credNameManagedIdentity, "authentication failed", resp, nil)
 }
 
 func (c *managedIdentityClient) createAccessToken(res *http.Response) (azcore.AccessToken, error) {
@@ -210,10 +210,10 @@ func (c *managedIdentityClient) createAccessToken(res *http.Response) (azcore.Ac
 		if expiresOn, err := strconv.Atoi(v); err == nil {
 			return azcore.AccessToken{Token: value.Token, ExpiresOn: time.Unix(int64(expiresOn), 0).UTC()}, nil
 		}
-		return azcore.AccessToken{}, newAuthenticationFailedError(credNameManagedIdentity, "unexpected expires_on value: "+v, res)
+		return azcore.AccessToken{}, newAuthenticationFailedError(credNameManagedIdentity, "unexpected expires_on value: "+v, res, nil)
 	default:
 		msg := fmt.Sprintf("unsupported type received in expires_on: %T, %v", v, v)
-		return azcore.AccessToken{}, newAuthenticationFailedError(credNameManagedIdentity, msg, res)
+		return azcore.AccessToken{}, newAuthenticationFailedError(credNameManagedIdentity, msg, res, nil)
 	}
 }
 
@@ -228,7 +228,7 @@ func (c *managedIdentityClient) createAuthRequest(ctx context.Context, id Manage
 		key, err := c.getAzureArcSecretKey(ctx, scopes)
 		if err != nil {
 			msg := fmt.Sprintf("failed to retreive secret key from the identity endpoint: %v", err)
-			return nil, newAuthenticationFailedError(credNameManagedIdentity, msg, nil)
+			return nil, newAuthenticationFailedError(credNameManagedIdentity, msg, nil, err)
 		}
 		return c.createAzureArcAuthRequest(ctx, id, scopes, key)
 	case msiTypeServiceFabric:
@@ -322,7 +322,7 @@ func (c *managedIdentityClient) getAzureArcSecretKey(ctx context.Context, resour
 	// of the secret key file. Any other status code indicates an error in the request.
 	if response.StatusCode != 401 {
 		msg := fmt.Sprintf("expected a 401 response, received %d", response.StatusCode)
-		return "", newAuthenticationFailedError(credNameManagedIdentity, msg, response)
+		return "", newAuthenticationFailedError(credNameManagedIdentity, msg, response, nil)
 	}
 	header := response.Header.Get("WWW-Authenticate")
 	if len(header) == 0 {