Added support for DEFAULT_IDENTITY_CLIENT_ID environment variable in Machine Learning Managed Identity

src/client/Microsoft.Identity.Client/ManagedIdentity/EnvironmentVariables.cs

@@ -14,5 +14,6 @@ internal class EnvironmentVariables
         public static string MsiEndpoint => Environment.GetEnvironmentVariable("MSI_ENDPOINT");
         public static string MsiSecret => Environment.GetEnvironmentVariable("MSI_SECRET");
         public static string IdentityServerThumbprint => Environment.GetEnvironmentVariable("IDENTITY_SERVER_THUMBPRINT");
+        public static string MachineLearningDefaultClientId => Environment.GetEnvironmentVariable("DEFAULT_IDENTITY_CLIENT_ID");
     }
 }

src/client/Microsoft.Identity.Client/ManagedIdentity/MachineLearningManagedIdentitySource.cs

@@ -10,12 +10,16 @@ namespace Microsoft.Identity.Client.ManagedIdentity
 {
     internal class MachineLearningManagedIdentitySource : AbstractManagedIdentity
     {
+        private const string MachineLearning = "Machine Learning";
+
         private const string MachineLearningMsiApiVersion = "2017-09-01";
         private const string SecretHeaderName = "secret";
 
         private readonly Uri _endpoint;
         private readonly string _secret;
 
+        public const string UnsupportedIdTypeError = "Only client id is supported for user-assigned managed identity in Machine Learning."; // referenced in unit test
+
         public static AbstractManagedIdentity Create(RequestContext requestContext)
         {
             requestContext.Logger.Info(() => "[Managed Identity] Machine learning managed identity is available.");
@@ -73,21 +77,45 @@ protected override ManagedIdentityRequest CreateRequest(string resource)
 
             switch (_requestContext.ServiceBundle.Config.ManagedIdentityId.IdType)
             {
+                case AppConfig.ManagedIdentityIdType.SystemAssigned:
+                    _requestContext.Logger.Info("[Managed Identity] Adding system assigned client id to the request.");
+
+                    // this environment variable is always set in an Azure Machine Learning source, but check if null just in case
+                    if (EnvironmentVariables.MachineLearningDefaultClientId == null)
+                    {
+                        throw MsalServiceExceptionFactory.CreateManagedIdentityException(
+                            MsalError.InvalidManagedIdentityIdType,
+                            "The DEFAULT_IDENTITY_CLIENT_ID environment variable is null.",
+                            null, // configuration error
+                            ManagedIdentitySource.MachineLearning,
+                            null); // statusCode is null in this case
+                    }
+
+                    // Use the new 2017 constant for older ML-based environment
+                    request.QueryParameters[Constants.ManagedIdentityClientId2017] = EnvironmentVariables.MachineLearningDefaultClientId;
+                    break;
+
                 case AppConfig.ManagedIdentityIdType.ClientId:
                     _requestContext.Logger.Info("[Managed Identity] Adding user assigned client id to the request.");
                     // Use the new 2017 constant for older ML-based environment
                     request.QueryParameters[Constants.ManagedIdentityClientId2017] = _requestContext.ServiceBundle.Config.ManagedIdentityId.UserAssignedId;
                     break;
 
-                case AppConfig.ManagedIdentityIdType.ResourceId:
-                    _requestContext.Logger.Info("[Managed Identity] Adding user assigned resource id to the request.");
-                    request.QueryParameters[Constants.ManagedIdentityResourceId] = _requestContext.ServiceBundle.Config.ManagedIdentityId.UserAssignedId;
-                    break;
-
-                case AppConfig.ManagedIdentityIdType.ObjectId:
-                    _requestContext.Logger.Info("[Managed Identity] Adding user assigned object id to the request.");
-                    request.QueryParameters[Constants.ManagedIdentityObjectId] = _requestContext.ServiceBundle.Config.ManagedIdentityId.UserAssignedId;
-                    break;
+                default:
+                    string errorMessage = string.Format(
+                        CultureInfo.InvariantCulture,
+                        MsalErrorMessage.ManagedIdentityInvalidIdType,
+                        "client id is",
+                        MachineLearning);
+
+                    var exception = MsalServiceExceptionFactory.CreateManagedIdentityException(
+                        MsalError.InvalidManagedIdentityIdType,
+                        errorMessage,
+                        null, // configuration error
+                        ManagedIdentitySource.MachineLearning,
+                        null); // statusCode is null in this case
+
+                    throw exception;
             }
 
             return request;

src/client/Microsoft.Identity.Client/MsalError.cs

@@ -1105,6 +1105,16 @@ public static class MsalError
         /// </summary>
         public const string InvalidManagedIdentityResponse = "invalid_managed_identity_response";
 
+        /// <summary>
+        /// The managed identity's source does not select a specific id type.
+        /// </summary>
+        public const string InvalidManagedIdentityIdType = "invalid_managed_identity_id_type";
+
+        /// <summary>
+        /// The managed identity is missing a required environment variable.
+        /// </summary>
+        public const string MissingManagedIdentityEnvVar = "missing_managed_identity_env_var";
+
         /// <summary>
         /// Managed Identity error response was received.
         /// </summary>

src/client/Microsoft.Identity.Client/MsalErrorMessage.cs

@@ -415,6 +415,7 @@ public static string InvalidTokenProviderResponseValue(string invalidValueName)
 
         public const string ManagedIdentityNoResponseReceived = "[Managed Identity] Authentication unavailable. No response received from the managed identity endpoint.";
         public const string ManagedIdentityInvalidResponse = "[Managed Identity] Invalid response, the authentication response received did not contain the expected fields.";
+        public const string ManagedIdentityInvalidIdType = "Only {0} supported for user-assigned managed identity in {1}";
         public const string ManagedIdentityJsonParseFailure = "[Managed Identity] MSI returned 200 OK, but the response could not be parsed.";
         public const string ManagedIdentityUnexpectedResponse = "[Managed Identity] Unexpected exception occurred when parsing the response. See the inner exception for details.";
         public const string ManagedIdentityExactlyOneScopeExpected = "[Managed Identity] To acquire token for managed identity, exactly one scope must be passed.";

src/client/Microsoft.Identity.Client/PublicApi/net462/PublicAPI.Unshipped.txt

@@ -1 +1,2 @@
-
+const Microsoft.Identity.Client.MsalError.InvalidManagedIdentityIdType = "invalid_managed_identity_id_type" -> string
+const Microsoft.Identity.Client.MsalError.MissingManagedIdentityEnvVar = "missing_managed_identity_env_var" -> string

src/client/Microsoft.Identity.Client/PublicApi/net472/PublicAPI.Unshipped.txt

@@ -1 +1,2 @@
-
+const Microsoft.Identity.Client.MsalError.InvalidManagedIdentityIdType = "invalid_managed_identity_id_type" -> string
+const Microsoft.Identity.Client.MsalError.MissingManagedIdentityEnvVar = "missing_managed_identity_env_var" -> string

src/client/Microsoft.Identity.Client/PublicApi/net8.0-android/PublicAPI.Unshipped.txt

@@ -0,0 +1,2 @@
+const Microsoft.Identity.Client.MsalError.InvalidManagedIdentityIdType = "invalid_managed_identity_id_type" -> string
+const Microsoft.Identity.Client.MsalError.MissingManagedIdentityEnvVar = "missing_managed_identity_env_var" -> string

src/client/Microsoft.Identity.Client/PublicApi/net8.0-ios/PublicAPI.Unshipped.txt

@@ -0,0 +1,2 @@
+const Microsoft.Identity.Client.MsalError.InvalidManagedIdentityIdType = "invalid_managed_identity_id_type" -> string
+const Microsoft.Identity.Client.MsalError.MissingManagedIdentityEnvVar = "missing_managed_identity_env_var" -> string

src/client/Microsoft.Identity.Client/PublicApi/net8.0/PublicAPI.Unshipped.txt

@@ -1 +1,2 @@
-
+const Microsoft.Identity.Client.MsalError.InvalidManagedIdentityIdType = "invalid_managed_identity_id_type" -> string
+const Microsoft.Identity.Client.MsalError.MissingManagedIdentityEnvVar = "missing_managed_identity_env_var" -> string

src/client/Microsoft.Identity.Client/PublicApi/netstandard2.0/PublicAPI.Unshipped.txt

@@ -0,0 +1,2 @@
+const Microsoft.Identity.Client.MsalError.InvalidManagedIdentityIdType = "invalid_managed_identity_id_type" -> string
+const Microsoft.Identity.Client.MsalError.MissingManagedIdentityEnvVar = "missing_managed_identity_env_var" -> string

tests/Microsoft.Identity.Test.Common/Core/Helpers/ManagedIdentityTestUtil.cs

@@ -62,6 +62,7 @@ public static void SetEnvironmentVariables(ManagedIdentitySource managedIdentity
                 case ManagedIdentitySource.MachineLearning:
                     Environment.SetEnvironmentVariable("MSI_ENDPOINT", endpoint);
                     Environment.SetEnvironmentVariable("MSI_SECRET", secret);
+                    Environment.SetEnvironmentVariable("DEFAULT_IDENTITY_CLIENT_ID", "fake_DEFAULT_IDENTITY_CLIENT_ID");
                     break;
             }
         }

tests/Microsoft.Identity.Test.Common/Core/Mocks/MockHttpManagerExtensions.cs

@@ -359,7 +359,7 @@ public static void AddRegionDiscoveryMockHandler(
                     });
         }
 
