Lozensky/enable managed identity

src/Microsoft.Identity.Web.TokenAcquisition/DefaultAuthorizationHeaderProvider.cs

@@ -49,19 +49,9 @@ public async Task<string> CreateAuthorizationHeaderForAppAsync(string scopes, Au
 
         private static TokenAcquisitionOptions CreateTokenAcquisitionOptionsFromApiOptions(AuthorizationHeaderProviderOptions? downstreamApiOptions, CancellationToken cancellationToken)
         {
-            return new TokenAcquisitionOptions()
-            {
-                AuthenticationOptionsName = downstreamApiOptions?.AcquireTokenOptions.AuthenticationOptionsName,
-                CancellationToken = cancellationToken,
-                Claims = downstreamApiOptions?.AcquireTokenOptions.Claims,
-                CorrelationId = downstreamApiOptions?.AcquireTokenOptions.CorrelationId ?? Guid.Empty,
-                ExtraQueryParameters = downstreamApiOptions?.AcquireTokenOptions.ExtraQueryParameters,
-                ForceRefresh = downstreamApiOptions?.AcquireTokenOptions.ForceRefresh ?? false,
-                LongRunningWebApiSessionKey = downstreamApiOptions?.AcquireTokenOptions.LongRunningWebApiSessionKey,
-                Tenant = downstreamApiOptions?.AcquireTokenOptions.Tenant,
-                UserFlow = downstreamApiOptions?.AcquireTokenOptions.UserFlow,
-                PopPublicKey = downstreamApiOptions?.AcquireTokenOptions.PopPublicKey,
-            };
+            TokenAcquisitionOptions tokenAcquisitionOptions = TokenAcquisitionOptions.CloneFromBaseClass(downstreamApiOptions?.AcquireTokenOptions);
+            tokenAcquisitionOptions.CancellationToken = cancellationToken;
+            return tokenAcquisitionOptions;
         }
     }
 }

src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.ManagedIdentity.cs

