Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
agrawalradhika-cell wants to merge 11 commits into
base: main
Choose a base branch
from
+612 −59
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
226afd4
fix: address concurrency crashes, state desynchronization, and test c…
agrawalradhika-cell Sep 10, 2026
c9ec394
Update packages/google-auth/google/auth/aio/transport/sessions.py
agrawalradhika-cell Sep 10, 2026
2f7baa5
Update packages/google-auth/google/auth/aio/transport/sessions.py
agrawalradhika-cell Sep 10, 2026
0ffd85f
fix(auth): fix fingerprint comparison on empty cached cert, catch Typ…
agrawalradhika-cell Sep 10, 2026
3675858
fix(auth): make configure_mtls_channel idempotent for repeated defaul…
agrawalradhika-cell Sep 10, 2026
5aef5c3
test(auth): make test_401_retry_raises_timeout_before_subsequent_retr…
agrawalradhika-cell Sep 10, 2026
96ec622
fix(auth): allow mTLS retry when credentials raise NotImplementedErro…
agrawalradhika-cell Sep 10, 2026
688f76f
style(auth): format with black and fix flake8 style issues
agrawalradhika-cell Sep 10, 2026
54eb0bc
test(auth): add unit test for consecutive mTLS certificate rotations
agrawalradhika-cell Sep 10, 2026
118785b
test(auth): pass mock_auth_request in cancellation test to avoid unmo…
agrawalradhika-cell Sep 10, 2026
e046760
test(auth): allow timeout_guard error message in test_401_retry_raise…
agrawalradhika-cell Sep 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion packages/google-auth/google/auth/aio/transport/mtls.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,7 +211,7 @@ def _fetch_fingerprints():
cached_cert
)
else:
cached_fingerprint = None
cached_fingerprint = current_fingerprint
return cached_fingerprint, current_fingerprint

