Adding api to enable additional caching parameters

src/client/Microsoft.Identity.Client/ApiConfig/AbstractAcquireTokenParameterBuilder.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.ComponentModel;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.ApiConfig.Parameters;

src/client/Microsoft.Identity.Client/ApiConfig/Parameters/AcquireTokenCommonParameters.cs

@@ -29,6 +29,6 @@ internal class AcquireTokenCommonParameters
         public PoPAuthenticationConfiguration PopAuthenticationConfiguration { get; set; }
         public Func<OnBeforeTokenRequestData, Task> OnBeforeTokenRequestHandler { get; internal set; }
         public X509Certificate2 MtlsCertificate { get; internal set; }
-
+        public List<string> AdditionalCacheParameters { get; set; }
     }
 }

src/client/Microsoft.Identity.Client/AuthenticationResult.cs

@@ -166,7 +166,9 @@ internal AuthenticationResult(
             CorrelationId = correlationID;
             ApiEvent = apiEvent;
             AuthenticationResultMetadata = new AuthenticationResultMetadata(tokenSource);
-            AdditionalResponseParameters = additionalResponseParameters;
+            AdditionalResponseParameters = msalAccessTokenCacheItem.PersistedCacheParameters?.Count > 0 ?
+                                                                    (IReadOnlyDictionary<string, string>)msalAccessTokenCacheItem.PersistedCacheParameters : 
+                                                                    additionalResponseParameters;
             if (msalAccessTokenCacheItem != null)
             {
                 AccessToken = authenticationScheme.FormatAccessToken(msalAccessTokenCacheItem);

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

@@ -3,12 +3,14 @@
 
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using Microsoft.Identity.Client.AuthScheme;
 using Microsoft.Identity.Client.Cache.Keys;
 using Microsoft.Identity.Client.Internal;
 using Microsoft.Identity.Client.OAuth2;
 using Microsoft.Identity.Client.Utils;
 #if SUPPORTS_SYSTEM_TEXT_JSON
+using System.Text.Json;
 using JObject = System.Text.Json.Nodes.JsonObject;
 #else
 using Microsoft.Identity.Json.Linq;
@@ -28,7 +30,8 @@ internal MsalAccessTokenCacheItem(
             string tenantId,
             string homeAccountId,
             string keyId = null,
-            string oboCacheKey = null)
+            string oboCacheKey = null,
+            IEnumerable<string> additionalRequestRapameters = null)
             : this(
                 scopes: ScopeHelper.OrderScopesAlphabetically(response.Scope), // order scopes to avoid cache duplication. This is not in the hot path.
                 cachedAt: DateTimeOffset.UtcNow,
@@ -45,11 +48,26 @@ internal MsalAccessTokenCacheItem(
             RawClientInfo = response.ClientInfo;
             HomeAccountId = homeAccountId;
             OboCacheKey = oboCacheKey;
+            PersistedCacheParameters = AcquireCacheParametersFromResponse(additionalRequestRapameters, response.ExtensionData);
 
             InitCacheKey();
         }
 
-
+        private IDictionary<string, string> AcquireCacheParametersFromResponse(
+                                                    IEnumerable<string> additionalRequestRapameters,
+#if SUPPORTS_SYSTEM_TEXT_JSON
+                                                    IDictionary<string, JsonElement> extraDataFromResponse)
+#else
+                                                    IDictionary<string, JToken> extraDataFromResponse)
+#endif
+        {
+            var cacheParameters = extraDataFromResponse
+                                    .Where(x => additionalRequestRapameters.Contains(x.Key))
+                                    .ToDictionary(x => x.Key, x => x.Value.ToString());
+
+            return cacheParameters;
+        }
+
         internal /* for test */ MsalAccessTokenCacheItem(
             string preferredCacheEnv,
             string clientId,
@@ -215,6 +233,12 @@ internal string TenantId
 
         internal string CacheKey { get; private set; }
 
+        /// <summary>
+        /// Additional parameters that were requested in the token request and are stored in the cache.
+        /// These are acquired from the response and are stored in the cache for later use.
+        /// </summary>
+        internal IDictionary<string, string> PersistedCacheParameters { get; private set; }
+
         private Lazy<IiOSKey> iOSCacheKeyLazy;
         public IiOSKey iOSCacheKey => iOSCacheKeyLazy.Value;
 

src/client/Microsoft.Identity.Client/Extensibility/AbstractConfidentialClientAcquireTokenParameterBuilderExtension.cs

@@ -2,6 +2,8 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Collections.Generic;
+using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Threading.Tasks;
 
@@ -55,5 +57,37 @@ public static AbstractAcquireTokenParameterBuilder<T> WithProofOfPosessionKeyId<
 
             return builder;
         }
+
+        /// <summary>
+        /// Specifies additional parameters acquired from authentication responses to be cached with the access token that are normally not included in the cache object.
+        /// these values can be read from the <see cref="AuthenticationResult.AdditionalResponseParameters"/> parameter.
+        /// </summary>
+        /// <typeparam name="T"></typeparam>
+        /// <param name="builder">The builder to chain options to</param>
+        /// <param name="cacheParameters">Additional parameters to cache</param>
+        /// <returns></returns>
+        public static AbstractAcquireTokenParameterBuilder<T> WithAdditionalCacheParameters<T>(
+            this AbstractAcquireTokenParameterBuilder<T> builder, 
+            IEnumerable<string> cacheParameters)
+            where T : AbstractAcquireTokenParameterBuilder<T>
+        {
+            if (cacheParameters != null && cacheParameters.Count() == 0)
+            {
+                return builder;
+            }
+
+            builder.ValidateUseOfExperimentalFeature();
+
+            //Check if the cache parameters are already initialized, if so, add to the existing list
+            if (builder.CommonParameters.AdditionalCacheParameters != null)
+            {
+                builder.CommonParameters.AdditionalCacheParameters.AddRange(cacheParameters);
+            }
+            else
+            {
+                builder.CommonParameters.AdditionalCacheParameters = cacheParameters.ToList<string>();
+            }
+            return builder;
+        }
     }   
 }

