Allow retries in DefaultKeyResolver.CanCreateAuthenticatedEncryptor

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,71 @@ 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 (retriesRemaining > 0)
+                {
+                    if (exceptions is null)
+                    {
+                        // The first failure doesn't count as a retry
+                        exceptions = new();
+                    }
+                    else
+                    {
+                        retriesRemaining--;
+                    }
+
+                    exceptions.Add(ex);
+                }
+
+                if (retriesRemaining <= 0) // Guard against infinite loops from programmer errors
+                {
+                    _logger.KeyIsIneligibleToBeTheDefaultKeyBecauseItsMethodFailed(key.KeyId, nameof(IKey.CreateEncryptor), exceptions is null ? ex : 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;
+            // Reset the descriptor to allow for a retry
+            (key as Key)?.ResetDescriptor();
+
+            // Don't retry immediately
+            Thread.Sleep(_decryptRetryDelay);
         }
     }
 
@@ -94,12 +137,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 +168,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/Key.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Xml.Linq;
 using Microsoft.AspNetCore.DataProtection.AuthenticatedEncryption;
 using Microsoft.AspNetCore.DataProtection.AuthenticatedEncryption.ConfigurationModel;
@@ -16,7 +17,13 @@ namespace Microsoft.AspNetCore.DataProtection.KeyManagement;
 /// </summary>
 internal sealed class Key : IKey
 {
-    private readonly Lazy<IAuthenticatedEncryptorDescriptor> _lazyDescriptor;
+    private IAuthenticatedEncryptorDescriptor? _descriptor;
+
+    // If descriptor is available at construction time, these will remain null forever
+    private readonly object? _descriptorLock; // Protects _descriptor and _descriptorException
+    private readonly Func<IAuthenticatedEncryptorDescriptor>? _descriptorFactory; // May not be used
+    private Exception? _descriptorException;
+
     private readonly IEnumerable<IAuthenticatedEncryptorFactory> _encryptorFactories;
 
     private IAuthenticatedEncryptor? _encryptor;
@@ -36,8 +43,10 @@ public Key(
               creationDate,
               activationDate,
               expirationDate,
-              new Lazy<IAuthenticatedEncryptorDescriptor>(() => descriptor),
-              encryptorFactories)
+              encryptorFactories,
+              descriptor,
+              descriptorFactory: null,
+              descriptorException: null)
     {
     }
 
@@ -57,8 +66,29 @@ public Key(
               creationDate,
               activationDate,
               expirationDate,
-              new Lazy<IAuthenticatedEncryptorDescriptor>(GetLazyDescriptorDelegate(keyManager, keyElement)),
-              encryptorFactories)
+              encryptorFactories,
+              descriptor: null,
+              descriptorFactory: GetLazyDescriptorDelegate(keyManager, keyElement),
+              descriptorException: null)
+    {
+    }
+
+    // internal for testing
+    internal Key(
+        Guid keyId,
+        DateTimeOffset creationDate,
+        DateTimeOffset activationDate,
+        DateTimeOffset expirationDate,
+        IEnumerable<IAuthenticatedEncryptorFactory> encryptorFactories,
+        Func<IAuthenticatedEncryptorDescriptor>? descriptorFactory)
+        : this(keyId,
+              creationDate,
+              activationDate,
+              expirationDate,
+              encryptorFactories,
+              descriptor: null,
+              descriptorFactory: descriptorFactory,
+              descriptorException: null)
     {
     }
 
@@ -67,15 +97,20 @@ private Key(
         DateTimeOffset creationDate,
         DateTimeOffset activationDate,
         DateTimeOffset expirationDate,
-        Lazy<IAuthenticatedEncryptorDescriptor> lazyDescriptor,
-        IEnumerable<IAuthenticatedEncryptorFactory> encryptorFactories)
+        IEnumerable<IAuthenticatedEncryptorFactory> encryptorFactories,
+        IAuthenticatedEncryptorDescriptor? descriptor,
+        Func<IAuthenticatedEncryptorDescriptor>? descriptorFactory,
+        Exception? descriptorException)
     {
         KeyId = keyId;
         CreationDate = creationDate;
         ActivationDate = activationDate;
         ExpirationDate = expirationDate;
-        _lazyDescriptor = lazyDescriptor;
         _encryptorFactories = encryptorFactories;
+        _descriptor = descriptor;
+        _descriptorFactory = descriptorFactory;
+        _descriptorException = descriptorException;
+        _descriptorLock = descriptor is null ? new() : null;
     }
 
     public DateTimeOffset ActivationDate { get; }
@@ -92,7 +127,56 @@ public IAuthenticatedEncryptorDescriptor Descriptor
     {
         get
         {
-            return _lazyDescriptor.Value;
+            // We could check for _descriptorException here, but there's no reason to optimize that case
+            // (i.e. by avoiding taking the lock)
+
+            if (_descriptor is not null) // Can only become less null, so losing a race here doesn't matter
+            {
+                Debug.Assert(_descriptorException is null); // Mutually exclusive with _descriptor
+                return _descriptor;
+            }
+
+            lock (_descriptorLock!)
+            {
+                if (_descriptorException is not null)
+                {
+                    throw _descriptorException;
+                }
+
+                if (_descriptor is not null)
+                {
+                    return _descriptor;
+                }
+
+                Debug.Assert(_descriptorFactory is not null, "Key constructed without either descriptor or descriptor factory");
+
+                try
+                {
+                    _descriptor = _descriptorFactory();
+                    return _descriptor;
+                }
+                catch (Exception ex)
+                {
+                    _descriptorException = ex;
+                    throw;
+                }
+            }
+        }
+    }
+
+    internal void ResetDescriptor()
+    {
+        if (_descriptor is not null)
+        {
+            Debug.Fail("ResetDescriptor called with descriptor available");
+            Debug.Assert(_descriptorException is null); // Mutually exclusive with _descriptor
+            return;
+        }
+
+        lock (_descriptorLock!)
+        {
+            _descriptor = null;
+            _descriptorException = null;
         }
     }
 
@@ -121,13 +205,16 @@ internal void SetRevoked()
 
     internal Key Clone()
     {
+        // Note that we don't reuse _descriptorLock
         return new Key(
             keyId: KeyId,
             creationDate: CreationDate,
             activationDate: ActivationDate,
             expirationDate: ExpirationDate,
-            lazyDescriptor: _lazyDescriptor,
-            encryptorFactories: _encryptorFactories)
+            encryptorFactories: _encryptorFactories,
+            descriptor: _descriptor,
+            descriptorFactory: _descriptorFactory,
+            descriptorException: _descriptorException)
         {
             IsRevoked = IsRevoked,
         };

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);
 }