[DO NOT MERGE] Allow Azure Arc managed identity requests to forward user-assigned identity selectors

.Pipelines/pipeline-publish.yml

@@ -36,6 +36,7 @@ pr: none
 variables:
   CDP_DEFINITION_BUILD_COUNT: $[counter('', 0)]
   system.debug: ${{ parameters.debug }}
+  WindowsContainerImage: 'onebranch.azurecr.io/windows/ltsc2022/vse2022:latest' # https://aka.ms/obpipelines/containers
 
 resources:
   repositories:
@@ -64,6 +65,7 @@ extends:
         tsaEnabled: true
     featureFlags:
       EnableCDPxPAT: false
+      WindowsHostVersion: '1ESWindows2022'
 
     stages:
 
@@ -82,34 +84,26 @@ extends:
       condition: eq(dependencies.E2ETests.result, 'Succeeded')
       jobs:
       - job: BuildDist
-        displayName: 'Build sdist + wheel (Python 3.12)'
+        displayName: 'Build sdist + wheel'
         pool:
-          type: linux
-          isCustom: true
-          vmImage: 'ubuntu-22.04'
+          type: windows
         variables:
           ob_outputDirectory: '$(Build.ArtifactStagingDirectory)'
         steps:
-        - task: UsePythonVersion@0
-          inputs:
-            versionSpec: '3.12'
-          displayName: 'Use Python 3.12'
+        - script: python --version
+          displayName: 'Check Python version'
 
-        - script: |
-            python -m pip install --upgrade pip build twine
-          displayName: 'Install build toolchain'
-
-        - script: |
-            python -m build
+        - script: python setup.py sdist bdist_wheel
           displayName: 'Build sdist and wheel'
 
-        - script: |
-            python -m twine check dist/*
-          displayName: 'Verify distribution (twine check)'
-
-        - script: |
-            cp -r dist/* $(ob_outputDirectory)/
+        - task: CopyFiles@2
           displayName: 'Copy dist to output directory'
+          inputs:
+            SourceFolder: '$(Build.SourcesDirectory)\dist'
+            Contents: |
+              *.whl
+              *.tar.gz
+            TargetFolder: '$(ob_outputDirectory)'
 
     # ========================================================================
     # Stage 5a - Publish to test.pypi.org (Preview / RC)
@@ -135,8 +129,9 @@ extends:
         - task: DownloadPipelineArtifact@2
           displayName: 'Download build artifact'
           inputs:
-            artifactName: drop_Build_BuildDist
-            targetPath: $(Pipeline.Workspace)/python-dist
+            buildType: 'current'
+            artifactName: 'drop_Build_BuildDist'
+            targetPath: $(Pipeline.Workspace)\python-dist
 
         - task: UsePythonVersion@0
           inputs:
@@ -211,8 +206,9 @@ extends:
         - task: DownloadPipelineArtifact@2
           displayName: 'Download python-dist artifact'
           inputs:
-            artifactName: drop_Build_BuildDist
-            targetPath: $(Pipeline.Workspace)/python-dist
+            buildType: 'current'
+            artifactName: 'drop_Build_BuildDist'
+            targetPath: $(Pipeline.Workspace)\python-dist
 
         - task: EsrpRelease@12
           displayName: 'Publish to PyPI via ESRP'
@@ -225,9 +221,9 @@ extends:
             intent: 'PackageDistribution'
             contenttype: 'PyPi'
             contentsource: 'Folder'
-            folderlocation: '$(Pipeline.Workspace)/python-dist'
+            folderlocation: '$(Pipeline.Workspace)\python-dist'
             waitforreleasecompletion: true
-            owners: 'ryauld@microsoft.com,avdunn@microsoft.com'
+            owners: 'avdunn@microsoft.com,bogavril@microsoft.com'
             approvers: 'avdunn@microsoft.com,bogavril@microsoft.com'
             serviceendpointurl: 'https://api.esrp.microsoft.com'
             mainpublisher: 'ESRPRELPACMAN'

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:

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.0rc2"
 SKU = "MSAL.Python"

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)
-