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();
         }
 
@@ -43,11 +46,14 @@ public async Task<string> CreateAuthorizationHeaderForAppAsync(
             AuthorizationHeaderProviderOptions? downstreamApiOptions = null,
             CancellationToken cancellationToken = default)
         {
+            var newTokenAcquisitionOptions = CreateTokenAcquisitionOptionsFromApiOptions(downstreamApiOptions, cancellationToken);
             var result = await _tokenAcquisition.GetAuthenticationResultForAppAsync(
                 scopes,
                 downstreamApiOptions?.AcquireTokenOptions.AuthenticationOptionsName,
                 downstreamApiOptions?.AcquireTokenOptions.Tenant,
-                CreateTokenAcquisitionOptionsFromApiOptions(downstreamApiOptions, cancellationToken)).ConfigureAwait(false);
+                newTokenAcquisitionOptions).ConfigureAwait(false);
+
+            UpdateOriginalTokenAcquisitionOptions(downstreamApiOptions?.AcquireTokenOptions, newTokenAcquisitionOptions);
             return result.CreateAuthorizationHeader();
         }
 
@@ -59,6 +65,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 +82,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 +92,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 +121,14 @@ 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)
+        {
+            acquireTokenOptions!.LongRunningWebApiSessionKey = newTokenAcquisitionOptions.LongRunningWebApiSessionKey;
+        }
     }
 }

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

@@ -0,0 +1,126 @@
+// 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 LongRunningSessionTest()
+        {
+            // Arrange
+            var tokenAcquirerFactory = InitTokenAcquirerFactoryForTest();
+            IServiceProvider serviceProvider = tokenAcquirerFactory.Build();
+
+            IAuthorizationHeaderProvider authorizationHeaderProvider =
+                serviceProvider.GetRequiredService<IAuthorizationHeaderProvider>();
+            var mockHttpClient = serviceProvider.GetRequiredService<IMsalHttpClientFactory>() as MockHttpClientFactory;
+
+            mockHttpClient!.AddMockHandler(MockHttpCreator.CreateClientCredentialTokenHandler());
+
+            // 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 = TokenAcquisitionOptions.LongRunningWebApiSessionKeyAuto
+                }
+            };
+
+            // Act & Assert
+
+            // Step 3: First call with ClaimsPrincipal to initiate LR session
+            var scopes = new[] { "https://graph.microsoft.com/.default" };
+            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
+            var result2 = await authorizationHeaderProvider.CreateAuthorizationHeaderForUserAsync(
+                                scopes,
+                                options,
+                                claimsPrincipal: null);
+
+            Assert.NotNull(result);
+            Assert.NotEqual(options.AcquireTokenOptions.LongRunningWebApiSessionKey, TokenAcquisitionOptions.LongRunningWebApiSessionKeyAuto);
+            Assert.Equal(key1, 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;
+        }
+    }
+}