Skip to content

fix(auth): address concurrency crashes, state desync, and test coverage in async mTLS sessions V2 - #18355

Open
andyrzhao wants to merge 19 commits into
googleapis:mainfrom
andyrzhao:main
Open

fix(auth): address concurrency crashes, state desync, and test coverage in async mTLS sessions V2#18355
andyrzhao wants to merge 19 commits into
googleapis:mainfrom
andyrzhao:main

Conversation

@andyrzhao

@andyrzhao andyrzhao commented Sep 12, 2026

Copy link
Copy Markdown
Contributor

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:

  • 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 Add cert rotation handling for aiohttp (Async HTTP) #18227

Overview

This PR resolves concurrency crashes, state desynchronization, and authentication recovery regressions in AsyncAuthorizedSession during asynchronous mTLS initialization and certificate rotation.
It builds upon and resolves open gaps identified in #18224 and #18334:

  1. Task Shielding & In-Flight Concurrency:
    • In request(), shielded self._mtls_init_task with asyncio.shield so that external caller cancellation does not cancel background session setup. Propagates asyncio.CancelledError cleanly without leaving the session in a corrupt state.
    • In configure_mtls_channel(), awaits any active, in-flight _mtls_init_task before creating a new one, eliminating race conditions that could clobber self._auth_request.
  2. Idempotency & Reconfiguration Controls:
    • Correctly handles transitions from a custom client_cert_callback back to default ADC (client_cert_callback=None).
    • Added a force: bool = False parameter to configure_mtls_channel() to allow explicit reconfiguration when dynamic callback providers rotate credentials without changing the callback reference.
  3. Lock Contention Deduplication on Failure:
    • In _recover_auth_state(), self._mtls_check_counter is now incremented unconditionally at the end of the check block. If parameter checks fail (e.g. ClientCertError or MutualTLSChannelError), subsequent concurrent tasks waiting on the lock properly detect that the check has already run and bypass redundant validation spikes.
  4. Accurate 401 Credential Recovery Semantics:
    • Fixed NotImplementedError handling so that non-mTLS requests with unrefreshable credentials return 401 immediately rather than triggering a futile retry.
    • Separated exceptions.InvalidOperation handling so that static credentials (e.g., CBA / static tokens) that cannot be refreshed will still retry if certificate rotation successfully reconfigured the mTLS channel (channel_reconfigured is True).
  5. Clean Session Teardown:
    • In close(), safely drains and closes _old_auth_requests and the active _auth_request in finally blocks, and cancels any pending initialization tasks without swallowing caller exceptions.
  6. Default Fingerprint Symmetry:
    • In mtls.check_parameters_for_unauthorized_response(), correctly defaults cached_fingerprint = current_fingerprint when cached_cert is None, ensuring equality checks skip redundant reconfigurations when appropriate.

agrawalradhika-cell and others added 13 commits September 10, 2026 18:24
…overage in async mTLS sessions

- Wrap await self._mtls_init_task in asyncio.shield in configure_mtls_channel to prevent external cancellations from destroying the init task
- Catch asyncio.CancelledError in request() when _mtls_init_task was cancelled
- Force credential refresh in _recover_auth_state when certificate rotation occurs, and conditionally increment _mtls_check_counter on check success
- Support reconfiguration in configure_mtls_channel when task is None or done without unsafe task variable resets
- Update close() to safely drain and close _old_auth_requests and _auth_request in a robust try...finally structure
- Atomically update self._is_mtls and self._cached_cert upon _auth_request swap
- Ensure 401 response is closed on timeout and update test assertions
- Restore 5 unit tests for configure_mtls_channel and add e2e rotation test
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
…eError in cert check, and reference InvalidOperation directly

- Match synchronous behavior in mtls.py by setting cached_fingerprint = current_fingerprint when cached_cert is falsy to prevent spurious mTLS reconfigurations on 401
- Catch TypeError in _recover_auth_state parameter checks and update warning log to reflect fallback to credential refresh and retry
- Reference exceptions.InvalidOperation directly in refresh exception handler instead of getattr fallback
- Add unit tests for TypeError fallback, empty cached cert comparison, and InvalidOperation handling
…t calls and preserve metadata state on error

- Distinguish between idempotent default calls and explicit reconfigurations in configure_mtls_channel
- Preserve self._is_mtls and self._cached_cert on configuration failure to keep metadata in sync with the active _auth_request
- Add unit tests verifying configure_mtls_channel idempotency and state preservation
…r on refresh

- In AsyncAuthorizedSession.request()'s _recover_auth_state(), do not return 401 response on NotImplementedError so that mTLS reconfiguration can fall through to retry
- Add unit test test_cert_rotation_credential_refresh_not_implemented_retries
@andyrzhao
andyrzhao requested review from a team as code owners September 12, 2026 05:10

@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 enhances the asynchronous mTLS channel configuration and 401 retry/rotation handling in google-auth. Key changes include adding a force parameter to configure_mtls_channel, shielding the initialization task, improving credential refresh logic during reconfiguration, and adding comprehensive unit tests. The review feedback highlights a potential AttributeError when checking cancellation on a None task, suggests cancelling the previous initialization task during reconfiguration to avoid blocking, and recommends mocking time.monotonic locally rather than globally to comply with the repository's style guide.

Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py
Comment thread packages/google-auth/tests/transport/aio/test_sessions_mtls.py Outdated
Comment thread packages/google-auth/tests/transport/aio/test_sessions_mtls.py Outdated
andyrzhao and others added 6 commits September 12, 2026 05:22
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
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