Adding WithExtraBodyParameters api

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

@@ -15,6 +15,8 @@
 using Microsoft.Identity.Client.Utils;
 using Microsoft.Identity.Client.Extensibility;
 using Microsoft.Identity.Client.OAuth2;
+using System.Security.Cryptography;
+using System.Text;
 
 namespace Microsoft.Identity.Client
 {
@@ -44,9 +46,9 @@ internal static AcquireTokenForClientParameterBuilder Create(
 
             if (!string.IsNullOrEmpty(confidentialClientApplicationExecutor.ServiceBundle.Config.CertificateIdToAssociateWithToken))
             {
-                builder.WithAdditionalCacheKeyComponents(new SortedList<string, string>
+                builder.WithAdditionalCacheKeyComponents(new SortedList<string, Func<CancellationToken, Task<string>>>
                 {
-                    { Constants.CertSerialNumber, confidentialClientApplicationExecutor.ServiceBundle.Config.CertificateIdToAssociateWithToken }
+                    { Constants.CertSerialNumber, (CancellationToken ct) => { return Task.FromResult(confidentialClientApplicationExecutor.ServiceBundle.Config.CertificateIdToAssociateWithToken); } }
                 });
             }
 
@@ -143,9 +145,9 @@ public AcquireTokenForClientParameterBuilder WithFmiPath(string pathSuffix)
                 throw new ArgumentNullException(nameof(pathSuffix));
             }
 
-            var cacheKey = new SortedList<string, string>
+            var cacheKey = new SortedList<string, Func<CancellationToken, Task<string>>>
             { 
-                { OAuth2Parameter.FmiPath, pathSuffix } 
+                { OAuth2Parameter.FmiPath, (CancellationToken ct) => {return Task.FromResult(pathSuffix);} } 
             };
 
             this.WithAdditionalCacheKeyComponents(cacheKey);

src/client/Microsoft.Identity.Client/ApiConfig/Executors/ClientApplicationBaseExecutor.cs

@@ -33,7 +33,8 @@ public async Task<AuthenticationResult> ExecuteAsync(
             var requestParameters = await _clientApplicationBase.CreateRequestParametersAsync(
                 commonParameters,
                 requestContext,
-                _clientApplicationBase.UserTokenCacheInternal).ConfigureAwait(false);
+                _clientApplicationBase.UserTokenCacheInternal,
+                cancellationToken).ConfigureAwait(false);
 
             requestParameters.SendX5C = silentParameters.SendX5C ?? false;
 
@@ -59,7 +60,8 @@ public async Task<AuthenticationResult> ExecuteAsync(
             var requestParameters = await _clientApplicationBase.CreateRequestParametersAsync(
                 commonParameters,
                 requestContext,
-                _clientApplicationBase.UserTokenCacheInternal).ConfigureAwait(false);
+                _clientApplicationBase.UserTokenCacheInternal,
+                cancellationToken).ConfigureAwait(false);
 
             requestContext.Logger.Info(() => LogMessages.UsingXScopesForRefreshTokenRequest(commonParameters.Scopes.Count()));
 

src/client/Microsoft.Identity.Client/ApiConfig/Executors/ConfidentialClientExecutor.cs

@@ -39,7 +39,8 @@ public async Task<AuthenticationResult> ExecuteAsync(
             AuthenticationRequestParameters requestParams = await _confidentialClientApplication.CreateRequestParametersAsync(
                 commonParameters,
                 requestContext,
-                _confidentialClientApplication.UserTokenCacheInternal).ConfigureAwait(false);
+                _confidentialClientApplication.UserTokenCacheInternal,
+                cancellationToken).ConfigureAwait(false);
             requestParams.SendX5C = authorizationCodeParameters.SendX5C ?? false;
 
             var handler = new ConfidentialAuthCodeRequest(
@@ -60,7 +61,8 @@ public async Task<AuthenticationResult> ExecuteAsync(
             AuthenticationRequestParameters requestParams = await _confidentialClientApplication.CreateRequestParametersAsync(
                 commonParameters,
                 requestContext,
-                _confidentialClientApplication.AppTokenCacheInternal).ConfigureAwait(false);
+                _confidentialClientApplication.AppTokenCacheInternal,
+                cancellationToken).ConfigureAwait(false);
 
             requestParams.SendX5C = clientParameters.SendX5C ?? false;
 
@@ -82,7 +84,8 @@ public async Task<AuthenticationResult> ExecuteAsync(
             AuthenticationRequestParameters requestParams = await _confidentialClientApplication.CreateRequestParametersAsync(
                 commonParameters,
                 requestContext,
-                _confidentialClientApplication.UserTokenCacheInternal).ConfigureAwait(false);
+                _confidentialClientApplication.UserTokenCacheInternal,
+                cancellationToken).ConfigureAwait(false);
 
             requestParams.SendX5C = onBehalfOfParameters.SendX5C ?? false;
             requestParams.UserAssertion = onBehalfOfParameters.UserAssertion;
@@ -106,7 +109,8 @@ public async Task<Uri> ExecuteAsync(
             AuthenticationRequestParameters requestParameters = await _confidentialClientApplication.CreateRequestParametersAsync(
                 commonParameters,
                 requestContext,
-                _confidentialClientApplication.UserTokenCacheInternal).ConfigureAwait(false);
+                _confidentialClientApplication.UserTokenCacheInternal,
+                cancellationToken).ConfigureAwait(false);
 
             requestParameters.Account = authorizationRequestUrlParameters.Account;
             requestParameters.LoginHint = authorizationRequestUrlParameters.LoginHint;
@@ -142,7 +146,8 @@ public async Task<AuthenticationResult> ExecuteAsync(
             AuthenticationRequestParameters requestParams = await _confidentialClientApplication.CreateRequestParametersAsync(
                 commonParameters,
                 requestContext,
-                _confidentialClientApplication.UserTokenCacheInternal).ConfigureAwait(false);
+                _confidentialClientApplication.UserTokenCacheInternal,
+                cancellationToken).ConfigureAwait(false);
 
             requestParams.SendX5C = usernamePasswordParameters.SendX5C ?? false;
 

src/client/Microsoft.Identity.Client/ApiConfig/Executors/ManagedIdentityExecutor.cs

@@ -38,7 +38,8 @@ public async Task<AuthenticationResult> ExecuteAsync(
             var requestParams = await _managedIdentityApplication.CreateRequestParametersAsync(
                 commonParameters,
                 requestContext,
-                _managedIdentityApplication.AppTokenCacheInternal).ConfigureAwait(false);
+                _managedIdentityApplication.AppTokenCacheInternal,
+                cancellationToken).ConfigureAwait(false);
 
             var handler = new ManagedIdentityAuthRequest(
                 ServiceBundle,

src/client/Microsoft.Identity.Client/ApiConfig/Executors/PublicClientExecutor.cs

@@ -31,7 +31,8 @@ public async Task<AuthenticationResult> ExecuteAsync(
             AuthenticationRequestParameters requestParams = await _publicClientApplication.CreateRequestParametersAsync(
                 commonParameters,
                 requestContext,
-                _publicClientApplication.UserTokenCacheInternal).ConfigureAwait(false);
+                _publicClientApplication.UserTokenCacheInternal,
+                cancellationToken).ConfigureAwait(false);
 
             requestParams.LoginHint = interactiveParameters.LoginHint;
             requestParams.Account = interactiveParameters.Account;
@@ -52,7 +53,8 @@ public async Task<AuthenticationResult> ExecuteAsync(
             var requestParams = await _publicClientApplication.CreateRequestParametersAsync(
                 commonParameters,
                 requestContext,
-                _publicClientApplication.UserTokenCacheInternal).ConfigureAwait(false);
+                _publicClientApplication.UserTokenCacheInternal,
+                cancellationToken).ConfigureAwait(false);
 
             var handler = new DeviceCodeRequest(
                 ServiceBundle,
@@ -72,7 +74,8 @@ public async Task<AuthenticationResult> ExecuteAsync(
             var requestParams = await _publicClientApplication.CreateRequestParametersAsync(
                 commonParameters,
                 requestContext,
-                _publicClientApplication.UserTokenCacheInternal).ConfigureAwait(false);
+                _publicClientApplication.UserTokenCacheInternal,
+                cancellationToken).ConfigureAwait(false);
 
             var handler = new IntegratedWindowsAuthRequest(
                 ServiceBundle,
@@ -92,7 +95,8 @@ public async Task<AuthenticationResult> ExecuteAsync(
             var requestParams = await _publicClientApplication.CreateRequestParametersAsync(
                 commonParameters,
                 requestContext,
-                _publicClientApplication.UserTokenCacheInternal).ConfigureAwait(false);
+                _publicClientApplication.UserTokenCacheInternal,
+                cancellationToken).ConfigureAwait(false);
 
             var handler = new UsernamePasswordRequest(
                 ServiceBundle,

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

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Security.Cryptography.X509Certificates;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.AppConfig;
 using Microsoft.Identity.Client.AuthScheme;
@@ -27,10 +28,10 @@ internal class AcquireTokenCommonParameters
         public IAuthenticationOperation AuthenticationOperation { get; set; } = new BearerAuthenticationOperation();
         public IDictionary<string, string> ExtraHttpHeaders { get; set; }
         public PoPAuthenticationConfiguration PopAuthenticationConfiguration { get; set; }
-        public Func<OnBeforeTokenRequestData, Task> OnBeforeTokenRequestHandler { get; internal set; }
+        public IList<Func<OnBeforeTokenRequestData, Task>> OnBeforeTokenRequestHandler { get; internal set; }
         public X509Certificate2 MtlsCertificate { get; internal set; }
         public List<string> AdditionalCacheParameters { get; set; }
-        public SortedList<string, string> CacheKeyComponents { get; internal set; }
+        public SortedList<string, Func<CancellationToken, Task<string>>> CacheKeyComponents { get; internal set; }
         public string FmiPathSuffix { get; internal set; }
         public string ClientAssertionFmiPath { get; internal set; }
     }

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

@@ -2,6 +2,8 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.ApiConfig.Parameters;
 using Microsoft.Identity.Client.Internal;
@@ -27,18 +29,41 @@ internal ApplicationBase(ApplicationConfiguration config)
         internal virtual async Task<AuthenticationRequestParameters> CreateRequestParametersAsync(
             AcquireTokenCommonParameters commonParameters,
             RequestContext requestContext,
-            ITokenCacheInternal cache)
+            ITokenCacheInternal cache,
+            CancellationToken cancellationToken)
         {
             Instance.Authority authority = await Instance.Authority.CreateAuthorityForRequestAsync(
                requestContext,
                commonParameters.AuthorityOverride).ConfigureAwait(false);
 
+            var cacheKeyComponents = await InitializeCacheKeyComponentsAsync(commonParameters.CacheKeyComponents, cancellationToken).ConfigureAwait(false);
+
             return new AuthenticationRequestParameters(
                 ServiceBundle,
                 cache,
                 commonParameters,
                 requestContext,
-                authority);
+                authority,
+                cacheKeyComponents: cacheKeyComponents);
+        }
+
+        internal async Task<SortedList<string, string>> InitializeCacheKeyComponentsAsync(SortedList<string, Func<CancellationToken, Task<string>>> cacheKeyComponents, CancellationToken cancellationToken)
+        {
+            if (cacheKeyComponents != null && cacheKeyComponents.Count > 0)
+            {
+                var initializedCacheKeyComponents = new SortedList<string, string>();
+
+                foreach (var kvp in cacheKeyComponents)
+                {
+                    if (kvp.Value != null)
+                    {
+                        initializedCacheKeyComponents.Add(kvp.Key, await kvp.Value.Invoke(cancellationToken).ConfigureAwait(false));
+                    }
+                }
+                return initializedCacheKeyComponents;
+            }
+
+            return null;
         }
 
         internal static void GuardMobileFrameworks()

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

@@ -207,9 +207,10 @@ AcquireTokenByRefreshTokenParameterBuilder IByRefreshToken.AcquireTokenByRefresh
         internal override async Task<AuthenticationRequestParameters> CreateRequestParametersAsync(
             AcquireTokenCommonParameters commonParameters,
             RequestContext requestContext,
-            ITokenCacheInternal cache)
+            ITokenCacheInternal cache,
+            CancellationToken cancellationToken)
         {
-            AuthenticationRequestParameters requestParams = await base.CreateRequestParametersAsync(commonParameters, requestContext, cache).ConfigureAwait(false);
+            AuthenticationRequestParameters requestParams = await base.CreateRequestParametersAsync(commonParameters, requestContext, cache, cancellationToken).ConfigureAwait(false);
             return requestParams;
         }
     }

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

@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Runtime.CompilerServices;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.OAuth2;
 
@@ -29,12 +30,14 @@ public static AbstractAcquireTokenParameterBuilder<T> OnBeforeTokenRequest<T>(
             Func<OnBeforeTokenRequestData, Task> onBeforeTokenRequestHandler) 
             where T : AbstractAcquireTokenParameterBuilder<T>
         {
-            if (builder.CommonParameters.OnBeforeTokenRequestHandler != null && onBeforeTokenRequestHandler != null)
+            if (builder.CommonParameters.OnBeforeTokenRequestHandler == null)
             {
-                throw new InvalidOperationException("Cannot set OnBeforeTokenRequest handler twice.");
+                builder.CommonParameters.OnBeforeTokenRequestHandler = new List<Func<OnBeforeTokenRequestData, Task>> { onBeforeTokenRequestHandler };
+            }
+            else
+            {
+                builder.CommonParameters.OnBeforeTokenRequestHandler.Add(onBeforeTokenRequestHandler);
             }
-
-            builder.CommonParameters.OnBeforeTokenRequestHandler = onBeforeTokenRequestHandler;
 
             return builder;
         }
@@ -75,13 +78,18 @@ public static AbstractAcquireTokenParameterBuilder<T> WithAuthenticationExtensio
            MsalAuthenticationExtension authenticationExtension)
             where T : AbstractAcquireTokenParameterBuilder<T>
         {
-            if (builder.CommonParameters.OnBeforeTokenRequestHandler != null && authenticationExtension.OnBeforeTokenRequestHandler != null)
+            if (authenticationExtension.OnBeforeTokenRequestHandler != null)
             {
-                throw new InvalidOperationException("Cannot set both an AuthenticaitonOperation and an OnBeforeTokenRequestHandler");
+                if (builder.CommonParameters.OnBeforeTokenRequestHandler == null)
+                {
+                    builder.CommonParameters.OnBeforeTokenRequestHandler = new List<Func<OnBeforeTokenRequestData, Task>> { authenticationExtension.OnBeforeTokenRequestHandler };
+                }
+                else
+                {
+                    builder.CommonParameters.OnBeforeTokenRequestHandler.Add(authenticationExtension.OnBeforeTokenRequestHandler);
+                }
             }
 
-            builder.CommonParameters.OnBeforeTokenRequestHandler = authenticationExtension.OnBeforeTokenRequestHandler;
-
             if (authenticationExtension.AuthenticationOperation != null)
                 builder.WithAuthenticationOperation(authenticationExtension.AuthenticationOperation);
 
@@ -137,7 +145,7 @@ public static AbstractAcquireTokenParameterBuilder<T> WithAdditionalCacheParamet
         /// </remarks>
         internal static AbstractAcquireTokenParameterBuilder<T> WithAdditionalCacheKeyComponents<T>(
             this AbstractAcquireTokenParameterBuilder<T> builder,
-            IDictionary<string, string> cacheKeyComponents)
+            IDictionary<string, Func<CancellationToken, Task<string>>> cacheKeyComponents)
             where T : AbstractAcquireTokenParameterBuilder<T>
         {
             if (cacheKeyComponents == null || cacheKeyComponents.Count == 0)
@@ -148,7 +156,7 @@ internal static AbstractAcquireTokenParameterBuilder<T> WithAdditionalCacheKeyCo
 
             if (builder.CommonParameters.CacheKeyComponents == null)
             {
-                builder.CommonParameters.CacheKeyComponents = new SortedList<string, string>(cacheKeyComponents);
+                builder.CommonParameters.CacheKeyComponents = new SortedList<string, Func<CancellationToken, Task<string>>>(cacheKeyComponents);
             }
             else
             {
@@ -187,9 +195,9 @@ public static AbstractAcquireTokenParameterBuilder<T> WithFmiPathForClientAssert
             builder.CommonParameters.ClientAssertionFmiPath = fmiPath;
 
             // Add the fmi_path to the cache key so that it is used for cache lookups
-            var cacheKey = new SortedList<string, string>
+            var cacheKey = new SortedList<string, Func<CancellationToken, Task<string>>>
             {
-                { "credential_fmi_path", fmiPath }
+                { "credential_fmi_path", (CancellationToken ct) => Task.FromResult(fmiPath) }
             };
 
             WithAdditionalCacheKeyComponents(builder, cacheKey);

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

@@ -6,6 +6,8 @@
 using System.ComponentModel;
 using System.Diagnostics.CodeAnalysis;
 using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.Identity.Client.Cache;
 using Microsoft.Identity.Client.OAuth2;
 
@@ -40,5 +42,35 @@ public static AcquireTokenForClientParameterBuilder WithProofOfPosessionKeyId(
 
             return builder;
         }
+
+        /// <summary>
+        /// Add extra body parameters to the token request. These parameters are added to the cache key to associate these parameters with the acquired token.
+        /// </summary>
+        /// <param name="builder"></param>
+        /// <param name="extrabodyparams">List of additional body parameters</param>
+        /// <returns></returns>
+        public static AcquireTokenForClientParameterBuilder WithExtraBodyParameters(
+            this AcquireTokenForClientParameterBuilder builder, 
+            Dictionary<string, Func<CancellationToken, Task<string>>> extrabodyparams)
+        {
+            builder.ValidateUseOfExperimentalFeature();
+            if (extrabodyparams == null || extrabodyparams.Count == 0)
+            {
+                return builder;
+            }
+            builder.OnBeforeTokenRequest(async (data) =>
+            {
+                foreach (var param in extrabodyparams)
+                {
+                    if (param.Value != null)
+                    {
+                        data.BodyParameters.Add(param.Key, await param.Value(data.CancellationToken).ConfigureAwait(false));
+                    }
+                }
+            });
+
+            builder.WithAdditionalCacheKeyComponents(extrabodyparams);
+            return builder;
+        }
     }
 }