Adding api to enable additional caching parameters (#4919)

* Adding api to enable additional caching parameters * clean up * Clean up, Refactoring, Updating tests * Fixing test issue * Resolving build issue * Test fix * Adding support for arrays and objects * Refactoring. Test updates --------- Co-authored-by: trwalke <[email protected]>
AzureAD · Sep 13, 2024 · b5177a4 · b5177a4
1 parent 4cb3b57
commit b5177a4
Show file tree

Hide file tree

Showing 13 changed files with 304 additions and 16 deletions.
diff --git a/build/platform_and_feature_flags.props b/build/platform_and_feature_flags.props
@@ -23,4 +23,7 @@
   <PropertyGroup Condition="'$(TargetFramework)' == '$(TargetFrameworkNetStandard)'">
     <DefineConstants>$(DefineConstants);NETSTANDARD;SUPPORTS_CONFIDENTIAL_CLIENT;SUPPORTS_BROKER;SUPPORTS_CUSTOM_CACHE;SUPPORTS_WIN32;</DefineConstants>
   </PropertyGroup>
+  <PropertyGroup Condition="'$(TargetFramework)' == '$(TargetFrameworkNet6Android)' or '$(TargetFramework)' == '$(TargetFrameworkNet6Ios)'">
+    <DefineConstants>$(DefineConstants);MOBILE</DefineConstants>
+  </PropertyGroup>
 </Project>
diff --git a/src/client/Microsoft.Identity.Client/ApiConfig/AbstractAcquireTokenParameterBuilder.cs b/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;

diff --git a/src/client/Microsoft.Identity.Client/ApiConfig/Parameters/AcquireTokenCommonParameters.cs b/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; }
     }
 }
diff --git a/src/client/Microsoft.Identity.Client/AuthenticationResult.cs b/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);

diff --git a/src/client/Microsoft.Identity.Client/Cache/Items/MsalAccessTokenCacheItem.cs b/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> persistedCacheParameters = 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,39 @@ internal MsalAccessTokenCacheItem(
             RawClientInfo = response.ClientInfo;
             HomeAccountId = homeAccountId;
             OboCacheKey = oboCacheKey;
-
+#if !MOBILE
+            PersistedCacheParameters = AcquireCacheParametersFromResponse(persistedCacheParameters, response.ExtensionData);
+#endif
             InitCacheKey();
         }
 
-
+#if !MOBILE
+        private IDictionary<string, string> AcquireCacheParametersFromResponse(
+                                                    IEnumerable<string> persistedCacheParameters,
+#if SUPPORTS_SYSTEM_TEXT_JSON
+                                                    Dictionary<string, JsonElement> extraDataFromResponse)
+#else
+                                                    Dictionary<string, JToken> extraDataFromResponse)
+#endif
+        {
+            if (persistedCacheParameters == null || !persistedCacheParameters.Any())
+            {
+                return null;
+            }
+
+            var cacheParameters = extraDataFromResponse
+                                    .Where(x => persistedCacheParameters.Contains(x.Key, StringComparer.InvariantCultureIgnoreCase))
+#if SUPPORTS_SYSTEM_TEXT_JSON
+                                    .ToDictionary(x => x.Key, x => x.Value.ToString());
+#else
+                                    //Avoid formatting arrays because it adds new lines after every element
+                                    .ToDictionary(x => x.Key, x => x.Value.Type == JTokenType.Array || x.Value.Type == JTokenType.Object ?
+                                                                                    x.Value.ToString(Json.Formatting.None) :
+                                                                                    x.Value.ToString());
+#endif
+            return cacheParameters;
+        }
+#endif
         internal /* for test */ MsalAccessTokenCacheItem(
             string preferredCacheEnv,
             string clientId,
@@ -215,6 +246,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;
 

diff --git a/...y.Client/Extensibility/AbstractConfidentialClientAcquireTokenParameterBuilderExtension.cs b/...y.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,39 @@ public static AbstractAcquireTokenParameterBuilder<T> WithProofOfPosessionKeyId<
 
             return builder;
         }
