Add HybridCache for IMDSv2 Attestation Token Caching #5539

Robbie-Microsoft · 2025-10-21T21:08:53Z

Summary

Implements a new hybrid caching strategy for IMDSv2 attestation tokens with both in-memory and file-based persistence.

Changes

Added IHybridCache.cs - Interface for hybrid cache operations
Added HybridCache.cs - Hybrid cache implementation with:
- In-memory cache for fast access within process
- File-based cache for persistence across app restarts
- Cross-process synchronization using named OS mutex
- Automatic expired token cleanup
- Atomic file operations to prevent corruption
Modified ImdsV2ManagedIdentitySource.cs - Updated to use HybridCache to cache MAA tokens

Benefits

Improved performance with dual-layer caching
Cross-process token sharing
Robust error handling with graceful fallback
Comprehensive logging for diagnostics

TODO: Testing

Comprehensive unit tests covering all cache scenarios
Thread safety and concurrency validation
File persistence and cross-instance testing

…store MAA tokens

bgavrilMS · 2025-10-22T09:28:33Z

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

+    /// For persistent implementations, cross-process synchronization should be considered.
+    /// Hybrid implementations should synchronize between memory and persistent storage when possible.
+    /// </remarks>
+    internal interface IHybridCache


I am not sure about the name, as it conflicts with https://learn.microsoft.com/en-us/aspnet/core/performance/caching/hybrid?view=aspnetcore-9.0

bgavrilMS · 2025-10-22T09:44:24Z

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

+        /// <exception cref="OperationCanceledException">
+        /// Thrown when the operation is canceled via the cancellation token.
+        /// </exception>
+        Task RemoveAsync(long key, CancellationToken ct);


nit: I believe it's customary to return Task to indicate if an entry was actually removed.

bgavrilMS · 2025-10-22T09:45:01Z

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

+        /// <exception cref="OperationCanceledException">
+        /// Thrown when the operation is canceled via the cancellation token.
+        /// </exception>
+        Task SetAsync(long key, string token, DateTimeOffset expiresOnUtc, CancellationToken ct);


Why don't you set a AttestationTokenResponse ?

bgavrilMS · 2025-10-22T10:01:12Z

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

+        /// 
+        /// This entry structure is used for both in-memory and file cache storage.
+        /// </remarks>
+        private class CacheEntry


Why do you need this object when you have AttestationTokenResponse ?

bgavrilMS · 2025-10-22T11:34:59Z

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

+                    rsaCng.Key.Handle != null &&
+                    !rsaCng.Key.Handle.IsInvalid)
+                {
+                    return rsaCng.Key.Handle.DangerousGetHandle().ToInt64();


I think you should use a key thumbprint instead. When I asked copilot it said that "The handle does not represent the key’s content, just its memory address in the current process."

A key thumbprint (e.g., SHA-256 hash of the public key) would be a stable, cryptographically unique identifier for the key material, regardless of process or machine. This is a standard approach for caching or identifying cryptographic keys.

rsaCng.Key.Export(CngKeyBlobFormat.GenericPublicBlob))
Hash the blob (e.g., using SHA-256)
Use the hash as the cache key.

Thoughts?

We can also put the thumbprint in the ManagedIdentityKeyInfo (via the ctor) so as to not compute it everytime.

gladjohn · 2025-10-22T13:36:15Z

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

+
+            // Create named mutex for cross-process synchronization
+            // Using a deterministic name based on cache file path to ensure same mutex across processes
+            _mutexName = $"Global\\MSAL_AttestationCache_{_cacheFilePath.GetHashCode():X8}";


not sure if this will work, as different processes may compute different mutex names?

gladjohn · 2025-10-22T13:38:21Z

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

+
+            _requestContext.Logger.Verbose(() => $"[ImdsV2] GetAttestationJwtAsync called for cache key: {cacheKey}");
+
+            var cache = new HybridCache(_requestContext.Logger);


HybridCache implements IDisposable but is not disposed? should this be wrapped in using for short‑lived use?