src/client/Microsoft.Identity.Client/Internal/Requests/AuthenticationRequestParameters.cs

@@ -115,6 +115,8 @@ public string Claims
 
         public IAuthenticationScheme AuthenticationScheme => _commonParameters.AuthenticationScheme;
 
+        public IEnumerable<string> AdditionalCacheParameters => _commonParameters.AdditionalCacheParameters;
+
         #region TODO REMOVE FROM HERE AND USE FROM SPECIFIC REQUEST PARAMETERS
         // TODO: ideally, these can come from the particular request instance and not be in RequestBase since it's not valid for all requests.
 

src/client/Microsoft.Identity.Client/OAuth2/MsalTokenResponse.cs

@@ -97,6 +97,7 @@ public IReadOnlyDictionary<string, string> CreateExtensionDataStringMap()
                    item.Value.ValueKind == JsonValueKind.Number ||
                    item.Value.ValueKind == JsonValueKind.True ||
                    item.Value.ValueKind == JsonValueKind.False ||
+                   item.Value.ValueKind == JsonValueKind.Array ||
                    item.Value.ValueKind == JsonValueKind.Null)
                 {
                     stringExtensionData.Add(item.Key, item.Value.ToString());
@@ -113,6 +114,7 @@ public IReadOnlyDictionary<string, string> CreateExtensionDataStringMap()
                    item.Value.Type == JTokenType.Guid ||
                    item.Value.Type == JTokenType.Integer ||
                    item.Value.Type == JTokenType.TimeSpan ||
+                   item.Value.Type == JTokenType.Array ||
                    item.Value.Type == JTokenType.Null)
                 {
                     stringExtensionData.Add(item.Key, item.Value.ToString());

src/client/Microsoft.Identity.Client/TokenCache.ITokenCacheInternal.cs

@@ -79,7 +79,8 @@ async Task<Tuple<MsalAccessTokenCacheItem, MsalIdTokenCacheItem, Account>> IToke
                         tenantId,
                         homeAccountId,
                         requestParams.AuthenticationScheme.KeyId,
-                        CacheKeyFactory.GetOboKey(requestParams.LongRunningOboCacheKey, requestParams.UserAssertion));
+                        CacheKeyFactory.GetOboKey(requestParams.LongRunningOboCacheKey, requestParams.UserAssertion),
+                        requestParams.AdditionalCacheParameters);
             }
 
             if (!string.IsNullOrEmpty(response.RefreshToken))

