chore: Add support for other error types in async/mtls and take care of race conditions

Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>

Author: agrawalradhika-cell

google/auth/aio/transport/mtls.py

@@ -70,20 +70,19 @@ def make_client_cert_ssl_context(
     Raises:
         google.auth.exceptions.TransportError: If there is an error loading the certificate.
     """
-    try:
-        context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
-
-        # Write cert and key to temp files because ssl.load_cert_chain requires paths
-        with _create_temp_file(cert_bytes) as cert_path:
-            with _create_temp_file(key_bytes) as key_path:
-                context.load_cert_chain(
-                    certfile=cert_path, keyfile=key_path, password=passphrase
-                )
-        return context
-    except (ssl.SSLError, OSError) as exc:
-        raise exceptions.TransportError(
-            "Failed to load client certificate and key for mTLS."
-        ) from exc
+    with _create_temp_file(cert_bytes) as cert_path, _create_temp_file(
+        key_bytes
+    ) as key_path:
+        try:
+            context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
+            context.load_cert_chain(
+                certfile=cert_path, keyfile=key_path, password=passphrase
+            )
+            return context
+        except (ssl.SSLError, OSError, IOError, ValueError, RuntimeError) as exc:
+            raise exceptions.TransportError(
+                "Failed to load client certificate and key for mTLS."
+            ) from exc
 
 
 async def _run_in_executor(func, *args):
@@ -104,7 +103,7 @@ def default_client_cert_source():
     """Get a callback which returns the default client SSL credentials.
 
     Returns:
-        Awaitable[Callable[[], [bytes, bytes]]]: A callback which returns the default
+        Awaitable[Callable[[], Tuple[bytes, bytes]]]: A callback which returns the default
             client certificate bytes and private key bytes, both in PEM format.
 
     Raises:

google/auth/aio/transport/sessions.py

@@ -143,6 +143,8 @@ async def configure_mtls_channel(self, client_cert_callback=None):
         successfully obtained (from the given client_cert_callback or from application
         default SSL credentials), the underlying transport will be reconfigured
         to use mTLS.
+        Note: This function does nothing if the `aiohttp` library is not
+        installed.
 
         Args:
             client_cert_callback (Optional[Callable[[], (bytes, bytes)]]):
@@ -180,8 +182,9 @@ async def configure_mtls_channel(self, client_cert_callback=None):
                 if AIOHTTP_INSTALLED and isinstance(self._auth_request, AiohttpRequest):
                     connector = aiohttp.TCPConnector(ssl=ssl_context)
                     new_session = aiohttp.ClientSession(connector=connector)
-                    await self._auth_request.close()
+                    old_auth_request = self._auth_request
                     self._auth_request = AiohttpRequest(session=new_session)
+                    await old_auth_request.close()
 
         except (
             exceptions.ClientCertError,

tests/transport/aio/test_sessions_mtls.py

@@ -46,22 +46,23 @@ async def test_configure_mtls_channel(self):
         ), mock.patch(
             "google.auth.aio.transport.mtls.get_client_cert_and_key"
         ) as mock_helper, mock.patch(
-            "ssl.create_default_context"
-        ) as mock_ssl:
+            "google.auth.aio.transport.mtls.make_client_cert_ssl_context"
+        ) as mock_make_context:
             mock_exists.return_value = True
             mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data")
 
             mock_context = mock.Mock(spec=ssl.SSLContext)
-            mock_ssl.return_value = mock_context
+            mock_make_context.return_value = mock_context
 
-            # Use AsyncMock for credentials to avoid "coroutine never awaited" warnings
             mock_creds = mock.AsyncMock(spec=credentials.Credentials)
             session = sessions.AsyncAuthorizedSession(mock_creds)
 
             await session.configure_mtls_channel()
 
             assert session._is_mtls is True
-            assert mock_context.load_cert_chain.called
+            mock_make_context.assert_called_once_with(
+                b"fake_cert_data", b"fake_key_data"
+            )
 
     @pytest.mark.asyncio
     async def test_configure_mtls_channel_disabled(self):

tests/transport/test_aio_mtls_helper.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import ssl
 from unittest import mock
 
 import pytest
@@ -24,6 +26,70 @@
 
 
 class TestMTLS:
+    @pytest.mark.asyncio
+    async def test__create_temp_file(self):
+        """Tests that _create_temp_file creates a file with correct content and deletes it."""
+        content = b"test cert data"
+
+        # Test file creation and content
+        with mtls._create_temp_file(content) as file_path:
+            assert os.path.exists(file_path)
+            # Verify file is not readable by others (mkstemp default)
+            if os.name == "posix":
+                assert (os.stat(file_path).st_mode & 0o777) == 0o600
+
+            with open(file_path, "rb") as f:
+                assert f.read() == content
+
+        # Test file deletion after context exit
+        assert not os.path.exists(file_path)
+
+    @pytest.mark.asyncio
+    async def test_make_client_cert_ssl_context_success(self):
+        """Tests successful creation of an SSLContext with client certificates."""
+        cert_bytes = b"cert_data"
+        key_bytes = b"key_data"
+        passphrase = b"password"
+
+        mock_context = mock.Mock(spec=ssl.SSLContext)
+
+        with mock.patch(
+            "ssl.create_default_context", return_value=mock_context
+        ) as mock_create:
+            context = mtls.make_client_cert_ssl_context(
+                cert_bytes, key_bytes, passphrase=passphrase
+            )
+
+            assert context == mock_context
+            mock_create.assert_called_once_with(ssl.Purpose.SERVER_AUTH)
+
+            # Verify load_cert_chain was called
+            assert mock_context.load_cert_chain.called
+            kwargs = mock_context.load_cert_chain.call_args.kwargs
+            assert "certfile" in kwargs
+            assert "keyfile" in kwargs
+            assert kwargs["password"] == passphrase
+
+            assert not os.path.exists(kwargs["certfile"])
+            assert not os.path.exists(kwargs["keyfile"])
+
+    @pytest.mark.asyncio
+    async def test_make_client_cert_ssl_context_error(self):
+        """Verifies that TransportError is raised when SSL loading fails."""
+        cert_bytes = b"cert_data"
+        key_bytes = b"key_data"
+
+        mock_context = mock.Mock(spec=ssl.SSLContext)
+        # Mocking an SSLError to trigger the exception handler in make_client_cert_ssl_context
+        mock_context.load_cert_chain.side_effect = ssl.SSLError("Mock SSL Error")
+
+        with mock.patch("ssl.create_default_context", return_value=mock_context):
+            with pytest.raises(exceptions.TransportError) as exc_info:
+                mtls.make_client_cert_ssl_context(cert_bytes, key_bytes)
+
+            assert "Failed to load client certificate" in str(exc_info.value)
+            assert isinstance(exc_info.value.__cause__, ssl.SSLError)
+
     @pytest.mark.asyncio
     @mock.patch(
         "google.auth.transport.mtls.has_default_client_cert_source", return_value=False