gladjohn · 2025-10-22T13:41:47Z

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

+            }
+            catch (Exception ex)
+            {
+                _logger.Warning($"[HybridCache] Error loading file cache from {_cacheFilePath}: {ex.Message}. Created a new file.");


but the new file isn't created here?

gladjohn · 2025-10-22T13:43:42Z

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

+
+                _logger.Verbose(() => $"[HybridCache] Loading cache from file: {_cacheFilePath}");
+
+                var json = await Task.Run(() => File.ReadAllText(_cacheFilePath)).ConfigureAwait(false);


can we not use ReadAllTextAsync?

gladjohn · 2025-10-22T13:44:09Z

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

+
+                // Write to temp file first, then move to avoid corruption during write
+                var tempFile = _cacheFilePath + ".tmp";
+                await Task.Run(() => File.WriteAllText(tempFile, json)).ConfigureAwait(false);


WriteAllTextAsync?

gladjohn · 2025-10-22T13:45:13Z

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

+                {
+                    if (!string.IsNullOrEmpty(memoryEntry.Token) && now + s_expirySkew < memoryEntry.ExpiresOnUtc)
+                    {
+                        _logger.Info(() => $"[HybridCache] Cache hit in memory for key: {key}");


gladjohn · 2025-10-22T13:45:18Z

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

+
+                    // Token expired in memory, remove it
+                    s_memoryCache.TryRemove(keyString, out _);
+                    _logger.Info(() => $"[HybridCache] Expired token removed from memory cache for key: {key}");


gladjohn · 2025-10-22T13:46:37Z

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

+                throw new ArgumentNullException(nameof(token));
+            }
+
+            _logger.Info(() => $"[HybridCache] SetAsync called for key: {key}, expires: {expiresOnUtc}");


gladjohn · 2025-10-22T13:46:50Z

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

+                throw new ObjectDisposedException(nameof(HybridCache));
+            }
+
+            _logger.Info(() => $"[HybridCache] RemoveAsync called for key: {key}");


gladjohn · 2025-10-22T13:46:57Z

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

+                    return Task.CompletedTask;
+                }, ct).ConfigureAwait(false);
+
+                _logger.Info(() => $"[HybridCache] Remove operation completed for key: {key}");


gladjohn · 2025-10-22T13:47:20Z

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

+                    {
+                        cache.Remove(key);
+                    }
+                    _logger.Info(() => $"[HybridCache] Cleaned up {keysToRemove.Count} expired entries from file cache");


gladjohn · 2025-10-22T13:58:18Z

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

+
+            _logger.Info(() => $"[HybridCache] Initializing cache with directory: {cacheDirectory}");
+
+            Directory.CreateDirectory(cacheDirectory);


can this operation fail?

Created interface and implementation for HybridCache, and used it to …

9ea94d3

…store MAA tokens

Robbie-Microsoft changed the title ~~Add HybridCache for IMDSv2 attestation token caching~~ Add HybridCache for IMDSv2 Attestation Token Caching Oct 21, 2025

bgavrilMS reviewed Oct 22, 2025

View reviewed changes

gladjohn reviewed Oct 22, 2025

View reviewed changes


		_requestContext.Logger.Verbose(() => $"[ImdsV2] GetAttestationJwtAsync called for cache key: {cacheKey}");

		var cache = new HybridCache(_requestContext.Logger);


		_logger.Verbose(() => $"[HybridCache] Loading cache from file: {_cacheFilePath}");

		var json = await Task.Run(() => File.ReadAllText(_cacheFilePath)).ConfigureAwait(false);


		_logger.Info(() => $"[HybridCache] Initializing cache with directory: {cacheDirectory}");

		Directory.CreateDirectory(cacheDirectory);

Add HybridCache for IMDSv2 Attestation Token Caching #5539

Are you sure you want to change the base?

Add HybridCache for IMDSv2 Attestation Token Caching #5539

Conversation

Robbie-Microsoft commented Oct 21, 2025

Summary

Changes

Benefits

TODO: Testing

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants