Add in-process MAA token caching to PopKeyAttestor#5887
Merged
Robbie-Microsoft merged 19 commits intomainfrom Apr 28, 2026
Merged
Add in-process MAA token caching to PopKeyAttestor#5887Robbie-Microsoft merged 19 commits intomainfrom
Robbie-Microsoft merged 19 commits intomainfrom
Conversation
Cache MAA (Microsoft Azure Attestation) JWT tokens in-process to
avoid redundant native DLL calls and network round-trips to the MAA
endpoint on every mTLS cert mint.
Changes:
- PopKeyAttestor: add ConcurrentDictionary<string, AttestationToken>
keyed by '{maaEndpoint}|{clientId}'. MAA cache is checked before any
native or network call. On MAA cache hit the key handle is not
required. MAA tokens within 5 min of expiry (matching MSAL's
AccessTokenExpirationBuffer) are considered stale and trigger a fresh
MAA attestation call. Key validation is deferred to cache-miss path.
Add ResetCacheForTest() and overridable s_expirationBuffer for tests.
- ImdsV2Tests: call ResetCacheForTest() in TestCleanup. Add three tests:
- MaaTokenCache_Hit_DoesNotCallAttestationProviderAgain
- MaaTokenCache_ExpiredToken_CallsAttestationProviderAgain
- MaaTokenCache_DifferentEndpoints_CachedSeparately
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds in-process caching of Microsoft Azure Attestation (MAA) tokens in PopKeyAttestor to avoid repeated native DLL calls and network round-trips during mTLS certificate minting, and introduces unit tests to validate cache behavior.
Changes:
- Added a static in-memory token cache keyed by
{maaEndpoint}|{clientId}with an expiration buffer. - Cached-token fast path now returns without requiring a key handle; cache-miss path performs key validation and attestation.
- Added/updated
ImdsV2Teststo reset the cache and validate cache hit/expiry/separate endpoints behavior.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsV2Tests.cs | Adds cleanup to reset the new cache and introduces tests for cache hit/expiry/isolation by endpoint. |
| src/client/Microsoft.Identity.Client.KeyAttestation/PopKeyAttestor.cs | Implements token caching, cache lookup before native/network call, and test hooks to reset cache/change expiration buffer. |
- Add per-key SemaphoreSlim (s_keyLocks) to prevent concurrent in-flight attestation calls for the same cache key (thundering herd protection) - Double-check cache after acquiring lock so waiters skip the network call - Evict stale entries in TryGetCachedToken via TryRemove to prevent unbounded memory growth - Replace null with string.Empty in cache-hit AttestationResult message for consistency and null-safety - ResetCacheForTest() now also clears s_keyLocks - Add MaaTokenCache_ConcurrentCacheMiss_SingleFlightCallsProviderOnce test Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
gladjohn
reviewed
Mar 26, 2026
gladjohn
reviewed
Mar 26, 2026
Contributor
gladjohn
left a comment
gladjohn
left a comment
There was a problem hiding this comment.
Thanks Robbie, looks good overall. Are we able to use any of the cache helper pattern from the cert cache here?
bgavrilMS
reviewed
Mar 30, 2026
bgavrilMS
reviewed
Mar 30, 2026
bgavrilMS
reviewed
Mar 30, 2026
bgavrilMS
reviewed
Mar 30, 2026
…he key comment, new test - Move keyHandle validation before cache check (Bogdan) - Remove redundant ThrowIfCancellationRequested (Bogdan) - Add ILoggerAdapter param; log cache hit/miss/store (Bogdan + gladjohn) - Register PopKeyAttestor.ResetCacheForTest with ApplicationBase.ResetStateForTest via callback (gladjohn) - Add cache key design comment explaining why Key.ID is not included (Bogdan + gladjohn) - Update AttestationTokenProvider delegate type to thread logger from request context - Update MaaTokenCache_DifferentEndpoints test: always pass valid handles (no null on cache hit) - Add MaaTokenCache_CertCacheMiss_MaaTokenCacheStillFresh test (gladjohn) - Update XML doc comment on PopKeyAttestor class (gladjohn) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
Author
The cert cache uses a two-tier (in-memory + persistent) design with a full interface hierarchy (ICertificateCache, IMtlsCertificateCache, etc.) because certs need to survive process restarts and are encrypted at rest. The MAA token cache is intentionally simpler — in-memory only, with a ConcurrentDictionary + per-key semaphore for single-flight. Reusing the cert cache pattern would bring in persistence, encryption, and MtlsBindingCache orchestration that aren't warranted for a short-lived JWT. If the requirements change (e.g. persist across reboots), we can adopt that pattern then. |
…erflow, test fixes - Add ArgumentNullException guard to ApplicationBase.RegisterResetCallback - Replace hand-rolled ConcurrentDictionary<string,SemaphoreSlim> with KeyedSemaphorePool - Fix TryGetCachedToken: flip comparison to avoid DateTimeOffset.MinValue - buffer overflow - MaaTokenCache_Hit: add WithForceRefresh + mocks so second acquire exercises cert/MAA path - MaaTokenCache_Concurrent: release gate for all tasks.Length so regression fails fast instead of hanging - MaaTokenCache_CertCacheMiss: add TokenSource.IdentityProvider assertion to prove path was exercised Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…l is stateless) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Include the CNG key name (CngKey.KeyName) as a third component of the
MAA token cache key, changing the format from '{endpoint}|{clientId}'
to '{endpoint}|{clientId}|{keyId}'.
This prevents a stale MAA token for one CNG key from being returned
for a different key that shares the same endpoint+clientId (e.g. after
key rotation).
Changes:
- PopKeyAttestor: add optional keyId param to AttestCredentialGuardAsync,
update s_testAttestationProvider delegate type, update BuildCacheKey,
rewrite cache-key design comment
- AcquireTokenCommonParameters / AcquireTokenForManagedIdentityParameters:
update AttestationTokenProvider delegate type
- ImdsV2ManagedIdentitySource: extract rsaCng.Key.KeyName and forward
as keyId to the attestation delegate
- ManagedIdentityAttestationExtensions: lambda accepts and forwards keyId
- Test helpers (ManagedIdentityTestExtensions, TestAttestationProviders):
update all delegate signatures
- ImdsV2Tests: update existing lambdas and direct AttestCredentialGuardAsync
calls; add MaaTokenCache_DifferentKeyIds_CachedSeparately test
Addresses review feedback from bgavrilMS in PR #5887.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
gladjohn
approved these changes
Apr 8, 2026
…emaphore comment, test keyId fixes - Switch Info(string) to Info(Func<string>) in PopKeyAttestor to avoid eager string allocation - Add NormalizeEndpoint() helper that trims trailing slash before composing cache key - Add comment explaining KeyedSemaphorePool stays bounded (named keys only; ephemeral share one slot) - Fix MaaTokenCache_DifferentEndpoints_CachedSeparately: use const string keyIds (rsa.Key.KeyName is null for ephemeral keys) - Fix MaaTokenCache_ConcurrentCacheMiss_SingleFlightCallsProviderOnce: add explicit keyId argument 82/82 ImdsV2Tests pass on .NET 8. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
bgavrilMS
approved these changes
Apr 24, 2026
bgavrilMS
approved these changes
Apr 24, 2026
When a CNG key has no KeyName (ephemeral/non-KSP key), the previous code used a null keyId which collapsed all ephemeral keys at the same endpoint to a single cache slot. This could cause distinct key handles to share an attestation token minted for a different key — a correctness and security bug. Fix: - ImdsV2ManagedIdentitySource: derive keyId from SHA-256 of RSA public key (Modulus+Exponent) when KeyName is null. Each distinct key now gets its own MAA cache entry. - PopKeyAttestor: retain defensive bypass for any caller that still passes a null keyId. - ImdsV2Tests: share a single TestKeyGuardManagedIdentityKeyProvider across both acquires in the cert-cache-miss test, correctly simulating a persistent hardware key surviving a cert re-mint. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1. Update stale comments: keyId is no longer null/empty for ephemeral keys. ImdsV2ManagedIdentitySource always derives a non-empty keyId (KeyName or SHA-256 fingerprint), so the null/empty bypass is a defensive fallback, not the common ephemeral-key path. 2. Change cache hit/miss/store logging from Info to Verbose and mask the full cache key to avoid exposing endpoint URLs and key fingerprints in logs. Masking follows the same convention as InMemoryCertificateCache. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
gladjohn
approved these changes
Apr 27, 2026
- Rewrite MaaTokenCache_Hit_DoesNotCallAttestationProviderAgain to force a cert-cache miss before the second acquire (via ResetCertCacheForTest + ResetSourceForTest), ensuring the test actually exercises the MAA cache code path rather than the cert-reuse shortcut that bypasses attestation. A shared TestKeyGuardManagedIdentityKeyProvider instance ensures both acquires use the same CNG key fingerprint so the MAA cache key matches. - Add cancellation note to AttestCredentialGuardAsync XML doc clarifying that Task.Run is used for the native call and cancellation prevents starting the operation but cannot interrupt it once under way. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
… loop - Switch s_tokenCache from OrdinalIgnoreCase to Ordinal; update NormalizeEndpoint to lowercase the endpoint string so host comparisons remain case-insensitive while keeping keyId (CNG key name / SHA-256 fingerprint) case-sensitive. Matches InMemoryCertificateCache convention. - Wrap each callback in ResetStateForTest() in a try/catch so a single failing callback does not abort the rest of the reset sequence. Exceptions are collected and re-thrown as AggregateException after all callbacks have run. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
gladjohn
approved these changes
Apr 28, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
MAA (Microsoft Azure Attestation) tokens are acquired during mTLS binding certificate minting, requiring a native DLL call + network round-trip to the MAA endpoint. Since MAA tokens carry their
own JWT expiry and have a longer TTL than binding certs, they can safely be reused across cert re-mints.
This PR adds an in-memory cache for MAA tokens inside
PopKeyAttestorto avoid redundant native DLL and network calls.Changes
PopKeyAttestor.csConcurrentDictionary<string, AttestationToken>cache keyed by"{endpoint}|{keyId}"endpoint— different MAA regional endpoints have different trust domainskeyId—CngKey.KeyName; ensures key rotation does not return a stale token for the new keyclientIdintentionally omitted — MAA attests the CNG key, not the identity; SAMI and UAMI sharing the same key at the same endpoint correctly share the cached tokennull,IsInvalid,SafeNCryptKeyHandlecast) always runs before the cache checkAccessTokenExpirationBuffer); freshness check guards againstDateTimeOffset.MinValue - bufferoverflowKeyedSemaphorePoolprevents thundering-herd redundant attestation callsILoggerAdapterthreaded from request context throughAttestationTokenProviderdelegateResetCacheForTestwithApplicationBase.RegisterResetCallbacksoApplicationBase.ResetStateForTest()automatically clears the MAA cacheApplicationBase.csRegisterResetCallback(Action)+ConcurrentBag<Action> s_resetCallbacksResetStateForTest()invokes all registered callbacks after its own resetsAcquireTokenCommonParameters/AcquireTokenForManagedIdentityParametersAttestationTokenProviderdelegate updated to includeILoggerAdapterandstring keyIdImdsV2ManagedIdentitySource.csrsaCng.Key.KeyNameand forwards it askeyIdto the attestation delegateImdsV2Tests.cs[TestCleanup]callsPopKeyAttestor.ResetCacheForTest()(belt-and-suspenders alongside theApplicationBasecallback)MaaTokenCache_Hit_DoesNotCallAttestationProviderAgainMaaTokenCache_ExpiredToken_CallsAttestationProviderAgainMaaTokenCache_DifferentEndpoints_CachedSeparatelyMaaTokenCache_DifferentKeyIds_CachedSeparately— two different CNG keys at the same endpoint get separate cache entries (key rotation safety)MaaTokenCache_ConcurrentCacheMiss_SingleFlightCallsProviderOnceMaaTokenCache_CertCacheMiss_MaaTokenCacheStillFresh_DoesNotCallAttestationProviderAgainTest results
81/81
ImdsV2Testspassing on .NET 4.8 and .NET 8, no regressions.No public API changes
Caching is entirely internal to
PopKeyAttestor.WithAttestationSupport()is unchanged.