Add public API to search in cache for Initiate long-running OBO (#4135)

* Add public API to search in cache for Initiate long-running OBO method. * Fix. * Move to extensibility class. Update comments. * Update comments. --------- Co-authored-by: Gladwin Johnson <[email protected]>
AzureAD · May 22, 2023 · f17b6fb · f17b6fb
1 parent f2498bc
commit f17b6fb
Show file tree

Hide file tree

Showing 11 changed files with 174 additions and 24 deletions.
diff --git a/src/client/Microsoft.Identity.Client/ApiConfig/AcquireTokenOnBehalfOfParameterBuilder.cs b/src/client/Microsoft.Identity.Client/ApiConfig/AcquireTokenOnBehalfOfParameterBuilder.cs
@@ -24,7 +24,7 @@ namespace Microsoft.Identity.Client
     public sealed class AcquireTokenOnBehalfOfParameterBuilder :
         AbstractConfidentialClientAcquireTokenParameterBuilder<AcquireTokenOnBehalfOfParameterBuilder>
     {
-        private AcquireTokenOnBehalfOfParameters Parameters { get; } = new AcquireTokenOnBehalfOfParameters();
+        internal AcquireTokenOnBehalfOfParameters Parameters { get; } = new AcquireTokenOnBehalfOfParameters();
 
         /// <inheritdoc />
         internal AcquireTokenOnBehalfOfParameterBuilder(IConfidentialClientApplicationExecutor confidentialClientApplicationExecutor)
@@ -161,7 +161,7 @@ public AcquireTokenOnBehalfOfParameterBuilder WithCcsRoutingHint(string userName
 
             this.WithExtraHttpHeaders(ccsRoutingHeader);
             return this;
-        }
+        }
 
         /// <inheritdoc />
         internal override Task<AuthenticationResult> ExecuteInternalAsync(CancellationToken cancellationToken)

diff --git a/...client/Microsoft.Identity.Client/ApiConfig/Parameters/AcquireTokenOnBehalfOfParameters.cs b/...client/Microsoft.Identity.Client/ApiConfig/Parameters/AcquireTokenOnBehalfOfParameters.cs
@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using System.Collections.Generic;
 using System.Text;
 using Microsoft.Identity.Client.Core;
 
@@ -9,13 +10,22 @@ namespace Microsoft.Identity.Client.ApiConfig.Parameters
     internal class AcquireTokenOnBehalfOfParameters : AbstractAcquireTokenConfidentialClientParameters, IAcquireTokenParameters
     {
         /// <remarks>
-        /// User assertion is null when <see cref="ILongRunningWebApi.AcquireTokenInLongRunningProcess"/> is called.
+        /// User assertion is null when <see cref="ILongRunningWebApi.InitiateLongRunningProcessInWebApi(IEnumerable{string}, string, ref string)"/> is called.
         /// </remarks>
         public UserAssertion UserAssertion { get; set; }
         /// <summary>
         /// User-provided cache key for long-running OBO flow.
         /// </summary>
         public string LongRunningOboCacheKey { get; set; }
+
+        /// <summary>
+        /// Only affects <see cref="ILongRunningWebApi.InitiateLongRunningProcessInWebApi(IEnumerable{string}, string, ref string)"/>.
+        /// When enabled, mimics MSAL 4.50.0 and below behavior - checks in cache for cached tokens first, 
+        /// and if not found, then uses user assertion to request new tokens from AAD.
+        /// When disabled (default behavior), doesn't search in cache, but uses the user assertion to retrieve tokens from AAD.
+        /// </summary>
+        public bool SearchInCacheForLongRunningObo { get; set; }
+
         public bool ForceRefresh { get; set; }
 
         /// <inheritdoc />
@@ -28,6 +38,7 @@ public void LogParameters(ILoggerAdapter logger)
                 builder.AppendLine("SendX5C: " + SendX5C);
                 builder.AppendLine("ForceRefresh: " + ForceRefresh);
                 builder.AppendLine("UserAssertion set: " + (UserAssertion != null));
+                builder.AppendLine("SearchInCacheForLongRunningObo: " + SearchInCacheForLongRunningObo);
                 builder.AppendLine("LongRunningOboCacheKey set: " + !string.IsNullOrWhiteSpace(LongRunningOboCacheKey));
                 if (UserAssertion != null && !string.IsNullOrWhiteSpace(LongRunningOboCacheKey))
                 {

diff --git a/...crosoft.Identity.Client/Extensibility/AcquireTokenOnBehalfOfParameterBuilderExtensions.cs b/...crosoft.Identity.Client/Extensibility/AcquireTokenOnBehalfOfParameterBuilderExtensions.cs
@@ -0,0 +1,35 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.Collections.Generic;
+using System.ComponentModel;
+
+namespace Microsoft.Identity.Client.Extensibility
+{
+    /// <summary>
+    /// Extension methods for the <see cref="AcquireTokenOnBehalfOfParameterBuilder" />
+    /// </summary>
+    public static class AcquireTokenOnBehalfOfParameterBuilderExtensions
+    {
+        /// <summary>
+        /// Only affects <see cref="ILongRunningWebApi.InitiateLongRunningProcessInWebApi(IEnumerable{string}, string, ref string)"/>.
+        /// When enabled, mimics MSAL 4.50.0 and below behavior - checks in cache for cached tokens first, 
+        /// and if not found, then uses user assertion to request new tokens from AAD.
+        /// When disabled (default behavior), doesn't search in cache, but uses the user assertion to retrieve tokens from AAD.
+        /// </summary>
+        /// <remarks>
+        /// This method should only be used in specific cases for backwards compatibility. For most cases, rely on the default behavior
+        /// of <see cref="ILongRunningWebApi.InitiateLongRunningProcessInWebApi(IEnumerable{string}, string, ref string)"/> and
+        /// <see cref="ILongRunningWebApi.AcquireTokenInLongRunningProcess(IEnumerable{string}, string)"/> described in https://aka.ms/msal-net-long-running-obo .
+        /// </remarks>
+        /// <param name="builder"></param>
+        /// <param name="searchInCache">Whether to search in cache.</param>
+        /// <returns>The builder to chain the .With methods</returns>
+        [EditorBrowsable(EditorBrowsableState.Never)]
+        public static AcquireTokenOnBehalfOfParameterBuilder WithSearchInCacheForLongRunningProcess(this AcquireTokenOnBehalfOfParameterBuilder builder, bool searchInCache = true)
+        {
+            builder.Parameters.SearchInCacheForLongRunningObo = searchInCache;
+            return builder;
+        }
+    }
+}
diff --git a/src/client/Microsoft.Identity.Client/ILongRunningWebApi.cs b/src/client/Microsoft.Identity.Client/ILongRunningWebApi.cs
@@ -14,11 +14,14 @@ public interface ILongRunningWebApi
         /// <summary>
         /// Acquires an access token for this web API from the authority configured in the application,
         /// in order to access another downstream protected web API on behalf of a user using the OAuth 2.0 On-Behalf-Of flow.
-        /// See https://aka.ms/msal-net-long-running-obo .
+        /// See https://aka.ms/msal-net-long-running-obo.
         /// This confidential client application was itself called with a token which will be provided in the
         /// <paramref name="userToken">userToken</paramref> parameter.
         /// Use <seealso cref="ConfidentialClientApplicationExtensions.StopLongRunningProcessInWebApiAsync"/> to stop the long running process.
         /// </summary>
+        /// <remarks>
+        /// This method should be called once when the long running session is started.
+        /// </remarks>
         /// <param name="scopes">Scopes requested to access a protected API</param>
         /// <param name="userToken">A JSON Web Token which was used to call the web API and contains the credential information
         /// about the user on behalf of whom to get a token.</param>
@@ -34,7 +37,7 @@ public interface ILongRunningWebApi
         /// Use <seealso cref="ConfidentialClientApplicationExtensions.StopLongRunningProcessInWebApiAsync"/> to stop the long running process.
         /// </summary>
         /// <remarks>
-        /// This method is intended to be used in the long running processes inside of web APIs.
+        /// This method should be called during the long running session to retrieve the token from the cache
         /// </remarks>
         /// <param name="scopes">Scopes requested to access a protected API</param>
         /// <param name="longRunningProcessSessionKey">Key by which to look up the token in the cache</param>

diff --git a/src/client/Microsoft.Identity.Client/Internal/Requests/OnBehalfOfRequest.cs b/src/client/Microsoft.Identity.Client/Internal/Requests/OnBehalfOfRequest.cs
@@ -44,9 +44,9 @@ protected override async Task<AuthenticationResult> ExecuteAsync(CancellationTok
             CacheRefreshReason cacheInfoTelemetry = CacheRefreshReason.NotApplicable;
 
             //Check if initiating a long running process
-            if (AuthenticationRequestParameters.ApiId == ApiEvent.ApiIds.InitiateLongRunningObo)
+            if (AuthenticationRequestParameters.ApiId == ApiEvent.ApiIds.InitiateLongRunningObo && !_onBehalfOfParameters.SearchInCacheForLongRunningObo)
             {
-                //Long running process should not use cached tokens
+                //Long running OBO doesn't search in cache by default
                 logger.Info("[OBO Request] Initiating long running process. Fetching OBO token from ESTS.");
                 return await FetchNewAccessTokenAsync(cancellationToken).ConfigureAwait(false);
             }
@@ -143,7 +143,6 @@ private async Task<AuthenticationResult> RefreshRtOrFetchNewAccessTokenAsync(Can
             {
                 AuthenticationRequestParameters.RequestContext.Logger.Info("[OBO request] Long-running OBO flow, trying to refresh using a refresh token flow.");
 
-
                 // Look for a refresh token
                 MsalRefreshTokenCacheItem cachedRefreshToken = await CacheManager.FindRefreshTokenAsync().ConfigureAwait(false);
 
@@ -176,11 +175,11 @@ private async Task<AuthenticationResult> RefreshRtOrFetchNewAccessTokenAsync(Can
                     throw new MsalClientException(MsalError.OboCacheKeyNotInCacheError, MsalErrorMessage.OboCacheKeyNotInCache);
                 }
 
-                AuthenticationRequestParameters.RequestContext.Logger.Info("[OBO request] No Refresh Token was found in the cache. Fetching OBO token from ESTS");
+                AuthenticationRequestParameters.RequestContext.Logger.Info("[OBO request] No refresh token was found in the cache. Fetching OBO tokens from ESTS.");
             }
             else
             {
-                logger.Info("[OBO request] Normal OBO flow, skipping to fetching access token via OBO flow.");
+                logger.Info("[OBO request] Fetching tokens via normal OBO flow.");
             }
 
             return await FetchNewAccessTokenAsync(cancellationToken).ConfigureAwait(false);

diff --git a/src/client/Microsoft.Identity.Client/TokenCache.ITokenCacheInternal.cs b/src/client/Microsoft.Identity.Client/TokenCache.ITokenCacheInternal.cs
@@ -2,7 +2,6 @@
 // Licensed under the MIT License.
 
 using System;
-using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Globalization;
@@ -54,7 +53,7 @@ async Task<Tuple<MsalAccessTokenCacheItem, MsalIdTokenCacheItem, Account>> IToke
 
             string suggestedWebCacheKey = CacheKeyFactory.GetExternalCacheKeyFromResponse(requestParams, homeAccountId);
 
-            // token could be comming from a different cloud than the one configured
+            // token could be coming from a different cloud than the one configured
             if (requestParams.AppConfig.MultiCloudSupportEnabled && !string.IsNullOrEmpty(response.AuthorityUrl))
             {
                 var url = new Uri(response.AuthorityUrl);

diff --git a/tests/Microsoft.Identity.Test.Common/Core/Mocks/MockHelpers.cs b/tests/Microsoft.Identity.Test.Common/Core/Mocks/MockHelpers.cs
@@ -66,12 +66,12 @@ public static string GetTokenResponseWithNoOidClaim()
                                         "\"trace_id\":\"dd25f4fb-3e8d-458e-90e7-179524ce0000\",\"correlation_id\":" +
                                         "\"f11508ab-067f-40d4-83cb-ccc67bf57e45\"}";
 
-        public static string GetDefaultTokenResponse(string accessToken = TestConstants.ATSecret)
+        public static string GetDefaultTokenResponse(string accessToken = TestConstants.ATSecret, string refreshToken = TestConstants.RTSecret)
         {
               return
             "{\"token_type\":\"Bearer\",\"expires_in\":\"3599\",\"refresh_in\":\"2400\",\"scope\":" +
             "\"r1/scope1 r1/scope2\",\"access_token\":\"" + accessToken + "\"" +
-            ",\"refresh_token\":\"" + Guid.NewGuid() + "\",\"client_info\"" +
+            ",\"refresh_token\":\"" + refreshToken + "\",\"client_info\"" +
             ":\"" + CreateClientInfo() + "\",\"id_token\"" +
             ":\"" + CreateIdToken(TestConstants.UniqueId, TestConstants.DisplayableId) + "\"}";
         }
@@ -81,7 +81,7 @@ public static string GetPopTokenResponse()
             return
           "{\"token_type\":\"pop\",\"expires_in\":\"3599\",\"scope\":" +
           "\"r1/scope1 r1/scope2\",\"access_token\":\"" + TestConstants.ATSecret + "\"" +
-          ",\"refresh_token\":\"" + Guid.NewGuid() + "\",\"client_info\"" +
+          ",\"refresh_token\":\"" + TestConstants.RTSecret + "\",\"client_info\"" +
           ":\"" + CreateClientInfo() + "\",\"id_token\"" +
           ":\"" + CreateIdToken(TestConstants.UniqueId, TestConstants.DisplayableId) +
           "\",\"id_token_expires_in\":\"3600\"}";
@@ -92,7 +92,7 @@ public static string GetHybridSpaTokenResponse(string spaCode)
             return
             "{\"token_type\":\"Bearer\",\"expires_in\":\"3599\",\"refresh_in\":\"2400\",\"scope\":" +
             "\"r1/scope1 r1/scope2\",\"access_token\":\"" + TestConstants.ATSecret + "\"" +
-            ",\"refresh_token\":\"" + Guid.NewGuid() + "\",\"client_info\"" +
+            ",\"refresh_token\":\"" + TestConstants.RTSecret + "\",\"client_info\"" +
             ":\"" + CreateClientInfo() + "\",\"id_token\"" +
             ":\"" + CreateIdToken(TestConstants.UniqueId, TestConstants.DisplayableId) +
             "\",\"spa_code\":\"" + spaCode + "\"" +
@@ -104,7 +104,7 @@ public static string GetBridgedHybridSpaTokenResponse(string spaAccountId)
             return
             "{\"token_type\":\"Bearer\",\"expires_in\":\"3599\",\"refresh_in\":\"2400\",\"scope\":" +
             "\"r1/scope1 r1/scope2\",\"access_token\":\"" + TestConstants.ATSecret + "\"" +
-            ",\"refresh_token\":\"" + Guid.NewGuid() + "\",\"client_info\"" +
+            ",\"refresh_token\":\"" + TestConstants.RTSecret + "\",\"client_info\"" +
             ":\"" + CreateClientInfo() + "\",\"id_token\"" +
             ":\"" + CreateIdToken(TestConstants.UniqueId, TestConstants.DisplayableId) +
             "\",\"spa_accountId\":\"" + spaAccountId + "\"" +
@@ -205,10 +205,10 @@ public static HttpResponseMessage CreateSuccessTokenResponseMessage(
                 scopes, idToken, clientInfo));
         }
 
-        public static HttpResponseMessage CreateSuccessTokenResponseMessage(bool foci = false, string accessToken = TestConstants.ATSecret)
+        public static HttpResponseMessage CreateSuccessTokenResponseMessage(bool foci = false, string accessToken = TestConstants.ATSecret, string refreshToken = TestConstants.RTSecret)
         {
             return CreateSuccessResponseMessage(
-                foci ? GetFociTokenResponse() : GetDefaultTokenResponse(accessToken));
+                foci ? GetFociTokenResponse() : GetDefaultTokenResponse(accessToken, refreshToken));
         }
 
         public static HttpResponseMessage CreateSuccessTokenResponseMessageWithUid(

diff --git a/tests/Microsoft.Identity.Test.Unit/CacheTests/CacheSerializationTests.cs b/tests/Microsoft.Identity.Test.Unit/CacheTests/CacheSerializationTests.cs
@@ -608,18 +608,20 @@ public void TestJsonSerialization()
         [TestMethod]
         public void TestSerializeContainsNoNulls()
         {
+            var prefix = $"_SOMERANDOMPREFIX";
             var accessor = CreateTokenCacheAccessor();
 
             // Create a refresh token with a null family id in it
             var item = CreateRefreshTokenItem();
             item.FamilyId = null;
-            item.Environment += $"_SOMERANDOMPREFIX"; // ensure we get unique cache keys
+            item.Environment += prefix; // ensure we get unique cache keys
             accessor.SaveRefreshToken(item);
 
             var s1 = new TokenCacheJsonSerializer(accessor);
             byte[] bytes = s1.Serialize(null);
             string json = CoreHelpers.ByteArrayToString(bytes);
-            Console.WriteLine(json);
+
+            Assert.IsTrue(json.Contains(prefix));
             Assert.IsFalse(json.ToLowerInvariant().Contains("null"));
         }
 

diff --git a/tests/Microsoft.Identity.Test.Unit/PublicApiTests/ConfidentialClientApplicationTests.cs b/tests/Microsoft.Identity.Test.Unit/PublicApiTests/ConfidentialClientApplicationTests.cs
@@ -1644,7 +1644,8 @@ public void EnsurePublicApiSurfaceExistsOnInterface()
             var longRunningOboBuilder = ((ILongRunningWebApi)app).InitiateLongRunningProcessInWebApi(
                                TestConstants.s_scope.ToArray(),
                                TestConstants.DefaultClientAssertion,
-                               ref oboCacheKey);
+                               ref oboCacheKey)
+                .WithSearchInCacheForLongRunningProcess();
             PublicClientApplicationTests.CheckBuilderCommonMethods(longRunningOboBuilder);
 
             longRunningOboBuilder = ((ILongRunningWebApi)app).AcquireTokenInLongRunningProcess(