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 19 commits into
Conversation
feat: Add retry for cert rotation handling
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.
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.
| or hasattr(data, "read") | ||
| ) | ||
| if getattr(self, "is_mtls", False) and any( | ||
| prefix in url for prefix in MTLS_URL_PREFIXES |
There was a problem hiding this comment.
prefix in url searches the entire URL, so paths or query parameters can trigger certificate rotation on standard hosts. Parse and normalize urlsplit(url).hostname, then match only exact hostnames or subdomains of documented mTLS endpoints.
| else: | ||
| response.close() | ||
| try: | ||
| await self._credentials.refresh(self._auth_request) |
There was a problem hiding this comment.
Closing the 401 response before credentials refresh returns an unreadable response if RefreshError occurs. Also, StaticCredentials and AnonymousCredentials raise uncaught InvalidOperation errors during refresh. Close the response only after refresh succeeds. Return the open original response if refresh fails or is unsupported.
| ): | ||
| 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.
| "`auth_request` must either be configured or the external package `aiohttp` must be installed to use the default value." | ||
| ) | ||
| self._auth_request = _auth_request | ||
| self._mtls_rotation_lock = asyncio.Lock() |
There was a problem hiding this comment.
self._mtls_rotation_lock = asyncio.Lock() runs in synchronous __init__. In Python 3.8 and 3.9, calling asyncio.Lock() outside an active event loop grabs asyncio.get_event_loop(). If someone creates the session in a sync factory, background thread, or test fixture that runs before the loop starts, Python throws RuntimeError: There is no current event loop in thread. If the session is used across different test loops, it fails with RuntimeError: Task got Future attached to a different loop. Set self._mtls_rotation_lock = None in __init__, then create the lock lazily inside an async helper on first use.
| ) | ||
| except Exception as e: | ||
| _LOGGER.warning( | ||
| "Failed to check client certificate parameters: %s. Proceeding with original response.", |
There was a problem hiding this comment.
The warning log on line 362 says we proceed with the original response, but the block never actually returns it. Code falls straight through to line 400, refreshes credentials, and retries the request twice on the unrotated transport. Add return response inside this except Exception as e: block so parameter check errors stop execution immediately.
| "Failed to reconfigure mTLS channel: %s", e | ||
| ) | ||
| raise exceptions.MutualTLSChannelError( | ||
| "Failed to reconfigure mTLS channel" |
There was a problem hiding this comment.
Raising MutualTLSChannelError here escapes before the code reaches response.close() on line 394. The 401 response socket stays open and leaks on the rotation failure path. Close the response before raising the error.
| ) | ||
|
|
||
| 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.
| else: | ||
| response.close() | ||
| try: | ||
| await self._credentials.refresh(self._auth_request) |
There was a problem hiding this comment.
The refresh and retry logic on lines 399 to 417 sits outside the mTLS check, which ends on line 391. Every 401 response across all endpoints now triggers credentials refresh and retries up to two times. While this matches the behavior of the synchronous session, the async session did not do this before, and the PR description frames this change purely around mTLS certificate mismatch. If this behavior change is intentional, please mention it in the PR description and add a test confirming that non-mTLS 401s refresh and retry.
feat: [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads
Fixes #18227 #18227 🦕