AbstractConfidentialClientAcquireTokenParameterBuilderExtension.cs

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Identity.Client.Core;
using Microsoft.Identity.Client.Internal;
using Microsoft.Identity.Client.OAuth2;

namespace Microsoft.Identity.Client.Extensibility
{

    /// <summary>
    /// Extensions for all AcquireToken methods
    /// </summary>
    public static class AbstractConfidentialClientAcquireTokenParameterBuilderExtension
    {
        /// <summary>
        /// Specifies client-originated claims to include in the token request.
        /// Unlike <see cref="AbstractAcquireTokenParameterBuilder{T}.WithClaims"/> (for server-issued
        /// claims challenges), tokens acquired with client claims <b>are cached</b> and the cache entry
        /// is keyed on the normalized claims value. Different claims values produce separate cache entries.
        /// Use stable, non-dynamic values to avoid unbounded cache growth.
        /// </summary>
        /// <typeparam name="T">The concrete confidential client builder type.</typeparam>
        /// <param name="builder">The builder to chain options to.</param>
        /// <param name="claimsJson">A JSON string containing the client-originated claims. Must be valid JSON.</param>
        /// <returns>The builder to chain the .With methods.</returns>
        public static T WithClientClaims<T>(
            this AbstractConfidentialClientAcquireTokenParameterBuilder<T> builder,
            string claimsJson)
            where T : AbstractConfidentialClientAcquireTokenParameterBuilder<T>
        {
            if (string.IsNullOrWhiteSpace(claimsJson))
            {
                return (T)builder;
            }

            string normalized = ClaimsHelper.NormalizeClaimsJson(claimsJson);
            builder.CommonParameters.ClientClaims = normalized;

            var cacheKey = new SortedList<string, Func<CancellationToken, Task<string>>>
            {
                { "client_claims", _ => Task.FromResult(normalized) }
            };

            WithAdditionalCacheKeyComponents(builder, cacheKey);

            return (T)builder;
        }

        /// <summary>
        /// Intervenes in the request pipeline, by executing a user provided delegate before MSAL makes the token request.
        /// The delegate can modify the request payload by adding or removing  body parameters and headers. <see cref="OnBeforeTokenRequestData"/>
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="builder">The builder to chain options to</param>
        /// <param name="onBeforeTokenRequestHandler">An async delegate which gets invoked just before MSAL makes a token request</param>
        /// <returns>The builder to chain other options to.</returns>
        public static AbstractAcquireTokenParameterBuilder<T> OnBeforeTokenRequest<T>(
            this AbstractAcquireTokenParameterBuilder<T> builder, 
            Func<OnBeforeTokenRequestData, Task> onBeforeTokenRequestHandler) 
            where T : AbstractAcquireTokenParameterBuilder<T>
        {
            if (builder.CommonParameters.OnBeforeTokenRequestHandler == null)
            {
                builder.CommonParameters.OnBeforeTokenRequestHandler = new List<Func<OnBeforeTokenRequestData, Task>> { onBeforeTokenRequestHandler };
            }
            else
            {
                builder.CommonParameters.OnBeforeTokenRequestHandler.Add(onBeforeTokenRequestHandler);
            }

            return builder;
        }