@@ -0,0 +1,119 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.Collections.Concurrent;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Identity.Abstractions;
+using Microsoft.Identity.Client;
+using Microsoft.Identity.Client.AppConfig;
+using Microsoft.IdentityModel.Tokens;
+
+namespace Microsoft.Identity.Web
+{
+    /// <summary>
+    /// Portion of the TokenAcquisition class that handles logic unique to managed identity.
+    /// </summary>
+#if NETSTANDARD2_0 || NET462 || NET472
+        internal partial class TokenAcquisition : ITokenAcquisitionInternal
+#else
+    internal partial class TokenAcquisition
+    #endif
+    {
+        private readonly ConcurrentDictionary<string, IManagedIdentityApplication> _managedIdentityApplicationsByClientId = new();
+        private readonly SemaphoreSlim _managedIdSemaphore = new(1, 1);
+        private const string SystemAssignedManagedIdentityKey = "SYSTEM";
+
+        /// <summary>
+        /// Gets a cached ManagedIdentityApplication object or builds a new one if not found.
+        /// </summary>
+        /// <param name="mergedOptions">The configuration options for the app.</param>
+        /// <param name="managedIdentityOptions">The configuration specific to managed identity.</param>
+        /// <returns>The application object used to request a token with managed identity.</returns>
+        internal async Task<IManagedIdentityApplication> GetOrBuildManagedIdentityApplication(
+            MergedOptions mergedOptions,
+            ManagedIdentityOptions managedIdentityOptions)
+        {
+            string key = GetCacheKeyForManagedId(managedIdentityOptions);
+
+            // Check if the application is already built, if so return it without grabbing the lock
+            if (_managedIdentityApplicationsByClientId.TryGetValue(key, out IManagedIdentityApplication? application))
+            {
+                return application;
+            }
+
+            // Lock the potential write of the dictionary to prevent multiple threads from creating the same application.
+            await _managedIdSemaphore.WaitAsync();
+            try
+            {
+                // Check if the application is already built (could happen between previous check and obtaining the key)
+                if (_managedIdentityApplicationsByClientId.TryGetValue(key, out application))
+                {
+                    return application;
+                }
+
+                // Set managedIdentityId to the correct value for either system or user assigned
+                ManagedIdentityId managedIdentityId;
+                if (key == SystemAssignedManagedIdentityKey)
+                {
+                    managedIdentityId = ManagedIdentityId.SystemAssigned;
+                }
+                else
+                {
+                    managedIdentityId = ManagedIdentityId.WithUserAssignedClientId(key);
+                }
+
+                // Build the application
+                application = BuildManagedIdentityApplication(
+                    managedIdentityId,
+                    mergedOptions.ConfidentialClientApplicationOptions.EnablePiiLogging
+                );
+
+                // Add the application to the cache
+                _managedIdentityApplicationsByClientId.TryAdd(key, application);
+            }
+            finally
+            {
+                // Now that the dictionary is updated, release the semaphore
+                _managedIdSemaphore.Release();
+            }
+            return application;
+        }
+
+        /// <summary>
+        /// Creates a managed identity client application.
+        /// </summary>
+        /// <param name="managedIdentityId">Indicates if system-assigned or user-assigned managed identity is used.</param>
+        /// <param name="enablePiiLogging">Indicates if logging that may contain personally identifiable information is enabled.</param>
+        /// <returns>A managed identity application.</returns>
+        private IManagedIdentityApplication BuildManagedIdentityApplication(ManagedIdentityId managedIdentityId, bool enablePiiLogging)
+        {
+            return ManagedIdentityApplicationBuilder
+                .Create(managedIdentityId)
+                .WithHttpClientFactory(_httpClientFactory)
+                .WithLogging(
+                    Log,
+                    ConvertMicrosoftExtensionsLogLevelToMsal(_logger),
+                    enablePiiLogging: enablePiiLogging)
+                .Build();
+        }
+
+        /// <summary>
+        /// Gets the key value for the Managed Identity cache, the default key for system-assigned identity is used if there is
+        /// no clientId for a user-assigned identity specified. The method is internal rather than private for testing purposes.
+        /// </summary>
+        /// <param name="managedIdOptions">Holds the clientId for managed identity if none is present.</param>
+        /// <returns>A key value for the Managed Identity cache.</returns>
+        internal static string GetCacheKeyForManagedId(ManagedIdentityOptions managedIdOptions)
+        {
+            if (managedIdOptions.UserAssignedClientId.IsNullOrEmpty())
+            {
+                return SystemAssignedManagedIdentityKey;
+            }
+            else
+            {
+                return managedIdOptions.UserAssignedClientId!;
+            }
+        }
+    }
+}

src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.cs

@@ -19,11 +19,11 @@
 using Microsoft.Identity.Client;
 using Microsoft.Identity.Client.Advanced;
 using Microsoft.Identity.Client.Extensibility;
+using Microsoft.Identity.Web.Experimental;
 using Microsoft.Identity.Web.TokenCacheProviders;
 using Microsoft.Identity.Web.TokenCacheProviders.InMemory;
 using Microsoft.IdentityModel.JsonWebTokens;
 using Microsoft.IdentityModel.Tokens;
