Do not dispose SignatureProvider when AsymmetricAdapter faults. (#2682)

Move compacted SignatureProviders to new cache to ensure dispose is called. Add delegate to check if signature provider should be removed from cache. Separate expired from compaction Dispose SignatureProvider if it was never cached. Co-authored-by: id4s <[email protected]>
AzureAD · Jul 22, 2024 · 49a3a95 · 49a3a95
1 parent ecc2b08
commit 49a3a95
Show file tree

Hide file tree

Showing 9 changed files with 290 additions and 193 deletions.
diff --git a/buildpackLocal.bat b/buildpackLocal.bat
@@ -0,0 +1,3 @@
+dotnet clean Product.proj > clean.log
+dotnet build /r Product.proj
+dotnet pack --no-restore -o c:\localpackages --no-build Product.proj
diff --git a/src/Microsoft.IdentityModel.Tokens/AsymmetricSignatureProvider.cs b/src/Microsoft.IdentityModel.Tokens/AsymmetricSignatureProvider.cs
@@ -208,15 +208,13 @@ public override bool Sign(ReadOnlySpan<byte> input, Span<byte> signature, out in
             catch
             {
                 CryptoProviderCache?.TryRemove(this);
-                Dispose(true);
                 throw;
             }
             finally
             {
-                if (!_disposed)
+                if (asym != null)
                     _asymmetricAdapterObjectPool.Free(asym);
             }
-
         }
 #endif
 
@@ -248,12 +246,11 @@ public override byte[] Sign(byte[] input)
             catch
             {
                 CryptoProviderCache?.TryRemove(this);
-                Dispose(true);
                 throw;
             }
             finally
             {
-                if (!_disposed)
+                if (asym != null)
                     _asymmetricAdapterObjectPool.Free(asym);
             }
         }
@@ -279,12 +276,11 @@ public override byte[] Sign(byte[] input, int offset, int count)
             catch
             {
                 CryptoProviderCache?.TryRemove(this);
-                Dispose(true);
                 throw;
             }
             finally
             {
-                if (!_disposed)
+                if (asym != null)
                     _asymmetricAdapterObjectPool.Free(asym);
             }
         }
@@ -380,12 +376,11 @@ public override bool Verify(byte[] input, byte[] signature)
             catch
             {
                 CryptoProviderCache?.TryRemove(this);
-                Dispose(true);
                 throw;
             }
             finally
             {
-                if (!_disposed)
+                if (asym != null)
                     _asymmetricAdapterObjectPool.Free(asym);
             }
         }
@@ -474,15 +469,14 @@ public override bool Verify(byte[] input, int inputOffset, int inputLength, byte
             }
             catch
             {
-                Dispose(true);
+                CryptoProviderCache?.TryRemove(this);
                 throw;
             }
             finally
             {
-                if (!_disposed)
+                if (asym != null)
                     _asymmetricAdapterObjectPool.Free(asym);
             }
-
         }
 
         /// <summary>

diff --git a/src/Microsoft.IdentityModel.Tokens/CryptoProviderFactory.cs b/src/Microsoft.IdentityModel.Tokens/CryptoProviderFactory.cs
@@ -18,6 +18,8 @@ public class CryptoProviderFactory
         private static readonly ConcurrentDictionary<string, string> _typeToAlgorithmMap = new ConcurrentDictionary<string, string>();
         private static readonly object _cacheLock = new object();
         private static int _defaultSignatureProviderObjectPoolCacheSize = Environment.ProcessorCount * 4;
+        private static string _typeofAsymmetricSignatureProvider = typeof(AsymmetricSignatureProvider).ToString();
+        private static string _typeofSymmetricSignatureProvider = typeof(SymmetricSignatureProvider).ToString();
         private int _signatureProviderObjectPoolCacheSize = _defaultSignatureProviderObjectPoolCacheSize;
 
         /// <summary>
