Allow retries in DefaultKeyResolver.CanCreateAuthenticatedEncryptor

This code is trying to ensure that the selected key can be decrypted (i.e. is usable).  It may fail if, for example, Azure KeyVault is unreachable due to connectivity issues.  If it fails, there's a log message and then an immediately-activated key will be generated.  An immediately-activated key can cause problems for sessions making requests to multiple app instances and those problems won't obviously be connected to the (almost silent) failure in CanCreateAuthenticatedEncryptor.  Rather than effectively swallowing such errors, we should allow some retries.

Part of dotnet#52678, which is part of dotnet#36157

Author: amcasey

src/DataProtection/DataProtection/src/KeyManagement/DefaultKeyResolver.cs

@@ -3,7 +3,9 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
+using System.Threading;
 using Microsoft.AspNetCore.Cryptography;
 using Microsoft.AspNetCore.DataProtection.AuthenticatedEncryption;
 using Microsoft.AspNetCore.DataProtection.KeyManagement.Internal;
@@ -28,6 +30,9 @@ internal sealed class DefaultKeyResolver : IDefaultKeyResolver
     /// </remarks>
     private readonly TimeSpan _keyPropagationWindow;
 
+    private readonly int _maxDecryptRetries;
+    private readonly TimeSpan _decryptRetryDelay;
+
     private readonly ILogger _logger;
 
     /// <summary>
@@ -41,33 +46,65 @@ internal sealed class DefaultKeyResolver : IDefaultKeyResolver
     /// </remarks>
     private readonly TimeSpan _maxServerToServerClockSkew;
 
-    public DefaultKeyResolver()
-        : this(NullLoggerFactory.Instance)
+    public DefaultKeyResolver(IOptions<KeyManagementOptions> keyManagementOptions)
+        : this(keyManagementOptions, NullLoggerFactory.Instance)
     { }
 
-    public DefaultKeyResolver(ILoggerFactory loggerFactory)
+    public DefaultKeyResolver(IOptions<KeyManagementOptions> keyManagementOptions, ILoggerFactory loggerFactory)
     {
         _keyPropagationWindow = KeyManagementOptions.KeyPropagationWindow;
         _maxServerToServerClockSkew = KeyManagementOptions.MaxServerClockSkew;
+        _maxDecryptRetries = keyManagementOptions.Value.MaximumDefaultKeyResolverRetries;
+        _decryptRetryDelay = keyManagementOptions.Value.DefaultKeyResolverRetryDelay;
         _logger = loggerFactory.CreateLogger<DefaultKeyResolver>();
     }
 
-    private bool CanCreateAuthenticatedEncryptor(IKey key)
+    private bool CanCreateAuthenticatedEncryptor(IKey key, ref int retriesRemaining)
     {
-        try
+        List<Exception>? exceptions = null;
+        while (true)
         {
-            var encryptorInstance = key.CreateEncryptor();
-            if (encryptorInstance == null)
+            try
+            {
+                var encryptorInstance = key.CreateEncryptor();
+                if (encryptorInstance is null)
+                {
+                    // This generally indicates a key (descriptor) without a corresponding IAuthenticatedEncryptorFactory,
+                    // which is not something that can succeed on retry - return immediately
+                    _logger.KeyIsIneligibleToBeTheDefaultKeyBecauseItsMethodFailed(
+                        key.KeyId,
+                        nameof(IKey.CreateEncryptor),
+                        new InvalidOperationException($"{nameof(IKey.CreateEncryptor)} returned null."));
+                    return false;
+                }
+
+                return true;
+            }
+            catch (Exception ex)
             {
-                CryptoUtil.Fail<IAuthenticatedEncryptor>("CreateEncryptorInstance returned null.");
+                if (exceptions is null)
+                {
+                    // The first failure doesn't count as a retry
+                    exceptions = new();
+                }
+                else
+                {
+                    retriesRemaining--;
+                }
+
+                exceptions.Add(ex);
+
+                if (retriesRemaining == 0)
+                {
+                    _logger.KeyIsIneligibleToBeTheDefaultKeyBecauseItsMethodFailed(key.KeyId, nameof(IKey.CreateEncryptor), new AggregateException(exceptions));
+                    return false;
+                }
+
+                _logger.RetryingMethodOfKeyAfterFailure(key.KeyId, nameof(IKey.CreateEncryptor), ex);
             }
 
-            return true;
-        }
-        catch (Exception ex)
-        {
-            _logger.KeyIsIneligibleToBeTheDefaultKeyBecauseItsMethodFailed(key.KeyId, nameof(IKey.CreateEncryptor), ex);
-            return false;
+            // Don't retry immediately
+            Thread.Sleep(_decryptRetryDelay);
         }
     }
 
@@ -94,12 +131,14 @@ where key.CreationDate > propagationCutoff
                                                       orderby key.ActivationDate ascending, key.KeyId ascending
                                                       select key).FirstOrDefault();
 
