aci returns non-json response for probing request (#36193)

* aci returns non-json response for probing request

* update

* update release date

Author: xiangyan99

sdk/identity/azure-identity/CHANGELOG.md

@@ -1,14 +1,10 @@
 # Release History
 
-## 1.17.1 (Unreleased)
-
-### Features Added
-
-### Breaking Changes
+## 1.17.1 (2024-06-21)
 
 ### Bugs Fixed
 
-### Other Changes
+- Continue to attempt requesting token if the probing request receives non-json response. ([#36184](https://github.com/Azure/azure-sdk-for-python/pull/36184))
 
 ## 1.17.0 (2024-06-18)
 

sdk/identity/azure-identity/azure/identity/_credentials/imds.py

@@ -87,9 +87,6 @@ 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)
@@ -102,6 +99,9 @@ def _request_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
 
         try:
             token = self._client.request_token(*scopes, headers={"Metadata": "true"})
+        except CredentialUnavailableError:
+            # Response is not json, skip the IMDS credential
+            raise
         except HttpResponseError as ex:
             # 400 in response to a token request indicates managed identity is disabled,
             # or the identity with the specified client_id is not available

sdk/identity/azure-identity/azure/identity/aio/_credentials/imds.py

@@ -45,9 +45,6 @@ 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)
@@ -60,6 +57,9 @@ async def _request_token(self, *scopes: str, **kwargs: Any) -> AccessToken:  # p
 
         try:
             token = await self._client.request_token(*scopes, headers={"Metadata": "true"})
+        except CredentialUnavailableError:
+            # Response is not json, skip the IMDS credential
+            raise
         except HttpResponseError as ex:
             # 400 in response to a token request indicates managed identity is disabled,
             # or the identity with the specified client_id is not available

sdk/identity/azure-identity/tests/test_imds_credential.py

@@ -93,7 +93,7 @@ def test_retries():
 
     total_retries = PIPELINE_SETTINGS["retry_total"]
 
-    assert within_credential_chain.get() == False
+    assert not within_credential_chain.get()
     for status_code in (404, 410, 429, 500):
         mock_send.reset_mock()
         mock_response.status_code = status_code
@@ -251,3 +251,40 @@ def test_user_assigned_tenant_id(self):
         token = credential.get_token(self.scope, tenant_id="tenant_id")
         assert token.token
         assert isinstance(token.expires_on, int)
+
+    def test_managed_identity_aci_probe(self):
+        access_token = "****"
+        expires_on = 42
+        expected_token = AccessToken(access_token, expires_on)
+        scope = "scope"
+        transport = validating_transport(
+            requests=[
+                Request(base_url=IMDS_AUTHORITY + IMDS_TOKEN_PATH),
+                Request(
+                    base_url=IMDS_AUTHORITY + IMDS_TOKEN_PATH,
+                    method="GET",
+                    required_headers={"Metadata": "true"},
+                    required_params={"resource": scope},
+                ),
+            ],
+            responses=[
+                # probe receives error response
+                mock_response(status_code=400),
+                mock_response(
+                    json_payload={
+                        "access_token": access_token,
+                        "expires_in": 42,
+                        "expires_on": expires_on,
+                        "ext_expires_in": 42,
+                        "not_before": int(time.time()),
+                        "resource": scope,
+                        "token_type": "Bearer",
+                    }
+                ),
+            ],
+        )
+        within_credential_chain.set(True)
+        cred = ImdsCredential(transport=transport)
+        token = cred.get_token(scope)
+        assert token == expected_token
+        within_credential_chain.set(False)

sdk/identity/azure-identity/tests/test_imds_credential_async.py

@@ -14,6 +14,7 @@
 from azure.identity._credentials.imds import IMDS_AUTHORITY, IMDS_TOKEN_PATH
 from azure.identity._internal.user_agent import USER_AGENT
 from azure.identity.aio._credentials.imds import ImdsCredential, PIPELINE_SETTINGS
+from azure.identity._internal.utils import within_credential_chain
 import pytest
 
 from helpers import mock_response, Request
@@ -291,3 +292,40 @@ async def test_user_assigned_tenant_id(self):
         token = await credential.get_token(self.scope, tenant_id="tenant_id")
         assert token.token
         assert isinstance(token.expires_on, int)
+
+    @pytest.mark.asyncio
+    async def test_managed_identity_aci_probe(self):
+        access_token = "****"
+        expires_on = 42
+        expected_token = AccessToken(access_token, expires_on)
+        scope = "scope"
+        transport = async_validating_transport(
+            requests=[
+                Request(base_url=IMDS_AUTHORITY + IMDS_TOKEN_PATH),
+                Request(
+                    base_url=IMDS_AUTHORITY + IMDS_TOKEN_PATH,
+                    method="GET",
+                    required_headers={"Metadata": "true"},
+                    required_params={"resource": scope},
+                ),
+            ],
+            responses=[
+                mock_response(status_code=400),
+                mock_response(
+                    json_payload={
+                        "access_token": access_token,
+                        "expires_in": 42,
+                        "expires_on": expires_on,
+                        "ext_expires_in": 42,
+                        "not_before": int(time.time()),
+                        "resource": scope,
+                        "token_type": "Bearer",
+                    }
+                ),
+            ],
+        )
+        within_credential_chain.set(True)
+        cred = ImdsCredential(transport=transport)
+        token = await cred.get_token(scope)
+        assert token == expected_token
+        within_credential_chain.set(False)