feat(auth): [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads - #18224
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>
chore: Reorder response closing logic for clarity
chore: Handle additional exception during credential refresh
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
Added debug logging for unimplemented refresh method in credentials.
|
|
||
| old_auth_request = self._auth_request | ||
| self._auth_request = AiohttpRequest(session=new_session) | ||
| self._old_auth_requests.append(old_auth_request) |
There was a problem hiding this comment.
I think we can solve this cleanly without background tasks. Since rotations occur hours apart and _do_configure() is already an async coroutine, can we consider bounding self._old_auth_requests to two generations and awaiting close() in-line on the oldest entry when appending a new one?
This gives in-flight requests on the immediately displaced session ample time to complete, and bounds retention.
Added refresh lock and counter to manage concurrent credential refreshes.
Limit the number of old authentication requests to 2 and ensure proper closure of the oldest requests.
Added tests for handling 401 responses with timeout and cancellation scenarios.
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
…onse Reordered parameters in check_parameters_for_unauthorized_response function.
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
| return has_cert, cert, key | ||
|
|
||
|
|
||
| async def check_parameters_for_unauthorized_response(cached_cert, client_cert_callback): |
There was a problem hiding this comment.
Reordering the arguments to (cached_cert, client_cert_callback) dropped the =None default on client_cert_callback. Both _mtls_helper.check_parameters_for_unauthorized_response(cached_cert, client_cert_callback=None) and mtls.get_client_cert_and_key(client_cert_callback=None) default client_cert_callback to None. Calling await mtls.check_parameters_for_unauthorized_response(cached_cert) now raises a TypeError. Add =None to match: async def check_parameters_for_unauthorized_response(cached_cert, client_cert_callback=None):.
|
|
||
| mock_resp_401 = mock.Mock() | ||
| mock_resp_401.status_code = http_client.UNAUTHORIZED | ||
| mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) |
There was a problem hiding this comment.
mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) returns 401 on retry after rotation, which triggers a second recovery wave (mock_check.call_count == 2, mock_creds.refresh.call_count == 2) until hitting _auth_retry_count == 2. Because the retry fails, this test only asserts mock_conf.assert_called_once() and cannot verify that concurrent 401s deduplicate mock_check and mock_creds.refresh. Match test_cert_rotation_lock_contention_no_cert_change on line 590 by setting side_effect=[mock_resp_401] * 3 + [mock_resp_200] * 3 and asserting resp == mock_resp_200, mock_check.assert_called_once(), and mock_creds.refresh.call_count == 1.
|
|
||
| async def _recover_auth_state(): | ||
| is_mtls_endpoint = False | ||
| refresh_counter_at_error = self._refresh_counter |
There was a problem hiding this comment.
refresh_counter_at_error = self._refresh_counter is captured when _recover_auth_state() starts, after the 401 response arrives. If two concurrent requests go out with stale credentials and the first finishes recovery before the second receives its 401 network response, the second request captures the already incremented _refresh_counter and triggers an extra refresh. Capture refresh_counter_at_start = self._refresh_counter at the top of request() before before_request() on line 322 and compare against that value during recovery.
|
|
||
|
|
||
| async def check_parameters_for_unauthorized_response(cached_cert, client_cert_callback): | ||
| """Async helper to retrieve certs and compute fingerprints for mTLS rotation.""" |
There was a problem hiding this comment.
check_parameters_for_unauthorized_response only has a one-line summary docstring. Add Args: and Returns: sections to match _mtls_helper.py:812 and the other public functions in this file.
| cached_fp, | ||
| current_fp, | ||
| ) = await mtls.check_parameters_for_unauthorized_response( | ||
| cached_cert=CERT_BYTES, client_cert_callback=callback |
There was a problem hiding this comment.
test_check_parameters_cert_matched and test_check_parameters_cert_mismatch_rotation mock parse_certificate and get_cached_cert_fingerprint with fixed return values, but never check their call arguments. Add mock_parse.assert_called_once_with(CERT_BYTES) and mock_get_cached.assert_called_once_with(CERT_BYTES) to verify _fetch_fingerprints() passes the current and cached certs to the right functions.
6 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 🦕