-using Microsoft.Identity.Web.Experimental;
 
 namespace Microsoft.Identity.Web
 {
@@ -47,9 +47,9 @@ class OAuthConstants
         private readonly object _applicationSyncObj = new();
 
         /// <summary>
-        ///  Please call GetOrBuildConfidentialClientApplication instead of accessing this field directly.
+        ///  Please call GetOrBuildConfidentialClientApplication instead of accessing _applicationsByAuthorityClientId directly.
         /// </summary>
-        private readonly ConcurrentDictionary<string, IConfidentialClientApplication?> _applicationsByAuthorityClientId = new ConcurrentDictionary<string, IConfidentialClientApplication?>();
+        private readonly ConcurrentDictionary<string, IConfidentialClientApplication?> _applicationsByAuthorityClientId = new();
         private bool _retryClientCertificate;
         protected readonly IMsalHttpClientFactory _httpClientFactory;
         protected readonly ILogger _logger;
@@ -115,7 +115,7 @@ public async Task<AcquireTokenResult> AddAccountToCacheFromAuthorizationCodeAsyn
             _ = Throws.IfNull(authCodeRedemptionParameters.Scopes);
             MergedOptions mergedOptions = _tokenAcquisitionHost.GetOptions(authCodeRedemptionParameters.AuthenticationScheme, out string effectiveAuthenticationScheme);
 
-            IConfidentialClientApplication? application=null;
+            IConfidentialClientApplication? application = null;
             try
             {
                 application = GetOrBuildConfidentialClientApplication(mergedOptions);
@@ -321,9 +321,10 @@ private void LogAuthResult(AuthenticationResult? authenticationResult)
 
         /// <summary>
         /// Acquires an authentication result from the authority configured in the app, for the confidential client itself (not on behalf of a user)
-        /// using the client credentials flow. See https://aka.ms/msal-net-client-credentials.
+        /// using either a client credentials or managed identity flow. See https://aka.ms/msal-net-client-credentials for client credentials or
+        /// https://aka.ms/Entra/ManagedIdentityOverview for managed identity.
         /// </summary>
-        /// <param name="scope">The scope requested to access a protected API. For this flow (client credentials), the scope
+        /// <param name="scope">The scope requested to access a protected API. For these flows (client credentials or managed identity), the scope
         /// should be of the form "{ResourceIdUri/.default}" for instance <c>https://management.azure.net/.default</c> or, for Microsoft
         /// Graph, <c>https://graph.microsoft.com/.default</c> as the requested scopes are defined statically with the application registration
         /// in the portal, and cannot be overridden in the application, as you can request a token for only one resource at a time (use
@@ -358,10 +359,28 @@ public async Task<AuthenticationResult> GetAuthenticationResultForAppAsync(
                 throw new ArgumentException(IDWebErrorMessage.ClientCredentialTenantShouldBeTenanted, nameof(tenant));
             }
 
+            // If using managed identity 
+            if (tokenAcquisitionOptions != null && tokenAcquisitionOptions.ManagedIdentity != null)
+            {
+                try
+                {
+                    IManagedIdentityApplication managedIdApp = await GetOrBuildManagedIdentityApplication(
+                        mergedOptions,
+                        tokenAcquisitionOptions.ManagedIdentity
+                    );
+                    return await managedIdApp.AcquireTokenForManagedIdentity(scope).ExecuteAsync().ConfigureAwait(false);
+                }
+                catch (Exception ex)
+                {
+                    Logger.TokenAcquisitionError(_logger, ex.Message, ex);
+                    throw;
+                }
+            }
+
             // Use MSAL to get the right token to call the API
-            var application = GetOrBuildConfidentialClientApplication(mergedOptions);
+            IConfidentialClientApplication application = GetOrBuildConfidentialClientApplication(mergedOptions);
 
-            var builder = application
+            AcquireTokenForClientParameterBuilder builder = application
                    .AcquireTokenForClient(new[] { scope }.Except(_scopesRequestedByMsal))
                    .WithSendX5C(mergedOptions.SendX5C);
 
@@ -585,7 +604,6 @@ private bool IsInvalidClientCertificateOrSignedAssertionError(MsalServiceExcepti
                     _applicationsByAuthorityClientId.TryAdd(GetApplicationKey(mergedOptions), application);
                 }
             }
-
             return application;
         }
 