+        var decryptRetriesRemaining = _maxDecryptRetries;
+
         if (preferredDefaultKey != null)
         {
             _logger.ConsideringKeyWithExpirationDateAsDefaultKey(preferredDefaultKey.KeyId, preferredDefaultKey.ExpirationDate);
 
             // if the key has been revoked or is expired, it is no longer a candidate
-            if (preferredDefaultKey.IsRevoked || preferredDefaultKey.IsExpired(now) || !CanCreateAuthenticatedEncryptor(preferredDefaultKey))
+            if (preferredDefaultKey.IsRevoked || preferredDefaultKey.IsExpired(now) || !CanCreateAuthenticatedEncryptor(preferredDefaultKey, ref decryptRetriesRemaining))
             {
                 _logger.KeyIsNoLongerUnderConsiderationAsDefault(preferredDefaultKey.KeyId);
             }
@@ -123,13 +162,14 @@ where key.CreationDate > propagationCutoff
         // keep trying until we find one that works.
         var unrevokedKeys = allKeys.Where(key => !key.IsRevoked);
         fallbackKey = (from key in (from key in unrevokedKeys
+                                    where !ReferenceEquals(key, preferredDefaultKey) // Don't reconsider it as a fallback
                                     where key.CreationDate <= propagationCutoff
                                     orderby key.CreationDate descending
                                     select key).Concat(from key in unrevokedKeys
                                                        where key.CreationDate > propagationCutoff
                                                        orderby key.CreationDate ascending
                                                        select key)
-                       where CanCreateAuthenticatedEncryptor(key)
+                       where CanCreateAuthenticatedEncryptor(key, ref decryptRetriesRemaining)
                        select key).FirstOrDefault();
 
         _logger.RepositoryContainsNoViableDefaultKey();

src/DataProtection/DataProtection/src/KeyManagement/KeyManagementOptions.cs

@@ -94,6 +94,25 @@ internal static TimeSpan MaxServerClockSkew
         }
     }
 
+    /// <summary>
+    /// During each default key resolution, if a key decryption attempt fails,
+    /// it can be retried, as long as the total number of retries across all keys
+    /// does not exceed this value.
+    /// </summary>
+    /// <remarks>
+    /// Settable for testing.
+    /// </remarks>
+    internal int MaximumDefaultKeyResolverRetries { get; set; } = 10;
+
+    /// <summary>
+    /// Wait this long before each default key resolution decryption retry.
+    /// <seealso cref="MaximumDefaultKeyResolverRetries"/>
+    /// </summary>
+    /// <remarks>
+    /// Settable for testing.
+    /// </remarks>
+    internal TimeSpan DefaultKeyResolverRetryDelay { get; set; } = TimeSpan.FromMilliseconds(200);
+
     /// <summary>
     /// Controls the lifetime (number of days before expiration)
     /// for newly-generated keys.

src/DataProtection/DataProtection/src/LoggingExtensions.cs

@@ -252,4 +252,7 @@ private static bool IsLogLevelEnabledCore([NotNullWhen(true)] ILogger? logger, L
 
     [LoggerMessage(72, LogLevel.Error, "The key ring's default data protection key {KeyId:B} has been revoked.", EventName = "KeyRingDefaultKeyIsRevoked")]
     public static partial void KeyRingDefaultKeyIsRevoked(this ILogger logger, Guid keyId);
+
+    [LoggerMessage(73, LogLevel.Debug, "Key {KeyId:B} method {MethodName} failed. Retrying.", EventName = "RetryingMethodOfKeyAfterFailure")]
+    public static partial void RetryingMethodOfKeyAfterFailure(this ILogger logger, Guid keyId, string methodName, Exception exception);
 }

src/DataProtection/DataProtection/test/Microsoft.AspNetCore.DataProtection.Tests/KeyManagement/DefaultKeyResolverTests.cs

@@ -291,9 +291,113 @@ public void ResolveDefaultKeyPolicy_OlderUnpropagatedKeyPreferred()
         Assert.False(resolution.ShouldGenerateNewKey);
     }
 
