Introdduce a LKG configuration cache to store each valid base configu…

…ration

instead of a single entry of configuration.

src/Microsoft.IdentityModel.JsonWebTokens/JsonWebTokenHandler.cs

@@ -1258,14 +1258,22 @@ private async Task<TokenValidationResult> ValidateTokenAsync(JsonWebToken jsonWe
                         }
                     }
 
-                    if (TokenUtilities.IsRecoverableConfiguration(validationParameters, currentConfiguration, out currentConfiguration))
+                    if (TokenUtilities.ShouldValidateWithLKG(validationParameters))
                     {
                         validationParameters.RefreshBeforeValidation = false;
                         validationParameters.ValidateWithLKG = true;
-                        tokenValidationResult = await ValidateTokenAsync(jsonWebToken, validationParameters, currentConfiguration).ConfigureAwait(false);
+                        var recoverableException = tokenValidationResult.Exception;
 
-                        if (tokenValidationResult.IsValid)
-                            return tokenValidationResult;
+                        foreach (BaseConfiguration lkgConfiguration in validationParameters.ConfigurationManager.GetValidLkgConfiguraitonFromCache())
+                        {
+                            if (!lkgConfiguration.Equals(currentConfiguration) && TokenUtilities.IsRecoverableConfiguration(jsonWebToken.Kid, currentConfiguration, lkgConfiguration, recoverableException))
+                            {
+                                tokenValidationResult = await ValidateTokenAsync(jsonWebToken, validationParameters, lkgConfiguration).ConfigureAwait(false);
+
+                                if (tokenValidationResult.IsValid)
+                                    return tokenValidationResult;
+                            }
+                        }
                     }
                 }
             }

src/Microsoft.IdentityModel.Tokens/BaseConfigurationComparer.cs

@@ -0,0 +1,40 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.Collections.Generic;
+using System.Linq;
+
+namespace Microsoft.IdentityModel.Tokens
+{
+    /// <summary>
+    /// Comparison class for a <see cref="BaseConfiguration"/>.
+    /// </summary>
+    internal class BaseConfigurationComparer : IEqualityComparer<BaseConfiguration>
+    {
+        public bool Equals(BaseConfiguration config1, BaseConfiguration config2)
+        {
+            if (config1 == null && config2 == null)
+                return true;
+            else if (config1 == null || config2 == null)
+                return false;
+            else if (config1.Issuer == config2.Issuer && config1.SigningKeys.Count == config2.SigningKeys.Count
+                     && !config1.SigningKeys.Select(x => x.InternalId).Except(config1.SigningKeys.Select(x => x.InternalId)).Any())
+                return true;
+            else
+                return false;
+        }
+
+        public int GetHashCode(BaseConfiguration config)
+        {
+            int defaultHash = string.Empty.GetHashCode();
+            int hashCode = defaultHash;
+            hashCode ^= string.IsNullOrEmpty(config.Issuer) ? defaultHash : config.Issuer.GetHashCode();
+            foreach(string internalId in config.SigningKeys.Select(x => x.InternalId))
+            {
+                hashCode ^= internalId.GetHashCode();
+            }
+
+            return hashCode;
+        }
+    }
+}

src/Microsoft.IdentityModel.Tokens/BaseConfigurationManager.cs

@@ -3,6 +3,8 @@
 
 
 using System;
+using System.Collections.Generic;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.IdentityModel.Logging;
@@ -20,6 +22,9 @@ public abstract class BaseConfigurationManager
         private BaseConfiguration _lastKnownGoodConfiguration;
         private DateTime? _lastKnownGoodConfigFirstUse = null;
 
+        private Dictionary<BaseConfiguration, DateTime> _lkgConfigurationCache = null;
+        private IEqualityComparer<BaseConfiguration> _baseConfigurationComparer = new BaseConfigurationComparer();
+
         /// <summary>
         /// Gets or sets the <see cref="TimeSpan"/> that controls how often an automatic metadata refresh should occur.
         /// </summary>
@@ -35,6 +40,18 @@ public TimeSpan AutomaticRefreshInterval
             }
         }
 
