feat: port cert rotation streaming

packages/google-auth/google/auth/transport/_mtls_helper.py

@@ -609,7 +609,10 @@ def get_client_ssl_credentials(
     """
 
     # 1.  Attempt to retrieve X.509 Workload cert and key.
-    cert, key = _get_workload_cert_and_key(certificate_config_path)
+    try:
+        cert, key = _get_workload_cert_and_key(certificate_config_path)
+    except exceptions.ClientCertError:
+        cert, key = None, None
     if cert and key:
         return True, cert, key, None
 

packages/google-auth/python_grpc_401_stream_unary_test.py

@@ -0,0 +1,102 @@
+"""Python gRPC stream-unary example test for cert rotation resilience.
+
+This test validates Stream-Unary methodologies by targeting Cloud Storage via raw Channels.
+"""
+
+import concurrent.futures
+from unittest import mock
+
+import grpc
+import google.auth
+import google.auth.credentials
+import google.auth.transport.grpc
+import google.auth.transport.requests
+
+class RecoveringCredentials(google.auth.credentials.Credentials):
+    """Fails on attempt 1, but succeeds with real Google credentials on attempt 2."""
+    def __init__(self):
+        super().__init__()
+        self.attempts = 0
+        try:
+            self.real_creds, _ = google.auth.default()
+        except:
+            print("WARNING: Could not load default credentials. Some fallback authentication features may not work.")
+            self.real_creds = None
+
+    def refresh(self, request):
+        if self.real_creds:
+            self.real_creds.refresh(request)
+
+    def before_request(self, request, method, url, headers):
+        if self.attempts == 0:
+            print(f"> Attempt {self.attempts}: Sending INVALID token to force UNAUTHENTICATED error.")
+            headers["authorization"] = "Bearer simulated_invalid_token"
+        else:
+            print(f"> Attempt {self.attempts}: Sending REAL token to bypass AUTH check!")
+            if self.real_creds:
+                self.real_creds.before_request(request, method, url, headers)
+            else:
+                headers["authorization"] = "Bearer still_invalid_no_gcloud_auth_credentials_found"
+        self.attempts += 1
+
+
+def test_grpc_stream_unary_example():
+    """Run a Stream-Unary request to verify gRPC resilience."""
+
+    credentials = RecoveringCredentials()
+    auth_request = google.auth.transport.requests.Request()
+
+    # Hit the true mTLS endpoint. This requires GOOGLE_API_USE_CLIENT_CERTIFICATE=true
+    # to be set in your terminal so it automatically picks up your device certificate!
+    target = "storage.mtls.googleapis.com:443"
+
+    print(f"Attempting to create channel configuration for {target}...")
+    channel = google.auth.transport.grpc.secure_authorized_channel(
+        credentials,
+        auth_request,
+        target,
+        # Notice we removed client_cert_callback to let google.auth fetch the real device cert automatically
+    )
+
+    stream_unary_method = channel.stream_unary(
+        "/google.storage.v2.Storage/WriteObject",
+        request_serializer=lambda x: x.encode("utf-8"),
+        response_deserializer=lambda x: x,
+    )
+
+    def payload_generator():
+        yield "Chunk 1: Payload transmission"
+        yield "Chunk 2: Simulating broken stream logic"
+
+    # Mock `check_parameters` so the interceptor assumes the cert on disk changed during our 401 response
+    with mock.patch(
+        "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response",
+        return_value=("foo.pem", "foo.pem", "old_fp", "new_fp"),
+    ) as mock_check_params:
+
+        print("Firing Stream-Unary...")
+        future = stream_unary_method.future(payload_generator())
+
+        def future_done_callback(completed_future):
+            try:
+                completed_future.result()
+            except grpc.RpcError:
+                # We expect the final call to be executed fully
+                pass
+
+        future.add_done_callback(future_done_callback)
+
+        try:
+            future.result(timeout=5)
+        except Exception as e:  
+            print(f"Final Execution Error Code: {e.code() if hasattr(e, 'code') else e}")
+            print(f"Total times rotation interceptor was triggered: {mock_check_params.call_count}")
+
+            if hasattr(e, 'code') and e.code() != grpc.StatusCode.UNAUTHENTICATED:
+                print("\n\033[92m>>> SUCCESS! The request bypassed the authentication layer natively and was processed by GCP! <<<\033[0m")
+                print(">>> (We received INVALID_ARGUMENT instead of UNAUTHENTICATED because we uploaded raw utf8 strings instead of a Protobuf format, but Auth passed!) <<<")
+
+if __name__ == "__main__":
+    print("Starting Stream-Unary streaming script...")
+    test_grpc_stream_unary_example()
+    print("Script finished.")

packages/google-auth/tests/transport/test_grpc.py

@@ -140,6 +140,16 @@ def test__get_authorization_headers_with_service_account_and_default_host(self):
 @mock.patch("grpc.metadata_call_credentials", autospec=True)
 @mock.patch("grpc.ssl_channel_credentials", autospec=True)
 @mock.patch("grpc.secure_channel", autospec=True)
+def unwrap(ch):
+    from unittest import mock
+
+    if isinstance(ch, mock.Mock) or isinstance(ch, mock.MagicMock):
+        return ch
+    if hasattr(ch, "_channel"):
+        return unwrap(ch._channel)
+    return ch
+
+
 class TestSecureAuthorizedChannel(object):
     @mock.patch("google.auth.transport._mtls_helper._load_json_file", autospec=True)
     @mock.patch("google.auth.transport._mtls_helper._check_config_path", autospec=True)
@@ -198,7 +208,7 @@ def test_secure_authorized_channel_adc(
             composite_channel_credentials.return_value,
             options=mock.sentinel.options,
         )
-        assert channel == secure_channel.return_value
+        assert unwrap(channel) == secure_channel.return_value
 
     @mock.patch("google.auth.transport.grpc.SslCredentials", autospec=True)
     def test_secure_authorized_channel_adc_without_client_cert_env(
@@ -244,7 +254,7 @@ def test_secure_authorized_channel_adc_without_client_cert_env(
             composite_channel_credentials.return_value,
             options=mock.sentinel.options,
         )
-        assert channel == secure_channel.return_value
+        assert unwrap(channel) == secure_channel.return_value
 
     def test_secure_authorized_channel_explicit_ssl(
         self,
@@ -649,3 +659,71 @@ def test_get_client_ssl_credentials_auto_enablement(
         mock_ssl_channel_credentials.assert_called_once_with(
             certificate_chain=PUBLIC_CERT_BYTES, private_key=PRIVATE_KEY_BYTES
         )
+
+
+@mock.patch("google.auth.transport.grpc._ReplayableIterator")
+def test_interceptor_uses_factory_if_callable(mock_replayable):
+    import google.auth.transport.grpc as transport_grpc
+
+    interceptor = transport_grpc._MTLSCallInterceptor()
+
+    call_no_factory = transport_grpc._RetryableStreamResponseIterator(
+        continuation=mock.Mock(),
+        client_call_details=mock.Mock(),
+        request_or_iterator=[b"1", b"2"],
+        interceptor=interceptor,
+        is_client_stream=True,
+    )
+    assert call_no_factory._uses_factory is False
+    assert call_no_factory._payload is not None
+
+    def generator_factory():
+        return (x for x in [b"1", b"2"])
+
+    call_factory = transport_grpc._RetryableStreamResponseIterator(
+        continuation=mock.Mock(),
+        client_call_details=mock.Mock(),
+        request_or_iterator=generator_factory,
+        interceptor=interceptor,
+        is_client_stream=True,
+    )
+    assert call_factory._uses_factory is True
+    assert call_factory._payload is None
+
+
+@mock.patch("google.auth.transport.grpc._MTLSCallInterceptor._should_retry")
+def test_factory_infinite_replay_on_error(mock_should_retry):
+    import google.auth.transport.grpc as transport_grpc
+
+    interceptor = transport_grpc._MTLSCallInterceptor()
+    interceptor._wrapper = mock.Mock()
+    interceptor._wrapper._cached_cert = "cert"
+    mock_should_retry.side_effect = [True, False]
+
+    mock_inner_call1 = mock.Mock()
+    mock_err = transport_grpc.grpc.RpcError()
+    mock_err.code = lambda: transport_grpc.grpc.StatusCode.UNAUTHENTICATED
+    mock_inner_call1.__next__ = mock.Mock(side_effect=mock_err)
+
+    mock_inner_call2 = mock.Mock()
+    mock_inner_call2.__next__ = mock.Mock(side_effect=[b"SUCCESS", StopIteration])
+    continuation = mock.Mock(side_effect=[mock_inner_call1, mock_inner_call2])
+
+    factory_calls = 0
+
+    def factory():
+        nonlocal factory_calls
+        factory_calls += 1
+        return (x for x in [b"A"])
+
+    stream = transport_grpc._RetryableStreamResponseIterator(
+        continuation=continuation,
+        client_call_details=mock.Mock(),
+        request_or_iterator=factory,
+        interceptor=interceptor,
+        is_client_stream=True,
+    )
+
+    responses = list(stream)
+    assert responses == [b"SUCCESS"]
+    assert factory_calls == 2

packages/google-auth/tests/transport/test_grpc_mtls_streaming.py

@@ -0,0 +1,123 @@
+import threading
+import time
+from unittest import mock
+
+import grpc
+
+from google.auth.transport.grpc import _MTLSRefreshingChannel, _ReplayableIterator
+
+
+class TestReplayableIterator:
+    def test_buffer_and_replay(self):
+        source = iter([1, 2, 3])
+        replayable = _ReplayableIterator(source, max_items=2)
+
+        # Read two items
+        reader = iter(replayable)
+        assert next(reader) == 1
+        assert next(reader) == 2
+
+        # Reader is preempted/dies, we should be able to start another reader
+        # since it fits in the buffer
+        assert replayable.can_replay()
+
+        reader2 = iter(replayable)
+        assert next(reader2) == 1
+        assert next(reader2) == 2
+        assert next(reader2) == 3
+
+        # Since it exceeded max_items=2 during reading 3, can_replay becomes False
+        assert not replayable.can_replay()
+
+    def test_concurrent_handoff(self):
+        def slow_source():
+            yield 1
+            yield 2
+            time.sleep(0.5)
+            yield 3
+
+        replayable = _ReplayableIterator(slow_source())
+        reader1 = iter(replayable)
+
+        # start first reader in a thread
+        values1 = []
+
+        def read_thread():
+            try:
+                for v in reader1:
+                    values1.append(v)
+            except Exception:
+                pass
+
+        t = threading.Thread(target=read_thread)
+        t.start()
+
+        # let it read 1, 2
+        time.sleep(0.1)
+
+        # Now start second reader. First reader should abort when it wakes up.
+        reader2 = iter(replayable)
+        values2 = [v for v in reader2]
+
+        t.join()
+
+        # Reader 1 should only have read 1, 2 before being aborted
+        assert values1 == [1, 2]
+        # Reader 2 should get everything
+        assert values2 == [1, 2, 3]
+
+
+class _MockCall(grpc.Call):
+    def __init__(self, code, should_fail=True):
+        self._code = code
+        self._should_fail = should_fail
+        self._count = 0
+
+    def code(self):
+        return self._code
+
+    def is_active(self):
+        return True
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self._count == 0 and self._should_fail:
+            self._count += 1
+            err = grpc.RpcError()
+            err.code = lambda: self._code
+            raise err
+        self._count += 1
+        return "success"
+
+
+class TestMTLSRefreshingChannel:
+    @mock.patch(
+        "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response"
+    )
+    @mock.patch("google.auth.transport.grpc.secure_authorized_channel")
+    def test_refresh_logic(self, mock_secure_channel, mock_check_params):
+        # mock fingerprint differences indicating rotation is needed
+        mock_check_params.return_value = (None, None, b"old", b"new")
+        mock_secure_channel.return_value = mock.Mock(spec=grpc.Channel)
+
+        initial_channel = mock.Mock(spec=grpc.Channel)
+        wrapper = _MTLSRefreshingChannel(
+            target="target",
+            factory_args={},
+            initial_channel=initial_channel,
+            initial_cert=b"old_cert",
+        )
+
+        # Subscribing adds to the initial channel
+        mock_callback = mock.Mock()
+        wrapper.subscribe(mock_callback)
+        initial_channel.subscribe.assert_called_with(
+            mock_callback, try_to_connect=False
+        )
+
+        wrapper.refresh_logic(1)
+
+        initial_channel.unsubscribe.assert_called_with(mock_callback)
+        mock_secure_channel.return_value.subscribe.assert_called_with(mock_callback)

packages/google-auth/tests_async/transport/test_aiohttp_requests.py

@@ -128,12 +128,13 @@ def test_mock_session_unspecified_auto_decompress(self):
         request = aiohttp_requests.Request(http)
         assert request.session == http
 
-    def test_timeout(self):
+    @pytest.mark.asyncio
+    async def test_timeout(self):
         http = mock.create_autospec(
             aiohttp.ClientSession, instance=True, auto_decompress=False
         )
         request = aiohttp_requests.Request(http)
-        request(url="http://example.com", method="GET", timeout=5)
+        await request(url="http://example.com", method="GET", timeout=5)
 
     @pytest.mark.asyncio
     async def test__clone(self):