Adding fix for OBO cache key error in long running OBO with DefaultAuthorizationHeaderProvider

src/Microsoft.Identity.Web.TokenAcquisition/DefaultAuthorizationHeaderProvider.cs

@@ -27,13 +27,16 @@ public async Task<string> CreateAuthorizationHeaderForUserAsync(
             ClaimsPrincipal? claimsPrincipal = null,
             CancellationToken cancellationToken = default)
         {
+            var newTokenAcquisitionOptions = CreateTokenAcquisitionOptionsFromApiOptions(downstreamApiOptions, cancellationToken);
             var result = await _tokenAcquisition.GetAuthenticationResultForUserAsync(
                 scopes,
                 downstreamApiOptions?.AcquireTokenOptions.AuthenticationOptionsName,
                 downstreamApiOptions?.AcquireTokenOptions.Tenant,
                 downstreamApiOptions?.AcquireTokenOptions.UserFlow,
                 claimsPrincipal,
-                CreateTokenAcquisitionOptionsFromApiOptions(downstreamApiOptions, cancellationToken)).ConfigureAwait(false);
+                newTokenAcquisitionOptions).ConfigureAwait(false);
+
+            UpdateOriginalTokenAcquisitionOptions(downstreamApiOptions?.AcquireTokenOptions, newTokenAcquisitionOptions);
             return result.CreateAuthorizationHeader();
         }
 
@@ -48,6 +51,7 @@ public async Task<string> CreateAuthorizationHeaderForAppAsync(
                 downstreamApiOptions?.AcquireTokenOptions.AuthenticationOptionsName,
                 downstreamApiOptions?.AcquireTokenOptions.Tenant,
                 CreateTokenAcquisitionOptionsFromApiOptions(downstreamApiOptions, cancellationToken)).ConfigureAwait(false);
+
             return result.CreateAuthorizationHeader();
         }
 
@@ -59,6 +63,7 @@ public async Task<string> CreateAuthorizationHeaderAsync(
             CancellationToken cancellationToken = default)
         {
             Client.AuthenticationResult result;
+            var newTokenAcquisitionOptions = CreateTokenAcquisitionOptionsFromApiOptions(downstreamApiOptions, cancellationToken);
 
             // Previously, with the API name we were able to distinguish between app and user token acquisition
             // This context is missing in the new API, so can we enforce that downstreamApiOptions.RequestAppToken
@@ -75,8 +80,7 @@ public async Task<string> CreateAuthorizationHeaderAsync(
                     scopes.FirstOrDefault()!,
                     downstreamApiOptions?.AcquireTokenOptions.AuthenticationOptionsName,
                     downstreamApiOptions?.AcquireTokenOptions.Tenant,
-                    CreateTokenAcquisitionOptionsFromApiOptions(downstreamApiOptions, cancellationToken)).ConfigureAwait(false);
-                return result.CreateAuthorizationHeader();
+                    newTokenAcquisitionOptions).ConfigureAwait(false);
             }
             else
             {
@@ -86,9 +90,11 @@ public async Task<string> CreateAuthorizationHeaderAsync(
                     downstreamApiOptions?.AcquireTokenOptions?.Tenant,
                     downstreamApiOptions?.AcquireTokenOptions?.UserFlow,
                     claimsPrincipal,
-                    CreateTokenAcquisitionOptionsFromApiOptions(downstreamApiOptions, cancellationToken)).ConfigureAwait(false);
-                return result.CreateAuthorizationHeader();
+                    newTokenAcquisitionOptions).ConfigureAwait(false);
             }
+
+            UpdateOriginalTokenAcquisitionOptions(downstreamApiOptions?.AcquireTokenOptions, newTokenAcquisitionOptions);
+            return result.CreateAuthorizationHeader();
         }
 
         private static TokenAcquisitionOptions CreateTokenAcquisitionOptionsFromApiOptions(
@@ -113,5 +119,17 @@ private static TokenAcquisitionOptions CreateTokenAcquisitionOptionsFromApiOptio
                 FmiPath = downstreamApiOptions?.AcquireTokenOptions.FmiPath,
             };
         }
+
+        /// <summary>
+        /// Since AcquireTokenOptions is recreated, we need to update the original TokenAcquisitionOptions wth the parameters that were
+        /// updated in the new TokenAcquisitionOptions.
+        /// </summary>
+        private void UpdateOriginalTokenAcquisitionOptions(AcquireTokenOptions? acquireTokenOptions, TokenAcquisitionOptions newTokenAcquisitionOptions)
+        {
+            if (acquireTokenOptions is not null && newTokenAcquisitionOptions is not null)
+            {
+                acquireTokenOptions.LongRunningWebApiSessionKey = newTokenAcquisitionOptions.LongRunningWebApiSessionKey;
+            }
+        }
     }
 }