-        public static void AddManagedIdentityMockHandler(
+        public static MockHttpMessageHandler AddManagedIdentityMockHandler(
             this MockHttpManager httpManager,
             string expectedUrl,
             string resource,
@@ -383,18 +383,17 @@ public static void AddManagedIdentityMockHandler(
 
             MockHttpMessageHandler httpMessageHandler = BuildMockHandlerForManagedIdentitySource(managedIdentitySourceType, resource);
 
-            if (userAssignedIdentityId == UserAssignedIdentityId.ClientId)
+            if (managedIdentitySourceType == ManagedIdentitySource.MachineLearning)
             {
-                if (managedIdentitySourceType == ManagedIdentitySource.MachineLearning)
-                {
-                    // For Machine Learning (App Service 2017), the param is "clientid"
-                    httpMessageHandler.ExpectedQueryParams.Add(Constants.ManagedIdentityClientId2017, userAssignedId);
-                }
-                else
-                {
-                    // For App Service 2019, Azure Arc, IMDS, etc., the param is "client_id"
-                    httpMessageHandler.ExpectedQueryParams.Add(Constants.ManagedIdentityClientId, userAssignedId);
-                }
+                // For Machine Learning (App Service 2017), the client id param is "clientid"
+                // it will always be a query parameter, no matter the source type
+                // use env var for SAMI, passed-in userAssignedId for UAMI
+                httpMessageHandler.ExpectedQueryParams.Add(Constants.ManagedIdentityClientId2017,
+                    userAssignedId == null ?
+                        EnvironmentVariables.MachineLearningDefaultClientId : userAssignedId);
+            } else if (userAssignedIdentityId == UserAssignedIdentityId.ClientId) {
+                // For App Service 2019, Azure Arc, IMDS, etc., the param is "client_id"
+                httpMessageHandler.ExpectedQueryParams.Add(Constants.ManagedIdentityClientId, userAssignedId);
             }
 
             if (userAssignedIdentityId == UserAssignedIdentityId.ResourceId)
@@ -414,6 +413,8 @@ public static void AddManagedIdentityMockHandler(
             httpMessageHandler.ExpectedUrl = expectedUrl;
 
             httpManager.AddMockHandler(httpMessageHandler);
+
+            return httpMessageHandler;
         }
 
         private static MockHttpMessageHandler BuildMockHandlerForManagedIdentitySource(ManagedIdentitySource managedIdentitySourceType, string resource)

tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/MachineLearningTests.cs

@@ -3,12 +3,11 @@
 
 using System;
 using System.Globalization;
-using System.Net;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client;
 using Microsoft.Identity.Client.AppConfig;
+using Microsoft.Identity.Client.Internal;
 using Microsoft.Identity.Client.ManagedIdentity;
-using Microsoft.Identity.Test.Common;
 using Microsoft.Identity.Test.Common.Core.Helpers;
 using Microsoft.Identity.Test.Common.Core.Mocks;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
@@ -20,6 +19,82 @@ namespace Microsoft.Identity.Test.Unit.ManagedIdentityTests
     public class MachineLearningTests : TestBase
     {
         private const string MachineLearning = "Machine learning";
+        private const string MachineLearningEndpoint = "http://localhost:7071/msi/token";
+        internal const string Resource = "https://management.azure.com";
+
+        [DataTestMethod]
+        [DataRow(null, null)]                                              // SAMI
+        [DataRow(TestConstants.ClientId, UserAssignedIdentityId.ClientId)] // UAMI
+        public async Task MachineLearningUserAssignedHappyPathAndHasCorrectClientIdQueryParameterAsync(
+            string userAssignedId,
+            UserAssignedIdentityId userAssignedIdentityId)
+        {
+            using (new EnvVariableContext())
+            using (var httpManager = new MockHttpManager())
+            {
+                SetEnvironmentVariables(ManagedIdentitySource.MachineLearning, MachineLearningEndpoint);
+
+                ManagedIdentityId managedIdentityId = userAssignedId == null
+                    ? ManagedIdentityId.SystemAssigned
+                    : ManagedIdentityId.WithUserAssignedClientId(userAssignedId);
+                ManagedIdentityApplicationBuilder miBuilder = ManagedIdentityApplicationBuilder.Create(managedIdentityId)
+                    .WithHttpManager(httpManager);
+
+                IManagedIdentityApplication mi = miBuilder.Build();
+
+                // Disabling shared cache options to avoid cross test pollution.
+                miBuilder.Config.AccessorOptions = null;
+
+                MockHttpMessageHandler mockHandler = httpManager.AddManagedIdentityMockHandler(
+                    MachineLearningEndpoint,
+                    Resource,
+                    MockHelpers.GetMsiSuccessfulResponse(),
+                    ManagedIdentitySource.MachineLearning,
+                    userAssignedId: userAssignedId,
+                    userAssignedIdentityId);
+
+                AuthenticationResult result = await mi.AcquireTokenForManagedIdentity(Resource).ExecuteAsync().ConfigureAwait(false);
+
+                Assert.IsNotNull(result);
+                Assert.IsNotNull(result.AccessToken);
+
+                // Verify query parameter is "clientid" and not "client_id"
+                Assert.IsTrue(mockHandler.ExpectedQueryParams.ContainsKey(Constants.ManagedIdentityClientId2017), "Query parameter should use 'clientid' and not 'client_id'");
+
+                // Verify the clientid value based on identity type
+                string expectedClientId = userAssignedId ?? EnvironmentVariables.MachineLearningDefaultClientId;
+                Assert.AreEqual(expectedClientId, mockHandler.ExpectedQueryParams[Constants.ManagedIdentityClientId2017],
+                    "Clientid value should match the provided user assigned ID for UAMI or environment variable for SAMI");
+            }
+        }
+
+        [DataTestMethod]
+        [DataRow(TestConstants.MiResourceId, UserAssignedIdentityId.ResourceId)]
+        [DataRow(TestConstants.MiResourceId, UserAssignedIdentityId.ObjectId)]
+        public async Task MachineLearningUserAssignedNonClientIdThrowsAsync(
+            string userAssignedId,
+            UserAssignedIdentityId userAssignedIdentityId)
+        {
+            using (new EnvVariableContext())
+            using (var httpManager = new MockHttpManager())
+            {
+                SetEnvironmentVariables(ManagedIdentitySource.MachineLearning, MachineLearningEndpoint);
+
+                ManagedIdentityApplicationBuilder miBuilder = CreateMIABuilder(userAssignedId, userAssignedIdentityId);
+                miBuilder
+                    .WithHttpManager(httpManager);
+
+                IManagedIdentityApplication mi = miBuilder.Build();
+
+                MsalServiceException ex = await Assert.ThrowsExceptionAsync<MsalServiceException>(async () =>
+                    await mi.AcquireTokenForManagedIdentity(Resource)
+                    .ExecuteAsync().ConfigureAwait(false)).ConfigureAwait(false);
+
+                Assert.IsNotNull(ex);
+                Assert.AreEqual(ManagedIdentitySource.MachineLearning.ToString(), ex.AdditionalExceptionData[MsalException.ManagedIdentitySource]);
+                Assert.AreEqual(MsalError.InvalidManagedIdentityIdType, ex.ErrorCode);
+            }
+        }
 
         [TestMethod]
         public async Task MachineLearningTestsInvalidEndpointAsync()
@@ -29,13 +104,13 @@ public async Task MachineLearningTestsInvalidEndpointAsync()
             {
                 SetEnvironmentVariables(ManagedIdentitySource.MachineLearning, "127.0.0.1:41564/msi/token");
 
-                var miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned)
+                ManagedIdentityApplicationBuilder miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned)
                     .WithHttpManager(httpManager);
 
                 // Disabling shared cache options to avoid cross test pollution.
                 miBuilder.Config.AccessorOptions = null;
 
-                var mi = miBuilder.Build();
+                IManagedIdentityApplication mi = miBuilder.Build();
 
                 MsalServiceException ex = await Assert.ThrowsExceptionAsync<MsalServiceException>(async () =>
                     await mi.AcquireTokenForManagedIdentity(ManagedIdentityTests.Resource)

tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ManagedIdentityTests.cs

@@ -127,8 +127,6 @@ public async Task ManagedIdentityHappyPathAsync(
         [DataRow(ServiceFabricEndpoint, ManagedIdentitySource.ServiceFabric, TestConstants.MiResourceId, UserAssignedIdentityId .ResourceId)]
         [DataRow(ServiceFabricEndpoint, ManagedIdentitySource.ServiceFabric, TestConstants.MiResourceId, UserAssignedIdentityId.ObjectId)]
         [DataRow(MachineLearningEndpoint, ManagedIdentitySource.MachineLearning, TestConstants.ClientId, UserAssignedIdentityId.ClientId)]
-        [DataRow(MachineLearningEndpoint, ManagedIdentitySource.MachineLearning, TestConstants.MiResourceId, UserAssignedIdentityId.ResourceId)]
-        [DataRow(MachineLearningEndpoint, ManagedIdentitySource.MachineLearning, TestConstants.MiResourceId, UserAssignedIdentityId.ObjectId)]
         public async Task ManagedIdentityUserAssignedHappyPathAsync(
             string endpoint,
             ManagedIdentitySource managedIdentitySource,