Allow Azure Arc managed identity requests to forward user-assigned identity selectors

msal/managed_identity.py

@@ -452,13 +452,8 @@ def _obtain_token(
         )
     arc_endpoint = _get_arc_endpoint()
     if arc_endpoint:
-        if ManagedIdentity.is_user_assigned(managed_identity):
-            raise ManagedIdentityError(  # Note: Azure Identity for Python raised exception too
-                "Invalid managed_identity parameter. "
-                "Azure Arc supports only system-assigned managed identity, "
-                "See also "
-                "https://learn.microsoft.com/en-us/azure/service-fabric/configure-existing-cluster-enable-managed-identity-token-service")
-        return _obtain_token_on_arc(http_client, arc_endpoint, resource)
+        return _obtain_token_on_arc(
+            http_client, arc_endpoint, resource, managed_identity)
     return _obtain_token_on_azure_vm(http_client, managed_identity, resource)
 
 
@@ -643,12 +638,45 @@ def _obtain_token_on_service_fabric(
 class ArcPlatformNotSupportedError(ManagedIdentityError):
     pass
 
-def _obtain_token_on_arc(http_client, endpoint, resource):
+def _raise_if_arc_did_not_honor_user_assigned_identity(managed_identity, payload):
+    """Fail closed when a user-assigned identity was requested but Azure Arc did not confirm it.
+
+    A legacy Azure Arc agent ignores the client_id / object_id / msi_res_id selector and silently
+    returns the machine's system-assigned identity. An agent that supports user-assigned managed
+    identity echoes the identity it used in the token response. When that echo is missing or does
+    not match the requested selector, MSAL must not hand back a token for a different identity than
+    the one that was requested.
+    """
+    if not ManagedIdentity.is_user_assigned(managed_identity):
+        return  # System-assigned: there is no requested identity to confirm
+    requested = managed_identity.get(ManagedIdentity.ID)
+    echoed = {
+        ManagedIdentity.CLIENT_ID: payload.get("client_id"),
+        ManagedIdentity.OBJECT_ID: payload.get("object_id"),
+        # Azure Arc echoes msi_res_id; accept the mi_res_id spelling too as a safety net
+        ManagedIdentity.RESOURCE_ID: payload.get("msi_res_id") or payload.get("mi_res_id"),
+    }.get(managed_identity.get(ManagedIdentity.ID_TYPE))
+    # Compare case-insensitively: client_id / object_id are GUIDs, and an ARM resource id
+    # (msi_res_id) can legitimately differ in segment casing.
+    if not echoed or str(echoed).lower() != str(requested).lower():
+        raise ManagedIdentityError(
+            "Azure Arc did not confirm the requested user-assigned managed identity "
+            "in the token response. The agent likely does not support user-assigned "
+            "managed identities and returned the system-assigned identity.")
+
+def _obtain_token_on_arc(http_client, endpoint, resource, managed_identity=None):
     # https://learn.microsoft.com/en-us/azure/azure-arc/servers/managed-identity-authentication
     logger.debug("Obtaining token via managed identity on Azure Arc")
+    params = {"api-version": "2020-06-01", "resource": resource}
+    if managed_identity:
+        _adjust_param(params, managed_identity, types_mapping={
+            ManagedIdentity.CLIENT_ID: "client_id",
+            ManagedIdentity.RESOURCE_ID: "msi_res_id",  # Azure Arc honors the IMDS msi_res_id spelling; mi_res_id is ignored and returns the system-assigned identity
+            ManagedIdentity.OBJECT_ID: "object_id",
+        })
     resp = http_client.get(
         endpoint,
-        params={"api-version": "2020-06-01", "resource": resource},
+        params=params.copy(),
         headers={"Metadata": "true"},
         )
     www_auth = "www-authenticate"  # Header in lower case
@@ -674,13 +702,15 @@ def _obtain_token_on_arc(http_client, endpoint, resource):
         secret = f.read()
     response = http_client.get(
         endpoint,
-        params={"api-version": "2020-06-01", "resource": resource},
+        params=params.copy(),
         headers={"Metadata": "true", "Authorization": "Basic {}".format(secret)},
         )
     try:
         payload = json.loads(response.text)
         if payload.get("access_token") and payload.get("expires_in"):
             # Example: https://learn.microsoft.com/en-us/azure/azure-arc/servers/media/managed-identity-authentication/bash-token-output-example.png
+            _raise_if_arc_did_not_honor_user_assigned_identity(
+                managed_identity, payload)
             return {
                 "access_token": payload["access_token"],
                 "expires_in": int(payload["expires_in"]),

msal/sku.py

@@ -2,5 +2,5 @@
 """
 
 # The __init__.py will import this. Not the other way around.
-__version__ = "1.37.0"
+__version__ = "1.38.0"
 SKU = "MSAL.Python"

tests/test_mi.py

@@ -477,6 +477,72 @@ def test_arc_error_should_be_normalized(self, mocked_stat):
                 if sys.platform in _supported_arc_platforms_and_their_prefixes:
                     self.fail("Should not raise ArcPlatformNotSupportedError")
 
+    def _assert_user_assigned_selector(self, managed_identity, selector_name, selector_value):
+        app = ManagedIdentityClient(managed_identity, http_client=requests.Session())
+        with patch.object(app._http_client, "get", side_effect=[
+            self.challenge,
+            MinimalResponse(
+                status_code=200,
+                # A compliant Azure Arc agent echoes the identity it used; MSAL verifies it (fail closed).
+                text='{"access_token": "AT", "expires_in": "1234", "resource": "R", "%s": "%s"}' % (
+                    selector_name, selector_value),
+            ),
+        ]) as mocked_method:
+            try:
+                result = app.acquire_token_for_client(resource="R")
+                self.assertEqual("AT", result["access_token"])
+                expected_params = {
+                    "api-version": "2020-06-01",
+                    "resource": "R",
+                    selector_name: selector_value,
+                }
+                self.assertEqual(expected_params, mocked_method.call_args_list[0].kwargs["params"])
+                self.assertEqual(expected_params, mocked_method.call_args_list[1].kwargs["params"])
+            except ArcPlatformNotSupportedError:
+                if sys.platform in _supported_arc_platforms_and_their_prefixes:
+                    self.fail("Should not raise ArcPlatformNotSupportedError")
+
+    def test_arc_user_assigned_client_id_should_be_forwarded(self, mocked_stat):
+        self._assert_user_assigned_selector(
+            UserAssignedManagedIdentity(client_id="client-id"),
+            "client_id",
+            "client-id",
+        )
+
+    def test_arc_user_assigned_resource_id_should_be_forwarded_as_msi_res_id(self, mocked_stat):
+        self._assert_user_assigned_selector(
+            UserAssignedManagedIdentity(resource_id="resource-id"),
+            "msi_res_id",
+            "resource-id",
+        )
+
+    def test_arc_user_assigned_object_id_should_be_forwarded(self, mocked_stat):
+        self._assert_user_assigned_selector(
+            UserAssignedManagedIdentity(object_id="object-id"),
+            "object_id",
+            "object-id",
+        )
+
+    def test_arc_user_assigned_identity_not_confirmed_should_fail_closed(self, mocked_stat):
+        # A legacy Azure Arc agent ignores the selector and returns the system-assigned identity:
+        # the token response echoes a different identity than the one requested. MSAL must fail
+        # closed rather than hand back a token for a different identity than requested.
+        app = ManagedIdentityClient(
+            UserAssignedManagedIdentity(client_id="client-id"),
+            http_client=requests.Session())
+        with patch.object(app._http_client, "get", side_effect=[
+            self.challenge,
+            MinimalResponse(
+                status_code=200,
+                text='{"access_token": "AT", "expires_in": "1234", "resource": "R", "client_id": "a-different-id"}',
+            ),
+        ]):
+            with self.assertRaises(ManagedIdentityError):
+                app.acquire_token_for_client(resource="R")
+        self.assertEqual(
+            {}, app._token_cache._cache,
+            "An unconfirmed identity token must not be cached")
+
 
 class GetManagedIdentitySourceTestCase(unittest.TestCase):
 
@@ -531,4 +597,3 @@ def test_cloud_shell(self):
 
     def test_default_to_vm(self):
         self.assertEqual(get_managed_identity_source(), DEFAULT_TO_VM)
-