+        /// <summary>
+        /// Gets or sets the BaseConfgiurationComparer that to compare <see cref="BaseConfiguration"/>.
+        /// </summary>
+        public IEqualityComparer<BaseConfiguration> BaseConfigurationComparer
+        {
+            get { return _baseConfigurationComparer; }
+            set
+            {
+                _baseConfigurationComparer = value ?? throw LogHelper.LogExceptionMessage(new ArgumentNullException(nameof(value)));
+            }
+        }
+
         /// <summary>
         /// 12 hours is the default time interval that afterwards will obtain new configuration.
         /// </summary>
@@ -63,6 +80,25 @@ public virtual Task<BaseConfiguration> GetBaseConfigurationAsync(CancellationTok
             throw new NotImplementedException();
         }
 
+        /// <summary>
+        /// Gets all valid last known good configurations from the cache.
+        /// </summary>
+        /// <returns>A collection of all valid last known good configurations.</returns>
+        internal ICollection<BaseConfiguration> GetValidLkgConfiguraitonFromCache()
+        {
+            if (_lkgConfigurationCache == null)
+                return null;
+
+            var expiredLkgConfiguration = _lkgConfigurationCache.Where(x => x.Value < DateTime.UtcNow).ToArray();
+
+            foreach (KeyValuePair<BaseConfiguration, DateTime> lkgConfiguration in expiredLkgConfiguration)
+            {
+                _lkgConfigurationCache.Remove(lkgConfiguration.Key);
+            }
+
+            return _lkgConfigurationCache.Any() ? _lkgConfigurationCache.Keys : null;
+        }
+
         /// <summary>
         /// The last known good configuration or LKG (a configuration retrieved in the past that we were able to successfully validate a token against).
         /// </summary>
@@ -76,6 +112,19 @@ public BaseConfiguration LastKnownGoodConfiguration
             {
                 _lastKnownGoodConfiguration = value ?? throw LogHelper.LogArgumentNullException(nameof(value));
                 _lastKnownGoodConfigFirstUse = DateTime.UtcNow;
+
+                if (_lkgConfigurationCache == null)
+                    _lkgConfigurationCache = new Dictionary<BaseConfiguration, DateTime>(BaseConfigurationComparer);
+
+                _lkgConfigurationCache[_lastKnownGoodConfiguration] = DateTime.UtcNow + LastKnownGoodLifetime;
+
+                //remove expired configuration to avoid memory leak 
+                var expiredLkgConfiguration = _lkgConfigurationCache.Where(x => x.Value < DateTime.UtcNow).ToArray();
+
+                foreach (KeyValuePair<BaseConfiguration, DateTime> lkgConfiguration in expiredLkgConfiguration)
+                {
+                    _lkgConfigurationCache.Remove(lkgConfiguration.Key);
+                }
             }
         }
 

src/Microsoft.IdentityModel.Tokens/TokenUtilities.cs

@@ -190,7 +190,6 @@ internal static bool IsRecoverableException(Exception exception)
             // using 'GetType()' instead of 'is' as SecurityTokenUnableToValidException (and others) extend SecurityTokenInvalidSignatureException
             // we want to make sure that the clause for SecurityTokenUnableToValidateException is hit so that the ValidationFailure is checked
             return exception.GetType().Equals(typeof(SecurityTokenInvalidSignatureException))
-                   || exception is SecurityTokenInvalidSigningKeyException
                    || exception is SecurityTokenInvalidIssuerException
                    // we should not try to revalidate with the LKG or request a refresh if the token has an invalid lifetime
                    || (exception as SecurityTokenUnableToValidateException)?.ValidationFailure != ValidationFailure.InvalidLifetime
@@ -200,30 +199,50 @@ internal static bool IsRecoverableException(Exception exception)
         /// <summary>
         /// Check whether the given configuration is recoverable by LKG.
         /// </summary>