tests/Microsoft.Identity.Web.Test.Common/Mocks/MockHttpCreator.cs

@@ -21,6 +21,38 @@ private static HttpResponseMessage CreateSuccessfulClientCredentialTokenResponse
                 "{\"token_type\":\"" + tokenType + "\",\"expires_in\":" + expiry + ",\"client_info\":\"" + CreateClientInfo() + "\",\"access_token\":\"" + token + "\"}");
         }
 
+        public static HttpResponseMessage GetLrOboTokenResponse(string scopes, string accessToken = "header.payload.signature", string refreshToken = "header.payload.signatureRt")
+        {
+            return CreateSuccessResponseMessage(
+          "{\"token_type\":\"Bearer\",\"expires_in\":\"3599\",\"refresh_in\":\"2400\",\"scope\":" +
+          "\"" + scopes + "\",\"access_token\":\"" + accessToken + "\"" +
+          ",\"refresh_token\":\"" + refreshToken + "\",\"client_info\"" +
+          ":\"" + CreateClientInfo() + "\",\"id_token\"" +
+          ":\"" + CreateIdToken("UniqueId", "DisplayableId") + "\"}");
+        }
+
+        public static string CreateIdToken(string uniqueId, string displayableId)
+        {
+            return CreateIdToken(uniqueId, displayableId, TestConstants.Utid);
+        }
+
+        public static string CreateIdToken(string uniqueId, string displayableId, string tenantId)
+        {
+            string id = "{\"aud\": \"e854a4a7-6c34-449c-b237-fc7a28093d84\"," +
+                        "\"iss\": \"https://login.microsoftonline.com/6c3d51dd-f0e5-4959-b4ea-a80c4e36fe5e/v2.0/\"," +
+                        "\"iat\": 1455833828," +
+                        "\"nbf\": 1455833828," +
+                        "\"exp\": 1455837728," +
+                        "\"ipaddr\": \"131.107.159.117\"," +
+                        "\"name\": \"Marrrrrio Bossy\"," +
+                        "\"oid\": \"" + uniqueId + "\"," +
+                        "\"preferred_username\": \"" + displayableId + "\"," +
+                        "\"sub\": \"K4_SGGxKqW1SxUAmhg6C1F6VPiFzcx-Qd80ehIEdFus\"," +
+                        "\"tid\": \"" + tenantId + "\"," +
+                        "\"ver\": \"2.0\"}";
+            return string.Format(CultureInfo.InvariantCulture, "someheader.{0}.somesignature", Base64UrlHelpers.Encode(id));
+        }
+
         public static HttpResponseMessage CreateSuccessResponseMessage(string successResponse)
         {
             HttpResponseMessage responseMessage = new HttpResponseMessage(HttpStatusCode.OK);
@@ -62,6 +94,18 @@ public static MockHttpMessageHandler CreateClientCredentialTokenHandler(
             return handler;
         }
 
+        public static MockHttpMessageHandler CreateLrOboTokenHandler(
+            string scopes, string accessToken = "header.payload.signature", string refreshToken = "header.payload.signatureRt")
+        {
+            var handler = new MockHttpMessageHandler()
+            {
+                ExpectedMethod = HttpMethod.Post,
+                ResponseMessage = GetLrOboTokenResponse(scopes, accessToken, refreshToken),
+            };
+
+            return handler;
+        }
+
         public static MockHttpMessageHandler CreateHandlerToValidatePostData(
             HttpMethod expectedMethod,
             IDictionary<string, string> expectedPostData)

tests/Microsoft.Identity.Web.Test/AuthorizationHeaderProviderTests.cs

@@ -0,0 +1,232 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Security.Claims;
+using System.Text;
+using System.Threading.Tasks;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Identity.Abstractions;
+using Microsoft.Identity.Client;
+using Microsoft.Identity.Web.Test.Common.Mocks;
+using Microsoft.Identity.Web.TestOnly;
+using Microsoft.IdentityModel.Tokens;
+using Xunit;
+
+namespace Microsoft.Identity.Web.Test
+{
+    public class AuthorizationHeaderProviderTests
+    {
+        [Fact]
+        public async Task LongRunningSessionForDefaultAuthProviderForUserDefaultKeyTest()
+        {
+            // Arrange
+            var tokenAcquirerFactory = InitTokenAcquirerFactoryForTest();
+            IServiceProvider serviceProvider = tokenAcquirerFactory.Build();
+
+            IAuthorizationHeaderProvider authorizationHeaderProvider =
+                serviceProvider.GetRequiredService<IAuthorizationHeaderProvider>();
+            var mockHttpClient = serviceProvider.GetRequiredService<IMsalHttpClientFactory>() as MockHttpClientFactory;
+
+            using (mockHttpClient)
+            {
+                // Create a test ClaimsPrincipal
+                var claims = new List<Claim>
+                            {
+                                new Claim(ClaimTypes.Name, "[email protected]")
+                            };
+
+                var identity = new CaseSensitiveClaimsIdentity(claims, "TestAuth");
+                identity.BootstrapContext = CreateTestJwt();
+                var claimsPrincipal = new ClaimsPrincipal(identity);
+
+                // Create options with LongRunningWebApiSessionKey
+                var options = new AuthorizationHeaderProviderOptions
+                {
+                    AcquireTokenOptions = new AcquireTokenOptions
+                    {
+                        //When this is set to Auto, the first call will set the LongRunningWebApiSessionKey in the options to
+                        //the hash of the ClaimsPrincipal's access token.
+                        LongRunningWebApiSessionKey = TokenAcquisitionOptions.LongRunningWebApiSessionKeyAuto
+                    }
+                };
+
+                // Act & Assert
+
+                // Step 3: First call with ClaimsPrincipal to initiate LR session
+                var scopes = new[] { "User.Read" };
+                mockHttpClient!.AddMockHandler(MockHttpCreator.CreateLrOboTokenHandler("User.Read"));
+                var result = await authorizationHeaderProvider.CreateAuthorizationHeaderForUserAsync(
+                    scopes,
+                    options,
+                    claimsPrincipal);
+
+                Assert.NotNull(result);
+                Assert.NotEqual(options.AcquireTokenOptions.LongRunningWebApiSessionKey, TokenAcquisitionOptions.LongRunningWebApiSessionKeyAuto);
+                string key1 = options.AcquireTokenOptions.LongRunningWebApiSessionKey;
+
+                // Step 4: Second call without ClaimsPrincipal should return the token from cache
+                result = await authorizationHeaderProvider.CreateAuthorizationHeaderForUserAsync(
+                                    scopes,
+                                    options);
+
+                Assert.NotNull(result);
+                Assert.NotEqual(options.AcquireTokenOptions.LongRunningWebApiSessionKey, TokenAcquisitionOptions.LongRunningWebApiSessionKeyAuto);
+                Assert.Equal(key1, options.AcquireTokenOptions.LongRunningWebApiSessionKey);
+
+                // Step 5: First call with ClaimsPrincipal to initiate LR session for CreateAuthorizationHeaderAsync
+                scopes = new[] { "User.Write" };
+                mockHttpClient!.AddMockHandler(MockHttpCreator.CreateLrOboTokenHandler("User.Write"));
+                result = await authorizationHeaderProvider.CreateAuthorizationHeaderAsync(
+                    scopes,
+                    options,
+                    claimsPrincipal);
+
+                Assert.NotNull(result);
+                Assert.NotEqual(options.AcquireTokenOptions.LongRunningWebApiSessionKey, TokenAcquisitionOptions.LongRunningWebApiSessionKeyAuto);
+                key1 = options.AcquireTokenOptions.LongRunningWebApiSessionKey;
+
+                // Step 6: Second call without ClaimsPrincipal should return the token from cache for CreateAuthorizationHeaderAsync
+                result = await authorizationHeaderProvider.CreateAuthorizationHeaderAsync(
+                                    scopes,
+                                    options);
+
+                Assert.NotNull(result);
+                Assert.NotEqual(options.AcquireTokenOptions.LongRunningWebApiSessionKey, TokenAcquisitionOptions.LongRunningWebApiSessionKeyAuto);
+                Assert.Equal(key1, options.AcquireTokenOptions.LongRunningWebApiSessionKey);
+            }
+        }
+
+        [Fact]
+        public async Task LongRunningSessionForDefaultAuthProviderForUserTest()
+        {
+            // Arrange
+            var tokenAcquirerFactory = InitTokenAcquirerFactoryForTest();
+            IServiceProvider serviceProvider = tokenAcquirerFactory.Build();
+
+            IAuthorizationHeaderProvider authorizationHeaderProvider =
+                serviceProvider.GetRequiredService<IAuthorizationHeaderProvider>();
+            var mockHttpClient = serviceProvider.GetRequiredService<IMsalHttpClientFactory>() as MockHttpClientFactory;
+
+            using (mockHttpClient)
+            {
+                // Create a test ClaimsPrincipal
+                var claims = new List<Claim>
+                            {
+                                new Claim(ClaimTypes.Name, "[email protected]")
+                            };
+
+                var identity = new CaseSensitiveClaimsIdentity(claims, "TestAuth");
+                identity.BootstrapContext = CreateTestJwt();
+                var claimsPrincipal = new ClaimsPrincipal(identity);
+
+                // Create options with LongRunningWebApiSessionKey
+                var options = new AuthorizationHeaderProviderOptions
+                {
+                    AcquireTokenOptions = new AcquireTokenOptions
+                    {
+                        LongRunningWebApiSessionKey = "oboKey1"
+                    }
+                };
+
+                // Act & Assert
+
+                // Step 3: First call with ClaimsPrincipal to initiate LR session
+                var scopes = new[] { "User.Read" };
+                mockHttpClient!.AddMockHandler(MockHttpCreator.CreateLrOboTokenHandler("User.Read"));
+                var result = await authorizationHeaderProvider.CreateAuthorizationHeaderForUserAsync(
+                    scopes,
+                    options,
+                    claimsPrincipal);
+
+                Assert.NotNull(result);
+                Assert.Equal("oboKey1", options.AcquireTokenOptions.LongRunningWebApiSessionKey);
+
+                // Step 4: Second call without ClaimsPrincipal should return the token from cache
+                result = await authorizationHeaderProvider.CreateAuthorizationHeaderForUserAsync(
+                                    scopes,
+                                    options);
+
+                Assert.NotNull(result);
+                Assert.Equal("oboKey1", options.AcquireTokenOptions.LongRunningWebApiSessionKey);
+
+                options = new AuthorizationHeaderProviderOptions
+                {
+                    AcquireTokenOptions = new AcquireTokenOptions
+                    {
+                        LongRunningWebApiSessionKey = "oboKey2"
+                    }
+                };
+
+                // Step 5: First call with ClaimsPrincipal to initiate LR session for CreateAuthorizationHeaderAsync
+                scopes = new[] { "User.Write" };
+                mockHttpClient!.AddMockHandler(MockHttpCreator.CreateLrOboTokenHandler("User.Write"));
+                result = await authorizationHeaderProvider.CreateAuthorizationHeaderAsync(
+                    scopes,
+                    options,
+                    claimsPrincipal);
+
+                Assert.NotNull(result);
+                Assert.Equal("oboKey2", options.AcquireTokenOptions.LongRunningWebApiSessionKey);
+
+                // Step 6: Second call without ClaimsPrincipal should return the token from cache for CreateAuthorizationHeaderAsync
+                result = await authorizationHeaderProvider.CreateAuthorizationHeaderAsync(
+                                    scopes,
+                                    options);
+
+                Assert.NotNull(result);
+                Assert.Equal("oboKey2", options.AcquireTokenOptions.LongRunningWebApiSessionKey);
+            }
+        }
+
+        private static string CreateTestJwt()
+        {
+            var header = new Dictionary<string, object>
+                {
+                    { "alg", "HS256" },
+                    { "typ", "JWT" }
+                };
+
+            var payload = new Dictionary<string, object>
+                {
+                    { "iss", "https://login.microsoftonline.com/test-tenant-id/v2.0" }
+                };
+
+            string headerJson = System.Text.Json.JsonSerializer.Serialize(header);
+            string payloadJson = System.Text.Json.JsonSerializer.Serialize(payload);
+
+            string headerBase64 = Base64UrlEncoder.Encode(headerJson);
+            string payloadBase64 = Base64UrlEncoder.Encode(payloadJson);
+
+            // For testing purposes, we're using a fixed signature
+            const string signature = "test_signature";
+            string signatureBase64 = Base64UrlEncoder.Encode(signature);
+
+            return $"{headerBase64}.{payloadBase64}.{signatureBase64}";
+        }
+
+        private TokenAcquirerFactory InitTokenAcquirerFactoryForTest()
+        {
+            TokenAcquirerFactoryTesting.ResetTokenAcquirerFactoryInTest();
+            TokenAcquirerFactory tokenAcquirerFactory = TokenAcquirerFactory.GetDefaultInstance();
+            tokenAcquirerFactory.Services.Configure<MicrosoftIdentityApplicationOptions>(options =>
+            {
+                options.Instance = "https://login.microsoftonline.com/";
+                options.TenantId = "testTenantId";
+                options.ClientId = "testClientId";
+                options.ClientCredentials = [ new CredentialDescription() {
+                                                SourceType = CredentialSource.ClientSecret,
+                                                ClientSecret = "test-secret"
+                                            }];
+            });
+
+            // Add required services
+            tokenAcquirerFactory.Services.AddSingleton<IMsalHttpClientFactory, MockHttpClientFactory>();
+            tokenAcquirerFactory.Services.AddScoped<IAuthorizationHeaderProvider, DefaultAuthorizationHeaderProvider>();
+
+            return tokenAcquirerFactory;
+        }
+    }
+}