Make the key ring cache expire shortly after an immediately-activated key is generated

src/DataProtection/DataProtection/src/KeyManagement/Internal/CacheableKeyRing.cs

@@ -21,16 +21,24 @@ internal CacheableKeyRing(CancellationToken expirationToken, DateTimeOffset expi
     }
 
     internal CacheableKeyRing(CancellationToken expirationToken, DateTimeOffset expirationTime, IKeyRing keyRing)
+        : this(expirationToken, expirationTime, hasShortRefreshPeriod: false, keyRing)
+    {
+    }
+
+    internal CacheableKeyRing(CancellationToken expirationToken, DateTimeOffset expirationTime, bool hasShortRefreshPeriod, IKeyRing keyRing)
     {
         _expirationToken = expirationToken;
         ExpirationTimeUtc = expirationTime.UtcDateTime;
         KeyRing = keyRing;
+        HasShortRefreshPeriod = hasShortRefreshPeriod;
     }
 
     internal DateTime ExpirationTimeUtc { get; }
 
     internal IKeyRing KeyRing { get; }
 
+    internal bool HasShortRefreshPeriod { get; }
+
     internal static bool IsValid([NotNullWhen(true)] CacheableKeyRing? keyRing, DateTime utcNow)
     {
         return keyRing != null
@@ -46,6 +54,6 @@ internal static bool IsValid([NotNullWhen(true)] CacheableKeyRing? keyRing, Date
     internal CacheableKeyRing WithTemporaryExtendedLifetime(DateTimeOffset now)
     {
         var extension = TimeSpan.FromMinutes(2);
-        return new CacheableKeyRing(CancellationToken.None, now + extension, KeyRing);
+        return new CacheableKeyRing(CancellationToken.None, now + extension, HasShortRefreshPeriod, KeyRing);
     }
 }

src/DataProtection/DataProtection/src/KeyManagement/Internal/ICacheableKeyRingProvider.cs

@@ -10,6 +10,13 @@ namespace Microsoft.AspNetCore.DataProtection.KeyManagement.Internal;
 /// This API supports infrastructure and is not intended to be used
 /// directly from your code. This API may change or be removed in future releases.
 /// </summary>
+/// <remarks>
+/// This interface is retained for the benefit of consumers who were casting <see cref="IKeyRingProvider"/>
+/// (also pseudo-internal) to <see cref="ICacheableKeyRingProvider"/> so that they could invoke
+/// <see cref="GetCacheableKeyRing"/>.  That method returns an object with no public properties and doesn't
+/// update the state of the <see cref="IKeyRingProvider"/>, but it does trigger calls to the backing
+/// <see cref="IKeyManager"/>, which may have observable effects.  It is no longer used as a test hook.
+/// </remarks>
 public interface ICacheableKeyRingProvider
 {
     /// <summary>

src/DataProtection/DataProtection/src/KeyManagement/Internal/IInternalCacheableKeyRingProvider.cs

@@ -0,0 +1,24 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Diagnostics.CodeAnalysis;
+
+namespace Microsoft.AspNetCore.DataProtection.KeyManagement.Internal;
+
+/// <summary>
+/// This API supports infrastructure and is not intended to be used
+/// directly from your code. This API may change or be removed in future releases.
+/// </summary>
+/// <remarks>
+/// Replaces <see cref="ICacheableKeyRingProvider"/> as a test hook for validating
+/// <see cref="KeyRingProvider"/>.
+/// </remarks>
+internal interface IInternalCacheableKeyRingProvider
+{
+    /// <summary>
+    /// This API supports infrastructure and is not intended to be used
+    /// directly from your code. This API may change or be removed in future releases.
+    /// </summary>
+    CacheableKeyRing GetCacheableKeyRing(DateTimeOffset now, bool allowShortRefreshPeriod = true);
+}

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

@@ -76,6 +76,20 @@ internal static TimeSpan KeyRingRefreshPeriod
         }
     }
 