-        /// <param name="validationParameters">The <see cref="TokenValidationParameters"/> to be used for validation.</param>
-        /// <param name="configuration">The <see cref="BaseConfiguration"/> to check.</param>
-        /// <param name="currentConfiguration">The updated <see cref="BaseConfiguration"/>.</param>
+        /// <param name="kid">The kid from token."/></param>
+        /// <param name="currentConfiguration">The <see cref="BaseConfiguration"/> to check.</param>
+        /// <param name="lkgConfiguration">The LKG exception to check.</param>
+        /// <param name="currentException">The exception to check.</param>
         /// <returns><c>true</c> if the configuration is recoverable otherwise, <c>false</c>.</returns>
-        internal static bool IsRecoverableConfiguration(TokenValidationParameters validationParameters, BaseConfiguration configuration, out BaseConfiguration currentConfiguration)
+        internal static bool IsRecoverableConfiguration(string kid, BaseConfiguration currentConfiguration, BaseConfiguration lkgConfiguration, Exception currentException)
         {
-            bool isRecoverableConfiguration = (validationParameters.ConfigurationManager.UseLastKnownGoodConfiguration
-                && validationParameters.ConfigurationManager.LastKnownGoodConfiguration != null
-                && !ReferenceEquals(configuration, validationParameters.ConfigurationManager.LastKnownGoodConfiguration));
+            Lazy<bool> isRecoverableIssuer = new Lazy<bool>(() => currentConfiguration.Issuer != lkgConfiguration.Issuer);
+            Lazy<bool> isRecoverableSigningKey = new Lazy<bool>(() => lkgConfiguration.SigningKeys.Any(signingKey => signingKey.KeyId == kid));
 
-            currentConfiguration = configuration;
-            if (isRecoverableConfiguration)
+            if (currentException is SecurityTokenInvalidIssuerException)
             {
-                // Inform the user that the LKG is expired.
-                if (!validationParameters.ConfigurationManager.IsLastKnownGoodValid)
-                {
-                    LogHelper.LogInformation(TokenLogMessages.IDX10263);
-                    return false;
-                }
-                else
-                    currentConfiguration = validationParameters.ConfigurationManager.LastKnownGoodConfiguration;
+                return isRecoverableIssuer.Value;
             }
+            else if (currentException is SecurityTokenSignatureKeyNotFoundException)
+            {
+                return isRecoverableSigningKey.Value;
+            }
+            else if ((currentException as SecurityTokenUnableToValidateException)?.ValidationFailure == ValidationFailure.InvalidIssuer)
+            {
+                return isRecoverableIssuer.Value && isRecoverableSigningKey.Value;
+            }
+            else if (currentException.GetType().Equals(typeof(SecurityTokenInvalidSignatureException)))
+            {
+                SecurityKey currentSigningKey = currentConfiguration.SigningKeys.FirstOrDefault(x => x.KeyId == kid);
+                if (currentSigningKey == null)
+                    return isRecoverableSigningKey.Value;
 
-            return isRecoverableConfiguration;
+                SecurityKey lkgSigningKey = lkgConfiguration.SigningKeys.FirstOrDefault(signingKey => signingKey.KeyId == kid);
+                return lkgSigningKey != null && currentSigningKey.InternalId != lkgSigningKey.InternalId;
+            }
+
+            return false;
+        }
+
+        /// <summary>
+        /// Check if the token can be validated by LKG.
+        /// </summary>
+        /// <param name="validationParameters">The <see cref="TokenValidationParameters"/> to be used for validation.</param>
+        /// <returns><c>true</c> if the token can be validated by LKG, <c>false</c>.</returns>
+        internal static bool ShouldValidateWithLKG(TokenValidationParameters validationParameters)
+        {
+            return  (validationParameters.ConfigurationManager.UseLastKnownGoodConfiguration
+                && validationParameters.ConfigurationManager.GetValidLkgConfiguraitonFromCache() != null);
         }
     }
 }

src/System.IdentityModel.Tokens.Jwt/JwtSecurityTokenHandler.cs

@@ -929,15 +929,25 @@ private ClaimsPrincipal ValidateToken(string token, JwtSecurityToken outerToken,
                         }
                     }
 
