Allow Azure Arc managed identity selectors

Author: Copilot

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,19 @@ def _obtain_token_on_service_fabric(
 class ArcPlatformNotSupportedError(ManagedIdentityError):
     pass
 
-def _obtain_token_on_arc(http_client, endpoint, resource):
+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: "mi_res_id",
+            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,7 +676,7 @@ 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:

tests/test_mi.py

@@ -477,6 +477,50 @@ 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,
+                text='{"access_token": "AT", "expires_in": "1234", "resource": "R"}',
+            ),
+        ]) 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_mi_res_id(self, mocked_stat):
+        self._assert_user_assigned_selector(
+            UserAssignedManagedIdentity(resource_id="resource-id"),
+            "mi_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",
+        )
+
 
 class GetManagedIdentitySourceTestCase(unittest.TestCase):
 
@@ -531,4 +575,3 @@ def test_cloud_shell(self):
 
     def test_default_to_vm(self):
         self.assertEqual(get_managed_identity_source(), DEFAULT_TO_VM)
-