+
+#if !MOBILE
+        /// <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.Any())
+            {
+                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;
+        }
+#endif
     }   
 }
diff --git a/src/client/Microsoft.Identity.Client/Internal/Requests/AuthenticationRequestParameters.cs b/src/client/Microsoft.Identity.Client/Internal/Requests/AuthenticationRequestParameters.cs
@@ -115,6 +115,8 @@ public string Claims
 
         public IAuthenticationScheme AuthenticationScheme => _commonParameters.AuthenticationScheme;
 
+        public IEnumerable<string> PersistedCacheParameters => _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.
 

diff --git a/src/client/Microsoft.Identity.Client/OAuth2/MsalTokenResponse.cs b/src/client/Microsoft.Identity.Client/OAuth2/MsalTokenResponse.cs
@@ -93,11 +93,8 @@ public IReadOnlyDictionary<string, string> CreateExtensionDataStringMap()
 #if SUPPORTS_SYSTEM_TEXT_JSON
             foreach (KeyValuePair<string, JsonElement> item in ExtensionData)
             {
-                if (item.Value.ValueKind == JsonValueKind.String ||
-                   item.Value.ValueKind == JsonValueKind.Number ||
-                   item.Value.ValueKind == JsonValueKind.True ||
-                   item.Value.ValueKind == JsonValueKind.False ||
-                   item.Value.ValueKind == JsonValueKind.Null)
+                if (item.Value.ValueKind != JsonValueKind.Undefined ||
+                    item.Value.ValueKind != JsonValueKind.Null)
                 {
                     stringExtensionData.Add(item.Key, item.Value.ToString());
                 }
@@ -117,6 +114,10 @@ public IReadOnlyDictionary<string, string> CreateExtensionDataStringMap()
                 {
                     stringExtensionData.Add(item.Key, item.Value.ToString());
                 }
+                else if (item.Value.Type == JTokenType.Array || item.Value.Type == JTokenType.Object)
+                {
+                    stringExtensionData.Add(item.Key, item.Value.ToString(Formatting.None));
+                }
             }
 #endif
             return stringExtensionData;

diff --git a/src/client/Microsoft.Identity.Client/TokenCache.ITokenCacheInternal.cs b/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.PersistedCacheParameters);
             }
 
             if (!string.IsNullOrEmpty(response.RefreshToken))

diff --git a/tests/Microsoft.Identity.Test.Common/Core/Mocks/MockHelpers.cs b/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\"],\"additional_param5\":{\"value5json\":\"value5\"}"
+            )
+        {
+            return CreateSuccessResponseMessage(
+                "{\"token_type\":\"" + tokenType + "\",\"expires_in\":\"" + expiry + "\",\"access_token\":\"" + token + "\"" + additionalparams + "}");
         }
 
         public static HttpResponseMessage CreateSuccessTokenResponseMessage(

diff --git a/tests/Microsoft.Identity.Test.Common/Core/Mocks/MockHttpManagerExtensions.cs b/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\"],\"additional_param5\":{\"value5json\":\"value5\"}")
+        {
+            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)
         {

diff --git a/tests/Microsoft.Identity.Test.Unit/PublicApiTests/AuthenticationResultTests.cs b/tests/Microsoft.Identity.Test.Unit/PublicApiTests/AuthenticationResultTests.cs
@@ -180,10 +180,6 @@ public async Task MsalTokenResponseParseTestAsync()
                 Assert.IsFalse(extMap.ContainsKey("id_token"));
                 Assert.IsFalse(extMap.ContainsKey("client_info"));
 
-                // only scalar properties should be in the map
-                Assert.IsFalse(extMap.ContainsKey("object_extension"));
-                Assert.IsFalse(extMap.ContainsKey("array_extension"));
-
                 // all other properties should be in the map
                 Assert.AreEqual("1209599", extMap["number_extension"]);
                 Assert.AreEqual("True", extMap["true_extension"]);