-                    if (TokenUtilities.IsRecoverableConfiguration(validationParameters, currentConfiguration, out currentConfiguration))
+                    if (TokenUtilities.ShouldValidateWithLKG(validationParameters))
                     {
-                        validationParameters.ValidateWithLKG = true;
                         validationParameters.RefreshBeforeValidation = false;
-                        claimsPrincipal = outerToken != null ? ValidateJWE(token, outerToken, validationParameters, currentConfiguration, out signatureValidatedToken, out exceptionThrown) :
-                            ValidateJWS(token, validationParameters, currentConfiguration, out signatureValidatedToken, out exceptionThrown);
+                        validationParameters.ValidateWithLKG = true;
+                        var recoverableException = exceptionThrown.SourceException;
+                        string kid = outerToken != null ? outerToken.Header.Kid :
+                            (ValidateSignatureUsingDelegates(token, validationParameters, null) ?? GetJwtSecurityTokenFromToken(token, validationParameters)).Header.Kid;
 
-                        if (claimsPrincipal != null)
-                            return claimsPrincipal;
+                        foreach (BaseConfiguration lkgConfiguration in validationParameters.ConfigurationManager.GetValidLkgConfiguraitonFromCache())
+                        {
+                            if (!lkgConfiguration.Equals(currentConfiguration) && TokenUtilities.IsRecoverableConfiguration(kid, currentConfiguration, lkgConfiguration, recoverableException))
+                            {
+                                claimsPrincipal = outerToken != null ? ValidateJWE(token, outerToken, validationParameters, lkgConfiguration, out signatureValidatedToken, out exceptionThrown) :
+                                    ValidateJWS(token, validationParameters, lkgConfiguration, out signatureValidatedToken, out exceptionThrown);
+
+                                if (claimsPrincipal != null)
+                                    return claimsPrincipal;
+                            }
+                        }
                     }
                 }
             }

test/Microsoft.IdentityModel.Protocols.OpenIdConnect.Tests/ConfigurationManagerTests.cs

@@ -105,11 +105,12 @@ public void GetSets()
         {
             TestUtilities.WriteHeader($"{this}.GetSets", "GetSets", true);
 
+            int ExpectedPropertyCount = 8;
             var configManager = new ConfigurationManager<OpenIdConnectConfiguration>("OpenIdConnectMetadata.json", new OpenIdConnectConfigurationRetriever(), new FileDocumentRetriever());
             Type type = typeof(ConfigurationManager<OpenIdConnectConfiguration>);
             PropertyInfo[] properties = type.GetProperties();
-            if (properties.Length != 7)
-                Assert.True(false, "Number of properties has changed from 7 to: " + properties.Length + ", adjust tests");
+            if (properties.Length != ExpectedPropertyCount)
+                Assert.True(false, "Number of properties has changed from {ExpectedPropertyCount} to: " + properties.Length + ", adjust tests");
 
             var defaultAutomaticRefreshInterval = ConfigurationManager<OpenIdConnectConfiguration>.DefaultAutomaticRefreshInterval;
             var defaultRefreshInterval = ConfigurationManager<OpenIdConnectConfiguration>.DefaultRefreshInterval;
@@ -330,6 +331,43 @@ public void ResetLastKnownGoodLifetime()
             TestUtilities.AssertFailIfErrors(context);
         }
 
