Adding support for MSAL telemetry cache data. (#2256)

* Adding support for MSAL telemetry cache data.

* Refactoring telemetryData to remove it from public api

* Adding more tests
adding implementation for in memory cache

* Test update

* Resolving test static cache issues

* Resolving release build issue

* Resolving test issues

* Updating instance discovery intercept logic

* Clean up

* Removing unneeded use statements
resolving nullable warnings

---------

Co-authored-by: trwalke <trwalke@microsoft.com>
Co-authored-by: Bogdan Gavril <bogavril@microsoft.com>

src/Microsoft.Identity.Web.TokenCache/CacheSerializerHints.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Threading;
+using Microsoft.Identity.Client.TelemetryCore.TelemetryClient;
 
 namespace Microsoft.Identity.Web.TokenCacheProviders
 {
@@ -21,5 +22,10 @@ public class CacheSerializerHints
         /// with the app token cache.
         /// </summary>
         public DateTimeOffset? SuggestedCacheExpiry { get; set; }
+
+        /// <summary>
+        /// Stores details to log to MSAL's telemetry client
+        /// </summary>
+        internal TelemetryData? TelemetryData { get; set; }
     }
 }

src/Microsoft.Identity.Web.TokenCache/Distributed/MsalDistributedTokenCacheAdapter.cs

@@ -147,6 +147,7 @@ await L2OperationWithRetryOnFailureAsync(
         {
             const string read = "Read";
             byte[]? result = null;
+            var telemetryData = cacheSerializerHints.TelemetryData;
 
             if (_memoryCache != null)
             {
@@ -169,6 +170,11 @@ await L2OperationWithRetryOnFailureAsync(
                 }, cacheSerializerHints.CancellationToken).Measure().ConfigureAwait(false);
 #pragma warning restore CA1062 // Validate arguments of public methods
 
+                if (result != null && telemetryData != null)
+                {
+                    telemetryData.CacheLevel = Client.Cache.CacheLevel.L2Cache;
+                }
+
                 Logger.DistributedCacheReadTime(_logger, _distributedCacheType, read, measure.MilliSeconds);
 
                 if (_memoryCache != null)
@@ -195,6 +201,11 @@ await L2OperationWithRetryOnFailureAsync(
                        (cacheKey) => _distributedCache.RefreshAsync(cacheKey, cacheSerializerHints.CancellationToken),
                        cacheKey,
                        result!).ConfigureAwait(false);
+
+                if (telemetryData != null)
+                {
+                    telemetryData.CacheLevel = Client.Cache.CacheLevel.L1Cache;
+                }
             }
 
 #pragma warning disable CS8603 // Possible null reference return.

src/Microsoft.Identity.Web.TokenCache/InMemory/MsalMemoryTokenCacheProvider.cs

@@ -62,6 +62,24 @@ protected override Task RemoveKeyAsync(string cacheKey)
             return Task.FromResult(tokenCacheBytes);
         }
 
+        /// <summary>
+        /// Method to be overridden by concrete cache serializers to Read the cache bytes.
+        /// </summary>
+        /// <param name="cacheKey">Cache key.</param>
+        /// <param name="cacheSerializerHints">Hints for the cache serialization implementation optimization.</param>
+        /// <returns>Read bytes.</returns>
+        protected override Task<byte[]?> ReadCacheBytesAsync(string cacheKey, CacheSerializerHints cacheSerializerHints)
+        {
+            byte[]? tokenCacheBytes = (byte[]?)_memoryCache.Get(cacheKey);
+
+            if (tokenCacheBytes != null && cacheSerializerHints.TelemetryData != null)
+            {
+                cacheSerializerHints.TelemetryData.CacheLevel = Client.Cache.CacheLevel.L1Cache;
+            }
+
+            return Task.FromResult(tokenCacheBytes);
+        }
+
         /// <summary>
         /// Writes a token cache blob to the serialization cache (identified by its key).
         /// </summary>

src/Microsoft.Identity.Web.TokenCache/MsalAbstractTokenCacheProvider.cs

@@ -101,7 +101,7 @@ private byte[] ProtectBytes(byte[] msalBytes)
         }
 
         private static CacheSerializerHints CreateHintsFromArgs(TokenCacheNotificationArgs args) => new CacheSerializerHints 
