Fix exception-contract and cache-key regressions vs #882

Blocker 1 — Exception contract:
RuntimeError from msal-key-attestation (DLL load, attestation call) now
gets caught and wrapped as MsiV2Error at both the provider call site in
msi_v2.py and the outer boundary in managed_identity.py.  Only
MsiV2Error (or its subclasses) can escape to the caller.

Blocker 2 — Stable attestation cache key:
The provider callback signature is expanded from (endpoint, key_handle,
client_id) to (endpoint, key_handle, client_id, cache_key).  MSAL now
passes the stable per-boot key name as cache_key, which
get_attestation_jwt() uses for its MAA token cache instead of falling
back to the less cache-friendly numeric handle.

Tests: 59 passed (44 core + 15 attestation), including new tests for
RuntimeError wrapping and cache_key forwarding.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: 2 people

msal-key-attestation/msal_key_attestation/attestation.py

@@ -338,17 +338,21 @@ def get_attestation_jwt(
 
 # ---------------------------------------------------------------------------
 # Public factory — matches the callback signature MSAL expects:
-#   (endpoint: str, key_handle: int, client_id: str) -> str
+#   (endpoint: str, key_handle: int, client_id: str, cache_key: str) -> str
 # ---------------------------------------------------------------------------
 
-def create_attestation_provider() -> Callable[[str, int, str], str]:
+def create_attestation_provider() -> Callable[[str, int, str, str], str]:
     """
     Create an attestation token provider callable for MSAL MSI v2.
 
     The returned callable has signature::
 
         provider(attestation_endpoint: str, key_handle: int,
-                 client_id: str) -> str
+                 client_id: str, cache_key: str) -> str
+
+    ``cache_key`` should be the stable per-boot key name.  Using the key
+    name (rather than the numeric handle) maximizes MAA-token cache hits
+    across key re-opens.
 
     It wraps :func:`get_attestation_jwt` with caching support.
 
@@ -372,10 +376,12 @@ def _provider(
         attestation_endpoint: str,
         key_handle: int,
         client_id: str,
+        cache_key: str = "",
     ) -> str:
         return get_attestation_jwt(
             attestation_endpoint=attestation_endpoint,
             client_id=client_id,
             key_handle=key_handle,
+            cache_key=cache_key or None,
         )
     return _provider

msal-key-attestation/tests/test_attestation.py

@@ -109,12 +109,26 @@ def test_returns_callable(self):
     def test_provider_calls_get_attestation_jwt(self, mock_get):
         mock_get.return_value = "fake.attestation.jwt"
         provider = create_attestation_provider()
-        result = provider("https://attest.example.com", 12345, "client-id")
+        result = provider(
+            "https://attest.example.com", 12345, "client-id", "my-key-name")
         self.assertEqual(result, "fake.attestation.jwt")
         mock_get.assert_called_once_with(
             attestation_endpoint="https://attest.example.com",
             client_id="client-id",
             key_handle=12345,
+            cache_key="my-key-name",
+        )
+
+    @patch("msal_key_attestation.attestation.get_attestation_jwt")
+    def test_provider_forwards_empty_cache_key_as_none(self, mock_get):
+        mock_get.return_value = "fake.jwt"
+        provider = create_attestation_provider()
+        provider("https://ep", 1, "cid", "")
+        mock_get.assert_called_once_with(
+            attestation_endpoint="https://ep",
+            client_id="cid",
+            key_handle=1,
+            cache_key=None,
         )
 
 

msal/managed_identity.py

@@ -343,11 +343,17 @@ def acquire_token_for_client(
                     "Install it with: pip install msal-key-attestation")
 
             from .msi_v2 import obtain_token as _obtain_token_v2
-            result = _obtain_token_v2(
-                self._http_client, self._managed_identity, resource,
-                attestation_enabled=True,
-                attestation_token_provider=attestation_token_provider,
-            )
+            try:
+                result = _obtain_token_v2(
+                    self._http_client, self._managed_identity, resource,
+                    attestation_enabled=True,
+                    attestation_token_provider=attestation_token_provider,
+                )
+            except MsiV2Error:
+                raise
+            except Exception as exc:
+                raise MsiV2Error(
+                    f"[msi_v2] Unexpected failure: {exc}") from exc
             if "access_token" in result and "error" not in result:
                 result[self._TOKEN_SOURCE] = self._TOKEN_SOURCE_IDP
             return result

msal/msi_v2.py

@@ -1232,9 +1232,10 @@ def _acquire_token_mtls_schannel(
 # Public API
 # ---------------------------------------------------------------------------
 
-# Type alias for attestation provider callback matching .NET delegate:
-#   (endpoint: str, key_handle: int, client_id: str) -> str (JWT)
-AttestationTokenProvider = Callable[[str, int, str], str]
+# Type alias for attestation provider callback.
+# Signature: (endpoint, key_handle, client_id, cache_key) -> JWT string.
+# cache_key is the stable per-boot key name for optimal caching.
+AttestationTokenProvider = Callable[[str, int, str, str], str]
 
 
 def obtain_token(
@@ -1263,8 +1264,10 @@ def obtain_token(
         resource: Resource URI for token acquisition
         attestation_enabled: Whether attestation is enabled
         attestation_token_provider: Callback (endpoint, key_handle,
-            client_id) -> JWT string.  Provided by msal-key-attestation
-            package.  None means non-attested flow.
+            client_id, cache_key) -> JWT string.  Provided by
+            msal-key-attestation package.  cache_key is the stable
+            per-boot key name for optimal caching.  None means
+            non-attested flow.
 
     Returns:
         Token response dict with access_token, expires_in, token_type,
@@ -1335,10 +1338,18 @@ def obtain_token(
                 if not attestation_endpoint:
                     raise MsiV2Error(
                         "[msi_v2] attestationEndpoint missing from metadata.")
-                att_jwt = attestation_token_provider(
-                    str(attestation_endpoint),
-                    int(key.value),
-                    str(client_id))
+                try:
+                    att_jwt = attestation_token_provider(
+                        str(attestation_endpoint),
+                        int(key.value),
+                        str(client_id),
+                        str(key_name))
+                except MsiV2Error:
+                    raise
+                except Exception as exc:
+                    raise MsiV2Error(
+                        f"[msi_v2] Attestation provider failed: {exc}"
+                    ) from exc
                 if not att_jwt or not str(att_jwt).strip():
                     raise MsiV2Error(
                         "[msi_v2] Attestation provider returned empty JWT.")

tests/test_msi_v2.py

@@ -403,7 +403,7 @@ def test_v2_called_when_both_flags_true(self, mock_v2, _):
         with patch.dict("sys.modules", {
             "msal_key_attestation": MagicMock(
                 create_attestation_provider=MagicMock(
-                    return_value=lambda ep, kh, ci: "fake.jwt"))
+                    return_value=lambda ep, kh, ci, ck="": "fake.jwt"))
         }):
             res = client.acquire_token_for_client(
                 resource="https://mtlstb.graph.microsoft.com",
@@ -425,7 +425,7 @@ def test_strict_v2_failure_raises_no_v1_fallback(
         with patch.dict("sys.modules", {
             "msal_key_attestation": MagicMock(
                 create_attestation_provider=MagicMock(
-                    return_value=lambda ep, kh, ci: "fake.jwt"))
+                    return_value=lambda ep, kh, ci, ck="": "fake.jwt"))
         }):
             with self.assertRaises(MsiV2Error):
                 client.acquire_token_for_client(
@@ -434,6 +434,26 @@ def test_strict_v2_failure_raises_no_v1_fallback(
                     with_attestation_support=True)
         mock_v1.assert_not_called()
 
+    @patch("msal.msi_v2.obtain_token",
+           side_effect=RuntimeError("DLL load failed"))
+    @patch("msal.managed_identity._obtain_token")
+    def test_runtime_error_wrapped_as_msi_v2_error(
+            self, mock_v1, mock_v2):
+        """RuntimeError from provider/DLL must surface as MsiV2Error."""
+        client = self._make_client()
+        with patch.dict("sys.modules", {
+            "msal_key_attestation": MagicMock(
+                create_attestation_provider=MagicMock(
+                    return_value=lambda ep, kh, ci, ck="": "fake.jwt"))
+        }):
+            with self.assertRaises(MsiV2Error) as ctx:
+                client.acquire_token_for_client(
+                    resource="R",
+                    mtls_proof_of_possession=True,
+                    with_attestation_support=True)
+            self.assertIn("DLL load failed", str(ctx.exception))
+        mock_v1.assert_not_called()
+
     def test_missing_attestation_package_raises_clear_error(self):
         client = self._make_client()
         with patch.dict("sys.modules", {"msal_key_attestation": None}):