cached_fingerprint, current_cert_fingerprint = await _run_in_executor(
Expand Down
94 changes: 55 additions & 39 deletions packages/google-auth/google/auth/aio/transport/sessions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,7 +193,20 @@ async def configure_mtls_channel(self, client_cert_callback=None):
google.auth.exceptions.MutualTLSChannelError: If mutual TLS channel
creation failed for any reason.
"""
if self._mtls_init_task is None:
is_explicit_reconfig = (
client_cert_callback is not None
and client_cert_callback != self._client_cert_callback
)
task_failed = (
self._mtls_init_task is not None
and self._mtls_init_task.done()
and (
self._mtls_init_task.cancelled()
or self._mtls_init_task.exception() is not None
)
)

if self._mtls_init_task is None or is_explicit_reconfig or task_failed:
self._client_cert_callback = client_cert_callback

async def _do_configure():
Expand Down Expand Up @@ -224,21 +237,24 @@ async def _do_configure():

old_auth_request = self._auth_request
self._auth_request = AiohttpRequest(session=new_session)
self._is_mtls = True
self._cached_cert = cert
self._old_auth_requests.append(old_auth_request)

while len(self._old_auth_requests) > 2:
oldest_auth_request = self._old_auth_requests[0]
oldest_auth_request = self._old_auth_requests.pop(0)
try:
if hasattr(oldest_auth_request, "close"):
res = oldest_auth_request.close()
if inspect.isawaitable(res):
await res
except Exception:
pass
self._old_auth_requests.pop(0)

else:
is_mtls = False
self._is_mtls = False
self._cached_cert = None
warnings.warn(
"Attempted to establish mTLS, but a custom async transport was provided. "
"google-auth cannot automatically configure custom transports for mTLS. "
Expand All @@ -247,11 +263,8 @@ async def _do_configure():
"using Certificate-Bound Tokens.",
UserWarning,
)

self._is_mtls = is_mtls
if is_mtls:
self._cached_cert = cert
else:
self._is_mtls = False
self._cached_cert = None

except Exception as caught_exc:
Expand All @@ -260,11 +273,7 @@ async def _do_configure():

self._mtls_init_task = asyncio.create_task(_do_configure())

try:
return await self._mtls_init_task
except BaseException:
self._mtls_init_task = None
raise
return await asyncio.shield(self._mtls_init_task)

async def request(
self,
Expand Down Expand Up @@ -319,6 +328,11 @@ async def request(
# Suppress all exceptions from the background mTLS initialization task,
# allowing the request to fail naturally elsewhere.
pass
except asyncio.CancelledError:
if self._mtls_init_task.cancelled():
pass
else:
raise
retries = _exponential_backoff.AsyncExponentialBackoff(
total_attempts=total_attempts,
)
Expand Down Expand Up @@ -371,6 +385,7 @@ async def request(
)

async def _recover_auth_state():
channel_reconfigured = False
is_mtls_endpoint = False
if self._is_mtls:
hostname = urllib.parse.urlsplit(url).hostname
Expand All @@ -394,6 +409,7 @@ async def _recover_auth_state():
):
pass
else:
check_passed = False
try:
(
call_cert_bytes,
Expand All @@ -404,15 +420,18 @@ async def _recover_auth_state():
self._cached_cert,
self._client_cert_callback,
)
check_passed = True
except (
exceptions.ClientCertError,
exceptions.MutualTLSChannelError,
OSError,
ValueError,
TypeError,
ImportError,
) as e:
_LOGGER.warning(
"Failed to check client certificate parameters: %s. Proceeding with original response.",
"Failed to check client certificate parameters: %s. "
"Falling back to credential refresh and retry.",
e,
)
else:
Expand All @@ -429,21 +448,13 @@ async def _recover_auth_state():
"Client certificate has changed, reconfiguring mTLS "
"channel."
)
if self._mtls_init_task is not None:
if (
not self._mtls_init_task.done()
):
try:
await self._mtls_init_task
except Exception:
pass
self._mtls_init_task = None
await self.configure_mtls_channel(
lambda: (
call_cert_bytes,
call_key_bytes,
)
)
channel_reconfigured = True
Comment thread
agrawalradhika-cell marked this conversation as resolved.
except Exception as e:
_LOGGER.error(
"Failed to reconfigure mTLS channel: %s",
Expand All @@ -467,14 +478,17 @@ async def _recover_auth_state():
"Skipping reconfiguration of mTLS channel because the client"
" certificate has not changed."
)
# Always increment so waiting tasks skip the check block
self._mtls_check_counter += 1
if check_passed:
self._mtls_check_counter += 1
if self._refresh_lock is None:
self._refresh_lock = asyncio.Lock()

async with self._refresh_lock:
# Check if another task already refreshed credentials while we were waiting
if self._refresh_counter > refresh_counter_at_error:
if (
not channel_reconfigured
and self._refresh_counter > refresh_counter_at_error
):
_LOGGER.debug(
"Credentials were already refreshed by a concurrent task. Skipping duplicate refresh."
)
Expand All @@ -485,10 +499,9 @@ async def _recover_auth_state():
_LOGGER.debug(
"Credentials do not implement refresh()."
)
return response
except (
exceptions.RefreshError,
getattr(exceptions, "InvalidOperation", Exception),
exceptions.InvalidOperation,
) as e:
_LOGGER.debug(
"Credential refresh failed, returning 401 response. Error: %s",
Expand Down Expand Up @@ -821,24 +834,27 @@ async def close(self) -> None:
"""
Close the underlying auth request session.
"""
if self._mtls_init_task and not self._mtls_init_task.done():
self._mtls_init_task.cancel()
try:
await self._mtls_init_task
except asyncio.CancelledError:
pass
try:
if hasattr(self._auth_request, "close"):
res = self._auth_request.close()
if inspect.isawaitable(res):
await res
if self._mtls_init_task and not self._mtls_init_task.done():
self._mtls_init_task.cancel()
try:
await self._mtls_init_task
except (Exception, asyncio.CancelledError):
pass
finally:
for old_request in self._old_auth_requests:
while self._old_auth_requests:
old_request = self._old_auth_requests.pop(0)
try:
if hasattr(old_request, "close"):
res = old_request.close()
if inspect.isawaitable(res):
await res
except Exception:
pass
self._old_auth_requests.clear()
try:
if hasattr(self._auth_request, "close"):
res = self._auth_request.close()
if inspect.isawaitable(res):
await res
except Exception:
pass
4 changes: 2 additions & 2 deletions packages/google-auth/tests/transport/aio/test_mtls.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,9 +159,9 @@ def callback():

assert cert == CERT_BYTES
assert key == KEY_BYTES
assert cached_fp is None
assert cached_fp == "FINGERPRINT_CURRENT"
assert current_fp == "FINGERPRINT_CURRENT"
assert cached_fp != current_fp
assert cached_fp == current_fp
mock_get_cached.assert_not_called()


Expand Down
Loading
Loading