Add in-process MAA token caching to PopKeyAttestor (#5887)

Author: Robbie-Microsoft

src/client/Microsoft.Identity.Client.KeyAttestation/ManagedIdentityAttestationExtensions.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Threading.Tasks;
+using Microsoft.Identity.Client.Core;
 
 namespace Microsoft.Identity.Client.KeyAttestation
 {
@@ -25,13 +26,14 @@ public static AcquireTokenForManagedIdentityParameterBuilder WithAttestationSupp
                 throw new ArgumentNullException(nameof(builder));
             }
 
-            // Set the attestation provider delegate
-            builder.CommonParameters.AttestationTokenProvider = async (endpoint, keyHandle, clientId, ct) =>
+            builder.CommonParameters.AttestationTokenProvider = async (endpoint, keyHandle, clientId, keyId, logger, ct) =>
             {
                 var result = await PopKeyAttestor.AttestCredentialGuardAsync(
                     endpoint,
                     keyHandle,
                     clientId,
+                    keyId,
+                    logger,
                     ct).ConfigureAwait(false);
 
                 // Return JWT on success, null for non-attested flow on failure

src/client/Microsoft.Identity.Client.KeyAttestation/PopKeyAttestor.cs

@@ -14,6 +14,7 @@
 using Microsoft.Identity.Client.Extensibility;
 using Microsoft.Identity.Client.Internal;
 using Microsoft.Identity.Client.Internal.ClientCredential;
+using Microsoft.Identity.Client.Core;
 using Microsoft.Identity.Client.ManagedIdentity;
 using Microsoft.Identity.Client.TelemetryCore.Internal.Events;
 using Microsoft.Identity.Client.Utils;
@@ -47,8 +48,9 @@ internal class AcquireTokenCommonParameters
         /// Optional delegate for obtaining attestation JWT for Credential Guard keys.
         /// Set by the KeyAttestation package via .WithAttestationSupport().
         /// Returns null for non-attested flows.
+        /// Signature: (endpoint, keyHandle, clientId, keyId, logger, cancellationToken) → JWT or null.
         /// </summary>
-        public Func<string, SafeHandle, string, CancellationToken, Task<string>> AttestationTokenProvider { get; set; }
+        public Func<string, SafeHandle, string, string, ILoggerAdapter, CancellationToken, Task<string>> AttestationTokenProvider { get; set; }
 
         /// <summary>
         /// This tries to see if the token request should be done over mTLS or over normal HTTP 

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

@@ -31,8 +31,9 @@ internal class AcquireTokenForManagedIdentityParameters : IAcquireTokenParameter
         /// <summary>
         /// Optional delegate for obtaining attestation JWT for Credential Guard keys.
         /// Set by the KeyAttestation package via .WithAttestationSupport().
+        /// Signature: (endpoint, keyHandle, clientId, keyId, logger, cancellationToken) → JWT or null.
         /// </summary>
-        public Func<string, SafeHandle, string, CancellationToken, Task<string>> AttestationTokenProvider { get; set; }
+        public Func<string, SafeHandle, string, string, ILoggerAdapter, CancellationToken, Task<string>> AttestationTokenProvider { get; set; }
 
         public void LogParameters(ILoggerAdapter logger)
         {

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

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
@@ -28,6 +29,22 @@ public abstract class ApplicationBase : IApplicationBase
         /// </summary>
         internal const string DefaultAuthority = "https://login.microsoftonline.com/common/";
 
+        // Allows extension packages (e.g. KeyAttestation) to register their own static-cache reset logic
+        // without introducing a circular dependency. Invoked by ResetStateForTest().
+        private static readonly ConcurrentBag<Action> s_resetCallbacks = new ConcurrentBag<Action>();
+
+        /// <summary>
+        /// Registers a callback to be invoked by <see cref="ResetStateForTest"/>.
+        /// Intended for extension packages (e.g. KeyAttestation) that own static caches MSAL cannot reference directly.
+        /// </summary>
+        internal static void RegisterResetCallback(Action callback)
+        {
+            if (callback is null)
+                throw new ArgumentNullException(nameof(callback));
+
+            s_resetCallbacks.Add(callback);
+        }
+
         internal IServiceBundle ServiceBundle { get; }
 
         internal ApplicationBase(ApplicationConfiguration config)
@@ -101,6 +118,28 @@ public static void ResetStateForTest()
 
             InMemoryPartitionedAppTokenCacheAccessor.ClearStaticCacheForTest();
             InMemoryPartitionedUserTokenCacheAccessor.ClearStaticCacheForTest();
+
+            List<Exception> callbackExceptions = null;
+            foreach (var cb in s_resetCallbacks)
+            {
+                try
+                {
+                    cb();
+                }
+                catch (Exception ex)
+                {
+                    callbackExceptions ??= new List<Exception>();
+                    callbackExceptions.Add(new InvalidOperationException(
+                        "A registered reset callback threw during ResetStateForTest().", ex));
+                }
+            }
+
+            if (callbackExceptions != null)
+            {
+                throw new AggregateException(
+                    "One or more registered reset callbacks failed during ResetStateForTest().",
+                    callbackExceptions);
+            }
         }
     }
 }

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