@@ -599,7 +617,7 @@ private IConfidentialClientApplication BuildConfidentialClientApplication(Merged
 
             try
             {
-                var builder = ConfidentialClientApplicationBuilder
+                ConfidentialClientApplicationBuilder builder = ConfidentialClientApplicationBuilder
                         .CreateWithApplicationOptions(mergedOptions.ConfidentialClientApplicationOptions)
                         .WithHttpClientFactory(_httpClientFactory)
                         .WithLogging(
@@ -848,8 +866,10 @@ private static void CheckAssertionsForInjectionAttempt(string assertion, string
                 if (!assertion.IsNullOrEmpty() && assertion.Contains('&')) throw new ArgumentException(IDWebErrorMessage.InvalidAssertion, nameof(assertion));
                 if (!subAssertion.IsNullOrEmpty() && subAssertion.Contains('&')) throw new ArgumentException(IDWebErrorMessage.InvalidSubAssertion, nameof(subAssertion));
 #else
-                if (!assertion.IsNullOrEmpty() && assertion.Contains('&', StringComparison.InvariantCultureIgnoreCase)) throw new ArgumentException(IDWebErrorMessage.InvalidAssertion, nameof(assertion));
-                if (!subAssertion.IsNullOrEmpty() && subAssertion.Contains('&', StringComparison.InvariantCultureIgnoreCase)) throw new ArgumentException(IDWebErrorMessage.InvalidSubAssertion, nameof(subAssertion));
+                if (!assertion.IsNullOrEmpty() && assertion.Contains('&', StringComparison.InvariantCultureIgnoreCase))
+                    throw new ArgumentException(IDWebErrorMessage.InvalidAssertion, nameof(assertion));
+                if (!subAssertion.IsNullOrEmpty() && subAssertion.Contains('&', StringComparison.InvariantCultureIgnoreCase))
+                    throw new ArgumentException(IDWebErrorMessage.InvalidSubAssertion, nameof(subAssertion));
 #endif
             }
         }

src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisitionOptions.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using System;
 using System.Threading;
 using Microsoft.Identity.Abstractions;
 using Microsoft.Identity.Client.AppConfig;
@@ -25,6 +26,47 @@ public class TokenAcquisitionOptions : AcquireTokenOptions
         /// </summary>
         public CancellationToken CancellationToken { get; set; } = CancellationToken.None;
 
+        /// <summary>
+        /// Creates a new instance of TokenAcquisitionOptions from an instance of the base class AcquireTokenOptions.
+        /// </summary>
+        /// <param name="acquireTokenOptions">The instance of the base class to clone</param>
+        /// <returns>A new instance of TokenAquisitionOptions</returns>
+        public static TokenAcquisitionOptions CloneFromBaseClass(AcquireTokenOptions? acquireTokenOptions)
+        {
+            if (acquireTokenOptions == null)
+            {
+                return new TokenAcquisitionOptions();
+            }
+
+            // alternate option 1: serialize the base class then deserialize it back to the derived class, this is probably not
+            // ideal since it likely has small performance repercussions but it doesn't have to be done with json.
+
+            // alternate option 2: add the two extra fields of PoPConfiguration and CancellationToken to the base class in
+            // Abstractions and leave this child as a 1:1 wrapper with the base class to avoid breaking anyone.
+            // This is my personal preference as it is most performant since no conversion would be needed. The drawback is I'd
+            // need to replace all instances of TokenAcquisitionOptions with AcquireTokenOptions in the codebase.
+            //      2.1 - I spoke with JM and he said this would be a longer conversation as he wants to deprecate
+            //      TokenAcquisitionOptions entirely and drop the CancellationToken and PoPConfiguration fields as they
+            //      apparently don't belong in AcquireTokenOptions.
+
+            // This is how we have been doing it so far (in the DefaultAuthorizationHeaderProvider class) I put it here to make
+            // it easier to find if we ever add more fields to the base class, but it will still require manual intervention to
+            // add the new fields here.
+            return new TokenAcquisitionOptions
+            {
+                AuthenticationOptionsName = acquireTokenOptions.AuthenticationOptionsName,
+                Claims = acquireTokenOptions.Claims,
+                CorrelationId = acquireTokenOptions.CorrelationId ?? Guid.Empty,
+                ExtraQueryParameters = acquireTokenOptions.ExtraQueryParameters,
+                ForceRefresh = acquireTokenOptions.ForceRefresh,
+                LongRunningWebApiSessionKey = acquireTokenOptions.LongRunningWebApiSessionKey,
+                ManagedIdentity = acquireTokenOptions.ManagedIdentity,
+                Tenant = acquireTokenOptions.Tenant,
+                UserFlow = acquireTokenOptions.UserFlow,
+                PopPublicKey = acquireTokenOptions.PopPublicKey,
+            };
+        }
+
         /// <summary>
         /// Clone the options (to be able to override them).
         /// </summary>

tests/E2E Tests/TokenAcquirerTests/OnlyOnAzureDevopsFactAttribute.cs

@@ -0,0 +1,20 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using Xunit;
+
+namespace TokenAcquirerTests
+{
+    public sealed class OnlyOnAzureDevopsFactAttribute : FactAttribute
+    {
+        public OnlyOnAzureDevopsFactAttribute()
+        {
+            if (IgnoreOnAzureDevopsFactAttribute.IsRunningOnAzureDevOps())
+            {
+                return;
+            }
+            Skip = "Ignored when not on Azure DevOps";
+        }
+    }
+}

tests/E2E Tests/TokenAcquirerTests/TokenAcquirer.cs

@@ -3,9 +3,11 @@
 
 using System;
 using System.Linq;
+using System.Net.Http;
 using System.Security.Cryptography;
 using System.Security.Cryptography.X509Certificates;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Options;
 using Microsoft.Graph;
@@ -299,5 +301,46 @@ private static async Task CreateGraphClientAndAssert(TokenAcquirerFactory tokenA
             Assert.NotNull(result.AccessToken);
         }
     }
+
+    public class AcquireTokenManagedIdentity 
+    {
+        [OnlyOnAzureDevopsFact]
+        //[Fact]
+        public async Task AcquireTokenWithManagedIdentity_UserAssigned()
+        {
+            // Arrange
+            const string scope = "https://vault.azure.net/.default";
+            const string baseUrl = "https://vault.azure.net";
+            const string clientId = "9c5896db-a74a-4b1a-a259-74c5080a3a6a";
+            TokenAcquirerFactory tokenAcquirerFactory = TokenAcquirerFactory.GetDefaultInstance();
+            _ = tokenAcquirerFactory.Services;
+            IServiceProvider serviceProvider = tokenAcquirerFactory.Build();
+
+            // Act: Get the authorization header provider and add the options to tell it to use Managed Identity
+            IAuthorizationHeaderProvider? api = serviceProvider.GetRequiredService<IAuthorizationHeaderProvider>();
+            Assert.NotNull(api);
+            string result = await api.CreateAuthorizationHeaderForAppAsync(scope, GetAuthHeaderOptions_ManagedId(baseUrl, clientId));
+
+            // Assert: Make sure we got a token
+            Assert.False(string.IsNullOrEmpty(result));
+        }
+
+        private static AuthorizationHeaderProviderOptions GetAuthHeaderOptions_ManagedId(string baseUrl, string? userAssignedClientId=null) 
+        {
+            ManagedIdentityOptions managedIdentityOptions = new()
+            {
+                UserAssignedClientId = userAssignedClientId
+            };
+            AcquireTokenOptions aquireTokenOptions = new()
+            {
+                ManagedIdentity = managedIdentityOptions
+            };
+            return new AuthorizationHeaderProviderOptions()
+            {
+                BaseUrl = baseUrl,
+                AcquireTokenOptions = aquireTokenOptions
+            };
+        }
+    }
 #endif //FROM_GITHUB_ACTION
 }