@@ -513,7 +515,13 @@ private SignatureProvider CreateSignatureProvider(SecurityKey key, string algori
             {
                 signatureProvider = CustomCryptoProvider.Create(algorithm, key, willCreateSignatures) as SignatureProvider;
                 if (signatureProvider == null)
-                    throw LogHelper.LogExceptionMessage(new InvalidOperationException(LogHelper.FormatInvariant(LogMessages.IDX10646, LogHelper.MarkAsNonPII(algorithm), key, LogHelper.MarkAsNonPII(typeof(SignatureProvider)))));
+                    throw LogHelper.LogExceptionMessage(
+                        new InvalidOperationException(
+                            LogHelper.FormatInvariant(
+                                LogMessages.IDX10646,
+                                LogHelper.MarkAsNonPII(algorithm),
+                                key,
+                                LogHelper.MarkAsNonPII(typeof(SignatureProvider)))));
 
                 return signatureProvider;
             }
@@ -523,7 +531,7 @@ private SignatureProvider CreateSignatureProvider(SecurityKey key, string algori
             bool createAsymmetric = true;
             if (key is AsymmetricSecurityKey)
             {
-                typeofSignatureProvider = typeof(AsymmetricSignatureProvider).ToString();
+                typeofSignatureProvider = _typeofAsymmetricSignatureProvider;
             }
             else if (key is JsonWebKey jsonWebKey)
             {
@@ -533,22 +541,22 @@ private SignatureProvider CreateSignatureProvider(SecurityKey key, string algori
                     {
                         if (convertedSecurityKey is AsymmetricSecurityKey)
                         {
-                            typeofSignatureProvider = typeof(AsymmetricSignatureProvider).ToString();
+                            typeofSignatureProvider = _typeofAsymmetricSignatureProvider;
                         }
                         else if (convertedSecurityKey is SymmetricSecurityKey)
                         {
-                            typeofSignatureProvider = typeof(SymmetricSignatureProvider).ToString();
+                            typeofSignatureProvider = _typeofSymmetricSignatureProvider;
                             createAsymmetric = false;
                         }
                     }
                     // this code is simply to maintain the same exception thrown
                     else
                     {
                         if (jsonWebKey.Kty == JsonWebAlgorithmsKeyTypes.RSA || jsonWebKey.Kty == JsonWebAlgorithmsKeyTypes.EllipticCurve)
-                            typeofSignatureProvider = typeof(AsymmetricSignatureProvider).ToString();
+                            typeofSignatureProvider = _typeofAsymmetricSignatureProvider;
                         else if (jsonWebKey.Kty == JsonWebAlgorithmsKeyTypes.Octet)
                         {
-                            typeofSignatureProvider = typeof(SymmetricSignatureProvider).ToString();
+                            typeofSignatureProvider = _typeofSymmetricSignatureProvider;
                             createAsymmetric = false;
                         }
                     }
@@ -560,12 +568,20 @@ private SignatureProvider CreateSignatureProvider(SecurityKey key, string algori
             }
             else if (key is SymmetricSecurityKey)
             {
-                typeofSignatureProvider = typeof(SymmetricSignatureProvider).ToString();
+                typeofSignatureProvider = _typeofSymmetricSignatureProvider;
                 createAsymmetric = false;
             }
 
             if (typeofSignatureProvider == null)
-                throw LogHelper.LogExceptionMessage(new NotSupportedException(LogHelper.FormatInvariant(LogMessages.IDX10621, LogHelper.MarkAsNonPII(typeof(SymmetricSignatureProvider)), LogHelper.MarkAsNonPII(typeof(SecurityKey)), LogHelper.MarkAsNonPII(typeof(AsymmetricSecurityKey)), LogHelper.MarkAsNonPII(typeof(SymmetricSecurityKey)), LogHelper.MarkAsNonPII(key.GetType()))));
+                throw LogHelper.LogExceptionMessage(
+                    new NotSupportedException(
+                        LogHelper.FormatInvariant(
+                            LogMessages.IDX10621,
+                            LogHelper.MarkAsNonPII(typeof(SymmetricSignatureProvider)),
+                            LogHelper.MarkAsNonPII(typeof(SecurityKey)),
+                            LogHelper.MarkAsNonPII(typeof(AsymmetricSecurityKey)),
+                            LogHelper.MarkAsNonPII(typeof(SymmetricSecurityKey)),
+                            LogHelper.MarkAsNonPII(key.GetType()))));
 
             if (CacheSignatureProviders && cacheProvider)
             {
@@ -592,7 +608,7 @@ private SignatureProvider CreateSignatureProvider(SecurityKey key, string algori
                         signatureProvider = new SymmetricSignatureProvider(key, algorithm, willCreateSignatures);
 
                     if (ShouldCacheSignatureProvider(signatureProvider))
-                        CryptoProviderCache.TryAdd(signatureProvider);
+                        signatureProvider.IsCached = CryptoProviderCache.TryAdd(signatureProvider);
                 }
             }
             else
@@ -737,7 +753,7 @@ public virtual void ReleaseSignatureProvider(SignatureProvider signatureProvider
             signatureProvider.Release();
             if (CustomCryptoProvider != null && CustomCryptoProvider.IsSupportedAlgorithm(signatureProvider.Algorithm))
                 CustomCryptoProvider.Release(signatureProvider);
-            else if (signatureProvider.CryptoProviderCache == null && signatureProvider.RefCount == 0)
+            else if (signatureProvider.CryptoProviderCache == null && signatureProvider.RefCount == 0 && !signatureProvider.IsCached)
                 signatureProvider.Dispose();
         }
     }