+        [Fact]
+        public void TestConfigurationComparer()
+        {
+            TestUtilities.WriteHeader($"{this}.TestConfigurationComparer", "TestConfigurationComparer", true);
+            var context = new CompareContext();
+
+            var config1 = new OpenIdConnectConfiguration() { TokenEndpoint = Default.Issuer + "oauth/token", Issuer = Default.Issuer };
+            config1.SigningKeys.Add(KeyingMaterial.DefaultX509Key_2048);
+            config1.SigningKeys.Add(KeyingMaterial.DefaultRsaSecurityKey1);
+            config1.SigningKeys.Add(KeyingMaterial.DefaultRsaSecurityKey2);
+
+            var config2 = new OpenIdConnectConfiguration() { TokenEndpoint = Default.Issuer + "oauth/token", Issuer = Default.Issuer };
+            config2.SigningKeys.Add(KeyingMaterial.DefaultRsaSecurityKey1);
+            config2.SigningKeys.Add(KeyingMaterial.DefaultX509Key_2048);
+            config2.SigningKeys.Add(KeyingMaterial.DefaultRsaSecurityKey2);
+
+            var config3 = new OpenIdConnectConfiguration() { TokenEndpoint = Default.Issuer + "oauth/token", Issuer = Default.Issuer };
+            config3.SigningKeys.Add(KeyingMaterial.DefaultRsaSecurityKey1);
+
+            var config4 = new OpenIdConnectConfiguration() { TokenEndpoint = Default.Issuer + "oauth/token", Issuer = Default.Issuer + "1" };
+            config3.SigningKeys.Add(KeyingMaterial.DefaultRsaSecurityKey1);
+
+            var configurationManager = new MockConfigurationManager<OpenIdConnectConfiguration>(config1, config1);
+            IdentityComparer.AreEqual(configurationManager.GetValidLkgConfiguraitonFromCache().Count, 1, context);
+
+            configurationManager.LastKnownGoodConfiguration = config2;
+            IdentityComparer.AreEqual(configurationManager.GetValidLkgConfiguraitonFromCache().Count, 1, context);
+
+            configurationManager.LastKnownGoodConfiguration = config3;
+            IdentityComparer.AreEqual(configurationManager.GetValidLkgConfiguraitonFromCache().Count, 2, context);
+
+            configurationManager.LastKnownGoodConfiguration = config4;
+            IdentityComparer.AreEqual(configurationManager.GetValidLkgConfiguraitonFromCache().Count, 3, context);
+
+            TestUtilities.AssertFailIfErrors(context);
+        }
+
         [Theory, MemberData(nameof(ValidateOpenIdConnectConfigurationTestCases), DisableDiscoveryEnumeration = true)]
         public void ValidateOpenIdConnectConfigurationTests(ConfigurationManagerTheoryData<OpenIdConnectConfiguration> theoryData)
         {

test/Microsoft.IdentityModel.TestUtils/MockConfigurationManager.cs

@@ -74,6 +74,30 @@ public MockConfigurationManager(T configuration, T lkgConfiguration, T refreshed
             _refreshedConfiguration = refreshedConfiguration;
         }
 
+        /// <summary>
+        /// Initializes an new instance of <see cref="MockConfigurationManager{T}"/> with a Configuration instance and a LKG Configuration instance.
+        /// </summary>
+        /// <param name="configuration">Configuration of type OpenIdConnectConfiguration or WsFederationConfiguration.</param>
+        /// <param name="lkgLifetime">The LKG configuration lifetime.</param>
+        public MockConfigurationManager(T configuration, TimeSpan lkgLifetime) : this(configuration)
+        {
+            LastKnownGoodLifetime = lkgLifetime;
+        }
+
+        /// <summary>
+        /// Initializes an new instance of <see cref="MockConfigurationManager{T}"/> with a Configuration instance and a LKG Configuration instance.
+        /// </summary>
+        /// <param name="configuration">Configuration of type OpenIdConnectConfiguration or WsFederationConfiguration.</param>
+        /// <param name="lkgConfiguration">Configuration of type OpenIdConnectConfiguration or WsFederationConfiguration.</param>
+        /// <param name="lkgLifetime">The LKG configuration lifetime.</param>
+        public MockConfigurationManager(T configuration, T lkgConfiguration, TimeSpan lkgLifetime) : this(configuration, lkgLifetime)
+        {
+            if (lkgConfiguration == null)
+                throw LogHelper.LogExceptionMessage(new ArgumentNullException(nameof(lkgConfiguration)));
+
+            LastKnownGoodConfiguration = lkgConfiguration as BaseConfiguration;
+        }
+
         /// <summary>
         /// Obtains an updated version of Configuration.
         /// </summary>