Lazy load temporary access key

Author: terencefan

src/Microsoft.Azure.SignalR.Common/Auth/MicrosoftEntra/AccessKeySynchronizer.cs

@@ -38,27 +38,38 @@ internal AccessKeySynchronizer(ILoggerFactory loggerFactory, bool start)
         _logger = (loggerFactory ?? NullLoggerFactory.Instance).CreateLogger<AccessKeySynchronizer>();
     }
 
-    public void AddServiceEndpoint(ServiceEndpoint endpoint)
+    public Task AddServiceEndpoint(ServiceEndpoint endpoint)
     {
         if (endpoint.AccessKey is MicrosoftEntraAccessKey key)
         {
-            _ = InitializeAccessKeyTask(key);
             _keyMap.TryAdd(key, true);
+            return InitializeAccessKeyTask(key);
         }
+        return Task.CompletedTask;
     }
 
     public void Dispose() => _timer.Stop();
 
     public void UpdateServiceEndpoints(IEnumerable<ServiceEndpoint> endpoints)
     {
+        _keyMap.Clear();
         foreach (var endpoint in endpoints)
         {
             AddServiceEndpoint(endpoint);
         }
     }
 
+    /// <summary>
+    /// Test only
+    /// </summary>
+    /// <param name="e"></param>
+    /// <returns></returns>
     internal bool ContainsKey(ServiceEndpoint e) => _keyMap.ContainsKey(e.AccessKey as MicrosoftEntraAccessKey);
 
+    /// <summary>
+    /// Test only
+    /// </summary>
+    /// <returns></returns>
     internal int Count() => _keyMap.Count;
 
     private static async Task InitializeAccessKeyTask(MicrosoftEntraAccessKey key)

src/Microsoft.Azure.SignalR.Common/Auth/MicrosoftEntra/IAccessKeySynchronizer.cs

@@ -2,12 +2,13 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System.Collections.Generic;
+using System.Threading.Tasks;
 
 namespace Microsoft.Azure.SignalR;
 
 internal interface IAccessKeySynchronizer
 {
-    public void AddServiceEndpoint(ServiceEndpoint endpoint);
+    public Task AddServiceEndpoint(ServiceEndpoint endpoint);
 
     public void UpdateServiceEndpoints(IEnumerable<ServiceEndpoint> endpoints);
 }

test/Microsoft.Azure.SignalR.Common.Tests/Endpoints/AccessKeySynchronizerFacts.cs

@@ -1,6 +1,9 @@
 // Copyright (c) Microsoft. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
+using System;
+using System.Threading;
+using System.Threading.Tasks;
 using Azure.Identity;
 using Microsoft.Azure.SignalR.Tests.Common;
 using Microsoft.Extensions.Logging.Abstractions;
@@ -34,6 +37,54 @@ public void AddAndRemoveServiceEndpointsTest()
         Assert.Equal(0, synchronizer.Count());
     }
 
+    [Fact]
+    public async Task GenerateAccessTokenTriggerLoadedTaskTest()
+    {
+        var synchronizer = GetInstanceForTest();
+
+        var credential = new DefaultAzureCredential();
+        var endpoint = new TestServiceEndpoint(credential);
+
+        var initializeTask = synchronizer.AddServiceEndpoint(endpoint);
+
+        var source = new CancellationTokenSource(1000);
+
+        var key = Assert.IsType<MicrosoftEntraAccessKey>(endpoint.AccessKey);
+        Assert.False(key.LoadedTask.IsCompleted);
+
+        await Assert.ThrowsAsync<TaskCanceledException>(
+            async () => await key.GenerateAccessTokenAsync("https://localhost", [], TimeSpan.FromSeconds(1), AccessTokenAlgorithm.HS256, source.Token)
+        );
+        Assert.True(key.LoadedTask.IsCompletedSuccessfully);
+
+        await initializeTask;
+        Assert.False(key.Available);
+
+        await Assert.ThrowsAsync<AzureSignalRAccessTokenNotAuthorizedException>(
+            async () => await key.GenerateAccessTokenAsync("https://localhost", [], TimeSpan.FromSeconds(1), AccessTokenAlgorithm.HS256)
+        );
+    }
+
+    [Fact(Skip = "local test only")]
+    public async Task GenerateAccessTokenTest()
+    {
+        var synchronizer = GetInstanceForTest();
+        var credential = new VisualStudioCredential();
+        var endpoint = new ServiceEndpoint(new Uri("https://test-signalr-sdk.service.signalr.net"), credential);
+
+        var key = Assert.IsType<MicrosoftEntraAccessKey>(endpoint.AccessKey);
+        Assert.False(key.LoadedTask.IsCompleted);
+
+        var initializeTask = synchronizer.AddServiceEndpoint(endpoint);
+
+        var token = await key.GenerateAccessTokenAsync("https://localhost", [], TimeSpan.FromMinutes(1), AccessTokenAlgorithm.HS256);
+        Assert.NotNull(token);
+
+        Assert.True(initializeTask.IsCompleted);
+        Assert.True(key.LoadedTask.IsCompleted);
+        Assert.True(key.Available);
+    }
+
     private static AccessKeySynchronizer GetInstanceForTest()
     {
         return new AccessKeySynchronizer(NullLoggerFactory.Instance, false);

test/Microsoft.Azure.SignalR.Tests.Common/TestClasses/TestAccessKeySynchronizer.cs

@@ -1,4 +1,5 @@
 using System.Collections.Generic;
+using System.Threading.Tasks;
 
 namespace Microsoft.Azure.SignalR.Tests.Common;
 
@@ -10,7 +11,8 @@ public void UpdateServiceEndpoints(IEnumerable<ServiceEndpoint> endpoints)
     {
     }
 
-    public void AddServiceEndpoint(ServiceEndpoint endpoint)
+    public Task AddServiceEndpoint(ServiceEndpoint endpoint)
     {
+        return Task.CompletedTask;
     }
 }