refactor(auth): simplify fallback logic using custom exception and clean standard with block

Author: nbayati

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

@@ -70,6 +70,12 @@
 )
 
 
+class _MemfdCreationError(OSError):
+    """Raised when Linux in-memory virtual file creation (memfd) fails."""
+
+    pass
+
+
 @contextlib.contextmanager
 def secure_cert_key_paths(
     cert: Union[bytes, str, None],
@@ -114,13 +120,8 @@ def secure_cert_key_paths(
     # the bytes to anonymous in-memory files using memfd_create. This yields
     # /proc/self/fd/... paths, keeping the private key entirely in memory.
     if sys.platform == "linux" and hasattr(os, "memfd_create"):
-        cm = _memfd_cert_key_paths(cert_bytes, key_bytes)
         try:
-            cert_path, key_path = cm.__enter__()
-        except OSError:
-            pass  # Fallback to Tier 3 on failure.
-        else:
-            try:
+            with _memfd_cert_key_paths(cert_bytes, key_bytes) as (cert_path, key_path):
                 # Handle cases where path exists but might be restricted.
                 if (cert_path is None or os.path.exists(cert_path)) and (
                     key_path is None or os.path.exists(key_path)
@@ -129,9 +130,8 @@ def secure_cert_key_paths(
                         str, key_path or key
                     ), passphrase
                     return
-            finally:
-                cm.__exit__(*sys.exc_info())
-            # If verification failed, fall through to Tier 3.
+        except _MemfdCreationError:
+            pass  # Fallback to Tier 3 on failure.
 
     # Tier 3: Fallback Encrypted Temp Files. If in-memory files are not supported
     # (macOS/Windows), we write to disk. To protect the key, we encrypt plaintext
@@ -213,20 +213,24 @@ def _memfd_cert_key_paths(
             the active descriptors (e.g., '/proc/self/fd/3').
     """
     cleanup_fds = []
-    print("--- in memfd_cert_key_paths")
     paths = []
 
     try:
-        for data, name in [(cert_bytes, "mtls_cert"), (key_bytes, "mtls_key")]:
-            if data is not None:
-                # MFD_CLOEXEC prevents FD leaks to spawned subprocesses.
-                fd = os.memfd_create(name, os.MFD_CLOEXEC)  # type: ignore[attr-defined]
-                cleanup_fds.append(fd)
-                with os.fdopen(fd, "wb", closefd=False) as f:
-                    f.write(data)
-                paths.append(f"/proc/self/fd/{fd}")
-            else:
-                paths.append(None)
+        try:
+            for data, name in [(cert_bytes, "mtls_cert"), (key_bytes, "mtls_key")]:
+                if data is not None:
+                    # MFD_CLOEXEC prevents FD leaks to spawned subprocesses.
+                    fd = os.memfd_create(name, os.MFD_CLOEXEC)  # type: ignore[attr-defined]
+                    cleanup_fds.append(fd)
+                    with os.fdopen(fd, "wb", closefd=False) as f:
+                        f.write(data)
+                    paths.append(f"/proc/self/fd/{fd}")
+                else:
+                    paths.append(None)
+        except OSError as exc:
+            raise _MemfdCreationError(
+                "Failed to create in-memory virtual files"
+            ) from exc
 
         cert_path, key_path = paths
         yield cert_path, key_path
@@ -258,7 +262,6 @@ def _tempfile_cert_key_paths(
         else None
     )
     cleanup_files = []
-    print("--- in _tempfile_cert_key_paths")
     new_passphrase = passphrase
     cert_data = cert_bytes
     key_data = None

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

@@ -1069,7 +1069,9 @@ def test_tier2_fallback_to_tier3_on_oserror(
         self, mock_tempfile_cm, mock_memfd_cm, mock_memfd_create
     ):
         mock_memfd_ctx = mock.MagicMock()
-        mock_memfd_ctx.__enter__.side_effect = OSError("memfd failed")
+        mock_memfd_ctx.__enter__.side_effect = _mtls_helper._MemfdCreationError(
+            "memfd failed"
+        )
         mock_memfd_cm.return_value = mock_memfd_ctx
 
         mock_tempfile_ctx = mock.MagicMock()
@@ -1205,7 +1207,7 @@ def test_success_shm(
 
         with mock.patch.object(os, "remove") as mock_remove, mock.patch.object(
             os.path, "exists", return_value=True
-        ):
+        ), mock.patch.object(os, "access", return_value=True):
             with _mtls_helper._tempfile_cert_key_paths(b"cert", b"key", b"pass") as (
                 cert_path,
                 key_path,
@@ -1244,7 +1246,7 @@ def test_permission_error_loop_resilience(
 
         with mock.patch.object(os, "remove") as mock_remove, mock.patch.object(
             os.path, "exists", return_value=True
-        ):
+        ), mock.patch.object(os, "access", return_value=True):
             with _mtls_helper._tempfile_cert_key_paths(b"cert", b"key", b"pass"):
                 pass
             mock_remove.assert_called_once_with("/shm/cert")