diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index f2ced1280e50..232acdf2ba65 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -165,6 +165,7 @@ def __init__( self._auth_request = _auth_request self._mtls_rotation_lock: Optional[asyncio.Lock] = None self._mtls_check_counter = 0 + self._mtls_reconfig_counter = 0 self._refresh_lock: Optional[asyncio.Lock] = None self._refresh_counter = 0 @@ -326,6 +327,7 @@ async def request( start_time = time.monotonic() refresh_counter_at_error = self._refresh_counter check_counter_at_error = self._mtls_check_counter + reconfig_counter_at_error = self._mtls_reconfig_counter async with timeout_guard(max_allowed_time) as with_timeout: await with_timeout( # Note: before_request will attempt to refresh credentials if expired. @@ -444,6 +446,7 @@ async def _recover_auth_state(): call_key_bytes, ) ) + self._mtls_reconfig_counter += 1 except Exception as e: _LOGGER.error( "Failed to reconfigure mTLS channel: %s", @@ -485,7 +488,11 @@ async def _recover_auth_state(): _LOGGER.debug( "Credentials do not implement refresh()." ) - return response + if ( + self._mtls_reconfig_counter + <= reconfig_counter_at_error + ): + return response except ( exceptions.RefreshError, getattr(exceptions, "InvalidOperation", Exception), diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index f6f1185a660e..5b66f392c1d8 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -1126,3 +1126,184 @@ async def slow_mtls_init(): assert not session._mtls_init_task.cancelled() assert session._is_mtls is True await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_credential_refresh_not_implemented_retries(self): + """Validate credentials that raise NotImplementedError on refresh() + still trigger a retry after mTLS reconfiguration, not return the 401.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(side_effect=NotImplementedError) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + mock_resp_200.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + with mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check: + with mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf: + mock_check.return_value = ( + b"new_cert", + b"new_key", + b"old_fp", + b"new_fp", + ) + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + # Validate that the handler falls through to `return None` + # on NotImplementedError in order to signal retry. + assert resp == mock_resp_200 + 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() == (b"new_cert", b"new_key") + mock_creds.refresh.assert_called_once_with(mock_auth_req) + assert mock_auth_req.call_count == 2 + mock_resp_401.close.assert_called_once() + + await session.close() + + @pytest.mark.asyncio + async def test_credential_refresh_not_implemented_no_retry_on_non_mtls(self): + """Validate credentials raising NotImplementedError on refresh do NOT + retry when the request is on a non-mTLS endpoint.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(side_effect=NotImplementedError) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + + resp = await session.request("GET", "https://pubsub.googleapis.com/test") + + assert resp == mock_resp_401 + assert mock_auth_req.call_count == 1 + await session.close() + + @pytest.mark.asyncio + async def test_credential_refresh_not_implemented_no_retry_when_cert_not_rotated( + self, + ): + """Validate credentials raising NotImplementedError on refresh do NOT + retry when the mTLS certificate has not rotated.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(side_effect=NotImplementedError) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"current_cert" + + with mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check: + mock_check.return_value = ( + b"current_cert", + b"current_key", + b"same_fp", + b"same_fp", + ) + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp_401 + assert mock_auth_req.call_count == 1 + + await session.close() + + @pytest.mark.asyncio + async def test_credential_refresh_not_implemented_concurrent_rotation_retries( + self, + ): + """Validate that concurrent requests hitting 401 during rotation both retry + and succeed when credentials raise NotImplementedError.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(side_effect=NotImplementedError) + + mock_resp_401_a = mock.Mock( + status_code=http_client.UNAUTHORIZED, close=mock.AsyncMock() + ) + mock_resp_401_b = mock.Mock( + status_code=http_client.UNAUTHORIZED, close=mock.AsyncMock() + ) + mock_resp_200_a = mock.Mock(status_code=http_client.OK, close=mock.AsyncMock()) + mock_resp_200_b = mock.Mock(status_code=http_client.OK, close=mock.AsyncMock()) + + mock_auth_req = mock.AsyncMock( + side_effect=[ + mock_resp_401_a, + mock_resp_401_b, + mock_resp_200_a, + mock_resp_200_b, + ] + ) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + with mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check: + with mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ): + mock_check.return_value = ( + b"new_cert", + b"new_key", + b"old_fp", + b"new_fp", + ) + + resps = await asyncio.gather( + session.request("GET", "https://pubsub.mtls.googleapis.com/test"), + session.request("GET", "https://pubsub.mtls.googleapis.com/test"), + ) + assert resps == [mock_resp_200_a, mock_resp_200_b] + assert mock_auth_req.call_count == 4 + + await session.close()