From b80c4bb192c532d5e9e5c8f896bfbcc8c2d7ef32 Mon Sep 17 00:00:00 2001 From: Xiang Yan Date: Mon, 10 Jun 2024 16:56:31 -0700 Subject: [PATCH 1/5] Raise CredentialUnavailableError if the response is not json --- .../azure/identity/_credentials/imds.py | 3 +++ .../_internal/managed_identity_client.py | 2 ++ .../azure/identity/aio/_credentials/imds.py | 3 +++ .../tests/test_managed_identity.py | 24 ++++++++++++++----- .../tests/test_managed_identity_async.py | 20 ++++++++++++++++ 5 files changed, 46 insertions(+), 6 deletions(-) diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/imds.py b/sdk/identity/azure-identity/azure/identity/_credentials/imds.py index 1d775520143a..f5e407d45ca5 100644 --- a/sdk/identity/azure-identity/azure/identity/_credentials/imds.py +++ b/sdk/identity/azure-identity/azure/identity/_credentials/imds.py @@ -87,6 +87,9 @@ def _request_token(self, *scopes: str, **kwargs: Any) -> AccessToken: try: self._client.request_token(*scopes, connection_timeout=1, retry_total=0) self._endpoint_available = True + except CredentialUnavailableError: + # Response is not json, skip the IMDS credential + raise except HttpResponseError as ex: # IMDS responded _check_forbidden_response(ex) diff --git a/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py b/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py index dcce2524ce6c..7d1345ea94a7 100644 --- a/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py +++ b/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py @@ -13,6 +13,7 @@ from azure.core.pipeline.policies import ContentDecodePolicy from azure.core.pipeline import PipelineResponse from azure.core.pipeline.transport import HttpRequest +from .. import CredentialUnavailableError from .._internal import _scopes_to_resource from .._internal.pipeline import build_pipeline @@ -51,6 +52,7 @@ def _process_response(self, response: PipelineResponse, request_time: int) -> Ac message = "Failed to deserialize JSON from response" else: message = 'Unexpected content type "{}"'.format(response.http_response.content_type) + raise CredentialUnavailableError(message=message, response=response.http_response) from ex raise ClientAuthenticationError(message=message, response=response.http_response) from ex if not content: diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/imds.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/imds.py index 264a5fcbdf8b..027938c57647 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/imds.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/imds.py @@ -45,6 +45,9 @@ async def _request_token(self, *scopes: str, **kwargs: Any) -> AccessToken: # p try: await self._client.request_token(*scopes, connection_timeout=1, retry_total=0) self._endpoint_available = True + except CredentialUnavailableError: + # Response is not json, skip the IMDS credential + raise except HttpResponseError as ex: # IMDS responded _check_forbidden_response(ex) diff --git a/sdk/identity/azure-identity/tests/test_managed_identity.py b/sdk/identity/azure-identity/tests/test_managed_identity.py index b0be4b91ee5e..cfb7a9a12135 100644 --- a/sdk/identity/azure-identity/tests/test_managed_identity.py +++ b/sdk/identity/azure-identity/tests/test_managed_identity.py @@ -5,18 +5,15 @@ import os import sys import time - -try: - from unittest import mock -except ImportError: # python < 3.3 - import mock # type: ignore +from unittest import mock from azure.core.credentials import AccessToken from azure.core.exceptions import ClientAuthenticationError -from azure.identity import ManagedIdentityCredential +from azure.identity import ManagedIdentityCredential, CredentialUnavailableError from azure.identity._constants import EnvironmentVariables from azure.identity._credentials.imds import IMDS_AUTHORITY, IMDS_TOKEN_PATH from azure.identity._internal.user_agent import USER_AGENT +from azure.identity._internal import within_credential_chain import pytest from helpers import build_aad_response, validating_transport, mock_response, Request @@ -686,6 +683,21 @@ def test_imds_tenant_id(): assert token == expected_token +def test_imds_text_response(): + within_credential_chain.set(True) + response = mock.Mock( + text=lambda encoding=None: b"{This is a text response}", + headers={"content-type": "text/html; charset=UTF-8"}, + content_type="text/html; charset=UTF-8", + status_code=200, + ) + mock_send = mock.Mock(return_value=response) + credential = ManagedIdentityCredential(transport=mock.Mock(send=mock_send)) + with pytest.raises(CredentialUnavailableError): + token = credential.get_token("") + within_credential_chain.set(False) + + def test_client_id_none(): """the credential should ignore client_id=None""" diff --git a/sdk/identity/azure-identity/tests/test_managed_identity_async.py b/sdk/identity/azure-identity/tests/test_managed_identity_async.py index 79d254848744..f986dc9e7bf7 100644 --- a/sdk/identity/azure-identity/tests/test_managed_identity_async.py +++ b/sdk/identity/azure-identity/tests/test_managed_identity_async.py @@ -8,10 +8,12 @@ from azure.core.credentials import AccessToken from azure.core.exceptions import ClientAuthenticationError +from azure.identity import CredentialUnavailableError from azure.identity.aio import ManagedIdentityCredential from azure.identity._credentials.imds import IMDS_AUTHORITY, IMDS_TOKEN_PATH from azure.identity._constants import EnvironmentVariables from azure.identity._internal.user_agent import USER_AGENT +from azure.identity._internal import within_credential_chain import pytest @@ -716,6 +718,24 @@ async def test_imds_user_assigned_identity(): assert token == expected_token +@pytest.mark.asyncio +async def test_imds_text_response(): + async def send(request, **kwargs): + response = mock.Mock( + text=lambda encoding=None: b"{This is a text response}", + headers={"content-type": "text/html; charset=UTF-8"}, + content_type="text/html; charset=UTF-8", + status_code=200, + ) + return response + + within_credential_chain.set(True) + credential = ManagedIdentityCredential(transport=mock.Mock(send=send)) + with pytest.raises(CredentialUnavailableError): + token = await credential.get_token("") + within_credential_chain.set(False) + + @pytest.mark.asyncio async def test_service_fabric(): """Service Fabric 2019-07-01-preview""" From c7671fa5e90d653337882e4a3f25e852cd77d195 Mon Sep 17 00:00:00 2001 From: Xiang Yan Date: Mon, 10 Jun 2024 16:58:00 -0700 Subject: [PATCH 2/5] update changelog --- sdk/identity/azure-identity/CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sdk/identity/azure-identity/CHANGELOG.md b/sdk/identity/azure-identity/CHANGELOG.md index eb7dd155b7e5..b44909e67f2c 100644 --- a/sdk/identity/azure-identity/CHANGELOG.md +++ b/sdk/identity/azure-identity/CHANGELOG.md @@ -8,6 +8,8 @@ ### Bugs Fixed +- Raised `CredentialUnavailableError` in IMDS credential if the response is not json. ([#35938](https://github.com/Azure/azure-sdk-for-python/pull/35938)) + ### Other Changes ## 1.17.0b2 (2024-06-11) From ce08f3717cdb8dbe11872dcb1a51c1fb28e90896 Mon Sep 17 00:00:00 2001 From: Xiang Yan Date: Mon, 10 Jun 2024 16:59:36 -0700 Subject: [PATCH 3/5] update --- .../azure/identity/_internal/managed_identity_client.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py b/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py index 7d1345ea94a7..d77a3ff2ddae 100644 --- a/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py +++ b/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py @@ -50,10 +50,10 @@ def _process_response(self, response: PipelineResponse, request_time: int) -> Ac except DecodeError as ex: if response.http_response.content_type.startswith("application/json"): message = "Failed to deserialize JSON from response" + raise ClientAuthenticationError(message=message, response=response.http_response) from ex else: message = 'Unexpected content type "{}"'.format(response.http_response.content_type) - raise CredentialUnavailableError(message=message, response=response.http_response) from ex - raise ClientAuthenticationError(message=message, response=response.http_response) from ex + raise CredentialUnavailableError(message=message, response=response.http_response) from ex if not content: raise ClientAuthenticationError(message="No token received.", response=response.http_response) From 8f13e7a261ef13f9cca6d712cd6acfc3c28b858d Mon Sep 17 00:00:00 2001 From: Xiang Yan Date: Mon, 10 Jun 2024 17:49:03 -0700 Subject: [PATCH 4/5] update --- .../azure/identity/_internal/managed_identity_client.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py b/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py index d77a3ff2ddae..b44042811142 100644 --- a/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py +++ b/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py @@ -51,9 +51,8 @@ def _process_response(self, response: PipelineResponse, request_time: int) -> Ac if response.http_response.content_type.startswith("application/json"): message = "Failed to deserialize JSON from response" raise ClientAuthenticationError(message=message, response=response.http_response) from ex - else: - message = 'Unexpected content type "{}"'.format(response.http_response.content_type) - raise CredentialUnavailableError(message=message, response=response.http_response) from ex + message = 'Unexpected content type "{}"'.format(response.http_response.content_type) + raise CredentialUnavailableError(message=message, response=response.http_response) from ex if not content: raise ClientAuthenticationError(message="No token received.", response=response.http_response) From 60d7ec6fadfc28e56631a95f484e8b45e734a34b Mon Sep 17 00:00:00 2001 From: Xiang Yan Date: Tue, 11 Jun 2024 12:11:46 -0700 Subject: [PATCH 5/5] Update sdk/identity/azure-identity/CHANGELOG.md Co-authored-by: Paul Van Eck --- sdk/identity/azure-identity/CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/identity/azure-identity/CHANGELOG.md b/sdk/identity/azure-identity/CHANGELOG.md index b44909e67f2c..ef0ab1ab26f4 100644 --- a/sdk/identity/azure-identity/CHANGELOG.md +++ b/sdk/identity/azure-identity/CHANGELOG.md @@ -8,7 +8,7 @@ ### Bugs Fixed -- Raised `CredentialUnavailableError` in IMDS credential if the response is not json. ([#35938](https://github.com/Azure/azure-sdk-for-python/pull/35938)) +- Allow credential chains to continue when an IMDS probe request returns a non-JSON response in `ManagedIdentityCredential`. ([#36016](https://github.com/Azure/azure-sdk-for-python/pull/36016)) ### Other Changes