MSAL cache partition

cache_partition.cs

using Microsoft.Extensions.Caching.Memory;
using Microsoft.Extensions.Options;
using Microsoft.Identity.Client; // Microsoft.Identity.Client nuget package
using Microsoft.Identity.Web.TokenCacheProviders.InMemory; // Microsoft.Identity.Web.TokenCache nuget package

internal class Program
{
    private static async Task Main(string[] args)
    {
        TokenAcquirer tokenAcquirer = new TokenAcquirer();
        AuthenticationResult result;

        // All of these should come from the IdentityProvider (check TokenSource property)
        result = await tokenAcquirer.GetTokensAssociatedWithKey(cachePartition: "foo");
        result = await tokenAcquirer.GetTokensAssociatedWithKey(cachePartition: "bar");
        result = await tokenAcquirer.GetTokensAssociatedWithKey(cachePartition: null);

        // Now the same requests should come from the cache
        result = await tokenAcquirer.GetTokensAssociatedWithKey(cachePartition: "foo");
        result = await tokenAcquirer.GetTokensAssociatedWithKey(cachePartition: "bar");
        result = await tokenAcquirer.GetTokensAssociatedWithKey(cachePartition: null);
    }
}

public class TokenAcquirer
{
    private class PartitionedMsalTokenMemoryCacheProvider : MsalMemoryTokenCacheProvider // from Microsoft.Identity.Web.TokenCache nuget pacakge
    {
        private readonly string? _cacheKeySuffix;

        /// <summary>
        /// Ctor
        /// </summary>
        /// <param name="memoryCache">A memory cache which can be configured for max size etc.</param>
        /// <param name="cacheOptions">Additional cache options, which canbe ignored for app tokens.</param>
        /// <param name="cachePartition">An aditional partition key. If let null, the original cache scoping is used (clientID, tenantID). MSAL also looks for resource.</param>
        public PartitionedMsalTokenMemoryCacheProvider(
            IMemoryCache memoryCache,
            IOptions<MsalMemoryTokenCacheOptions> cacheOptions,
            string? cachePartition) : base(memoryCache, cacheOptions)
        {
            _cacheKeySuffix = cachePartition;
        }

        public override string GetSuggestedCacheKey(TokenCacheNotificationArgs args)
        {
            return base.GetSuggestedCacheKey(args) + (_cacheKeySuffix ?? "");
        }
    }

    private const int TokenCacheMemoryLimitInMb = 100;
    private static MemoryCache s_memoryCache = InitiatlizeMemoryCache(); // It's important that this remains static / shared

    private static MemoryCache InitiatlizeMemoryCache()
    {
        // For 100 MB limit ... ~2KB per token entry means 50,000 entries
        var options = Options.Create(new MemoryCacheOptions() { SizeLimit = (TokenCacheMemoryLimitInMb / 2) * 1000 });
        var cache = new MemoryCache(options);

        return cache;
    }

    /// <summary>
    /// Token cache for MSAL based on MemoryCache, which can be partitioned by an additional key.
    /// For app tokens, the default key is ClientID + TenantID (and MSAL also looks for resource).
    /// </summary>

    public async Task<AuthenticationResult> GetTokensAssociatedWithKey(string? cachePartition)
    {
        var confidentialApp = ConfidentialClientApplicationBuilder
                            .Create("ClientId")
                            .WithClientSecret("") // or assertion or certificate
                            .Build();

        var msalMemoryTokenCacheProvider =
            new PartitionedMsalTokenMemoryCacheProvider(
                s_memoryCache,
                Options.Create(new MsalMemoryTokenCacheOptions()),
                cachePartition: cachePartition);

        await msalMemoryTokenCacheProvider.InitializeAsync(confidentialApp.AppTokenCache).ConfigureAwait(false);

        AuthenticationResult result = await confidentialApp
            .AcquireTokenForClient(["https://graph.microsoft.com/.default"])
            // with extra parameter 
            .ExecuteAsync()
            .ConfigureAwait(false);

        // Monitor the token source!
        Console.WriteLine($"Token From: {result.AuthenticationResultMetadata.TokenSource}");

        return result;
    }
}