        /// <summary>
        /// Sends <paramref name="attributeTokens"/> as the <c>attribute_tokens</c> body parameter
        /// and includes them in the cache key. Null/empty/whitespace entries are ignored;
        /// a null or empty collection is a no-op.
        /// </summary>
        /// <typeparam name="T">The concrete confidential client builder type.</typeparam>
        /// <param name="builder">The builder to chain options to.</param>
        /// <param name="attributeTokens">Attribute tokens to include. Individual tokens must not contain whitespace.</param>
        /// <returns>The builder to chain method calls.</returns>
        /// <remarks>
        /// <para>
        /// For <c>AcquireTokenForClient</c> the <c>attribute_tokens</c> set is part of the cache
        /// partition key, so different sets are fully isolated.
        /// </para>
        /// <para>
        /// For <c>AcquireTokenOnBehalfOf</c> the partition key is the user assertion hash; the
        /// <c>attribute_tokens</c> set is stored as an item-level cache component within that
        /// partition. Multiple variants for the same assertion coexist and are disambiguated on
        /// read only when the caller supplies the same <c>WithAttributeTokens</c> set.
        /// <b>Mixing attributed and non-attributed reads against the same user assertion can
        /// return unintended cache entries or fail with <c>multiple_matching_tokens_detected</c>.</b>
        /// Be consistent: if you used <c>WithAttributeTokens</c> on a write, use it on every
        /// subsequent read for that assertion. Callers that require strict per-set cache isolation
        /// across different attribute-token sets should use separate <see cref="IConfidentialClientApplication"/>
        /// instances.
        /// </para>
        /// </remarks>
        /// <exception cref="ArgumentException">Thrown when any token contains embedded whitespace.</exception>
        /// <exception cref="MsalClientException">
        /// Thrown when the application was not configured to allow experimental features
        /// (this method transitively calls <see cref="WithExtraBodyParameters{T}"/>, which requires
        /// experimental features to be enabled via <c>WithExperimentalFeatures()</c> on the application builder).
        /// </exception>
        public static T WithAttributeTokens<T>(
            this AbstractConfidentialClientAcquireTokenParameterBuilder<T> builder,
            IEnumerable<string> attributeTokens)
            where T : AbstractConfidentialClientAcquireTokenParameterBuilder<T>
        {
            ILoggerAdapter logger = builder.ServiceBundle?.ApplicationLogger;

            if (attributeTokens is null)
            {
                logger?.Verbose(() => "[WithAttributeTokens] No attribute tokens passed.");
                return (T)builder;
            }

            var normalizedTokens = new List<string>();
            foreach (string token in attributeTokens)
            {
                if (!string.IsNullOrWhiteSpace(token))
                {
                    string trimmed = token.Trim();
                    if (trimmed.Any(char.IsWhiteSpace))
                    {
                        throw new ArgumentException(
                            $"Attribute tokens must not contain whitespace. Invalid token: '{trimmed}'",
                            nameof(attributeTokens));
                    }

                    normalizedTokens.Add(trimmed);
                }
            }

            if (normalizedTokens.Count == 0)
            {
                logger?.Verbose(() => "[WithAttributeTokens] collection contained no usable tokens.");
                return (T)builder;
            }

            // Deduplicate and sort so that callers passing the same logical set of tokens
            // (in any order, with possible duplicates) hit the same cache entry and
            // produce the same request body.
            normalizedTokens = normalizedTokens.Distinct(StringComparer.Ordinal).ToList();
            normalizedTokens.Sort(StringComparer.Ordinal);

            string joinedTokens = string.Join(" ", normalizedTokens);

            int count = normalizedTokens.Count;
            logger?.Info(() => $"[WithAttributeTokens] Attaching {count} attribute token(s) " +
                               "to the request body and including them in the cache key partition.");

            var extraBodyParams = new Dictionary<string, Func<CancellationToken, Task<string>>>
            {
                { OAuth2Parameter.AttributeTokens, _ => Task.FromResult(joinedTokens) }
            };

            return builder.WithExtraBodyParameters(extraBodyParams);
        }

        /// <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. Works for confidential client flows
        /// (AcquireTokenForClient, AcquireTokenOnBehalfOf, AcquireTokenByAuthorizationCode).
        /// </summary>
        /// <typeparam name="T">The concrete confidential client builder type.</typeparam>
        /// <param name="builder">The builder to chain options to.</param>
        /// <param name="extraBodyParams">List of additional body parameters.</param>
        /// <returns>The concrete builder to chain method calls.</returns>
        public static T WithExtraBodyParameters<T>(
            this AbstractConfidentialClientAcquireTokenParameterBuilder<T> builder,
            Dictionary<string, Func<CancellationToken, Task<string>>> extraBodyParams)
            where T : AbstractConfidentialClientAcquireTokenParameterBuilder<T>
        {
            builder.ValidateUseOfExperimentalFeature();

            if (extraBodyParams == null || extraBodyParams.Count == 0)
            {
                return (T)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 (T)builder;
        }

        /// <summary>
        /// Binds the token to a key in the cache.No cryptographic operations is performed on the token.
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="builder">The builder to chain options to</param>
        /// <param name="keyId">A key id to which the access token is associated. The token will not be retrieved from the cache unless the same key id is presented. Can be null.</param>
        /// <param name="expectedTokenTypeFromAad">AAD issues several types of bound tokens. MSAL checks the token type, which needs to match the value set by ESTS. Normal POP tokens have this as "pop"</param>
        /// <returns>the builder</returns>
        public static AbstractAcquireTokenParameterBuilder<T> WithProofOfPosessionKeyId<T>(
            this AbstractAcquireTokenParameterBuilder<T> builder,
            string keyId,
            string expectedTokenTypeFromAad = "Bearer")
            where T : AbstractAcquireTokenParameterBuilder<T>
        {
            if (string.IsNullOrEmpty(keyId))
            {
                throw new ArgumentNullException(nameof(keyId));
            }

            builder.ValidateUseOfExperimentalFeature();
            builder.CommonParameters.AuthenticationOperation = new ExternalBoundTokenScheme(keyId, expectedTokenTypeFromAad);

            return builder;
        }

        /// <summary>
        /// Enables client applications to provide a custom authentication operation to be used in the token acquisition request.
        /// </summary>
        /// <param name="builder">The builder to chain options to</param>
        /// <param name="authenticationExtension">The implementation of the authentication operation.</param>
        /// <returns></returns>
        public static AbstractAcquireTokenParameterBuilder<T> WithAuthenticationExtension<T>(
           this AbstractAcquireTokenParameterBuilder<T> builder,
           MsalAuthenticationExtension authenticationExtension)
            where T : AbstractAcquireTokenParameterBuilder<T>
        {
            if (authenticationExtension.OnBeforeTokenRequestHandler != null)
            {
                if (builder.CommonParameters.OnBeforeTokenRequestHandler == null)
                {
                    builder.CommonParameters.OnBeforeTokenRequestHandler = new List<Func<OnBeforeTokenRequestData, Task>> { authenticationExtension.OnBeforeTokenRequestHandler };
                }
                else
                {
                    builder.CommonParameters.OnBeforeTokenRequestHandler.Add(authenticationExtension.OnBeforeTokenRequestHandler);
                }
            }

            if (authenticationExtension.AuthenticationOperation != null)
                builder.WithAuthenticationOperation(authenticationExtension.AuthenticationOperation);

            if (authenticationExtension.AdditionalCacheParameters != null)
                builder.WithAdditionalCacheParameters(authenticationExtension.AdditionalCacheParameters);

            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.Any())
            {
                return builder;
            }

            //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;
        }