+    /// <summary>
+    /// If a key is required and none is available, one will be generated and immediately
+    /// activated.  Immediately-activated keys are problematic because they may be used
+    /// before they have propagated to other app instances.  What's more, in a scenario
+    /// where one instance needs to generate an immediately-activated key, it's likely that
+    /// other instances will need to do the same.  A given instance has no control over when
+    /// other instances pick up its immediately-activated keys, but it can do its best to
+    /// pick up theirs.  Therefore, every time it generates an immediately-activated key,
+    /// it assumes other instances may have as well and pre-emptively schedules a refresh
+    /// of its own key ring cache.  This property controls how long after an immediately-
+    /// activated key is generated the key ring cache will be refreshed.
+    /// </summary>
+    internal static TimeSpan ShortKeyRingRefreshPeriod { get; set; } = TimeSpan.FromSeconds(15);
+
     /// <summary>
     /// Specifies the maximum clock skew allowed between servers when reading
     /// keys from the key ring. The key ring may use a key which has not yet

src/DataProtection/DataProtection/src/KeyManagement/KeyRingProvider.cs

@@ -14,7 +14,7 @@
 
 namespace Microsoft.AspNetCore.DataProtection.KeyManagement;
 
-internal sealed class KeyRingProvider : ICacheableKeyRingProvider, IKeyRingProvider
+internal sealed class KeyRingProvider : ICacheableKeyRingProvider, IInternalCacheableKeyRingProvider, IKeyRingProvider
 {
     private CacheableKeyRing? _cacheableKeyRing;
     private readonly object _cacheableKeyRingLockObj = new object();
@@ -55,13 +55,16 @@ public KeyRingProvider(
     }
 
     // for testing
-    internal ICacheableKeyRingProvider CacheableKeyRingProvider { get; set; }
+    internal CacheableKeyRing? CacheableKeyRing => Volatile.Read(ref _cacheableKeyRing);
+
+    // for testing
+    internal IInternalCacheableKeyRingProvider CacheableKeyRingProvider { get; set; }
 
     internal DateTime AutoRefreshWindowEnd { get; set; }
 
     internal bool InAutoRefreshWindow() => DateTime.UtcNow < AutoRefreshWindowEnd;
 
-    private CacheableKeyRing CreateCacheableKeyRingCore(DateTimeOffset now, IKey? keyJustAdded)
+    private CacheableKeyRing CreateCacheableKeyRingCore(DateTimeOffset now, bool allowShortRefreshPeriod, IKey? keyJustAdded)
     {
         // Refresh the list of all keys
         var cacheExpirationToken = _keyManager.GetCacheExpirationToken();
@@ -80,7 +83,7 @@ private CacheableKeyRing CreateCacheableKeyRingCore(DateTimeOffset now, IKey? ke
         if (keyJustAdded != null)
         {
             var keyToUse = defaultKey ?? defaultKeyPolicy.FallbackKey ?? keyJustAdded;
-            return CreateCacheableKeyRingCoreStep2(now, cacheExpirationToken, keyToUse, allKeys);
+            return CreateCacheableKeyRingCoreStep2(now, cacheExpirationToken, keyToUse, keyJustAdded, allowShortRefreshPeriod, allKeys);
         }
 
         // Determine whether we need to generate a new key
@@ -107,7 +110,7 @@ private CacheableKeyRing CreateCacheableKeyRingCore(DateTimeOffset now, IKey? ke
         if (!shouldGenerateNewKey)
         {
             CryptoUtil.Assert(defaultKey != null, "Expected to see a default key.");
-            return CreateCacheableKeyRingCoreStep2(now, cacheExpirationToken, defaultKey, allKeys);
+            return CreateCacheableKeyRingCoreStep2(now, cacheExpirationToken, defaultKey, keyJustAdded, allowShortRefreshPeriod, allKeys);
         }
 
         _logger.PolicyResolutionStatesThatANewKeyShouldBeAddedToTheKeyRing();
@@ -127,7 +130,7 @@ private CacheableKeyRing CreateCacheableKeyRingCore(DateTimeOffset now, IKey? ke
             else
             {
                 _logger.UsingFallbackKeyWithExpirationAsDefaultKey(keyToUse.KeyId, keyToUse.ExpirationDate);
-                return CreateCacheableKeyRingCoreStep2(now, cacheExpirationToken, keyToUse, allKeys);
+                return CreateCacheableKeyRingCoreStep2(now, cacheExpirationToken, keyToUse, keyJustAdded, allowShortRefreshPeriod, allKeys);
             }
         }
 
@@ -139,19 +142,19 @@ private CacheableKeyRing CreateCacheableKeyRingCore(DateTimeOffset now, IKey? ke
             // The case where there's no default key is the easiest scenario, since it
             // means that we need to create a new key with immediate activation.
             var newKey = _keyManager.CreateNewKey(activationDate: now, expirationDate: now + _newKeyLifetime);
-            return CreateCacheableKeyRingCore(now, keyJustAdded: newKey); // recursively call
+            return CreateCacheableKeyRingCore(now, allowShortRefreshPeriod, keyJustAdded: newKey); // recursively call
         }
         else
         {
             // If there is a default key, then the new key we generate should become active upon
             // expiration of the default key. The new key lifetime is measured from the creation
             // date (now), not the activation date.
             var newKey = _keyManager.CreateNewKey(activationDate: defaultKey.ExpirationDate, expirationDate: now + _newKeyLifetime);
-            return CreateCacheableKeyRingCore(now, keyJustAdded: newKey); // recursively call
+            return CreateCacheableKeyRingCore(now, allowShortRefreshPeriod, keyJustAdded: newKey); // recursively call
         }
     }
 
-    private CacheableKeyRing CreateCacheableKeyRingCoreStep2(DateTimeOffset now, CancellationToken cacheExpirationToken, IKey defaultKey, IEnumerable<IKey> allKeys)
+    private CacheableKeyRing CreateCacheableKeyRingCoreStep2(DateTimeOffset now, CancellationToken cacheExpirationToken, IKey defaultKey, IKey? generatedKey, bool allowShortRefreshPeriod, IEnumerable<IKey> allKeys)
     {
         Debug.Assert(defaultKey != null);
 
@@ -167,7 +170,25 @@ private CacheableKeyRing CreateCacheableKeyRingCoreStep2(DateTimeOffset now, Can
 
         _logger.UsingKeyAsDefaultKey(defaultKey.KeyId);
 
-        var nextAutoRefreshTime = now + GetRefreshPeriodWithJitter(KeyManagementOptions.KeyRingRefreshPeriod);
+        // If we're in a scenario where we think other app instances may be generating keys that will be
+        // activated before they are propagated, we'll use a shorter refresh period to ensure that we pick
+        // up those keys so that we can select from the same pool of candidates as the other instances.
+        // 
+        // These not-yet-propagated keys are usually immediately-activated, but an app restarted after a period
+        // of downtime may discover that it has a valid, but soon-to-be-expired key. The replacement will not
+        // be immediately-activated, but may be activated before it has propagated.
+        var useShortRefreshPeriod = allowShortRefreshPeriod &&
+            (generatedKey is not null
+                // Either we had to generate a key and it will be activated before it has time to propagate
+                // (in which case, other instances may have done the same)
+                ? (generatedKey.ActivationDate < now + KeyManagementOptions.KeyPropagationWindow) // No clock skew on a key we generated
+                // Or we selected a key that has yet to propagate (presumably, from another instance)
+                : (defaultKey.CreationDate > now - KeyManagementOptions.KeyPropagationWindow));
+
+        var nextAutoRefreshTime = now + GetRefreshPeriodWithJitter(
+            useShortRefreshPeriod
+                ? KeyManagementOptions.ShortKeyRingRefreshPeriod
+                : KeyManagementOptions.KeyRingRefreshPeriod);
 
         // The cached keyring should expire at the earliest of (default key expiration, next auto-refresh time).
         // Since the refresh period and safety window are not user-settable, we can guarantee that there's at
@@ -179,8 +200,8 @@ private CacheableKeyRing CreateCacheableKeyRingCoreStep2(DateTimeOffset now, Can
         return new CacheableKeyRing(
             expirationToken: cacheExpirationToken,
             expirationTime: (defaultKey.ExpirationDate <= now) ? nextAutoRefreshTime : Min(defaultKey.ExpirationDate, nextAutoRefreshTime),
-            defaultKey: defaultKey,
-            allKeys: allKeys);
+            hasShortRefreshPeriod: useShortRefreshPeriod,
+            keyRing: new KeyRing(defaultKey, allKeys));
     }
 
     public IKeyRing GetCurrentKeyRing()
@@ -241,7 +262,7 @@ internal IKeyRing GetCurrentKeyRingCore(DateTime utcNow, bool forceRefresh = fal
 
                 try
                 {
-                    newCacheableKeyRing = CacheableKeyRingProvider.GetCacheableKeyRing(utcNow);
+                    newCacheableKeyRing = CacheableKeyRingProvider.GetCacheableKeyRing(utcNow, allowShortRefreshPeriod: existingCacheableKeyRing?.HasShortRefreshPeriod != true);
                 }
                 catch (Exception ex)
                 {
@@ -311,9 +332,18 @@ private static DateTimeOffset Min(DateTimeOffset a, DateTimeOffset b)
         return (a < b) ? a : b;
     }
 
+    // This is only here to support external consumers of our internal test hook.
+    // The object returned has no public properties, so we could probably return a dummy
+    // without doing any work, but maybe they're calling it for side-effects?
     CacheableKeyRing ICacheableKeyRingProvider.GetCacheableKeyRing(DateTimeOffset now)
     {
         // the entry point allows one recursive call
-        return CreateCacheableKeyRingCore(now, keyJustAdded: null);
+        return CreateCacheableKeyRingCore(now, allowShortRefreshPeriod: false, keyJustAdded: null);
+    }
+
+    CacheableKeyRing IInternalCacheableKeyRingProvider.GetCacheableKeyRing(DateTimeOffset now, bool allowShortRefreshPeriod)
+    {
+        // the entry point allows one recursive call
+        return CreateCacheableKeyRingCore(now, allowShortRefreshPeriod, keyJustAdded: null);
     }
 }

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

@@ -226,7 +226,7 @@ private static DateTime StringToDateTime(string input)
         return DateTimeOffset.ParseExact(input, "u", CultureInfo.InvariantCulture).UtcDateTime;
     }
 
-    private static KeyRingProvider CreateKeyRingProvider(ICacheableKeyRingProvider cacheableKeyRingProvider)
+    private static KeyRingProvider CreateKeyRingProvider(IInternalCacheableKeyRingProvider cacheableKeyRingProvider)
     {
         var mockEncryptorFactory = new Mock<IAuthenticatedEncryptorFactory>();
         mockEncryptorFactory.Setup(m => m.CreateEncryptorInstance(It.IsAny<IKey>())).Returns(new Mock<IAuthenticatedEncryptor>().Object);
@@ -345,16 +345,16 @@ public void Unprotect_KeyNotFound_RefreshOnce_CanFindKey()
         Assert.Empty(result);
     }
 
-    private class TestKeyRingProvider : ICacheableKeyRingProvider
+    private class TestKeyRingProvider : IInternalCacheableKeyRingProvider
     {
         private CacheableKeyRing _keyRing;
 
         public TestKeyRingProvider(CacheableKeyRing keys) => _keyRing = keys;
 
-        public CacheableKeyRing GetCacheableKeyRing(DateTimeOffset now) => _keyRing;
+        public CacheableKeyRing GetCacheableKeyRing(DateTimeOffset now, bool _allowShortRefreshPeriod) => _keyRing;
     }
 
-    private class RefreshTestKeyRingProvider : ICacheableKeyRingProvider
+    private class RefreshTestKeyRingProvider : IInternalCacheableKeyRingProvider
     {
         private CacheableKeyRing _keyRing;
         private CacheableKeyRing _refreshKeyRing;
@@ -366,7 +366,7 @@ public RefreshTestKeyRingProvider(CacheableKeyRing keys, CacheableKeyRing refres
             _refreshKeyRing = refreshKeys;
         }
 
-        public CacheableKeyRing GetCacheableKeyRing(DateTimeOffset now)
+        public CacheableKeyRing GetCacheableKeyRing(DateTimeOffset now, bool _allowShortRefreshPeriod)
         {
             if (!_called)
             {