feat: [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads - #18224
feat: [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads #18224agrawalradhika-cell wants to merge 25 commits into
Conversation
feat: Add retry for cert rotation handling
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces client certificate rotation handling for asynchronous authorized sessions when encountering an unauthorized response under mTLS. The review feedback highlights a violation of the repository style guide regarding exception contract compliance, suggesting that the certificate parameter check should be wrapped in a try-except block to gracefully fall back to the original response rather than crashing. Additionally, the feedback recommends updating the corresponding unit tests to assert this resilient fallback behavior.
agrawalradhika-cell
changed the title
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Handle exceptions during mTLS reconfiguration with warnings instead of errors.
…logs Updated test logic to assert response instead of expecting an error.
…sync executor Refactor unauthorized response handling to use async executor for MTLS parameter checks.
chore: Reset mTLS init task upon client certificate change
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
…eck after 401 check chore: Refactor mTLS channel reconfiguration logic for adding mTLS check after 401 check
Implement mTLS rotation lock to prevent race conditions during certificate reconfiguration.
chore: Change warning to error log for mTLS channel reconfiguration failure.
chore: Refactor mTLS handling for unauthorized responses
Remove unnecessary continue statement after mTLS configuration.
Refactor tests for certificate rotation and error handling in AsyncAuthorizedSession. Update test names for clarity and ensure proper logging of errors.
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
Handle RefreshError during credential refresh to prevent unhandled exceptions.
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
| url, | ||
| data=data, | ||
| headers=headers, | ||
| max_allowed_time=max_allowed_time, |
There was a problem hiding this comment.
max_allowed_time=max_allowed_time resets the timer on each retry instead of limiting total runtime.
| ): | ||
| self._mtls_init_task = None | ||
| await self.configure_mtls_channel( | ||
| lambda: (call_cert_bytes, call_key_bytes) |
There was a problem hiding this comment.
Initial setup can use client_cert_callback, but check_parameters_for_unauthorized_response always checks application default credentials. This mismatch can replace custom certificates or skip rotation. Save the initial certificate source, and reuse it for all fingerprint checks and reconfiguration.
| and self._mtls_init_task.done() | ||
| ): | ||
| self._mtls_init_task = None | ||
| await self.configure_mtls_channel( |
There was a problem hiding this comment.
configure_mtls_channel closes the active ClientSession, which aborts in-flight concurrent requests. Keep old sessions open until their requests finish, or close them in AsyncAuthorizedSession.close(). Add a concurrency regression test for this case.
| ) | ||
|
|
||
| assert resp == mock_resp | ||
| assert mock_check.call_count >= 1 |
There was a problem hiding this comment.
assert mock_check.call_count >= 1 masks the fallthrough retry bug. If checking parameters fails, the session should abort on the first attempt without calling credentials refresh or retrying the request. Change this to mock_check.assert_called_once(), and assert that mock_creds.refresh was not called.
| with mock.patch( | ||
| "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response" | ||
| ) as mock_check, mock.patch.object( | ||
| session, "configure_mtls_channel", new_callable=mock.AsyncMock |
There was a problem hiding this comment.
None of the new tests check lock contention under load. The mocked configure_mtls_channel never updates _cached_cert, so the double-checked lock skip check on line 346, if self._cached_cert != stale_cert:, is never exercised. Add a test with asyncio.gather where multiple concurrent requests receive 401s at the same time, confirming the channel is reconfigured once and following tasks take the skip branch.
chore: Reorder response closing logic for clarity
chore: Handle additional exception during credential refresh
Change _mtls_rotation_lock initialization to None and update its usage.
fix: Handle response closure in mTLS error handling
Fix: Fix improperly falling through to the credential refresh logic.
3 participants
feat: [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads
Changes included:
401 Unauthorizedresponses (not just mTLS).Fixes #18227 #18227 🦕