-        { CancellationToken = args.CancellationToken, SuggestedCacheExpiry = args.SuggestedCacheExpiry };
+        { CancellationToken = args.CancellationToken, SuggestedCacheExpiry = args.SuggestedCacheExpiry, TelemetryData = args.TelemetryData };
 
         private async Task OnBeforeAccessAsync(TokenCacheNotificationArgs args)
         {

src/Microsoft.Identity.Web.TokenCache/Properties/InternalsVisibleTo.cs

@@ -5,3 +5,4 @@
 
 // Allow this assembly to be serviced when run on desktop CLR
 [assembly: InternalsVisibleTo("Microsoft.Identity.Web.Test,          PublicKey=00240000048000009400000006020000002400005253413100040000010001002D96616729B54F6D013D71559A017F50AA4861487226C523959D1579B93F3FDF71C08B980FD3130062B03D3DE115C4B84E7AC46AEF5E192A40E7457D5F3A08F66CEAB71143807F2C3CB0DA5E23B38F0559769978406F6E5D30CEADD7985FC73A5A609A8B74A1DF0A29399074A003A226C943D480FEC96DBEC7106A87896539AD")]
+[assembly: InternalsVisibleTo("Microsoft.Identity.Web.Test.Common,          PublicKey=00240000048000009400000006020000002400005253413100040000010001002D96616729B54F6D013D71559A017F50AA4861487226C523959D1579B93F3FDF71C08B980FD3130062B03D3DE115C4B84E7AC46AEF5E192A40E7457D5F3A08F66CEAB71143807F2C3CB0DA5E23B38F0559769978406F6E5D30CEADD7985FC73A5A609A8B74A1DF0A29399074A003A226C943D480FEC96DBEC7106A87896539AD")]

src/Microsoft.Identity.Web.TokenCache/TokenCacheExtensions.cs

@@ -19,7 +19,8 @@ namespace Microsoft.Identity.Web
     /// </summary>
     public static class TokenCacheExtensions
     {
-        private static readonly ConcurrentDictionary<MethodInfo, IServiceProvider> s_serviceProviderFromAction
+        //internal for testing only
+        internal static readonly ConcurrentDictionary<MethodInfo, IServiceProvider> s_serviceProviderFromAction
             = new ConcurrentDictionary<MethodInfo, IServiceProvider>();
 
         /// <summary>

tests/Microsoft.Identity.Web.Test.Common/Mocks/MockHttpMessageHandler.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Net;
 using System.Net.Http;
 using System.Threading;
 using System.Threading.Tasks;
@@ -11,6 +12,8 @@ namespace Microsoft.Identity.Web.Test.Common.Mocks
 {
     public class MockHttpMessageHandler : HttpMessageHandler
     {
+        public Func<MockHttpMessageHandler, MockHttpMessageHandler> ReplaceMockHttpMessageHandler;
+
 #pragma warning disable CS8618 // Non-nullable field must contain a non-null value when exiting constructor. Consider declaring as nullable.
         public MockHttpMessageHandler()
 #pragma warning restore CS8618 // Non-nullable field must contain a non-null value when exiting constructor. Consider declaring as nullable.
@@ -41,6 +44,25 @@ protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage reques
 
             var uri = request.RequestUri;
             Assert.NotNull(uri);
+
+            //Intercept instance discovery requests and serve a response. 
+            //Also, requeue the current mock handler for MSAL's next request.
+#if NET6_0_OR_GREATER
+            if (uri.AbsoluteUri.Contains("/discovery/instance", StringComparison.OrdinalIgnoreCase))
+#else
+            if (uri.AbsoluteUri.Contains("/discovery/instance"))
+#endif
+            {
+                ReplaceMockHttpMessageHandler(this);
+
+                var responseMessage = new HttpResponseMessage(HttpStatusCode.OK)
+                {
+                    Content = new StringContent(TestConstants.DiscoveryJsonResponse),
+                };
+
+                return Task.FromResult(responseMessage);
+            }
+
             if (!string.IsNullOrEmpty(ExpectedUrl))
             {
                 Assert.Equal(
@@ -59,7 +81,7 @@ protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage reques
                 string postData = request.Content.ReadAsStringAsync().Result;
             }
 
-            return new TaskFactory().StartNew(() => ResponseMessage, cancellationToken);
+            return Task.FromResult(ResponseMessage);
         }
     }
 }

tests/Microsoft.Identity.Web.Test.Common/TestHelpers/TestMsalDistributedTokenCacheAdapter.cs

@@ -6,6 +6,7 @@
 using Microsoft.Extensions.Caching.Distributed;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
+using Microsoft.Identity.Client.TelemetryCore.TelemetryClient;
 using Microsoft.Identity.Web.TokenCacheProviders;
 using Microsoft.Identity.Web.TokenCacheProviders.Distributed;
 
@@ -42,5 +43,10 @@ public async Task TestWriteCacheBytesAsync(string cacheKey, byte[] bytes, CacheS
         {
             return await ReadCacheBytesAsync(cacheKey).ConfigureAwait(false);
         }
+
+        public async Task<byte[]?> TestReadCacheBytesAsync(string cacheKey, TelemetryData telemetryData)
+        {
+            return await ReadCacheBytesAsync(cacheKey, new CacheSerializerHints() { TelemetryData = telemetryData }).ConfigureAwait(false);
+        }
     }
 }