        /// <summary>
        /// Specifies additional cache key components to use when caching and retrieving tokens.
        /// </summary>
        /// <param name="cacheKeyComponents">The list of additional cache key components.</param>
        /// <param name="builder"></param>
        /// <returns>The builder.</returns>
        /// <remarks>
        /// <list type="bullet">
        /// <item><description>This api can be used to associate certificate key identifiers along with other keys with a particular token.</description></item>
        /// <item><description>In order for the tokens to be successfully retrieved from the cache, all components used to cache the token must be provided.</description></item>
        /// </list>
        /// </remarks>
        internal static AbstractAcquireTokenParameterBuilder<T> WithAdditionalCacheKeyComponents<T>(
            this AbstractAcquireTokenParameterBuilder<T> builder,
            IDictionary<string, Func<CancellationToken, Task<string>>> cacheKeyComponents)
            where T : AbstractAcquireTokenParameterBuilder<T>
        {
            if (cacheKeyComponents == null || cacheKeyComponents.Count == 0)
            {
                //no-op
                return builder;
            }

            if (builder.CommonParameters.CacheKeyComponents == null)
            {
                builder.CommonParameters.CacheKeyComponents = new SortedList<string, Func<CancellationToken, Task<string>>>(cacheKeyComponents);
            }
            else
            {
                foreach (var kvp in cacheKeyComponents)
                {
                    // Key conflicts are not allowed, it is expected for this method to fail.
                    builder.CommonParameters.CacheKeyComponents.Add(kvp.Key, kvp.Value);
                }
            }

            return builder;
        }

        /// <summary>
        /// Specifies an FMI path to be used for the client assertion. This lets higher level APIs like Id.Web 
        /// provide credentials which are FMI sensitive.
        /// Important: tokens are associated with the credential FMI path, which impacts cache lookups
        /// This is an extensibility API and should not be used by applications.
        /// </summary>
        /// <param name="builder">The builder.</param>
        /// <param name="fmiPath">The FMI path to use for client assertion.</param>
        /// <returns>The builder to chain the .With methods</returns>
        /// <exception cref="ArgumentNullException">Thrown when fmiPath is null or whitespace.</exception>
        public static AbstractAcquireTokenParameterBuilder<T> WithFmiPathForClientAssertion<T>(
            this AbstractAcquireTokenParameterBuilder<T> builder,
            string fmiPath)
            where T : AbstractAcquireTokenParameterBuilder<T>
        {
            builder.ValidateUseOfExperimentalFeature();

            if (string.IsNullOrWhiteSpace(fmiPath))
            {
                throw new ArgumentNullException(nameof(fmiPath));
            }

            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, Func<CancellationToken, Task<string>>>
            {
                { "credential_fmi_path", (CancellationToken ct) => Task.FromResult(fmiPath) }
            };

            WithAdditionalCacheKeyComponents(builder, cacheKey);

            return builder;
        }

        /// <summary>
        /// Specifies extra claims to be included in the client assertion. 
        /// These claims will be merged with default claims when the client assertion is generated.
        /// This lets higher level APIs like Microsoft.Identity.Web provide additional claims for the client assertion.
        /// Important: tokens are associated with the extra client assertion claims, which impacts cache lookups.
        /// This is an extensibility API and should not be used by applications directly.
        /// </summary>
        /// <param name="builder">The builder to chain options to</param>
        /// <param name="clientAssertionClaims">Additional claims in JSON format to be signed in the client assertion.</param>
        /// <returns>The builder to chain the .With methods</returns>
        /// <exception cref="ArgumentNullException">Thrown when clientAssertionClaims is null or whitespace.</exception>
        public static AbstractAcquireTokenParameterBuilder<T> WithExtraClientAssertionClaims<T>(
            this AbstractAcquireTokenParameterBuilder<T> builder,
            string clientAssertionClaims)
            where T : AbstractAcquireTokenParameterBuilder<T>
        {
            if (string.IsNullOrWhiteSpace(clientAssertionClaims))
            {
                throw new ArgumentNullException(nameof(clientAssertionClaims));
            }

            builder.CommonParameters.ExtraClientAssertionClaims = clientAssertionClaims;

            // Add the extra claims to the cache key so different claims result in different cache entries
            var cacheKey = new SortedList<string, Func<CancellationToken, Task<string>>>
            {
                { "extra_client_assertion_claims", (CancellationToken ct) => Task.FromResult(clientAssertionClaims) }
            };

            WithAdditionalCacheKeyComponents(builder, cacheKey);

            return builder;
        }
    }   
}