Skip to content

Commit

Permalink
aci returns non-json response for probing request (#36193)
Browse files Browse the repository at this point in the history
* aci returns non-json response for probing request

* update

* update release date
  • Loading branch information
xiangyan99 authored Jun 22, 2024
1 parent 6632e13 commit 4defbc7
Show file tree
Hide file tree
Showing 5 changed files with 84 additions and 13 deletions.
8 changes: 2 additions & 6 deletions sdk/identity/azure-identity/CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -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)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand Down
39 changes: 38 additions & 1 deletion sdk/identity/azure-identity/tests/test_imds_credential.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
38 changes: 38 additions & 0 deletions sdk/identity/azure-identity/tests/test_imds_credential_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

0 comments on commit 4defbc7

Please sign in to comment.