Ensure additional cache parameters are persisted in cache serialization (#5262)

* Ensure additional cache parameters are persisted in cache serialization

* Updating test attributes

---------

Co-authored-by: trwalke <[email protected]>

Author: trwalke

src/client/Microsoft.Identity.Client/Cache/Items/MsalAccessTokenCacheItem.cs

@@ -321,6 +321,7 @@ internal static MsalAccessTokenCacheItem FromJObject(JObject j)
             string tokenType = JsonHelper.ExtractExistingOrDefault<string>(j, StorageJsonKeys.TokenType) ?? StorageJsonValues.TokenTypeBearer;
             string scopes = JsonHelper.ExtractExistingOrEmptyString(j, StorageJsonKeys.Target);
             var additionalCacheKeyComponents = JsonHelper.ExtractInnerJsonAsDictionary(j, StorageJsonKeys.CacheExtensions);
+            var persistedCacheParameters = JsonHelper.ExtractInnerJsonAsDictionary(j, StorageJsonKeys.PersistedCacheParameters);
 
             var item = new MsalAccessTokenCacheItem(
                 scopes: scopes,
@@ -338,6 +339,7 @@ internal static MsalAccessTokenCacheItem FromJObject(JObject j)
                 item.CredentialType = StorageJsonValues.CredentialTypeAccessTokenExtended;
             }
 
+            item.PersistedCacheParameters = persistedCacheParameters;
             item.OboCacheKey = oboCacheKey;
             item.PopulateFieldsFromJObject(j);
 
@@ -367,22 +369,29 @@ internal override JObject ToJObject()
             // previous versions of MSAL used "ext_expires_on" instead of the correct "extended_expires_on".
             // this is here for back compatibility
             SetItemIfValueNotNull(json, StorageJsonKeys.ExtendedExpiresOn_MsalCompat, extExpiresUnixTimestamp);
-            if (AdditionalCacheKeyComponents != null)
+
+            StoreDictionaryInJson(json, StorageJsonKeys.CacheExtensions, AdditionalCacheKeyComponents);
+            StoreDictionaryInJson(json, StorageJsonKeys.PersistedCacheParameters, PersistedCacheParameters);
+            return json;
+        }
+
+        private void StoreDictionaryInJson(JObject json, string key, IDictionary<string, string> values)
+        {
+            if (values != null)
             {
 #if SUPPORTS_SYSTEM_TEXT_JSON
                 var obj = new JsonObject();
 
-                foreach (KeyValuePair<string, string> accId in AdditionalCacheKeyComponents)
+                foreach (KeyValuePair<string, string> value in values)
                 {
-                    obj[accId.Key] = accId.Value;
+                    obj[value.Key] = value.Value;
                 }
 
-                json[StorageJsonKeys.CacheExtensions] = obj;
+                json[key] = obj;
 #else
-                SetItemIfValueNotNull(json, StorageJsonKeys.CacheExtensions, JObject.FromObject(AdditionalCacheKeyComponents));
+                SetItemIfValueNotNull(json, key, JObject.FromObject(values));
 #endif
             }
-            return json;
         }
 
         internal string ToJsonString()

src/client/Microsoft.Identity.Client/Cache/StorageJsonKeys.cs

@@ -43,5 +43,6 @@ internal static class StorageJsonKeys
         public const string ExtendedExpiresOn_MsalCompat = "ext_expires_on";
 
         public const string CacheExtensions = "ext";
+        public const string PersistedCacheParameters = "persisted_cache_parameters";
     }
 }

tests/Microsoft.Identity.Test.Unit/AuthExtension/AuthenticationOperationTests.cs

@@ -18,6 +18,7 @@ namespace Microsoft.Identity.Test.Unit
     public class AuthenticationOperationTests : TestBase
     {
         private const string ProtectedUrl = "https://www.contoso.com/path1/path2?queryParam1=a&queryParam2=b";
+        private byte[] _serializedCache;
 
         [TestMethod]
         public async Task Should_UseCustomRequestHeaders_And_StoreAdditionalParameters()
@@ -66,8 +67,10 @@ public async Task Should_UseCustomRequestHeaders_And_StoreAdditionalParameters()
             }
         }
 
-        [TestMethod]
-        public async Task Should_UseCustomRequestHeaders_And_StoreAdditionalParametersWithCaching()
+        [DataTestMethod] // Fix for regression https://github.com/AzureAD/microsoft-authentication-library-for-dotnet/issues/5261
+        [DataRow(false)]
+        [DataRow(true)]
+        public async Task Should_UseCustomRequestHeaders_And_StoreAdditionalParametersWithCaching(bool useSerializedCache)
         {
             using (var httpManager = new MockHttpManager())
             {
@@ -100,6 +103,12 @@ public async Task Should_UseCustomRequestHeaders_And_StoreAdditionalParametersWi
                     AdditionalCacheParameters = new[] { "additional_param1", "additional_param2" }
                 };
 
+                if (useSerializedCache)
+                {
+                    app.AppTokenCache.SetBeforeAccess(BeforeCacheAccess);
+                    app.AppTokenCache.SetAfterAccess(AfterCacheAccess);
+                }
+
                 // Act
                 var result = await app.AcquireTokenForClient(TestConstants.s_scope.ToArray())
                     .WithTenantId(TestConstants.Utid)
@@ -130,6 +139,16 @@ public async Task Should_UseCustomRequestHeaders_And_StoreAdditionalParametersWi
             }
         }
 
+        private void BeforeCacheAccess(TokenCacheNotificationArgs args)
+        {
+            args.TokenCache.DeserializeMsalV3(_serializedCache);
+        }
+
+        private void AfterCacheAccess(TokenCacheNotificationArgs args)
+        {
+            _serializedCache = args.TokenCache.SerializeMsalV3();
+        }
+
         [TestMethod]
         public async Task Should_UseEmptyExtension_And_Parameters()
         {