Skip to content

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

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

feat(auth): [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads #18224
agrawalradhika-cell wants to merge 89 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

Changes included:

  • Adds mTLS reconfiguration logic when a certificate mismatch occurs.
  • Aligns async session behavior with the synchronous session by triggering a credential refresh and request retry on all 401 Unauthorized responses (not just mTLS).
  • 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
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/tests/transport/aio/test_sessions_mtls.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
…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>
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
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
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/mtls.py Outdated
Comment thread packages/google-auth/tests/transport/aio/test_sessions_mtls.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/mtls.py Outdated
Comment thread packages/google-auth/tests/transport/aio/test_mtls.py
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>
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py
Comment thread packages/google-auth/google/auth/transport/_mtls_helper.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/mtls.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/tests/transport/aio/test_sessions_mtls.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
Comment thread packages/google-auth/google/auth/aio/transport/mtls.py
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.
Refactor mTLS configuration tests to improve clarity and correctness. Update assertions and mock behaviors for better test coverage.
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
Refactor mTLS channel tests by removing outdated tests and updating assertions to ensure proper handling of exceptions and retries.
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
" certificate has not changed."
)
# Always increment so waiting tasks skip the check block
self._mtls_check_counter += 1

@nbayati nbayati Sep 8, 2026

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.

self._mtls_check_counter += 1 is nested inside the else: block at line 461, so it only executes when check_parameters_for_unauthorized_response() completes without raising. When handled exceptions occur (such as ClientCertError, MutualTLSChannelError, or OSError), the except block catches them and skips else:, leaving the counter unchanged. Any coroutines queued on _mtls_rotation_lock will then evaluate the skip guard as false and redundantly repeat the failing certificate check.

Note: moving the increment into a finally: block would advance the counter during asyncio.CancelledError unwinding or channel reconfiguration failures.

Consider placing the counter increment immediately after the try...except...else block.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Had placed self._mtls_check_counter += 1 is nested inside the else: block because of earlier comment that led to assumption - The counter only increments if the certificate check completely succeeds (and reconfiguration succeeds or wasn't needed). If Task A hits a transient error, the else block is skipped, and the counter stays the same. When Task B wakes up, it correctly re-runs the certificate check. This is much safer, as it forces concurrent tasks to retry the check rather than falsely assuming a failed check was resolved.

Please let me know if you still want me to alter it?

creation failed for any reason.
"""
if self._mtls_init_task is None:
self._client_cert_callback = client_cert_callback

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 _do_configure() encounters an exception or is cancelled, self._mtls_init_task is never reset to None, permanently blocking subsequent retries in configure_mtls_channel(). Additionally, because asyncio.CancelledError inherits from BaseException in Python 3.8+, the except Exception: handler in request() (line 314) fails to catch background task cancellation, causing subsequent calls to session.request() to crash with an unhandled CancelledError.

Note that catching CancelledError with pass in request() would swallow outer cancellation of the calling request; checking not self._mtls_init_task.cancelled() ensures caller cancellation is preserved.

Consider resetting self._mtls_init_task = None on failure in configure_mtls_channel(), and guarding the await in request():

# In configure_mtls_channel():
try:
    return await self._mtls_init_task
except BaseException:
    self._mtls_init_task = None
    raise

# In request() (lines 311-317):
if self._mtls_init_task and not self._mtls_init_task.done():
    try:
        await self._mtls_init_task
    except asyncio.CancelledError:
        if not self._mtls_init_task.cancelled():
            raise
    except Exception:
        pass

saved_callback
)
else:
_LOGGER.info(

@nbayati nbayati Sep 8, 2026

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.

nit: Consider distinguishing between a missing certificate and an unchanged certificate when logging:

new_callable=mock.AsyncMock,
) as mock_check,
mock.patch.object(
session, "configure_mtls_channel", new_callable=mock.AsyncMock

@nbayati nbayati Sep 8, 2026

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.

Because session._mtls_init_task is None in these rotation tests, the newly added _mtls_init_task.done() reset branch in sessions.py remains unexercised. The new pre-refresh and pre-retry TimeoutError branches also lack unit tests. Consider adding test cases to cover these paths, such as pre-populating _mtls_init_task with a completed task and simulating elapsed time on 401 retries.

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)

6 participants