From 6ce60b971f9c859e25dbfbbae4173df380e57593 Mon Sep 17 00:00:00 2001
From: Amer Jusupovic <32405726+amerjusupovic@users.noreply.github.com>
Date: Wed, 21 Feb 2024 13:23:21 -0800
Subject: [PATCH] Handle exceptions thrown during Key Vault reference
 resolution during startup (#518)

* treat keyvaultreferenceexception as failoverable when inner exception is failoverable

* check for aggregateexception, make isfailoverable for keyvault exception
---
 .../AzureAppConfigurationProvider.cs          | 34 ++++++++++++
 .../FailoverTests.cs                          | 53 +++++++++++++++++++
 2 files changed, 87 insertions(+)

diff --git a/src/Microsoft.Extensions.Configuration.AzureAppConfiguration/AzureAppConfigurationProvider.cs b/src/Microsoft.Extensions.Configuration.AzureAppConfiguration/AzureAppConfigurationProvider.cs
index 5090487d..06e7c155 100644
--- a/src/Microsoft.Extensions.Configuration.AzureAppConfiguration/AzureAppConfigurationProvider.cs
+++ b/src/Microsoft.Extensions.Configuration.AzureAppConfiguration/AzureAppConfigurationProvider.cs
@@ -630,6 +630,17 @@ private async Task<bool> TryInitializeAsync(IEnumerable<ConfigurationClient> cli
 
                 throw;
             }
+            catch (KeyVaultReferenceException exception)
+            {
+                if (IsFailOverable(exception))
+                {
+                    startupExceptions.Add(exception);
+
+                    return false;
+                }
+
+                throw;
+            }
             catch (AggregateException exception)
             {
                 if (exception.InnerExceptions?.Any(e => e is OperationCanceledException) ?? false)
@@ -980,6 +991,15 @@ private async Task<T> ExecuteWithFailOverPolicyAsync<T>(
                         throw;
                     }
                 }
+                catch (KeyVaultReferenceException kvre)
+                {
+                    if (!IsFailOverable(kvre) || !clientEnumerator.MoveNext())
+                    {
+                        backoffAllClients = true;
+
+                        throw;
+                    }
+                }
                 catch (AggregateException ae)
                 {
                     if (!IsFailOverable(ae) || !clientEnumerator.MoveNext())
@@ -1066,6 +1086,20 @@ innerException is SocketException ||
                    innerException is IOException;
         }
 
+        private bool IsFailOverable(KeyVaultReferenceException kvre)
+        {
+            if (kvre.InnerException is RequestFailedException rfe && IsFailOverable(rfe))
+            {
+                return true;
+            }
+            else if (kvre.InnerException is AggregateException ae && IsFailOverable(ae))
+            {
+                return true;
+            }
+
+            return false;
+        }
+
         private async Task<Dictionary<string, ConfigurationSetting>> MapConfigurationSettings(Dictionary<string, ConfigurationSetting> data)
         {
             Dictionary<string, ConfigurationSetting> mappedData = new Dictionary<string, ConfigurationSetting>(StringComparer.OrdinalIgnoreCase);
diff --git a/tests/Tests.AzureAppConfiguration/FailoverTests.cs b/tests/Tests.AzureAppConfiguration/FailoverTests.cs
index 78d1e858..32b4cd50 100644
--- a/tests/Tests.AzureAppConfiguration/FailoverTests.cs
+++ b/tests/Tests.AzureAppConfiguration/FailoverTests.cs
@@ -257,5 +257,58 @@ public void FailOverTests_BackoffStateIsUpdatedOnSuccessfulRequest()
             // The client enumerator should return 2 clients for the third time.
             Assert.Equal(2, configClientManager.GetAvailableClients(DateTimeOffset.UtcNow).Count());
         }
+
+        [Fact]
+        public void FailOverTests_FailOverOnKeyVaultReferenceException()
+        {
+            // Arrange
+            IConfigurationRefresher refresher = null;
+            var mockResponse = new Mock<Response>();
+
+            var mockClient1 = new Mock<ConfigurationClient>();
+            mockClient1.Setup(c => c.GetConfigurationSettingsAsync(It.IsAny<SettingSelector>(), It.IsAny<CancellationToken>()))
+                       .Throws(new KeyVaultReferenceException("Key vault reference failed.", new RequestFailedException(503, "Request failed.")));
+            mockClient1.Setup(c => c.GetConfigurationSettingAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
+                       .Throws(new KeyVaultReferenceException("Key vault reference failed.", new RequestFailedException(503, "Request failed.")));
+            mockClient1.Setup(c => c.GetConfigurationSettingAsync(It.IsAny<ConfigurationSetting>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()))
+                       .Throws(new KeyVaultReferenceException("Key vault reference failed.", new RequestFailedException(503, "Request failed.")));
+            mockClient1.Setup(c => c.Equals(mockClient1)).Returns(true);
+
+            var mockClient2 = new Mock<ConfigurationClient>();
+            mockClient2.Setup(c => c.GetConfigurationSettingsAsync(It.IsAny<SettingSelector>(), It.IsAny<CancellationToken>()))
+                       .Returns(new MockAsyncPageable(Enumerable.Empty<ConfigurationSetting>().ToList()));
+            mockClient2.Setup(c => c.GetConfigurationSettingAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
+                       .Returns(Task.FromResult(Response.FromValue<ConfigurationSetting>(kv, mockResponse.Object)));
+            mockClient2.Setup(c => c.GetConfigurationSettingAsync(It.IsAny<ConfigurationSetting>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()))
+                       .Returns(Task.FromResult(Response.FromValue<ConfigurationSetting>(kv, mockResponse.Object)));
+            mockClient2.Setup(c => c.Equals(mockClient2)).Returns(true);
+
+            ConfigurationClientWrapper cw1 = new ConfigurationClientWrapper(TestHelpers.PrimaryConfigStoreEndpoint, mockClient1.Object);
+            ConfigurationClientWrapper cw2 = new ConfigurationClientWrapper(TestHelpers.SecondaryConfigStoreEndpoint, mockClient2.Object);
+
+            var clientList = new List<ConfigurationClientWrapper>() { cw1, cw2 };
+            var configClientManager = new ConfigurationClientManager(clientList);
+
+            var config = new ConfigurationBuilder()
+                .AddAzureAppConfiguration(options =>
+                {
+                    options.ClientManager = configClientManager;
+                    options.Select("TestKey*");
+                    options.ConfigureRefresh(refreshOptions =>
+                    {
+                        refreshOptions.Register("TestKey1", "label")
+                            .SetCacheExpiration(TimeSpan.FromSeconds(1));
+                    });
+
+                    refresher = options.GetRefresher();
+                })
+                .Build();
+
+            // The client enumerator should return just 1 client.
+            Assert.Single(configClientManager.GetAvailableClients(DateTimeOffset.UtcNow));
+
+            // The build should be successful since one client was backed off and it failed over to the second client.
+            Assert.Equal("TestValue1", config["TestKey1"]);
+        }
     }
 }