Skip to content

feat: [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads - #18224

Open
agrawalradhika-cell wants to merge 11 commits into
mainfrom
cert-rotation-aiohttp
Open

feat: [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads #18224
agrawalradhika-cell wants to merge 11 commits into
mainfrom
cert-rotation-aiohttp

Conversation

@agrawalradhika-cell

@agrawalradhika-cell agrawalradhika-cell commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

feat: [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads

  • Make sure to open an issue as a bug/issue before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea
  • Ensure the tests and linter pass
  • Code coverage does not decrease (if any source code was changed)
  • Appropriate docs were updated (if necessary)

Fixes #18227 #18227 🦕

@agrawalradhika-cell
agrawalradhika-cell requested review from a team as code owners August 26, 2026 04:00

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/tests/transport/aio/test_sessions_mtls.py
@agrawalradhika-cell agrawalradhika-cell changed the title feat: [aiohttp] Add reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads feat: [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads Aug 26, 2026
agrawalradhika-cell and others added 2 commits August 26, 2026 10:34
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>
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
await self.configure_mtls_channel(
lambda: (call_cert_bytes, call_key_bytes)
)
continue

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the initial request was a streaming upload (e.g., data was passed as an AsyncIterable or generator), executing continue here routes the exhausted generator back into self._auth_request for the next retry attempt, uploading a zero-byte body or crashing.

Per the HLD (go/sdk-mds-bound-token ), we cannot safely retry streaming calls automatically. If we detect a streaming payload, we should still allow the mTLS rotation block to execute so the channel is rebuilt, but we must explicitly skip the continue and return the 401 response to the caller so they can safely reconstruct the stream and retry on the new channel.

I realized that this was a bug in the sync http as well. I've opened this bug to track the sync http fix separately: #18238

Comment thread packages/google-auth/google/auth/aio/transport/sessions.py
await self.configure_mtls_channel(
lambda: (call_cert_bytes, call_key_bytes)
)
continue

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Another possible issue with continue:
When we rotate the mTLS certificate and execute continue to retry, we're currently re-sending the exact same headers dictionary that contains the old access token.

Presenting the old token over a newly established mTLS connection will result in an immediate 401 rejection. We need to ensure the credentials are explicitly refreshed and the headers are updated with the new token before we retry the request.

Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
await self.configure_mtls_channel(
lambda: (call_cert_bytes, call_key_bytes)
)
continue

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Executing continue here routes the 401 retry through the AsyncExponentialBackoff loop, which introduces several unintended side effects:

  • Artificial Sleep Delay: It forces an await asyncio.sleep() delay (starting at ~1 second) before retrying locally. Credential and mTLS rotations should be retried immediately.
  • Shared Budget Exhaustion: It consumes one of the finite total_attempts (default 3) intended for transient 5xx server errors.
  • Lost Retry Edge Case: If the 401 occurs on the final iteration of the backoff loop, continue will raise StopAsyncIteration. The loop exits and returns the 401 to the user without ever executing the retry on the newly configured channel.
  • Socket Leaks: Executing continue without await response.close() leaves the unread 401 response open, which leaks underlying aiohttp connections.

We should handle 401 auth retries explicitly (e.g., via recursion like the sync requests.py transport does, or a dedicated outer loop) rather than injecting them into the exponential backoff loop.

Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like we're missing an end-to-end happy-path test for the 401 certificate rotation flow.

The current test suite covers the failure paths (test_cert_rotation_failure_logs, etc.) and the no-op path where the cert hasn't changed. However, there doesn't appear to be a test verifying the core success path: receiving a 401 -> detecting a cert change -> successfully executing configure_mtls_channel -> retrying the request -> returning a 200 OK.

Adding a full lifecycle mock test for this execution flow is helps ensure that the retry logic actually works end-to-end without leaking state or raising unexpected errors.

mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp")
mock_conf.side_effect = Exception("Failed to reconfigure")

resp = await session.request("GET", "http://example.com")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just a heads-up on the test URLs used here. These newly added rotation tests are currently issuing requests to "http://example.com".

Once we implement the URL prefix check (so we don't attempt mTLS rotations on non-mTLS domains), these tests will break because "http://example.com" will correctly bypass the certificate checking block.

To future-proof these tests, we should update the request URLs to use a valid mTLS domain (e.g., "https://pubsub.mtls.googleapis.com/test"). It would also be highly valuable to add a dedicated test verifying the inverse: that requests to non-mTLS URLs successfully bypass the rotation logic and just return the 401 immediately.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just a few test hygiene cleanups to ensure these newly added tests are robust:

  1. The tests test_cert_rotation_failure_logs and test_cert_rotation_check_params_fails don't actually assert that the warnings are logged. You can add pytest's caplog fixture to the test method signature and assert against caplog.text to verify the logging behavior.
  2. In test_no_cert_rotation_when_cert_match_and_mTLS_enabled, the return value of await session.request(...) is discarded. We should capture it and assert resp == mock_resp just like the other tests do.
  3. None of the three newly added tests call await session.close() at the end. This leaves aiohttp session resources unclosed and diverges from the cleanup pattern used in the rest of this test file.
  4. import http.client as http_client is declared repeatedly inside each test body. Let's put this to the top of the file!

…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add cert rotation handling for aiohttp (Async HTTP)

2 participants