Fix for #5375 - Add a FormatResultAsync to IAuthenticationOperation

src/client/Microsoft.Identity.Client/AuthScheme/Bearer/BearerAuthenticationOperation.cs

@@ -2,13 +2,15 @@
 // Licensed under the MIT License.
 
 using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.Identity.Client.Cache.Items;
 using Microsoft.Identity.Client.Internal;
 using Microsoft.Identity.Client.Utils;
 
 namespace Microsoft.Identity.Client.AuthScheme.Bearer
 {
-    internal class BearerAuthenticationOperation : IAuthenticationOperation
+    internal class BearerAuthenticationOperation : IAuthenticationOperation2
     {
         internal const string BearerTokenType = "bearer";
 
@@ -25,6 +27,12 @@ public void FormatResult(AuthenticationResult authenticationResult)
             // no-op
         }
 
+        public Task FormatResultAsync(AuthenticationResult authenticationResult, CancellationToken cancellationToken = default)
+        {
+            // no-op, return completed task
+            return Task.CompletedTask;
+        }
+
         public IReadOnlyDictionary<string, string> GetTokenRequestParams()
         {
             // ESTS issues Bearer tokens by default, no need for any extra params

src/client/Microsoft.Identity.Client/AuthScheme/IAuthenticationOperation.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System.Collections.Generic;
+using System.Threading.Tasks;
 using Microsoft.Identity.Client.Cache.Items;
 
 namespace Microsoft.Identity.Client.AuthScheme
@@ -54,7 +55,7 @@ public interface IAuthenticationOperation
         /// Creates the access token that goes into an Authorization HTTP header. 
         /// </summary>
         void FormatResult(AuthenticationResult authenticationResult);
-
+        
         /// <summary>
         /// Expected to match the token_type parameter returned by ESTS. Used to disambiguate
         /// between ATs of different types (e.g. Bearer and PoP) when loading from cache etc.

src/client/Microsoft.Identity.Client/AuthScheme/IAuthenticationOperation2.cs

@@ -0,0 +1,21 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Identity.Client.AuthScheme
+{
+    /// <summary>
+    /// This is an extensibility API and should only be used by SDKs.
+    /// Enhanced version of IAuthenticationOperation that supports asynchronous token formatting.
+    /// Used to modify the experience depending on the type of token asked with async capabilities.
+    /// </summary>
+    public interface IAuthenticationOperation2 : IAuthenticationOperation
+    {
+        /// <summary>
+        /// Will be invoked instead of IAuthenticationOperation.FormatResult
+        /// </summary>
+        Task FormatResultAsync(AuthenticationResult authenticationResult, CancellationToken cancellationToken = default);
+    }
+}

src/client/Microsoft.Identity.Client/AuthScheme/PoP/PopAuthenticationOperation.cs

@@ -6,6 +6,8 @@
 using System.Globalization;
 using System.Security.Cryptography;
 using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.Identity.Client.AppConfig;
 using Microsoft.Identity.Client.Cache.Items;
 using Microsoft.Identity.Client.Internal;
@@ -21,7 +23,7 @@
 
 namespace Microsoft.Identity.Client.AuthScheme.PoP
 {
-    internal class PopAuthenticationOperation : IAuthenticationOperation
+    internal class PopAuthenticationOperation : IAuthenticationOperation2
     {
         private readonly PoPAuthenticationConfiguration _popAuthenticationConfiguration;
         private readonly IPoPCryptoProvider _popCryptoProvider;
@@ -86,6 +88,14 @@ public void FormatResult(AuthenticationResult authenticationResult)
             authenticationResult.AccessToken = popToken;
         }
 
+        public Task FormatResultAsync(AuthenticationResult authenticationResult, CancellationToken cancellationToken = default)
+        {
+            // For now, PoP token creation is synchronous, so we wrap the sync method
+            // Future enhancement could make crypto operations truly async
+            FormatResult(authenticationResult);
+            return Task.CompletedTask;
+        }
+
         private JObject CreateBody(string accessToken)
         {
             var publicKeyJwk = JToken.Parse(_popCryptoProvider.CannonicalPublicKeyJwk);

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

@@ -6,6 +6,7 @@
 using System.ComponentModel;
 using System.Globalization;
 using System.Security.Claims;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.AuthScheme;
 using Microsoft.Identity.Client.Cache;
@@ -21,7 +22,7 @@ namespace Microsoft.Identity.Client
     /// </summary>
     public partial class AuthenticationResult
     {
-        private readonly IAuthenticationOperation _authenticationScheme;
+        private IAuthenticationOperation _authenticationScheme;
 
         /// <summary>
         /// Constructor meant to help application developers test their apps. Allows mocking of authentication flows.
@@ -126,19 +127,76 @@ public AuthenticationResult(
 
         }
 
-        internal AuthenticationResult(
+        /// <summary>
+        /// This method must be used by the product code to create an <see cref="AuthenticationResult"/> instance.
+        /// It calls IAuthenticationOperation.FormatResult or FormatResultAsync on the authentication scheme
+        /// </summary>
+        internal static async Task<AuthenticationResult> CreateAsync(
             MsalAccessTokenCacheItem msalAccessTokenCacheItem,
             MsalIdTokenCacheItem msalIdTokenCacheItem,
             IAuthenticationOperation authenticationScheme,
-            Guid correlationID,
+            Guid correlationId,
             TokenSource tokenSource,
             ApiEvent apiEvent,
             Account account,
             string spaAuthCode,
-            IReadOnlyDictionary<string, string> additionalResponseParameters)
+            IReadOnlyDictionary<string, string> additionalResponseParameters,
+            CancellationToken cancellationToken = default)
         {
-            _authenticationScheme = authenticationScheme ?? throw new ArgumentNullException(nameof(authenticationScheme));
+            if (authenticationScheme == null)
+            {
+                throw new ArgumentNullException(nameof(authenticationScheme));
+            }
+
+            // Create the AuthenticationResult without calling FormatResult in constructor
+            var result = new AuthenticationResult(
+                msalAccessTokenCacheItem,
+                msalIdTokenCacheItem,
+                correlationId,
+                tokenSource,
+                apiEvent,
+                account,
+                spaAuthCode,
+                additionalResponseParameters,
+                authenticationScheme);
+
+            // Apply token formatting (async if supported, sync otherwise)
+            var measuredResultDuration = await StopwatchService.MeasureCodeBlockAsync(async () =>
+            {
+                if (authenticationScheme is IAuthenticationOperation2 asyncAuthScheme)
+                {
+                    await asyncAuthScheme.FormatResultAsync(result, cancellationToken).ConfigureAwait(false);
+                }
+                else
+                {
+                    authenticationScheme.FormatResult(result);
+                }
+            }).ConfigureAwait(false);
+
+            // Update telemetry metadata
+            result.AuthenticationResultMetadata.DurationCreatingExtendedTokenInUs = measuredResultDuration.Microseconds;
+            result.AuthenticationResultMetadata.TelemetryTokenType = authenticationScheme.TelemetryTokenType;
+
+            return result;
+        }
+
+        //Default constructor for testing
+        internal AuthenticationResult() { }
 
+        /// <summary>
+        /// This is to help CreateAsync
+        /// </summary>
+        private AuthenticationResult(
+            MsalAccessTokenCacheItem msalAccessTokenCacheItem,
+            MsalIdTokenCacheItem msalIdTokenCacheItem,
+            Guid correlationID,
+            TokenSource tokenSource,
+            ApiEvent apiEvent,
+            Account account,
+            string spaAuthCode,
+            IReadOnlyDictionary<string, string> additionalResponseParameters,
+            IAuthenticationOperation authenticationScheme)
+        {
             string homeAccountId =
                 msalAccessTokenCacheItem?.HomeAccountId ??
                 msalIdTokenCacheItem?.HomeAccountId;
@@ -163,7 +221,7 @@ internal AuthenticationResult(
             TenantId = msalIdTokenCacheItem?.IdToken?.TenantId;
             IdToken = msalIdTokenCacheItem?.Secret;
             SpaAuthCode = spaAuthCode;
-
+            _authenticationScheme = authenticationScheme;
             CorrelationId = correlationID;
             ApiEvent = apiEvent;
             AuthenticationResultMetadata = new AuthenticationResultMetadata(tokenSource);
@@ -189,20 +247,8 @@ internal AuthenticationResult(
 
                 AccessToken = msalAccessTokenCacheItem.Secret;
             }
-
-            var measuredResultDuration = StopwatchService.MeasureCodeBlock(() =>
-            {
-                //Important: only call this at the end
-                authenticationScheme.FormatResult(this);
-            });
-
-            AuthenticationResultMetadata.DurationCreatingExtendedTokenInUs = measuredResultDuration.Microseconds;
-            AuthenticationResultMetadata.TelemetryTokenType = authenticationScheme.TelemetryTokenType;
         }
 
-        //Default constructor for testing
-        internal AuthenticationResult() { }
-
         /// <summary>
         /// Access Token that can be used as a bearer token to access protected web APIs
         /// </summary>

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

@@ -20,7 +20,7 @@ public class MsalAuthenticationExtension
         public Func<OnBeforeTokenRequestData, Task> OnBeforeTokenRequestHandler { get; set; }
 
         /// <summary>
-        /// Enables the developer to provide a custom authentication extension.
+        /// Important: IAuthenticationOperation2 exists, for asynchronous token formatting.
         /// </summary>
         public IAuthenticationOperation AuthenticationOperation { get; set; }
 

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

@@ -37,7 +37,7 @@ protected override async Task<AuthenticationResult> ExecuteAsync(CancellationTok
                 throw new MsalServiceException(msalTokenResponse.Error, msalTokenResponse.ErrorDescription, null);
             }
 
-            return await CacheTokenResponseAndCreateAuthenticationResultAsync(msalTokenResponse).ConfigureAwait(false);
+            return await CacheTokenResponseAndCreateAuthenticationResultAsync(msalTokenResponse, cancellationToken).ConfigureAwait(false);
         }
 
         private static Dictionary<string, string> GetBodyParameters(string refreshTokenSecret)

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

@@ -76,7 +76,7 @@ protected override async Task<AuthenticationResult> ExecuteAsync(CancellationTok
             // No access token or cached access token needs to be refreshed 
             if (cachedAccessTokenItem != null)
             {
-                authResult = CreateAuthenticationResultFromCache(cachedAccessTokenItem);
+                authResult = await CreateAuthenticationResultFromCacheAsync(cachedAccessTokenItem).ConfigureAwait(false);
 
                 try
                 {
@@ -101,7 +101,7 @@ protected override async Task<AuthenticationResult> ExecuteAsync(CancellationTok
                 }
                 catch (MsalServiceException e)
                 {
-                    return await HandleTokenRefreshErrorAsync(e, cachedAccessTokenItem).ConfigureAwait(false);
+                    return await HandleTokenRefreshErrorAsync(e, cachedAccessTokenItem, cancellationToken).ConfigureAwait(false);
                 }
             }
             else
@@ -128,7 +128,7 @@ private async Task<AuthenticationResult> GetAccessTokenAsync(
             if (ServiceBundle.Config.AppTokenProvider == null)
             {
                 MsalTokenResponse msalTokenResponse = await SendTokenRequestAsync(GetBodyParameters(), cancellationToken).ConfigureAwait(false);
-                return await CacheTokenResponseAndCreateAuthenticationResultAsync(msalTokenResponse).ConfigureAwait(false);
+                return await CacheTokenResponseAndCreateAuthenticationResultAsync(msalTokenResponse, cancellationToken).ConfigureAwait(false);
             }
 
             // Get a token from the app provider delegate
@@ -164,7 +164,7 @@ private async Task<AuthenticationResult> GetAccessTokenAsync(
                     else
                     {
                         logger.Verbose(() => "[ClientCredentialRequest] Checking for a cached access token.");
-                        authResult = CreateAuthenticationResultFromCache(cachedAccessTokenItem);
+                        authResult = await CreateAuthenticationResultFromCacheAsync(cachedAccessTokenItem).ConfigureAwait(false);
                     }
                 }
 
@@ -196,7 +196,7 @@ private async Task<AuthenticationResult> SendTokenRequestToAppTokenProviderAsync
             tokenResponse.Scope = appTokenProviderParameters.Scopes.AsSingleString();
             tokenResponse.CorrelationId = appTokenProviderParameters.CorrelationId;
 
-            AuthenticationResult authResult = await CacheTokenResponseAndCreateAuthenticationResultAsync(tokenResponse)
+            AuthenticationResult authResult = await CacheTokenResponseAndCreateAuthenticationResultAsync(tokenResponse, cancellationToken)
                 .ConfigureAwait(false);
 
             return authResult;
@@ -274,9 +274,9 @@ private void MarkAccessTokenAsCacheHit()
         /// </summary>
         /// <param name="cachedAccessTokenItem"></param>
         /// <returns></returns>
-        private AuthenticationResult CreateAuthenticationResultFromCache(MsalAccessTokenCacheItem cachedAccessTokenItem)
+        private Task<AuthenticationResult> CreateAuthenticationResultFromCacheAsync(MsalAccessTokenCacheItem cachedAccessTokenItem)
         {
-            AuthenticationResult authResult = new AuthenticationResult(
+            return AuthenticationResult.CreateAsync(
                                                             cachedAccessTokenItem,
                                                             null,
                                                             AuthenticationRequestParameters.AuthenticationScheme,
@@ -286,7 +286,6 @@ private AuthenticationResult CreateAuthenticationResultFromCache(MsalAccessToken
                                                             account: null,
                                                             spaAuthCode: null,
                                                             additionalResponseParameters: null);
-            return authResult;
         }
 
         /// <summary>

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

@@ -28,7 +28,7 @@ protected override async Task<AuthenticationResult> ExecuteAsync(CancellationTok
         {
             await ResolveAuthorityAsync().ConfigureAwait(false);
             var msalTokenResponse = await SendTokenRequestAsync(GetBodyParameters(), cancellationToken).ConfigureAwait(false);
-            return await CacheTokenResponseAndCreateAuthenticationResultAsync(msalTokenResponse).ConfigureAwait(false);
+            return await CacheTokenResponseAndCreateAuthenticationResultAsync(msalTokenResponse, cancellationToken).ConfigureAwait(false);
         }
 
         private Dictionary<string, string> GetBodyParameters()

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

@@ -58,7 +58,7 @@ protected override async Task<AuthenticationResult> ExecuteAsync(CancellationTok
             await _deviceCodeParameters.DeviceCodeResultCallback(deviceCodeResult).ConfigureAwait(false);
 
             var msalTokenResponse = await WaitForTokenResponseAsync(deviceCodeResult, cancellationToken).ConfigureAwait(false);
-            return await CacheTokenResponseAndCreateAuthenticationResultAsync(msalTokenResponse).ConfigureAwait(false);
+            return await CacheTokenResponseAndCreateAuthenticationResultAsync(msalTokenResponse, cancellationToken).ConfigureAwait(false);
         }
 
         private async Task<MsalTokenResponse> WaitForTokenResponseAsync(

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

@@ -64,7 +64,7 @@ protected override async Task<AuthenticationResult> ExecuteAsync(CancellationTok
                                                 GetAdditionalBodyParameters(userAssertion), cancellationToken)
                                             .ConfigureAwait(false);
 
-            return await CacheTokenResponseAndCreateAuthenticationResultAsync(msalTokenResponse).ConfigureAwait(false);
+            return await CacheTokenResponseAndCreateAuthenticationResultAsync(msalTokenResponse, cancellationToken).ConfigureAwait(false);
         }
 
         protected override KeyValuePair<string, string>? GetCcsHeader(IDictionary<string, string> additionalBodyParameters)

src/client/Microsoft.Identity.Client/Internal/Requests/Interactive/InteractiveRequest.cs

@@ -107,7 +107,7 @@ protected override async Task<AuthenticationResult> ExecuteAsync(
                 tokenResponse = await GetTokenResponseAsync(cancellationToken)
                     .ConfigureAwait(false);
             }
-            return await CacheTokenResponseAndCreateAuthenticationResultAsync(tokenResponse)
+            return await CacheTokenResponseAndCreateAuthenticationResultAsync(tokenResponse, cancellationToken)
                 .ConfigureAwait(false);
         }
         #endregion