-
Notifications
You must be signed in to change notification settings - Fork 3.3k
cosmos: handle HTTP 403/sub-status 5300 (AAD_REQUEST_NOT_AUTHORIZED) by refreshing bearer token and retrying #46167
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Draft
Copilot
wants to merge
11
commits into
main
Choose a base branch
from
copilot/fix-azure-cosmos-403-error
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Draft
Changes from all commits
Commits
Show all changes
11 commits
Select commit
Hold shift + click to select a range
32b88c1
Initial plan
Copilot 58c0cbf
Merge branch 'main' into copilot/fix-azure-cosmos-403-error
bambriz 792e84d
Fix 403 AAD token refresh in CosmosBearerTokenCredentialPolicy (sync …
Copilot de53008
Fix condition indentation in send() overrides
Copilot 8ef88db
Merge branch 'main' into copilot/fix-azure-cosmos-403-error
bambriz cb77da6
Rewrite auth policy tests using realistic Pipeline with MockTransport
Copilot 5ea1653
Merge branch 'main' into copilot/fix-azure-cosmos-403-error
bambriz a095d2c
Fix spelling: retriable -> retryable in test docstrings
Copilot 0277df9
Merge branch 'main' into copilot/fix-azure-cosmos-403-error
bambriz f2daec1
Update CHANGELOG with bug fix for HTTP 403/5300 AAD token refresh
Copilot 35c6bc6
Merge branch 'main' into copilot/fix-azure-cosmos-403-error
bambriz File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,219 @@ | ||
| # The MIT License (MIT) | ||
| # Copyright (c) Microsoft Corporation. All rights reserved. | ||
|
|
||
| """Unit tests for CosmosBearerTokenCredentialPolicy 403/AAD token refresh behavior. | ||
|
|
||
| Uses a realistic azure-core Pipeline with a mock transport that returns proper | ||
| requests.Response objects (including the x-ms-substatus header), and verifies | ||
| that the Authorization header is correctly set in the requests that reach the transport. | ||
| """ | ||
|
|
||
| import time | ||
| import unittest | ||
| from unittest.mock import Mock | ||
|
|
||
| from requests import Response | ||
|
|
||
| from azure.core.credentials import AccessToken | ||
| from azure.core.pipeline import Pipeline | ||
| from azure.core.pipeline.transport import HttpTransport, HttpRequest | ||
|
|
||
| from azure.cosmos._auth_policy import CosmosBearerTokenCredentialPolicy | ||
| from azure.cosmos.http_constants import HttpHeaders, SubStatusCodes | ||
|
|
||
| COSMOS_ACCOUNT_URL = "https://example.cosmos.azure.com" | ||
| ACCOUNT_SCOPE = "https://cosmos.azure.com/.default" | ||
| AAD_AUTH_PREFIX = "type=aad&ver=1.0&sig=" | ||
|
|
||
|
|
||
| def _make_response(status_code, sub_status=None): | ||
| """Create a requests.Response with optional x-ms-substatus header.""" | ||
| response = Response() | ||
| response.status_code = status_code | ||
| if sub_status is not None: | ||
| response.headers[HttpHeaders.SubStatus] = str(sub_status) | ||
| return response | ||
|
|
||
|
|
||
| def _make_credential(token_str="fake-token"): | ||
| """Create a sync credential mock that returns an AccessToken via get_token.""" | ||
| credential = Mock(spec_set=["get_token"]) | ||
| credential.get_token.return_value = AccessToken(token_str, int(time.time()) + 3600) | ||
| return credential | ||
|
|
||
|
|
||
| class MockTransport(HttpTransport): | ||
| """Minimal sync HTTP transport that replays a sequence of canned responses and | ||
| records each outgoing request so tests can inspect its headers.""" | ||
|
|
||
| def __init__(self, *responses): | ||
| self._responses = list(responses) | ||
| self.requests = [] | ||
|
|
||
| def open(self): | ||
| pass | ||
|
|
||
| def close(self): | ||
| pass | ||
|
|
||
| def __exit__(self, *args): | ||
| pass | ||
|
|
||
| def __enter__(self): | ||
| return self | ||
|
|
||
| def send(self, request, **kwargs): | ||
| self.requests.append(request) | ||
| return self._responses.pop(0) | ||
|
|
||
|
|
||
| class TestCosmosBearerTokenPolicySend(unittest.TestCase): | ||
|
|
||
| def _run(self, credential, *responses): | ||
| """Build a Pipeline with the Cosmos bearer policy and run a GET against it. | ||
|
|
||
| Returns (pipeline_response, transport) so callers can inspect both the | ||
| final response and the recorded outgoing requests. | ||
| """ | ||
| transport = MockTransport(*responses) | ||
| policy = CosmosBearerTokenCredentialPolicy(credential, ACCOUNT_SCOPE) | ||
| pipeline = Pipeline(transport=transport, policies=[policy]) | ||
| http_response = pipeline.run(HttpRequest("GET", f"{COSMOS_ACCOUNT_URL}/dbs")) | ||
| return http_response, transport | ||
|
|
||
| # ------------------------------------------------------------------ | ||
| # Pass-through cases — no retry expected | ||
| # ------------------------------------------------------------------ | ||
|
|
||
| def test_200_response_passes_through(self): | ||
| """A 200 response is forwarded to the caller with no retry.""" | ||
| credential = _make_credential() | ||
| _, transport = self._run(credential, _make_response(200)) | ||
|
|
||
| assert transport.requests[0].headers["Authorization"].startswith(AAD_AUTH_PREFIX) | ||
| assert len(transport.requests) == 1 | ||
|
|
||
| def test_403_without_substatus_no_retry(self): | ||
| """A 403 with no sub-status is not an AAD expiry — no retry should occur.""" | ||
| credential = _make_credential() | ||
| result, transport = self._run(credential, _make_response(403)) | ||
|
|
||
| assert result.http_response.status_code == 403 | ||
| assert len(transport.requests) == 1 | ||
|
|
||
| def test_403_write_forbidden_no_retry(self): | ||
| """403/WRITE_FORBIDDEN is a different error — no AAD-triggered retry.""" | ||
| credential = _make_credential() | ||
| result, transport = self._run( | ||
| credential, _make_response(403, sub_status=SubStatusCodes.WRITE_FORBIDDEN) | ||
| ) | ||
|
|
||
| assert result.http_response.status_code == 403 | ||
| assert len(transport.requests) == 1 | ||
|
|
||
| # ------------------------------------------------------------------ | ||
| # 403 / AAD_REQUEST_NOT_AUTHORIZED — retry expected | ||
| # ------------------------------------------------------------------ | ||
|
|
||
| def test_403_aad_expired_retries_and_succeeds(self): | ||
| """403/AAD_REQUEST_NOT_AUTHORIZED triggers a token refresh and one retry. | ||
|
|
||
| The retry must succeed with the fresh token, and both the initial request | ||
| and the retry must carry a properly-formatted Cosmos AAD Authorization header. | ||
| """ | ||
| credential = _make_credential("fresh-token") | ||
| result, transport = self._run( | ||
| credential, | ||
| _make_response(403, sub_status=SubStatusCodes.AAD_REQUEST_NOT_AUTHORIZED), | ||
| _make_response(200), | ||
| ) | ||
|
|
||
| assert result.http_response.status_code == 200 | ||
| assert len(transport.requests) == 2 | ||
|
|
||
| # Both requests must carry the Cosmos-specific AAD header format | ||
| for req in transport.requests: | ||
| assert req.headers["Authorization"].startswith(AAD_AUTH_PREFIX), ( | ||
| f"Expected Cosmos AAD header format, got: {req.headers.get('Authorization')}" | ||
| ) | ||
|
|
||
| def test_403_aad_expired_sends_fresh_token_on_retry(self): | ||
| """The retry request must use a freshly-acquired token, not the expired one. | ||
|
|
||
| We give the credential two different tokens: the first simulates an expired | ||
| cached token; the second is the fresh one returned after the cache is cleared. | ||
| """ | ||
| fresh_token = "brand-new-token" | ||
| expired_token = "old-expired-token" | ||
|
|
||
| call_count = [0] | ||
| tokens = [expired_token, fresh_token] | ||
|
|
||
| credential = Mock(spec_set=["get_token"]) | ||
|
|
||
| def rotating_get_token(*scopes, **kwargs): | ||
| token = tokens[min(call_count[0], len(tokens) - 1)] | ||
| call_count[0] += 1 | ||
| return AccessToken(token, int(time.time()) + 3600) | ||
|
|
||
| credential.get_token.side_effect = rotating_get_token | ||
|
|
||
| transport = MockTransport( | ||
| _make_response(403, sub_status=SubStatusCodes.AAD_REQUEST_NOT_AUTHORIZED), | ||
| _make_response(200), | ||
| ) | ||
| policy = CosmosBearerTokenCredentialPolicy(credential, ACCOUNT_SCOPE) | ||
| pipeline = Pipeline(transport=transport, policies=[policy]) | ||
| pipeline.run(HttpRequest("GET", f"{COSMOS_ACCOUNT_URL}/dbs")) | ||
|
|
||
| assert len(transport.requests) == 2 | ||
| retry_auth = transport.requests[1].headers["Authorization"] | ||
| assert fresh_token in retry_auth, ( | ||
| f"Expected fresh token '{fresh_token}' in retry Authorization header, got: {retry_auth}" | ||
| ) | ||
|
|
||
| def test_403_aad_expired_auth_header_cleared_before_retry(self): | ||
| """After 403/5300 the policy clears its cached token so the retry gets a new one. | ||
|
|
||
| We force the token cache to contain an expired-looking token and verify | ||
| that the Authorization header on the retry differs from the initial request. | ||
| """ | ||
| credential = _make_credential("fresh-token-after-expiry") | ||
| transport = MockTransport( | ||
| _make_response(403, sub_status=SubStatusCodes.AAD_REQUEST_NOT_AUTHORIZED), | ||
| _make_response(200), | ||
| ) | ||
| policy = CosmosBearerTokenCredentialPolicy(credential, ACCOUNT_SCOPE) | ||
| # Inject a "stale" token into the policy cache to simulate an expired token | ||
| policy._token = AccessToken("stale-token", int(time.time()) - 60) | ||
|
|
||
| pipeline = Pipeline(transport=transport, policies=[policy]) | ||
| pipeline.run(HttpRequest("GET", f"{COSMOS_ACCOUNT_URL}/dbs")) | ||
|
|
||
| assert len(transport.requests) == 2 | ||
| initial_auth = transport.requests[0].headers["Authorization"] | ||
| retry_auth = transport.requests[1].headers["Authorization"] | ||
| # The stale token must not appear in the retry request | ||
| assert "stale-token" not in retry_auth, ( | ||
| "Stale token should have been replaced before retry" | ||
| ) | ||
| # Both headers must still use the Cosmos-specific format | ||
| assert initial_auth.startswith(AAD_AUTH_PREFIX) | ||
| assert retry_auth.startswith(AAD_AUTH_PREFIX) | ||
|
|
||
| def test_403_aad_retry_still_fails_returns_second_response(self): | ||
| """If the retry also returns a non-retryable 403, that response is returned unchanged.""" | ||
| credential = _make_credential() | ||
| result, transport = self._run( | ||
| credential, | ||
| _make_response(403, sub_status=SubStatusCodes.AAD_REQUEST_NOT_AUTHORIZED), | ||
| _make_response(403, sub_status=SubStatusCodes.WRITE_FORBIDDEN), | ||
| ) | ||
|
|
||
| assert result.http_response.status_code == 403 | ||
| assert len(transport.requests) == 2 | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| unittest.main() | ||
|
|
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.