@@ -7,6 +7,7 @@
 using System.Net;
 using System.Net.Http;
 using System.Runtime.InteropServices;
+using System.Security.Cryptography;
 using System.Security.Cryptography.X509Certificates;
 using System.Threading;
 using System.Threading.Tasks;
@@ -27,7 +28,7 @@ internal class ImdsV2ManagedIdentitySource : AbstractManagedIdentity
         internal static readonly ICertificateCache s_mtlsCertificateCache = new InMemoryCertificateCache();
 
         private readonly IMtlsCertificateCache _mtlsCache;
-        private Func<string, SafeHandle, string, CancellationToken, Task<string>> _attestationTokenProvider;
+        private Func<string, SafeHandle, string, string, ILoggerAdapter, CancellationToken, Task<string>> _attestationTokenProvider;
 
         // used in unit tests
         public const string ApiVersionQueryParam = "cred-api-version";
@@ -470,11 +471,18 @@ private async Task<string> GetAttestationJwtAsync(
 
             try
             {
-                // Call the attestation token provider delegate
+                // Call the attestation token provider delegate.
+                // Prefer the CNG key name (stable for persisted/KSP keys).
+                // For ephemeral keys (no name), derive a stable identifier from the public key
+                // fingerprint so that the same key handle maps to the same cache entry while
+                // distinct ephemeral keys get distinct entries.
+                string keyId = rsaCng.Key.KeyName ?? GetPublicKeyFingerprint(rsaCng);
                 string attestationJwt = await _attestationTokenProvider(
                     attestationEndpoint.AbsoluteUri,
                     rsaCng.Key.Handle,
                     clientId,
+                    keyId,
+                    _requestContext.Logger,
                     cancellationToken).ConfigureAwait(false);
 
                 if (string.IsNullOrWhiteSpace(attestationJwt))
@@ -521,5 +529,22 @@ internal static void ResetCertCacheForTest()
                 s_mtlsCertificateCache.Clear();
             }
         }
+
+        /// <summary>
+        /// Computes a stable hex fingerprint of the RSA public key.
+        /// Used as a cache key for ephemeral CNG keys that have no key name.
+        /// Compatible with .NET Framework 4.6.2 and netstandard2.0.
+        /// </summary>
+        private static string GetPublicKeyFingerprint(RSA rsa)
+        {
+            RSAParameters p = rsa.ExportParameters(includePrivateParameters: false);
+            // Concatenate Modulus + Exponent as a stable, unique representation of the public key.
+            byte[] combined = new byte[p.Modulus.Length + p.Exponent.Length];
+            Buffer.BlockCopy(p.Modulus, 0, combined, 0, p.Modulus.Length);
+            Buffer.BlockCopy(p.Exponent, 0, combined, p.Modulus.Length, p.Exponent.Length);
+            using var sha256 = SHA256.Create();
+            byte[] hash = sha256.ComputeHash(combined);
+            return BitConverter.ToString(hash).Replace("-", string.Empty);
+        }
     }
 }