Introduces a lock while loading credentials from Credential Source (#…

…2438) * Introduces a lock while loading credentials from Credential Source * Make type nullable * Updates Abstractions version and CredentialLoader implementation * Remove unnecessary usings * Address feedback
AzureAD · Sep 7, 2023 · 050e3e0 · 050e3e0
1 parent 52cec44
commit 050e3e0
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 4 deletions.
diff --git a/Directory.Build.props b/Directory.Build.props
@@ -80,7 +80,7 @@
     <MicrosoftGraphVersion>4.34.0</MicrosoftGraphVersion>
     <MicrosoftGraphBetaVersion>4.50.0-preview</MicrosoftGraphBetaVersion>
     <MicrosoftExtensionsHttpVersion>3.1.3</MicrosoftExtensionsHttpVersion>
-    <MicrosoftIdentityAbstractions>4.0.0</MicrosoftIdentityAbstractions>
+    <MicrosoftIdentityAbstractions>4.1.0</MicrosoftIdentityAbstractions>
     <!--CVE-2021-24112-->
     <SystemDrawingCommon>4.7.2</SystemDrawingCommon>
   </PropertyGroup>

diff --git a/src/Microsoft.Identity.Web.Certificate/DefaultCredentialsLoader.cs b/src/Microsoft.Identity.Web.Certificate/DefaultCredentialsLoader.cs
@@ -1,8 +1,9 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
 using Microsoft.Identity.Abstractions;
@@ -15,6 +16,7 @@ namespace Microsoft.Identity.Web
     public class DefaultCredentialsLoader : ICredentialsLoader
     {
         ILogger<DefaultCredentialsLoader>? _logger;
+        private readonly ConcurrentDictionary<string, SemaphoreSlim> _loadingSemaphores = new ConcurrentDictionary<string, SemaphoreSlim>();
 
         /// <summary>
         /// Constructor with a logger
@@ -56,9 +58,24 @@ public async Task LoadCredentialsIfNeededAsync(CredentialDescription credentialD
 
             if (credentialDescription.CachedValue == null)
             {
-                if (CredentialSourceLoaders.TryGetValue(credentialDescription.SourceType, out ICredentialSourceLoader? loader))
+                // Get or create a semaphore for this credentialDescription
+                var semaphore = _loadingSemaphores.GetOrAdd(credentialDescription.Id, (v) => new SemaphoreSlim(1));
+
+                // Wait to acquire the semaphore
+                await semaphore.WaitAsync();
+
+                try
+                {
+                    if (credentialDescription.CachedValue == null)
+                    {
+                        if (CredentialSourceLoaders.TryGetValue(credentialDescription.SourceType, out ICredentialSourceLoader? loader))
+                            await loader.LoadIfNeededAsync(credentialDescription, parameters);
+                    }
+                }
+                finally
                 {
-                    await loader.LoadIfNeededAsync(credentialDescription, parameters);
+                    // Release the semaphore
+                    semaphore.Release();
                 }
             }
         }

diff --git a/tests/IntegrationTests/TokenAcquirerTests/TokenAcquirer.cs b/tests/IntegrationTests/TokenAcquirerTests/TokenAcquirer.cs
@@ -14,6 +14,7 @@
 using Microsoft.Identity.Web.TokenCacheProviders.InMemory;
 using Microsoft.IdentityModel.Tokens;
 using Xunit;
+using TaskStatus = System.Threading.Tasks.TaskStatus;
 
 namespace TokenAcquirerTests
 {
@@ -126,6 +127,46 @@ public async Task AcquireToken_WithFactoryAndAuthorityClientIdCert_ClientCredent
             Assert.False(string.IsNullOrEmpty(result.AccessToken));
         }
 
+        [IgnoreOnAzureDevopsFact]
+        //[Fact]
+        public async Task LoadCredentialsIfNeededAsync_MultipleThreads_WaitsForSemaphore()
+        {
+            TokenAcquirerFactory tokenAcquirerFactory = TokenAcquirerFactory.GetDefaultInstance();
+            IServiceCollection services = tokenAcquirerFactory.Services;
+
+            services.Configure<MicrosoftIdentityApplicationOptions>(s_optionName, option =>
+            {
+                option.Instance = "https://login.microsoftonline.com/";
+                option.TenantId = "msidentitysamplestesting.onmicrosoft.com";
+                option.ClientId = "6af093f3-b445-4b7a-beae-046864468ad6";
+                option.ClientCredentials = s_clientCredentials;
+            });
+
+            services.AddInMemoryTokenCaches();
+            var serviceProvider = tokenAcquirerFactory.Build();
+            var options = serviceProvider.GetRequiredService<IOptionsMonitor<MicrosoftIdentityApplicationOptions>>().Get(s_optionName);
+            var credentialsLoader = serviceProvider.GetRequiredService<ICredentialsLoader>();
+
+            var task1 = Task.Run(async () =>
+            {
+                await credentialsLoader.LoadCredentialsIfNeededAsync(options.ClientCredentials!.First());
+            });
+
+            var task2 = Task.Run(async () =>
+            {
+                await credentialsLoader.LoadCredentialsIfNeededAsync(options.ClientCredentials!.First());
+            });
+
+            // Run task1 and task2 concurrently
+            await Task.WhenAll(task1, task2);
+
+            var cert = options.ClientCredentials!.First().Certificate;
+
+            Assert.NotNull(cert);
+            Assert.Equal(TaskStatus.RanToCompletion, task1.Status);
+            Assert.Equal(TaskStatus.RanToCompletion, task2.Status);
+        }
+
         [IgnoreOnAzureDevopsFact]
         //[Fact]
         public async Task AcquireTokenWithPop_ClientCredentialsAsync()