tests/Microsoft.Identity.Test.Common/Core/Mocks/MockHelpers.cs

@@ -341,7 +341,18 @@ public static HttpResponseMessage CreateSuccessfulClientCredentialTokenResponseM
             string tokenType = "Bearer")
         {
             return CreateSuccessResponseMessage(
-                "{\"token_type\":\"" + tokenType + "\",\"expires_in\":\"" + expiry + "\",\"access_token\":\"" + token + "\"}");
+                "{\"token_type\":\"" + tokenType + "\",\"expires_in\":\"" + expiry + "\",\"access_token\":\"" + token + "\",\"additional_param1\":\"value1\",\"additional_param2\":\"value2\",\"additional_param3\":\"value3\"}");
+        }
+
+        public static HttpResponseMessage CreateSuccessfulClientCredentialTokenResponseWithAdditionalParamsMessage(
+            string token = "header.payload.signature",
+            string expiry = "3599",
+            string tokenType = "Bearer",
+            string additionalparams = ",\"additional_param1\":\"value1\",\"additional_param2\":\"value2\",\"additional_param3\":\"value3\",\"additional_param4\":[\"GUID\", \"GUID2\", \"GUID3\"]"
+            )
+        {
+            return CreateSuccessResponseMessage(
+                "{\"token_type\":\"" + tokenType + "\",\"expires_in\":\"" + expiry + "\",\"access_token\":\"" + token + "\"" + additionalparams + "}");
         }
 
         public static HttpResponseMessage CreateSuccessTokenResponseMessage(

tests/Microsoft.Identity.Test.Common/Core/Mocks/MockHttpManagerExtensions.cs

@@ -197,6 +197,26 @@ public static MockHttpMessageHandler AddMockHandlerSuccessfulClientCredentialTok
             return handler;
         }
 
+        public static MockHttpMessageHandler AddMockHandlerSuccessfulClientCredentialTokenResponseWithAdditionalParamsMessage(
+            this MockHttpManager httpManager,
+            string token = "header.payload.signature",
+            string expiresIn = "3599",
+            string tokenType = "Bearer",
+            IList<string> unexpectedHttpHeaders = null,
+            string additionalparams = ",\"additional_param1\":\"value1\",\"additional_param2\":\"value2\",\"additional_param3\":\"value3\",\"additional_param4\":[\"GUID\", \"GUID2\", \"GUID3\"]")
+        {
+            var handler = new MockHttpMessageHandler()
+            {
+                ExpectedMethod = HttpMethod.Post,
+                ResponseMessage = MockHelpers.CreateSuccessfulClientCredentialTokenResponseWithAdditionalParamsMessage(token, expiresIn, tokenType, additionalparams),
+                UnexpectedRequestHeaders = unexpectedHttpHeaders
+            };
+
+            httpManager.AddMockHandler(handler);
+
+            return handler;
+        }
+
         public static MockHttpMessageHandler AddMockHandlerForThrottledResponseMessage(
             this MockHttpManager httpManager)
         {

tests/Microsoft.Identity.Test.Unit/PublicApiTests/ExtensiblityTests.cs

@@ -218,6 +218,139 @@ public async Task ValidateAppTokenProviderAsync()
             }
         }
 
