Allow retries in DefaultKeyResolver.CanCreateAuthenticatedEncryptor (#54711)

amcasey · web-flow · commit 4d11def519f7 · 2024-04-19T21:13:15.000Z
* 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 #36157 * Roll our own Lazy that allows resets Retries against the actual `Key` type weren't working because the exception was getting cached in the key's lazy descriptor. Implement our own simple lazy and expose a method for clearing the cached value and exception.
diff --git a/src/DataProtection/DataProtection/src/KeyManagement/DefaultKeyResolver.cs b/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;
@@ -18,6 +20,13 @@ namespace Microsoft.AspNetCore.DataProtection.KeyManagement;
 /// </summary>
 internal sealed class DefaultKeyResolver : IDefaultKeyResolver
 {
+#if NETFRAMEWORK
+    // Sadly, numerical AppContext data support was not added until 4.7 and we target 4.6
+    private const string MaximumTotalDefaultKeyResolverRetriesSwitchKey = "Microsoft.AspNetCore.DataProtection.KeyManagement.DisableDefaultKeyResolverRetries";
+#else
+    private const string MaximumTotalDefaultKeyResolverRetriesDataKey = "Microsoft.AspNetCore.DataProtection.KeyManagement.MaximumTotalDefaultKeyResolverRetries";
+#endif
+
     /// <summary>
     /// The window of time before the key expires when a new key should be created
     /// and persisted to the keyring to ensure uninterrupted service.
@@ -28,6 +37,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 +53,95 @@ 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 = GetMaxDecryptRetriesFromAppContext() ?? keyManagementOptions.Value.MaximumTotalDefaultKeyResolverRetries;
+        _decryptRetryDelay = keyManagementOptions.Value.DefaultKeyResolverRetryDelay;
         _logger = loggerFactory.CreateLogger<DefaultKeyResolver>();
     }
 
-    private bool CanCreateAuthenticatedEncryptor(IKey key)
+    private static int? GetMaxDecryptRetriesFromAppContext()
     {
-        try
+#if NETFRAMEWORK
+        // Force to zero if retries are disabled.  Otherwise, use the configured value
+        return AppContext.TryGetSwitch(MaximumTotalDefaultKeyResolverRetriesSwitchKey, out var areRetriesDisabled) && areRetriesDisabled ? 0 : null;
+#else
+        var data = AppContext.GetData(MaximumTotalDefaultKeyResolverRetriesDataKey);
+
+        // Programmatically-configured values are usually ints
+        if (data is int count)
         {
-            var encryptorInstance = key.CreateEncryptor();
-            if (encryptorInstance == null)
-            {
-                CryptoUtil.Fail<IAuthenticatedEncryptor>("CreateEncryptorInstance returned null.");
-            }
+            return count;
+        }
 
-            return true;
+        // msbuild-configured values are usually strings
+        if (data is string countStr && int.TryParse(countStr, out var parsed))
+        {
+            return parsed;
         }
-        catch (Exception ex)
+
+        return null;
+#endif
+    }
+
+    private bool CanCreateAuthenticatedEncryptor(IKey key, ref int retriesRemaining)
+    {
+        List<Exception>? exceptions = null;
+        while (true)
         {
-            _logger.KeyIsIneligibleToBeTheDefaultKeyBecauseItsMethodFailed(key.KeyId, nameof(IKey.CreateEncryptor), ex);
-            return false;
+            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)
+            {
+                if (retriesRemaining > 0)
+                {
+                    if (exceptions is null)
+                    {
+                        // The first failure doesn't count as a retry
+                        exceptions = [];
+                    }
+                    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);
+            }
+
+            // Reset the descriptor to allow for a retry
+            (key as Key)?.ResetDescriptor();
+
+            // Don't retry immediately - allow a little time for the transient problem to clear
+            Thread.Sleep(_decryptRetryDelay);
         }
     }
 
@@ -94,12 +168,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 +199,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();
diff --git a/src/DataProtection/DataProtection/src/KeyManagement/Key.cs b/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 go from null to non-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,
         };
diff --git a/src/DataProtection/DataProtection/src/KeyManagement/KeyManagementOptions.cs b/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 MaximumTotalDefaultKeyResolverRetries { get; set; } = 10;
+
+    /// <summary>
+    /// Wait this long before each default key resolution decryption retry.
+    /// <seealso cref="MaximumTotalDefaultKeyResolverRetries"/>
+    /// </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.
diff --git a/src/DataProtection/DataProtection/src/LoggingExtensions.cs b/src/DataProtection/DataProtection/src/LoggingExtensions.cs
@@ -73,7 +73,7 @@ private static bool IsLogLevelEnabledCore([NotNullWhen(true)] ILogger? logger, L
     [LoggerMessage(11, LogLevel.Debug, "Using managed symmetric algorithm '{FullName}'.", EventName = "UsingManagedSymmetricAlgorithm")]
     public static partial void UsingManagedSymmetricAlgorithm(this ILogger logger, string fullName);
 
-    [LoggerMessage(12, LogLevel.Warning, "Key {KeyId:B} is ineligible to be the default key because its {MethodName} method failed.", EventName = "KeyIsIneligibleToBeTheDefaultKeyBecauseItsMethodFailed")]
+    [LoggerMessage(12, LogLevel.Warning, "Key {KeyId:B} is ineligible to be the default key because its {MethodName} method failed after the maximum number of retries.", EventName = "KeyIsIneligibleToBeTheDefaultKeyBecauseItsMethodFailed")]
     public static partial void KeyIsIneligibleToBeTheDefaultKeyBecauseItsMethodFailed(this ILogger logger, Guid keyId, string methodName, Exception exception);
 
     [LoggerMessage(13, LogLevel.Debug, "Considering key {KeyId:B} with expiration date {ExpirationDate:u} as default key.", EventName = "ConsideringKeyWithExpirationDateAsDefaultKey")]
@@ -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);
 }
diff --git a/src/DataProtection/DataProtection/test/Microsoft.AspNetCore.DataProtection.Tests/KeyManagement/DefaultKeyResolverTests.cs b/src/DataProtection/DataProtection/test/Microsoft.AspNetCore.DataProtection.Tests/KeyManagement/DefaultKeyResolverTests.cs
