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
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>
chore: Reorder response closing logic for clarity
chore: Handle additional exception during credential refresh
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>
Updated the 'check_parameters_for_unauthorized_response' function to include an optional client_cert_callback parameter and added detailed docstring for better understanding.
Store refresh counter at error for better tracking.
Updated mock patches to ensure correct assertions and added bound mocks for certificate parsing and caching.
Update mock authentication response handling in tests.
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
| ), mock.patch( | ||
| "aiohttp.ClientSession" | ||
| ) as mock_session: | ||
| """Tests that an exception in old_auth_request.close() does not abort configuration.""" |
There was a problem hiding this comment.
nit: In configure_mtls_channel(), old transports are buffered until len(self._old_auth_requests) >= 2 before closing. Because a single configuration call appends session._auth_request without calling close(), the eviction loop and exception suppression block in sessions.py remain unexercised here (the mock error is only caught at teardown in session.close()). Consider driving configure_mtls_channel() three times (clearing _mtls_init_task = None between calls) or pre-populating _old_auth_requests so the test exercises the eviction loop and verifies that close() exceptions do not abort configuration.
| ) | ||
|
|
||
| with pytest.raises(TimeoutError): | ||
| await session.request("GET", "https://example.com", max_allowed_time=0.01) |
There was a problem hiding this comment.
nit: max_allowed_time=0.01 (10ms) against real clock time risks flaking on loaded CI runners if event loop scheduling latency exceeds ~8.5ms during initial mock dispatch, which aborts before 401 recovery is entered and leaves mock_resp_401 unclosed. Consider setting max_allowed_time=0.1 (or 0.2) to provide ample headroom for mock dispatch while keeping unit test runtime brisk and reliably timing out inside slow_refresh (10s):
with pytest.raises(TimeoutError):
await session.request("GET", "https://example.com", max_allowed_time=0.1)|
|
||
| assert resp == mock_resp_200 | ||
| mock_check.assert_called_once() | ||
| mock_conf.assert_called_once_with(mock.ANY) |
There was a problem hiding this comment.
nit: In test_cert_rotation_success_and_retry and test_psc_endpoint_triggers_cert_rotation, asserting mock_conf.assert_called_once_with(mock.ANY) accepts any argument and does not verify that the newly fetched certificates are forwarded into configure_mtls_channel(). Consider capturing the callback and asserting its return value to ensure rotated credentials are sent:
mock_conf.assert_called_once()
cb = (
mock_conf.call_args.args[0]
if mock_conf.call_args.args
else mock_conf.call_args.kwargs["client_cert_callback"]
)
assert cb() == (new_cert, new_key)Refactor old authentication request handling in aiohttp transport.
Remove redundant return statement in certificate checking logic.
Update return type annotations for client certificate callback.
Refactor mTLS endpoint check and initialization task handling.
…for_unauthorized_response function and adjusted the implementation accordingly. Removed the client_cert_callback parameter from the check_parameters_for_unauthorized_response function and adjusted the implementation accordingly.
Added check_counter_at_error to track mTLS checks during errors.
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 🦕