+        [TestMethod]
+        public async Task ValidateAdditionalCacheParametersAreStored()
+        {
+            using (var httpManager = new MockHttpManager())
+            {
+                httpManager.AddInstanceDiscoveryMockHandler();
+
+                var app = ConfidentialClientApplicationBuilder.Create(TestConstants.ClientId)
+                              .WithAuthority("https://login.microsoftonline.com/tid/")
+                              .WithClientSecret(TestConstants.ClientSecret)
+                              .WithExperimentalFeatures(true)
+                              .WithHttpManager(httpManager)
+                              .BuildConcrete();
+
+                httpManager.AddMockHandlerSuccessfulClientCredentialTokenResponseWithAdditionalParamsMessage();
+
+                var result = await app.AcquireTokenForClient(TestConstants.s_scope.ToArray())
+                    .WithAdditionalCacheParameters(new List<string> { "additional_param1", "additional_param2", "additional_param3", "additional_param4" })
+                    .ExecuteAsync()
+                    .ConfigureAwait(false);
+
+                var parameters = app.AppTokenCacheInternal.Accessor.GetAllAccessTokens().Single().PersistedCacheParameters;
+                Assert.IsTrue(parameters.Count == 4);
+
+                parameters.TryGetValue("additional_param1", out string additionalParam1);
+                parameters.TryGetValue("additional_param2", out string additionalParam2);
+                parameters.TryGetValue("additional_param3", out string additionalParam3);
+                parameters.TryGetValue("additional_param4", out string additionalParam4);
+
+                Assert.AreEqual("value1", additionalParam1);
+                Assert.AreEqual("value2", additionalParam2);
+                Assert.AreEqual("value3", additionalParam3);
+                Assert.AreEqual("[\"GUID\", \"GUID2\", \"GUID3\"]", additionalParam4);
+
+                Assert.AreEqual("Bearer", result.TokenType);
+                Assert.AreEqual(TokenSource.IdentityProvider, result.AuthenticationResultMetadata.TokenSource);
+                //Validate that the additional parameters are reflected in the AuthenticationResult.AdditionalResponseParameters
+                Assert.AreEqual((IReadOnlyDictionary<string, string>)parameters, result.AdditionalResponseParameters);
+
+                //Verify cache parameters still exist
+                result = await app.AcquireTokenForClient(TestConstants.s_scope.ToArray())
+                    .WithAdditionalCacheParameters(new List<string> { "additional_param1", "additional_param2", "additional_param3" })
+                    .ExecuteAsync()
+                    .ConfigureAwait(false);
+
+                Assert.AreEqual("Bearer", result.TokenType);
+                Assert.AreEqual(TokenSource.Cache, result.AuthenticationResultMetadata.TokenSource);
+
+                parameters = app.AppTokenCacheInternal.Accessor.GetAllAccessTokens().Single().PersistedCacheParameters;
+                Assert.IsTrue(parameters.Count == 4);
+
+                parameters.TryGetValue("additional_param1", out additionalParam1);
+                parameters.TryGetValue("additional_param2", out additionalParam2);
+                parameters.TryGetValue("additional_param3", out additionalParam3);
+                parameters.TryGetValue("additional_param4", out additionalParam4);
+
+                Assert.AreEqual("value1", additionalParam1);
+                Assert.AreEqual("value2", additionalParam2);
+                Assert.AreEqual("value3", additionalParam3);
+                Assert.AreEqual("[\"GUID\", \"GUID2\", \"GUID3\"]", additionalParam4);
+
+                Assert.AreEqual((IReadOnlyDictionary<string, string>)parameters, result.AdditionalResponseParameters);
+
+                //Verify cache parameters still exist without using WithAdditionalCacheParameters
+                result = await app.AcquireTokenForClient(TestConstants.s_scope.ToArray())
+                                    .ExecuteAsync()
+                                    .ConfigureAwait(false);
+
+                Assert.AreEqual("Bearer", result.TokenType);
+                Assert.AreEqual(TokenSource.Cache, result.AuthenticationResultMetadata.TokenSource);
+
+                parameters = app.AppTokenCacheInternal.Accessor.GetAllAccessTokens().Single().PersistedCacheParameters;
+                Assert.IsTrue(parameters.Count == 4);
+
+                parameters.TryGetValue("additional_param1", out additionalParam1);
+                parameters.TryGetValue("additional_param2", out additionalParam2);
+                parameters.TryGetValue("additional_param3", out additionalParam3);
+                parameters.TryGetValue("additional_param4", out additionalParam4);
+
+                Assert.AreEqual("value1", additionalParam1);
+                Assert.AreEqual("value2", additionalParam2);
+                Assert.AreEqual("value3", additionalParam3);
+                Assert.AreEqual("[\"GUID\", \"GUID2\", \"GUID3\"]", additionalParam4);
+
+                Assert.AreEqual((IReadOnlyDictionary<string, string>)parameters, result.AdditionalResponseParameters);
+
+                //Ensure not all cache parameters are required
+                app = ConfidentialClientApplicationBuilder.Create(TestConstants.ClientId)
+                        .WithClientSecret(TestConstants.ClientSecret)
+                        .WithExperimentalFeatures(true)
+                        .WithAuthority("https://login.microsoftonline.com/tid/")
+                        .WithHttpManager(httpManager)
+                        .BuildConcrete();
+
+                httpManager.AddMockHandlerSuccessfulClientCredentialTokenResponseWithAdditionalParamsMessage();
+
+                result = await app.AcquireTokenForClient(TestConstants.s_scope.ToArray())
+                    .WithAdditionalCacheParameters(new List<string> { "additional_param1", "additional_param3" })
+                    .ExecuteAsync()
+                    .ConfigureAwait(false);
+
+                parameters = app.AppTokenCacheInternal.Accessor.GetAllAccessTokens().Single().PersistedCacheParameters;
+                Assert.IsTrue(parameters.Count == 2);
+
+                parameters.TryGetValue("additional_param1", out additionalParam1);
+                parameters.TryGetValue("additional_param3", out additionalParam3);
+
+                Assert.AreEqual("value1", additionalParam1);
+                Assert.AreEqual("value3", additionalParam3);
+
+                Assert.AreEqual((IReadOnlyDictionary<string, string>)parameters, result.AdditionalResponseParameters);
+
+                //Ensure missing cache parameters are not added
+                app = ConfidentialClientApplicationBuilder.Create(TestConstants.ClientId)
+                      .WithClientSecret(TestConstants.ClientSecret)
+                      .WithExperimentalFeatures(true)
+                      .WithAuthority("https://login.microsoftonline.com/tid/")
+                      .WithHttpManager(httpManager)
+                      .BuildConcrete();
+
+                httpManager.AddMockHandlerSuccessfulClientCredentialTokenResponseWithAdditionalParamsMessage();
+                result = await app.AcquireTokenForClient(TestConstants.s_scope.ToArray())
+                    .WithAdditionalCacheParameters(new List<string> { "additional_paramN" })
+                    .ExecuteAsync()
+                    .ConfigureAwait(false);
+
+                parameters = app.AppTokenCacheInternal.Accessor.GetAllAccessTokens().Single().PersistedCacheParameters;
+                parameters.TryGetValue("additional_param1", out string additionalParam);
+                Assert.IsNull(additionalParam);
+                Assert.IsTrue(result.AdditionalResponseParameters.Count == 4);
+            }
+        }
+
         private AppTokenProviderResult GetAppTokenProviderResult(string differentScopesForAt = "", long? refreshIn = 1000)
         {
             var token = new AppTokenProviderResult();