[MSI v2 - Feature] Add cert cache and Auth Operation to IMDS V2

src/client/Microsoft.Identity.Client/ApiConfig/Parameters/AcquireTokenForManagedIdentityParameters.cs

@@ -4,11 +4,13 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Security.Cryptography.X509Certificates;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.Core;
 using Microsoft.Identity.Client.ManagedIdentity;
+using Microsoft.Identity.Client.ManagedIdentity.V2;
 
 namespace Microsoft.Identity.Client.ApiConfig.Parameters
 {
@@ -24,6 +26,13 @@ internal class AcquireTokenForManagedIdentityParameters : IAcquireTokenParameter
 
         public bool IsMtlsPopRequested { get; set; }
 
+        // When the MI source produced / resolved an mTLS binding certificate, we attach it here
+        // so the request layer can apply a cache-correct IAuthenticationOperation.
+        public X509Certificate2 MtlsCertificate { get; set; }
+
+        // CSR response we get back when IMDSv2 minted the certificate.
+        internal CertificateRequestResponse CertificateRequestResponse { get; set; }
+
         internal Func<AttestationTokenInput, CancellationToken, Task<AttestationTokenResponse>> AttestationTokenProvider { get; set; }
 
         public void LogParameters(ILoggerAdapter logger)
@@ -38,6 +47,10 @@ public void LogParameters(ILoggerAdapter logger)
                      Claims: {!string.IsNullOrEmpty(Claims)}
                      RevokedTokenHash: {!string.IsNullOrEmpty(RevokedTokenHash)}
                      """);
+
+                logger.Info(() =>
+                    $"[AcquireTokenForManagedIdentityParameters] IsMtlsPopRequested={IsMtlsPopRequested}, " +
+                    $"MtlsCert={(MtlsCertificate != null ? MtlsCertificate.Thumbprint : "null")}");
             }
         }
     }

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

@@ -95,7 +95,7 @@ public static void ResetStateForTest()
             OidcRetrieverWithCache.ResetCacheForTest();
             AuthorityManager.ClearValidationCache();
             SingletonThrottlingManager.GetInstance().ResetCache();
-            ManagedIdentityClient.ResetSourceForTest();
+            ManagedIdentityClient.ResetSourceAndCertForTest();
             AuthorityManager.ClearValidationCache();
             PoPCryptoProviderFactory.Reset();
 

src/client/Microsoft.Identity.Client/Internal/Requests/AuthenticationRequestParameters.cs

@@ -115,6 +115,13 @@ public AuthenticationRequestParameters(
 
         public bool IsMtlsPopRequested => _commonParameters.IsMtlsPopRequested;
 
+        // Request‑scoped override for the authentication operation.
+        internal IAuthenticationOperation AuthenticationOperationOverride { get; set; }
+
+        // Effective operation for this request: prefer the override, otherwise the default.
+        public IAuthenticationOperation AuthenticationScheme =>
+            AuthenticationOperationOverride ?? _commonParameters.AuthenticationOperation;
+
         /// <summary>
         /// Indicates if the user configured claims via .WithClaims. Not affected by Client Capabilities
         /// </summary>
@@ -127,8 +134,6 @@ public string Claims
             }
         }
 
-        public IAuthenticationOperation AuthenticationScheme => _commonParameters.AuthenticationOperation;
-
         public IEnumerable<string> PersistedCacheParameters => _commonParameters.AdditionalCacheParameters;
 
         public SortedList<string, string> CacheKeyComponents {get; private set; }

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

@@ -43,7 +43,7 @@ protected AbstractManagedIdentity(RequestContext requestContext, ManagedIdentity
         }
 
         public virtual async Task<ManagedIdentityResponse> AuthenticateAsync(
-            AcquireTokenForManagedIdentityParameters parameters,
+            AcquireTokenForManagedIdentityParameters parameters, 
             CancellationToken cancellationToken)
         {
             if (cancellationToken.IsCancellationRequested)
@@ -54,14 +54,25 @@ public virtual async Task<ManagedIdentityResponse> AuthenticateAsync(
 
             HttpResponse response;
 
-            // Convert the scopes to a resource string.
             string resource = parameters.Resource;
-
             _isMtlsPopRequested = parameters.IsMtlsPopRequested;
 
             ManagedIdentityRequest request = await CreateRequestAsync(resource).ConfigureAwait(false);
 
-            // Automatically add claims / capabilities if this MI source supports them
+            // Bubble cert + CSR to parameters so upper layers can apply mtls_pop before caching
+            if (request != null)
+            {
+                if (request.MtlsCertificate != null)
+                {
+                    parameters.MtlsCertificate = request.MtlsCertificate;
+                }
+
+                if (request.CertificateRequestResponse != null)
+                {
+                    parameters.CertificateRequestResponse = request.CertificateRequestResponse;
+                }
+            }
+
             if (_sourceType.SupportsClaimsAndCapabilities())
             {
                 request.AddClaimsAndCapabilities(
@@ -110,7 +121,6 @@ public virtual async Task<ManagedIdentityResponse> AuthenticateAsync(
                             cancellationToken: cancellationToken,
                             retryPolicy: retryPolicy)
                         .ConfigureAwait(false);
-
                 }
 
                 return await HandleResponseAsync(parameters, response, cancellationToken).ConfigureAwait(false);

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

@@ -2,14 +2,16 @@
 // Licensed under the MIT License.
 
 using System;
-using System.Threading.Tasks;
+using System.Collections.Concurrent;
+using System.IO;
+using System.Security.Cryptography.X509Certificates;
 using System.Threading;
-using Microsoft.Identity.Client.Internal;
+using System.Threading.Tasks;
 using Microsoft.Identity.Client.ApiConfig.Parameters;
-using Microsoft.Identity.Client.PlatformsCommon.Shared;
-using System.IO;
 using Microsoft.Identity.Client.Core;
+using Microsoft.Identity.Client.Internal;
 using Microsoft.Identity.Client.ManagedIdentity.V2;
+using Microsoft.Identity.Client.PlatformsCommon.Shared;
 
 namespace Microsoft.Identity.Client.ManagedIdentity
 {
@@ -21,10 +23,47 @@ internal class ManagedIdentityClient
         private const string WindowsHimdsFilePath = "%Programfiles%\\AzureConnectedMachineAgent\\himds.exe";
         private const string LinuxHimdsFilePath = "/opt/azcmagent/bin/himds";
         internal static ManagedIdentitySource s_sourceName = ManagedIdentitySource.None;
+        private X509Certificate2 _mtlsBindingCertificate;
+
+        internal X509Certificate2 MtlsBindingCertificate
+        {
+            get => _mtlsBindingCertificate;
+            set
+            {
+                // Interlocked.Exchange returns the old value.
+                var old = Interlocked.Exchange(ref _mtlsBindingCertificate, value);
+
+                // Dispose the old cert if it is being replaced.
+                if (old != null && !ReferenceEquals(old, value))
+                {
+                    try
+                    { old.Dispose(); }
+                    catch { /* best effort */ }
+                }
+            }
+        }
+
+        // Identity‑scoped caches (per process). Key should uniquely represent the identity
+        // (e.g., UAMI clientId or system‑assigned, tenant, and CUID where available).
+        private static readonly ConcurrentDictionary<string, X509Certificate2> s_mtlsCertCache =
+            new ConcurrentDictionary<string, X509Certificate2>(StringComparer.Ordinal);
+
+        private static readonly ConcurrentDictionary<string, CertificateRequestResponse> s_certResponseCache =
+            new ConcurrentDictionary<string, CertificateRequestResponse>(StringComparer.Ordinal);
 
-        internal static void ResetSourceForTest()
+        internal static void ResetSourceAndCertForTest()
         {
             s_sourceName = ManagedIdentitySource.None;
+
+            foreach (var kvp in s_mtlsCertCache)
+            {
+                try
+                { kvp.Value?.Dispose(); }
+                catch { }
+            }
+
+            s_mtlsCertCache.Clear();
+            s_certResponseCache.Clear();
         }
 
         internal async Task<ManagedIdentityResponse> SendTokenRequestForManagedIdentityAsync(
@@ -157,5 +196,61 @@ private static bool ValidateAzureArcEnvironment(string identityEndpoint, string
             logger?.Verbose(() => "[Managed Identity] Azure Arc managed identity is not available.");
             return false;
         }
+
+        internal static bool TryGetCachedMtlsCertificate(
+            string identityKey,
+            out X509Certificate2 cert,
+            out CertificateRequestResponse response)
+        {
+            cert = null;
+            response = null;
+
+            if (string.IsNullOrEmpty(identityKey))
+            {
+                return false;
+            }
+
+            if (s_mtlsCertCache.TryGetValue(identityKey, out var c))
+            {
+                cert = c;
+            }
+
+            if (s_certResponseCache.TryGetValue(identityKey, out var r))
+            {
+                response = r;
+            }
+
+            return cert != null && response != null;
+        }
+
+        internal static void CacheMtlsCertificate(
+            string identityKey,
+            X509Certificate2 cert,
+            CertificateRequestResponse response)
+        {
+            if (string.IsNullOrEmpty(identityKey) || cert == null || response == null)
+            {
+                return;
+            }
+
+            s_mtlsCertCache[identityKey] = cert;
+            s_certResponseCache[identityKey] = response;
+        }
+
+        internal static bool IsMtlsCertExpiringSoon(string identityKey)
+        {
+            if (string.IsNullOrEmpty(identityKey))
+            {
+                return true;
+            }
+
+            if (!s_mtlsCertCache.TryGetValue(identityKey, out var cert) || cert == null)
+            {
+                return true;
+            }
+
+            var remaining = cert.NotAfter.ToUniversalTime() - DateTime.UtcNow;
+            return remaining <= TimeSpan.FromMinutes(1);
+        }
     }
 }

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

@@ -8,6 +8,7 @@
 using System.Security.Cryptography.X509Certificates;
 using Microsoft.Identity.Client.ApiConfig.Parameters;
 using Microsoft.Identity.Client.Core;
+using Microsoft.Identity.Client.ManagedIdentity.V2;
 using Microsoft.Identity.Client.OAuth2;
 using Microsoft.Identity.Client.Utils;
 
@@ -29,6 +30,8 @@ internal class ManagedIdentityRequest
 
         public X509Certificate2 MtlsCertificate { get; set; }
 
+        internal CertificateRequestResponse CertificateRequestResponse { get; set; }
+
         public ManagedIdentityRequest(
             HttpMethod method,
             Uri endpoint,

src/client/Microsoft.Identity.Client/ManagedIdentity/V2/ImdsV2ManagedIdentitySource.cs

@@ -6,6 +6,7 @@
 using System.Linq;
 using System.Net;
 using System.Net.Http;
+using System.Security.Cryptography.X509Certificates;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.Core;
@@ -284,29 +285,78 @@ protected override async Task<ManagedIdentityRequest> CreateRequestAsync(string
         {
             var csrMetadata = await GetCsrMetadataAsync(_requestContext, false).ConfigureAwait(false);
 
+            // Key the binding cache by the MSAL ClientId (UAMI) or SAMI’s constant client id.
+            string identityKey = _requestContext.ServiceBundle.Config.ClientId;
+
             IManagedIdentityKeyProvider keyProvider = _requestContext.ServiceBundle.PlatformProxy.ManagedIdentityKeyProvider;
 
+            // Try to reuse cached binding
+            X509Certificate2 cachedCert = null;
+            CertificateRequestResponse cachedResponse = null;
+            bool haveCached = !string.IsNullOrEmpty(identityKey) &&
+                              ManagedIdentityClient.TryGetCachedMtlsCertificate(identityKey, out cachedCert, out cachedResponse);
+            bool certFresh = haveCached && !ManagedIdentityClient.IsMtlsCertExpiringSoon(identityKey);
+
+            if (haveCached && certFresh)
+            {
+                string tokenType = _isMtlsPopRequested ? "mtls_pop" : "bearer";
+
+                var requestFromCache = BuildTokenRequest(resource, cachedResponse, tokenType);
+                requestFromCache.MtlsCertificate = cachedCert;              // ✅ attach cert even for Bearer
+                requestFromCache.CertificateRequestResponse = cachedResponse;
+
+                _requestContext.Logger.Info("[IMDSv2] Using cached mTLS binding (cert + metadata) for identity.");
+                return requestFromCache;
+            }
+
+            // Need to mint/rotate binding certificate
             ManagedIdentityKeyInfo keyInfo = await keyProvider
-                .GetOrCreateKeyAsync(
-                _requestContext.Logger, 
-                _requestContext.UserCancellationToken)
+                .GetOrCreateKeyAsync(_requestContext.Logger, _requestContext.UserCancellationToken)
                 .ConfigureAwait(false);
 
-            var (csr, privateKey) = _requestContext.ServiceBundle.Config.CsrFactory.Generate(keyInfo.Key, csrMetadata.ClientId, csrMetadata.TenantId, csrMetadata.CuId);
+            var (csr, privateKey) = _requestContext.ServiceBundle.Config.CsrFactory
+                .Generate(keyInfo.Key, csrMetadata.ClientId, csrMetadata.TenantId, csrMetadata.CuId);
 
             var certificateRequestResponse = await ExecuteCertificateRequestAsync(
                 csrMetadata.ClientId,
                 csrMetadata.AttestationEndpoint,
                 csr,
                 keyInfo).ConfigureAwait(false);
 
-            // transform certificateRequestResponse.Certificate to x509 with private key
-            var mtlsCertificate = CommonCryptographyManager.AttachPrivateKeyToCert(
+            // ✅ IMDS v2 requires client mTLS on the STS call for both Bearer & PoP – always attach & cache the cert
+            X509Certificate2 mtlsCertificate = CommonCryptographyManager.AttachPrivateKeyToCert(
                 certificateRequestResponse.Certificate,
                 privateKey);
 
-            ManagedIdentityRequest request = new(HttpMethod.Post, new Uri($"{certificateRequestResponse.MtlsAuthenticationEndpoint}/{certificateRequestResponse.TenantId}{AcquireEntraTokenPath}"));
+            if (!string.IsNullOrEmpty(identityKey))
+            {
+                ManagedIdentityClient.CacheMtlsCertificate(identityKey, mtlsCertificate, certificateRequestResponse);
+                _requestContext.Logger.Info("[IMDSv2] Minted and cached new mTLS certificate for identity.");
+            }
+            else
+            {
+                _requestContext.Logger.Warning("[IMDSv2] Skipping mTLS cache store due to missing identity key.");
+            }
+
+            string tokenTypeFinal = _isMtlsPopRequested ? "mtls_pop" : "bearer";
+            var request = BuildTokenRequest(resource, certificateRequestResponse, tokenTypeFinal);
+            request.MtlsCertificate = mtlsCertificate;                     // ✅ attach cert even for Bearer
+            request.CertificateRequestResponse = certificateRequestResponse;
 
+            return request;
+        }
+
+        private ManagedIdentityRequest BuildTokenRequest(string resource, CertificateRequestResponse resp, string tokenType)
+        {
+            // Build STS endpoint: {mtlsAuthEndpoint}/{tenantId}/oauth2/v2.0/token
+            var stsUri = new Uri($"{resp.MtlsAuthenticationEndpoint}/{resp.TenantId}{AcquireEntraTokenPath}");
+
+            var request = new ManagedIdentityRequest(HttpMethod.Post, stsUri)
+            {
+                RequestType = RequestType.STS
+            };
+
+            // Standard MSAL identification headers
             var idParams = MsalIdHelper.GetMsalIdParameters(_requestContext.Logger);
             foreach (var idParam in idParams)
             {
@@ -316,17 +366,12 @@ protected override async Task<ManagedIdentityRequest> CreateRequestAsync(string
             request.Headers.Add(ThrottleCommon.ThrottleRetryAfterHeaderName, ThrottleCommon.ThrottleRetryAfterHeaderValue);
             request.Headers.Add(OAuth2Header.RequestCorrelationIdInResponse, "true");
 
-            var tokenType = _isMtlsPopRequested ? "mtls_pop" : "bearer";
-
-            request.BodyParameters.Add("client_id", certificateRequestResponse.ClientId);
+            // Body
+            request.BodyParameters.Add("client_id", resp.ClientId);
             request.BodyParameters.Add("grant_type", OAuth2GrantType.ClientCredentials);
             request.BodyParameters.Add("scope", resource.TrimEnd('/') + "/.default");
             request.BodyParameters.Add("token_type", tokenType);
 
-            request.RequestType = RequestType.STS;
-
-            request.MtlsCertificate = mtlsCertificate;
-
             return request;
         }