[Core] Update on_challenge in credential policies (#41857)

We will now use authorize_request internally, which will cache the token
to maintain prior behavior of subclassing policies like ARM policies.

Signed-off-by: Paul Van Eck <paulvaneck@microsoft.com>

Author: pvaneck

sdk/core/azure-core/CHANGELOG.md

@@ -18,6 +18,7 @@
 
 - A timeout error when using the `aiohttp` transport (the default for async SDKs) will now be raised as a `azure.core.exceptions.ServiceResponseTimeoutError`, a subtype of the previously raised `ServiceResponseError`.
 - When using with `aiohttp` 3.10 or later, a connection timeout error will now be raised as a `azure.core.exceptions.ServiceRequestTimeoutError`, which can be retried.
+- The default implementation of `on_challenge` in `BearerTokenCredentialPolicy` and `AsyncBearerTokenCredentialPolicy` will now cache the retrieved token. #41857
 
 ## 1.34.0 (2025-05-01)
 

sdk/core/azure-core/azure/core/pipeline/policies/_authentication.py

@@ -204,9 +204,7 @@ def on_challenge(
                 padding_needed = -len(encoded_claims) % 4
                 claims = base64.urlsafe_b64decode(encoded_claims + "=" * padding_needed).decode("utf-8")
                 if claims:
-                    token = self._get_token(*self._scopes, claims=claims)
-                    bearer_token = cast(Union["AccessToken", "AccessTokenInfo"], token).token
-                    request.http_request.headers["Authorization"] = "Bearer " + bearer_token
+                    self.authorize_request(request, *self._scopes, claims=claims)
                     return True
             except Exception:  # pylint:disable=broad-except
                 return False

sdk/core/azure-core/azure/core/pipeline/policies/_authentication_async.py

@@ -149,9 +149,7 @@ async def on_challenge(
                 padding_needed = -len(encoded_claims) % 4
                 claims = base64.urlsafe_b64decode(encoded_claims + "=" * padding_needed).decode("utf-8")
                 if claims:
-                    token = await self._get_token(*self._scopes, claims=claims)
-                    bearer_token = cast(Union["AccessToken", "AccessTokenInfo"], token).token
-                    request.http_request.headers["Authorization"] = "Bearer " + bearer_token
+                    await self.authorize_request(request, *self._scopes, claims=claims)
                     return True
             except Exception:  # pylint:disable=broad-except
                 return False

sdk/core/azure-core/tests/async_tests/test_authentication_async.py

@@ -4,6 +4,7 @@
 # license information.
 # -------------------------------------------------------------------------
 import asyncio
+import base64
 import sys
 import time
 from unittest.mock import Mock, patch, AsyncMock, create_autospec
@@ -12,7 +13,7 @@
 from azure.core.credentials import AccessToken, AccessTokenInfo
 from azure.core.credentials_async import AsyncTokenCredential, AsyncSupportsTokenInfo
 from azure.core.exceptions import ServiceRequestError
-from azure.core.pipeline import AsyncPipeline, PipelineRequest, PipelineContext
+from azure.core.pipeline import AsyncPipeline, PipelineRequest, PipelineContext, PipelineResponse
 from azure.core.pipeline.policies import (
     AsyncBearerTokenCredentialPolicy,
     SansIOHTTPPolicy,
@@ -593,3 +594,49 @@ def test_async_token_credential_sync():
         # Ensure trio isn't in sys.modules (i.e. imported).
         sys.modules.pop("trio", None)
         AsyncBearerTokenCredentialPolicy(Mock(), "scope")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
+async def test_async_bearer_policy_on_challenge_caches_token(http_request):
+    """Test that async on_challenge caches the token when handling claims challenges"""
+    # Setup credentials that return different tokens for different calls
+    initial_token = AccessToken("initial_token", int(time.time()) + 3600)
+    claims_token = AccessToken("claims_token", int(time.time()) + 3600)
+
+    call_count = 0
+
+    async def mock_get_token_info(*scopes, options=None):
+        nonlocal call_count
+        call_count += 1
+        if options and "claims" in options:
+            return claims_token
+        return initial_token
+
+    fake_credential = Mock(spec_set=["get_token_info"], get_token_info=mock_get_token_info)
+    policy = AsyncBearerTokenCredentialPolicy(fake_credential, "scope")
+
+    # Create request and initial response
+    http_req = http_request("GET", "https://example.com")
+    request = PipelineRequest(http_req, PipelineContext(None))
+
+    # Create a 401 response with insufficient_claims challenge
+    test_claims = '{"access_token":{"foo":"bar"}}'
+    encoded_claims = base64.urlsafe_b64encode(test_claims.encode()).decode().rstrip("=")
+    challenge_header = f'Bearer error="insufficient_claims", claims="{encoded_claims}"'
+
+    response_mock = Mock(status_code=401, headers={"WWW-Authenticate": challenge_header})
+    response = PipelineResponse(request, response_mock, PipelineContext(None))
+
+    # Call on_challenge
+    result = await policy.on_challenge(request, response)
+
+    # Verify the challenge was handled successfully
+    assert result is True
+
+    # Verify the token was cached
+    assert policy._token is claims_token
+    assert policy._token.token == "claims_token"
+
+    # Verify the Authorization header was set correctly
+    assert request.http_request.headers["Authorization"] == "Bearer claims_token"

sdk/core/azure-core/tests/test_authentication.py

@@ -4,6 +4,7 @@
 # license information.
 # -------------------------------------------------------------------------
 from collections import namedtuple
+import base64
 import time
 from itertools import product
 from requests import Response
@@ -16,7 +17,7 @@
     AccessTokenInfo,
 )
 from azure.core.exceptions import ServiceRequestError
-from azure.core.pipeline import Pipeline, PipelineRequest, PipelineContext
+from azure.core.pipeline import Pipeline, PipelineRequest, PipelineContext, PipelineResponse
 from azure.core.pipeline.transport import HttpTransport, HttpRequest
 from azure.core.pipeline.policies import (
     BearerTokenCredentialPolicy,
@@ -791,3 +792,48 @@ def test_access_token_subscriptable():
     assert len(token) == 2
     assert token[0] == "token"
     assert token[1] == 42
+
+
+@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
+def test_bearer_policy_on_challenge_caches_token_with_claims(http_request):
+    """Test that on_challenge caches the token when handling claims challenges"""
+    # Setup credentials that return different tokens for different calls
+    initial_token = AccessToken("initial_token", int(time.time()) + 3600)
+    claims_token = AccessToken("claims_token", int(time.time()) + 3600)
+
+    call_count = 0
+
+    def mock_get_token_info(*scopes, options):
+        nonlocal call_count
+        call_count += 1
+        if options and "claims" in options:
+            return claims_token
+        return initial_token
+
+    fake_credential = Mock(spec_set=["get_token_info"], get_token_info=mock_get_token_info)
+    policy = BearerTokenCredentialPolicy(fake_credential, "scope")
+
+    # Create request and initial response
+    http_req = http_request("GET", "https://example.com")
+    request = PipelineRequest(
+        http_req, PipelineContext(None)
+    )  # Create a 401 response with insufficient_claims challenge
+    test_claims = '{"access_token":{"foo":"bar"}}'
+    encoded_claims = base64.urlsafe_b64encode(test_claims.encode()).decode().rstrip("=")
+    challenge_header = f'Bearer error="insufficient_claims", claims="{encoded_claims}"'
+
+    response_mock = Mock(status_code=401, headers={"WWW-Authenticate": challenge_header})
+    response = PipelineResponse(request, response_mock, PipelineContext(None))
+
+    # Call on_challenge
+    result = policy.on_challenge(request, response)
+
+    # Verify the challenge was handled successfully
+    assert result is True
+
+    # Verify the token was cached
+    assert policy._token is claims_token
+    assert policy._token.token == "claims_token"
+
+    # Verify the Authorization header was set correctly
+    assert request.http_request.headers["Authorization"] == "Bearer claims_token"