From 07668b0f3737a13223cd5fcc0b119c66c2f50ec5 Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Tue, 22 Jul 2025 02:15:34 +0000 Subject: [PATCH] [Identity] Fix issue with cae_cache not being used Signed-off-by: Paul Van Eck --- sdk/identity/azure-identity/CHANGELOG.md | 2 ++ .../_credentials/authorization_code.py | 2 +- .../aio/_credentials/authorization_code.py | 2 +- .../identity/aio/_credentials/on_behalf_of.py | 2 +- .../azure-identity/tests/test_auth_code.py | 21 +++++++++++++------ .../azure-identity/tests/test_obo_async.py | 12 ++++++++--- 6 files changed, 29 insertions(+), 12 deletions(-) diff --git a/sdk/identity/azure-identity/CHANGELOG.md b/sdk/identity/azure-identity/CHANGELOG.md index 2aff6f3d3791..121f72841073 100644 --- a/sdk/identity/azure-identity/CHANGELOG.md +++ b/sdk/identity/azure-identity/CHANGELOG.md @@ -8,6 +8,8 @@ ### Bugs Fixed +- Fixed an issue where CAE (Continuous Access Evaluation) caches were not properly used by `AuthorizationCodeCredential` and the asynchronous `OnBehalfOfCredential`. ([#42145](https://github.com/Azure/azure-sdk-for-python/pull/42145)) + ### Other Changes ## 1.24.0b1 (2025-07-17) diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/authorization_code.py b/sdk/identity/azure-identity/azure/identity/_credentials/authorization_code.py index 419c237457d8..1cd77ee4dcf8 100644 --- a/sdk/identity/azure-identity/azure/identity/_credentials/authorization_code.py +++ b/sdk/identity/azure-identity/azure/identity/_credentials/authorization_code.py @@ -127,7 +127,7 @@ def _request_token(self, *scopes: str, **kwargs) -> AccessTokenInfo: return token token = None - for refresh_token in self._client.get_cached_refresh_tokens(scopes): + for refresh_token in self._client.get_cached_refresh_tokens(scopes, **kwargs): if "secret" in refresh_token: token = self._client.obtain_token_by_refresh_token(scopes, refresh_token["secret"], **kwargs) if token: diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/authorization_code.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/authorization_code.py index 775b6ba6cfef..623a7b5ba095 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/authorization_code.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/authorization_code.py @@ -133,7 +133,7 @@ async def _request_token(self, *scopes: str, **kwargs: Any) -> AccessTokenInfo: return token token = cast(AccessTokenInfo, None) - for refresh_token in self._client.get_cached_refresh_tokens(scopes): + for refresh_token in self._client.get_cached_refresh_tokens(scopes, **kwargs): if "secret" in refresh_token: token = await self._client.obtain_token_by_refresh_token(scopes, refresh_token["secret"], **kwargs) if token: diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/on_behalf_of.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/on_behalf_of.py index c1399dfedde1..e3d1c8e47b44 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/on_behalf_of.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/on_behalf_of.py @@ -119,7 +119,7 @@ async def _request_token(self, *scopes: str, **kwargs: Any) -> AccessTokenInfo: # Note we assume the cache has tokens for one user only. That's okay because each instance of this class is # locked to a single user (assertion). This assumption will become unsafe if this class allows applications # to change an instance's assertion. - refresh_tokens = self._client.get_cached_refresh_tokens(scopes) + refresh_tokens = self._client.get_cached_refresh_tokens(scopes, **kwargs) if len(refresh_tokens) == 1: # there should be only one try: refresh_token = refresh_tokens[0]["secret"] diff --git a/sdk/identity/azure-identity/tests/test_auth_code.py b/sdk/identity/azure-identity/tests/test_auth_code.py index 0d2f5f4d8811..7b777bbaf68f 100644 --- a/sdk/identity/azure-identity/tests/test_auth_code.py +++ b/sdk/identity/azure-identity/tests/test_auth_code.py @@ -80,7 +80,8 @@ def test_tenant_id(get_token_method): @pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS) -def test_auth_code_credential(get_token_method): +@pytest.mark.parametrize("enable_cae", [True, False]) +def test_auth_code_credential(get_token_method, enable_cae): client_id = "client id" secret = "fake-client-secret" tenant_id = "tenant" @@ -118,6 +119,7 @@ def test_auth_code_credential(get_token_method): responses=[mock_response(json_payload=auth_response)] * 2, ) cache = msal.TokenCache() + cae_cache = msal.TokenCache() credential = AuthorizationCodeCredential( client_id=client_id, @@ -127,22 +129,29 @@ def test_auth_code_credential(get_token_method): redirect_uri=redirect_uri, transport=transport, cache=cache, + cae_cache=cae_cache, ) # first call should redeem the auth code - token = getattr(credential, get_token_method)(expected_scope) + kwargs = {"enable_cae": enable_cae} + if get_token_method == "get_token_info": + kwargs = {"options": kwargs} + token = getattr(credential, get_token_method)(expected_scope, **kwargs) assert token.token == expected_access_token assert transport.send.call_count == 1 # no auth code -> credential should return cached token - token = getattr(credential, get_token_method)(expected_scope) + token = getattr(credential, get_token_method)(expected_scope, **kwargs) assert token.token == expected_access_token assert transport.send.call_count == 1 # no auth code, no cached token -> credential should redeem refresh token - cached_access_token = list(cache.search(cache.CredentialType.ACCESS_TOKEN))[0] - cache.remove_at(cached_access_token) - token = getattr(credential, get_token_method)(expected_scope) + cache_being_used = cae_cache if enable_cae else cache + cached_tokens = list(cache_being_used.search(cache_being_used.CredentialType.ACCESS_TOKEN)) + assert cached_tokens + cached_access_token = cached_tokens[0] + cache_being_used.remove_at(cached_access_token) + token = getattr(credential, get_token_method)(expected_scope, **kwargs) assert token.token == expected_access_token assert transport.send.call_count == 2 diff --git a/sdk/identity/azure-identity/tests/test_obo_async.py b/sdk/identity/azure-identity/tests/test_obo_async.py index 5e90040a0af1..3232742d3e59 100644 --- a/sdk/identity/azure-identity/tests/test_obo_async.py +++ b/sdk/identity/azure-identity/tests/test_obo_async.py @@ -249,7 +249,8 @@ def test_invalid_cert(): @pytest.mark.asyncio @pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS) -async def test_refresh_token(get_token_method): +@pytest.mark.parametrize("enable_cae", [True, False]) +async def test_refresh_token(get_token_method, enable_cae): first_token = "***" second_token = first_token * 2 refresh_token = "refresh-token" @@ -274,10 +275,15 @@ async def send(request, **kwargs): credential = OnBehalfOfCredential( "tenant-id", "client-id", client_secret="secret", user_assertion="assertion", transport=Mock(send=send) ) - token = await getattr(credential, get_token_method)("scope") + + kwargs = {"enable_cae": enable_cae} + if get_token_method == "get_token_info": + kwargs = {"options": kwargs} + + token = await getattr(credential, get_token_method)("scope", **kwargs) assert token.token == first_token - token = await getattr(credential, get_token_method)("scope") + token = await getattr(credential, get_token_method)("scope", **kwargs) assert token.token == second_token assert requests == 2