+    [Fact]
+    public void CreateEncryptor_NoRetryOnNullReturn()
+    {
+        // Arrange
+        var resolver = CreateDefaultKeyResolver();
+
+        var now = ParseDateTimeOffset("2010-01-01 00:00:00Z");
+
+        var mockKey = new Mock<IKey>();
+        mockKey.Setup(o => o.KeyId).Returns(Guid.NewGuid());
+        mockKey.Setup(o => o.CreationDate).Returns(now.AddDays(-3)); // Propagated
+        mockKey.Setup(o => o.ActivationDate).Returns(now.AddDays(-1)); // Activated
+        mockKey.Setup(o => o.ExpirationDate).Returns(now.AddDays(14)); // Unexpired
+        mockKey.Setup(o => o.IsRevoked).Returns(false); // Unrevoked
+        mockKey.Setup(o => o.CreateEncryptor()).Returns((IAuthenticatedEncryptor)null); // Anomalous null return
+
+        // Act
+        var resolution = resolver.ResolveDefaultKeyPolicy(now, [mockKey.Object]);
+
+        // Assert
+        Assert.Null(resolution.DefaultKey);
+        Assert.Null(resolution.FallbackKey);
+        Assert.True(resolution.ShouldGenerateNewKey);
+
+        mockKey.Verify(o => o.CreateEncryptor(), Times.Once); // Not retried
+    }
+
+    [Fact]
+    public void CreateEncryptor_FirstAttemptIsNotARetry()
+    {
+        // Arrange
+        var options = Options.Create(new KeyManagementOptions()
+        {
+            MaximumDefaultKeyResolverRetries = 3,
+            DefaultKeyResolverRetryDelay = TimeSpan.Zero,
+        });
+
+        var resolver = new DefaultKeyResolver(options, NullLoggerFactory.Instance);
+
+        var now = ParseDateTimeOffset("2010-01-01 00:00:00Z");
+
+        var creation1 = now.AddDays(-3);
+        var activation1 = creation1.AddDays(2);
+        var expiration1 = creation1.AddDays(90);
+
+        // Newer but still propagated => preferred
+        var creation2 = creation1.AddHours(1);
+        var activation2 = activation1.AddHours(1);
+        var expiration2 = expiration1.AddHours(1);
+
+        var mockKey1 = CreateMockKey(activation1, expiration1, creation1, isRevoked: false, createEncryptorThrows: false);
+        var mockKey2 = CreateMockKey(activation2, expiration2, creation2, isRevoked: false, createEncryptorThrows: true); // Uses up all the retries
+
+        // Act
+        var resolution = resolver.ResolveDefaultKeyPolicy(now, [mockKey1.Object, mockKey2.Object]);
+
+        // Assert
+        Assert.Null(resolution.DefaultKey);
+        Assert.Same(mockKey1.Object, resolution.FallbackKey);
+        Assert.True(resolution.ShouldGenerateNewKey);
+
+        mockKey1.Verify(o => o.CreateEncryptor(), Times.Once); // 1 try
+        mockKey2.Verify(o => o.CreateEncryptor(), Times.Exactly(4)); // 1 try plus max (3) retries
+    }
+
+    [Fact]
+    public void CreateEncryptor_SucceedsOnRetry()
+    {
+        // Arrange
+        var options = Options.Create(new KeyManagementOptions()
+        {
+            MaximumDefaultKeyResolverRetries = 3,
+            DefaultKeyResolverRetryDelay = TimeSpan.Zero,
+        });
+
+        var resolver = new DefaultKeyResolver(options, NullLoggerFactory.Instance);
+
+        var now = ParseDateTimeOffset("2010-01-01 00:00:00Z");
+
+        var creation = now.AddDays(-3);
+        var activation = creation.AddDays(2);
+        var expiration = creation.AddDays(90);
+
+        var mockKey = CreateMockKey(activation, expiration, creation);
+        mockKey
+            .Setup(o => o.CreateEncryptor())
+            .Returns(() =>
+            {
+                // Added to list before the callback is called
+                if (mockKey.Invocations.Count(i => i.Method.Name == nameof(IKey.CreateEncryptor)) == 1)
+                {
+                    throw new Exception("This method fails.");
+                }
+                return new Mock<IAuthenticatedEncryptor>().Object;
+            });
+
+        // Act
+        var resolution = resolver.ResolveDefaultKeyPolicy(now, [mockKey.Object]);
+
+        // Assert
+        Assert.Same(mockKey.Object, resolution.DefaultKey);
+        mockKey.Verify(o => o.CreateEncryptor(), Times.Exactly(2)); // 1 try plus 1 retry
+    }
+
     private static IDefaultKeyResolver CreateDefaultKeyResolver()
     {
-        return new DefaultKeyResolver(NullLoggerFactory.Instance);
+        return new DefaultKeyResolver(Options.Create(new KeyManagementOptions()), NullLoggerFactory.Instance);
     }
 
     private static IKey CreateKey(string activationDate, string expirationDate, string creationDate = null, bool isRevoked = false, bool createEncryptorThrows = false)
@@ -302,6 +406,11 @@ private static IKey CreateKey(string activationDate, string expirationDate, stri
     }
 
     private static IKey CreateKey(DateTimeOffset activationDate, DateTimeOffset expirationDate, DateTimeOffset? creationDate = null, bool isRevoked = false, bool createEncryptorThrows = false)
+    {
+        return CreateMockKey(activationDate, expirationDate, creationDate, isRevoked, createEncryptorThrows).Object;
+    }
+
+    private static Mock<IKey> CreateMockKey(DateTimeOffset activationDate, DateTimeOffset expirationDate, DateTimeOffset? creationDate = null, bool isRevoked = false, bool createEncryptorThrows = false)
     {
         var mockKey = new Mock<IKey>();
         mockKey.Setup(o => o.KeyId).Returns(Guid.NewGuid());
@@ -318,7 +427,7 @@ private static IKey CreateKey(DateTimeOffset activationDate, DateTimeOffset expi
             mockKey.Setup(o => o.CreateEncryptor()).Returns(Mock.Of<IAuthenticatedEncryptor>());
         }
 
-        return mockKey.Object;
+        return mockKey;
     }
 
     private static DateTimeOffset ParseDateTimeOffset(string dto)