fix(auth): allow mTLS retry when credentials raise NotImplementedError on refresh

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,17 @@ async def _recover_auth_state():
                                     _LOGGER.debug(
                                         "Credentials do not implement refresh()."
                                     )
-                                    return response
+                                    # A retry only helps when an mTLS reconfiguration
+                                    # occurred for this mTLS endpoint. Short-circuit on
+                                    # non-mTLS endpoints first so that a concurrent
+                                    # rotation (which bumps the session-wide counter)
+                                    # cannot trigger a spurious retry here.
+                                    if (
+                                        not is_mtls_endpoint
+                                        or self._mtls_reconfig_counter
+                                        <= reconfig_counter_at_error
+                                    ):
+                                        return response
                                 except (
                                     exceptions.RefreshError,
                                     getattr(exceptions, "InvalidOperation", Exception),

packages/google-auth/tests/transport/aio/test_sessions_mtls.py

@@ -1126,3 +1126,241 @@ 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()
+
+    @pytest.mark.asyncio
+    async def test_credential_refresh_not_implemented_no_retry_on_non_mtls_during_concurrent_rotation(
+        self,
+    ):
+        """A non-mTLS request whose credentials raise NotImplementedError on
+        refresh() must NOT retry, even when a concurrent mTLS rotation has
+        advanced the session-wide reconfig counter. Guards against the race
+        where the shared counter alone would signal a spurious retry."""
+        mock_creds = mock.AsyncMock(spec=credentials.Credentials)
+        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()
+
+        # If a retry were (incorrectly) triggered, the 200 would be returned.
+        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"
+
+        # Advance the session-wide counter after this request snapshots it (the
+        # snapshot happens before before_request), simulating a concurrent mTLS
+        # request that reconfigured mTLS while this one was in flight.
+        async def advance_counter(*args, **kwargs):
+            session._mtls_reconfig_counter += 1
+
+        mock_creds.before_request = mock.AsyncMock(side_effect=advance_counter)
+
+        with (
+            mock.patch(
+                "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response",
+                new_callable=mock.AsyncMock,
+            ) as mock_check,
+            mock.patch.object(
+                session, "configure_mtls_channel", new_callable=mock.AsyncMock
+            ) as mock_conf,
+        ):
+            resp = await session.request("GET", "https://pubsub.googleapis.com/test")
+
+            # The 401 is returned as-is; no retry, no cert-rotation machinery
+            # for a non-mTLS endpoint.
+            assert resp == mock_resp_401
+            mock_check.assert_not_called()
+            mock_conf.assert_not_called()
+            mock_creds.refresh.assert_called_once_with(mock_auth_req)
+            assert mock_auth_req.call_count == 1
+
+        await session.close()