Improve alias caching

Author: Avery-Dunn

src/client/Microsoft.Identity.Client/Instance/Discovery/INetworkCacheMetadataProvider.cs

@@ -8,6 +8,13 @@ namespace Microsoft.Identity.Client.Instance.Discovery
     internal interface INetworkCacheMetadataProvider
     {
         void AddMetadata(string environment, InstanceDiscoveryMetadataEntry entry);
+
+        /// <summary>
+        /// Caches <paramref name="entry"/> under each host in <paramref name="entry"/>.Aliases.
+        /// If the entry has no aliases, falls back to <paramref name="fallbackEnvironment"/>.
+        /// </summary>
+        void AddMetadataWithAliases(InstanceDiscoveryMetadataEntry entry, string fallbackEnvironment);
+
         InstanceDiscoveryMetadataEntry GetMetadata(string environment, ILoggerAdapter logger);
     }
 }

src/client/Microsoft.Identity.Client/Instance/Discovery/InstanceDiscoveryManager.cs

@@ -212,7 +212,7 @@ private async Task<InstanceDiscoveryMetadataEntry> FetchNetworkMetadataOrFallbac
                     _knownMetadataProvider.GetMetadata(authorityUri.Host, Enumerable.Empty<string>(), requestContext.Logger)
                     ?? CreateEntryForSingleAuthority(authorityUri);
 
-                _networkCacheMetadataProvider.AddMetadata(authorityUri.Host, fallbackEntry);
+                _networkCacheMetadataProvider.AddMetadataWithAliases(fallbackEntry, authorityUri.Host);
 
                 return fallbackEntry;
             }

src/client/Microsoft.Identity.Client/Instance/Discovery/NetworkCacheMetadataProvider.cs

@@ -26,6 +26,23 @@ public void AddMetadata(string environment, InstanceDiscoveryMetadataEntry entry
             s_cache.AddOrUpdate(environment, entry, (_, _) => entry);
         }
 
+        public void AddMetadataWithAliases(InstanceDiscoveryMetadataEntry entry, string fallbackEnvironment)
+        {
+            string[] aliases = entry?.Aliases;
+            if (aliases != null && aliases.Length > 0)
+            {
+                foreach (string alias in aliases)
+                {
+                    if (!string.IsNullOrWhiteSpace(alias))
+                        s_cache.AddOrUpdate(alias, entry, (_, _) => entry);
+                }
+            }
+            else if (!string.IsNullOrWhiteSpace(fallbackEnvironment))
+            {
+                s_cache.AddOrUpdate(fallbackEnvironment, entry, (_, _) => entry);
+            }
+        }
+
         internal static void ResetStaticCacheForTest()
         {
             s_cache.Clear();

src/client/Microsoft.Identity.Client/Instance/Discovery/NetworkMetadataProvider.cs

@@ -67,10 +67,7 @@ private void CacheInstanceDiscoveryMetadata(InstanceDiscoveryResponse instanceDi
         {
             foreach (InstanceDiscoveryMetadataEntry entry in instanceDiscoveryResponse.Metadata ?? Enumerable.Empty<InstanceDiscoveryMetadataEntry>())
             {
-                foreach (string aliasedEnvironment in entry.Aliases ?? Enumerable.Empty<string>())
-                {
-                    _networkCacheMetadataProvider.AddMetadata(aliasedEnvironment, entry);
-                }
+                _networkCacheMetadataProvider.AddMetadataWithAliases(entry, fallbackEnvironment: null);
             }
         }
 

tests/Microsoft.Identity.Test.Unit/CoreTests/InstanceTests/InstanceDiscoveryManagerTests.cs

@@ -251,6 +251,52 @@ public async Task NetworkProviderFailures_AreIgnored_Async()
             _knownMetadataProvider.Received(1).GetMetadata("some_env.com", Enumerable.Empty<string>(), Arg.Any<ILoggerAdapter>());
         }
 
+        [TestMethod]
+        public async Task NetworkProviderFailures_FallbackEntry_AllAliasesCached_Async()
+        {
+            // Arrange — use a real cache so we can inspect what was written
+            var realNetworkCache = new NetworkCacheMetadataProvider();
+
+            var multiAliasEntry = new InstanceDiscoveryMetadataEntry()
+            {
+                Aliases = new[] { "some_env.com", "alias1.some_env.com", "alias2.some_env.com" },
+                PreferredCache = "some_env.com",
+                PreferredNetwork = "some_env.com"
+            };
+
+            _knownMetadataProvider
+                .GetMetadata("some_env.com", Enumerable.Empty<string>(), Arg.Any<ILoggerAdapter>())
+                .Returns(multiAliasEntry);
+
+            _networkMetadataProvider
+                .When(x => x.GetMetadataAsync(Arg.Any<Uri>(), _testRequestContext))
+                .Do(_ => throw new MsalServiceException("endpoint_busy", "simulated failure"));
+
+            _discoveryManager = new InstanceDiscoveryManager(
+                _harness.HttpManager,
+                null,
+                null,
+                _knownMetadataProvider,
+                realNetworkCache,
+                _networkMetadataProvider);
+
+            // Act
+            var result = await _discoveryManager.GetMetadataEntryAsync(
+                AuthorityInfo.FromAuthorityUri("https://some_env.com/tid", true),
+                _testRequestContext)
+                .ConfigureAwait(false);
+
+            Assert.AreSame(multiAliasEntry, result);
+
+            // Assert — fallback entry must be cached under every alias so subsequent requests
+            // using any alias won't retry the failing network call
+            ILoggerAdapter logger = _testRequestContext.Logger;
+            Assert.IsNotNull(realNetworkCache.GetMetadata("some_env.com", logger), "Primary host should be cached");
+            Assert.IsNotNull(realNetworkCache.GetMetadata("alias1.some_env.com", logger), "First alias should be cached");
+            Assert.IsNotNull(realNetworkCache.GetMetadata("alias2.some_env.com", logger), "Second alias should be cached");
+            Assert.IsNull(realNetworkCache.GetMetadata("unrelated_env.com", logger), "Unrelated host should not be in cache");
+        }
+
         [TestMethod]
         public async Task NetworkProviderFailures_WithNoKnownMetadata_ContinuesWithAuthority_Async()
         {