Adding MSAL token caching support to ManagedIdentityCredential (Azure#30432)

* updating to use native broker impl

* upgrade to latest msal

* adding conf client caching to managed identity cred

* disable manual test

* removing local nuget feed

* adding caching and refresh tests

Author: schaabs

sdk/identity/Azure.Identity/src/ManagedIdentityClient.cs

@@ -2,9 +2,13 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Linq;
+using System.Net.Http.Headers;
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
+using Microsoft.Identity.Client;
+using Microsoft.Identity.Client.Extensibility;
 
 namespace Azure.Identity
 {
@@ -14,6 +18,7 @@ internal class ManagedIdentityClient
             "ManagedIdentityCredential authentication unavailable. No Managed Identity endpoint found.";
 
         private Lazy<ManagedIdentitySource> _identitySource;
+        private MsalConfidentialClient _msal;
 
         protected ManagedIdentityClient()
         {
@@ -40,18 +45,35 @@ public ManagedIdentityClient(ManagedIdentityClientOptions options)
             ClientId = options.ClientId;
             Pipeline = options.Pipeline;
             _identitySource = new Lazy<ManagedIdentitySource>(() => SelectManagedIdentitySource(options));
+            _msal = new MsalConfidentialClient(Pipeline, "MANAGED-IDENTITY-RESOURCE-TENENT", ClientId ?? "SYSTEM-ASSIGNED-MANAGED-IDENTITY", AppTokenProviderImpl, options.Options);
         }
 
         internal CredentialPipeline Pipeline { get; }
 
         protected string ClientId { get; }
 
-        public virtual async ValueTask<AccessToken> AuthenticateAsync(bool async, TokenRequestContext context,
+        public async ValueTask<AccessToken> AuthenticateAsync(bool async, TokenRequestContext context, CancellationToken cancellationToken)
+        {
+            AuthenticationResult result = await _msal.AcquireTokenForClientAsync(context.Scopes, context.TenantId, async, cancellationToken).ConfigureAwait(false);
+
+            return new AccessToken(result.AccessToken, result.ExpiresOn);
+        }
+
+        public virtual async ValueTask<AccessToken> AuthenticateCoreAsync(bool async, TokenRequestContext context,
             CancellationToken cancellationToken)
         {
             return await _identitySource.Value.AuthenticateAsync(async, context, cancellationToken).ConfigureAwait(false);
         }
 
+        private async Task<AppTokenProviderResult> AppTokenProviderImpl(AppTokenProviderParameters parameters)
+        {
+            TokenRequestContext requestContext = new TokenRequestContext(parameters.Scopes.ToArray(), claims: parameters.Claims);
+
+            AccessToken token = await AuthenticateCoreAsync(true, requestContext, parameters.CancellationToken).ConfigureAwait(false);
+
+            return new AppTokenProviderResult() { AccessToken = token.Token, ExpiresInSeconds = Math.Max(Convert.ToInt64((token.ExpiresOn - DateTimeOffset.UtcNow).TotalSeconds), 1) };
+        }
+
         private static ManagedIdentitySource SelectManagedIdentitySource(ManagedIdentityClientOptions options)
         {
             return

sdk/identity/Azure.Identity/src/MsalConfidentialClient.cs

@@ -7,16 +7,19 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client;
+using Microsoft.Identity.Client.Extensibility;
 
 namespace Azure.Identity
 {
     internal class MsalConfidentialClient : MsalClientBase<IConfidentialClientApplication>
     {
+        private const string s_instanceMetadata = "{\"tenant_discovery_endpoint\":\"https://login.microsoftonline.com/common/v2.0/.well-known/openid-configuration\",\"api-version\":\"1.1\",\"metadata\":[{\"preferred_network\":\"login.microsoftonline.com\",\"preferred_cache\":\"login.windows.net\",\"aliases\":[\"login.microsoftonline.com\",\"login.windows.net\",\"login.microsoft.com\",\"sts.windows.net\"]}]}";
         internal readonly string _clientSecret;
         internal readonly bool _includeX5CClaimHeader;
         internal readonly IX509Certificate2Provider _certificateProvider;
         private readonly Func<string> _assertionCallback;
         private readonly Func<CancellationToken, Task<string>> _asyncAssertionCallback;
+        private readonly Func<AppTokenProviderParameters, Task<AppTokenProviderResult>> _appTokenProviderCallback;
 
         internal string RedirectUrl { get; }
 
@@ -52,15 +55,32 @@ public MsalConfidentialClient(CredentialPipeline pipeline, string tenantId, stri
             _asyncAssertionCallback = assertionCallback;
         }
 
+        public MsalConfidentialClient(CredentialPipeline pipeline, string tenantId, string clientId, Func<AppTokenProviderParameters, Task<AppTokenProviderResult>> appTokenProviderCallback, TokenCredentialOptions options)
+            : base(pipeline, tenantId, clientId, options)
+        {
+            _appTokenProviderCallback = appTokenProviderCallback;
+        }
+
         internal string RegionalAuthority { get; } = EnvironmentVariables.AzureRegionalAuthorityName;
 
         protected override async ValueTask<IConfidentialClientApplication> CreateClientAsync(bool async, CancellationToken cancellationToken)
         {
             ConfidentialClientApplicationBuilder confClientBuilder = ConfidentialClientApplicationBuilder.Create(ClientId)
-                .WithAuthority(Pipeline.AuthorityHost.AbsoluteUri, TenantId)
                 .WithHttpClientFactory(new HttpPipelineClientFactory(Pipeline.HttpPipeline))
                 .WithLogging(LogMsal, enablePiiLogging: IsPiiLoggingEnabled);
 
+            //special case for using appTokenProviderCallback, authority validation and instance metadata discovery should be disabled since we're not calling the STS
+            if (_appTokenProviderCallback != null)
+            {
+                confClientBuilder.WithAppTokenProvider(_appTokenProviderCallback)
+                    .WithAuthority(Pipeline.AuthorityHost.AbsoluteUri, TenantId, false)
+                    .WithInstanceDiscoveryMetadata(s_instanceMetadata);
+            }
+            else
+            {
+                confClientBuilder.WithAuthority(Pipeline.AuthorityHost.AbsoluteUri, TenantId);
+            }
+
             if (_clientSecret != null)
             {
                 confClientBuilder.WithClientSecret(_clientSecret);

sdk/identity/Azure.Identity/tests/ManagedIdentityCredentialTests.cs

@@ -29,6 +29,46 @@ public ManagedIdentityCredentialTests(bool isAsync) : base(isAsync)
 
         private const string ExpectedToken = "mock-msi-access-token";
 
+        [Test]
+        public async Task VerifyTokenCaching()
+        {
+            int callCount = 0;
+
+            var mockClient = new MockManagedIdentityClient(CredentialPipeline.GetInstance(null))
+            {
+                TokenFactory = () => { callCount++; return new AccessToken(Guid.NewGuid().ToString(), DateTimeOffset.UtcNow.AddHours(24)); }
+            };
+
+            var cred = InstrumentClient(new ManagedIdentityCredential(mockClient));
+
+            for (int i = 0; i < 5; i++)
+            {
+                await cred.GetTokenAsync(new TokenRequestContext(MockScopes.Default));
+            }
+
+            Assert.AreEqual(1, callCount);
+        }
+
+        [Test]
+        public async Task VerifyExpiringTokenRefresh()
+        {
+            int callCount = 0;
+
+            var mockClient = new MockManagedIdentityClient(CredentialPipeline.GetInstance(null))
+            {
+                TokenFactory = () => { callCount++; return new AccessToken(Guid.NewGuid().ToString(), DateTimeOffset.UtcNow.AddMinutes(2)); }
+            };
+
+            var cred = InstrumentClient(new ManagedIdentityCredential(mockClient));
+
+            for (int i = 0; i < 5; i++)
+            {
+                await cred.GetTokenAsync(new TokenRequestContext(MockScopes.Default));
+            }
+
+            Assert.AreEqual(5, callCount);
+        }
+
         [NonParallelizable]
         [Test]
         public async Task VerifyImdsRequestWithClientIdMockAsync()

sdk/identity/Azure.Identity/tests/Mock/MockManagedIdentityClient.cs

@@ -27,7 +27,7 @@ public MockManagedIdentityClient(CredentialPipeline pipeline, string clientId)
 
         public Func<AccessToken> TokenFactory { get; set; }
 
-        public override ValueTask<AccessToken> AuthenticateAsync(bool async, TokenRequestContext context, CancellationToken cancellationToken)
+        public override ValueTask<AccessToken> AuthenticateCoreAsync(bool async, TokenRequestContext context, CancellationToken cancellationToken)
               => TokenFactory != null ? new ValueTask<AccessToken>(TokenFactory()) : base.AuthenticateAsync(async, context, cancellationToken);
     }
 }