Add multi-user key store provider registration support

doc/snippets/Microsoft.Data.SqlClient/SqlColumnEncryptionKeyStoreProvider.xml

@@ -1,7 +1,11 @@
 <docs>
     <members name="SqlColumnEncryptionKeyStoreProvider">
         <SqlColumnEncryptionKeyStoreProvider>
-            <summary>Base class for all key store providers. A custom provider must derive from this class and override its member functions and then register it using SqlConnection.RegisterColumnEncryptionKeyStoreProviders(). For details see, <see href="https://docs.microsoft.com/sql/relational-databases/security/encryption/always-encrypted-database-engine"> Always Encrypted</see>.
+            <summary>Base class for all key store providers. A custom provider must derive from this class and override its member functions and then register it using 
+						<see cref="M:Microsoft.Data.SqlClient.SqlConnection.RegisterColumnEncryptionKeyStoreProviders()" />, 
+						<see cref="M:Microsoft.Data.SqlClient.SqlConnection.RegisterColumnEncryptionKeyStoreProvidersOnConnection()" /> or 
+						<see cref="M:Microsoft.Data.SqlClient.SqlCommand.RegisterColumnEncryptionKeyStoreProvidersOnCommand()" />.
+						For details see, <see href="https://docs.microsoft.com/sql/relational-databases/security/encryption/always-encrypted-database-engine"> Always Encrypted</see>.
             </summary>
             <remarks>To be added.</remarks>
         </SqlColumnEncryptionKeyStoreProvider>
@@ -55,5 +59,10 @@ The <xref:Microsoft.Data.SqlClient.SqlColumnEncryptionKeyStoreProvider.SignColum
             <returns>When implemented in a derived class, the method is expected to return true if the specified signature is valid, or false if the specified signature is not valid. The default implementation throws NotImplementedException.</returns>
             <remarks>To be added.</remarks>
         </VerifyColumnMasterKeyMetadata>
+        <ColumnEncryptionKeyCacheTtl>
+            <summary>Gets or sets the lifespan of the decrypted column encryption key in the cache. Once the timespan has elapsed, the decrypted column encryption key is discarded and must be revalidated.</summary>
+            <remarks>Internally, there is a cache of column encryption keys (once they are decrypted). This is useful for rapidly decrypting multiple data values. The default value is 2 hours. Setting this value to zero disables caching.
+        </remarks>
+        </ColumnEncryptionKeyCacheTtl>
     </members>
 </docs>

doc/snippets/Microsoft.Data.SqlClient/SqlCommand.xml

@@ -2772,6 +2772,25 @@ You must set the value for this property before the command is executed for it t
 ]]></format>
 			</remarks>
 		</Notification>
+		<RegisterColumnEncryptionKeyStoreProvidersOnCommand>
+			<param name="customProviders">Dictionary of custom column encryption key providers</param>
+			<summary>Registers the encryption key store providers on the <see cref="T:Microsoft.Data.SqlClient.SqlCommand" /> instance. If this function has been called, any providers registered using the <see cref="M:Microsoft.Data.SqlClient.SqlConnection.RegisterColumnEncryptionKeyStoreProviders()" /> or
+			<see cref="M:Microsoft.Data.SqlClient.SqlConnection.RegisterColumnEncryptionKeyStoreProvidersOnConnection()" /> methods will be ignored. This function can be called more than once. This does shallow copying of the dictionary so that the app cannot alter the custom provider list once it has been set.</summary>
+			<exception cref="T:System.ArgumentNullException">
+					A null dictionary was provided.
+
+					-or-
+
+					A string key in the dictionary was null or empty.
+
+					-or-
+
+					An EncryptionKeyStoreProvider value in the dictionary was null.
+			</exception>
+			<exception cref="T:System.ArgumentException">
+				A string key in the dictionary started with "MSSQL_". This prefix is reserved for system providers.
+			</exception>
+		</RegisterColumnEncryptionKeyStoreProvidersOnCommand>
 		<RetryLogicProvider>
 			<summary>
 				Gets or sets a value that specifies the

src/Microsoft.Data.SqlClient/add-ons/AzureKeyVaultProvider/Constants.cs

@@ -6,11 +6,6 @@ namespace Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider
 {
     internal static class Constants
     {
-        /// <summary>
-        /// Hashing algorithm used for signing
-        /// </summary>
-        internal const string HashingAlgorithm = @"RS256";
-
         /// <summary>
         /// Azure Key Vault Domain Name
         /// </summary>

src/Microsoft.Data.SqlClient/add-ons/AzureKeyVaultProvider/LocalCache.cs

@@ -0,0 +1,95 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using Microsoft.Extensions.Caching.Memory;
+using System;
+
+using static System.Math;
+
+namespace Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider
+{
+    /// <summary>
+    /// LocalCache is to reuse heavy objects.
+    /// When performing a heavy creation operation, we will save the result in our cache container.
+    /// The next time that we need that result, we will pull it from the cache container, instead of performing the heavy operation again.
+    /// </summary>
+    internal class LocalCache<TKey, TValue>
+    {
+        /// <summary>
+        /// A simple thread-safe implementation of an in-memory Cache.
+        /// When the process dies, the cache dies with it.
+        /// </summary>
+        private readonly MemoryCache cache;
+
+        private readonly int maxSize;
+
+        /// <summary>
+        /// Sets an absolute expiration time, relative to now.
+        /// </summary>
+        internal TimeSpan? TimeToLive { get; set; }
+
+        /// <summary>
+        /// Gets the count of the current entries for diagnostic purposes.
+        /// </summary>
+        internal int Count => cache.Count;
+
+        /// <summary>
+        /// Constructs a new <see cref="LocalCache{TKey, TValue}">LocalCache</see> object.
+        /// </summary>
+        internal LocalCache(int maxSizeLimit = int.MaxValue)
+        {
+            if(maxSizeLimit <= 0)
+            {
+                throw new ArgumentOutOfRangeException(nameof(maxSizeLimit));
+            }
+
+            maxSize = maxSizeLimit;
+            cache = new MemoryCache(new MemoryCacheOptions());
+        }
+
+        /// <summary>
+        /// Looks for the cache entry that maps to the <paramref name="key"/> value. If it exists (cache hit) it will simply be
+        /// returned. Otherwise, the <paramref name="createItem"/> delegate function will be invoked to create the value. 
+        /// It will then get stored it in the cache and set the time-to-live before getting returned.
+        /// </summary>
+        /// <param name="key">The encrypted key bytes.</param>
+        /// <param name="createItem">The delegate function that will decrypt the encrypted column encryption key.</param>
+        /// <returns>The decrypted column encryption key.</returns>
+        internal TValue GetOrCreate(TKey key, Func<TValue> createItem)
+        {
+            if (TimeToLive <= TimeSpan.Zero)
+            {
+                return createItem();
+            }
+
+            if (!cache.TryGetValue(key, out TValue cacheEntry))
+            {
+                if (cache.Count == maxSize)
+                {
+                    cache.Compact(Max(0.10, 1.0 / maxSize));
+                }
+
+                cacheEntry = createItem();
+                var cacheEntryOptions = new MemoryCacheEntryOptions
+                {
+                    AbsoluteExpirationRelativeToNow = TimeToLive
+                };
+
+                cache.Set(key, cacheEntry, cacheEntryOptions);
+            }
+
+            return cacheEntry;
+        }
+
+        /// <summary>
+        /// Determines whether the <see cref="LocalCache{TKey, TValue}">LocalCache</see> contains the specified key.
+        /// </summary>
+        /// <param name="key"></param>
+        /// <returns></returns>
+        internal bool Contains(TKey key)
+        {
+            return cache.TryGetValue(key, out _);
+        }
+    }
+}

src/Microsoft.Data.SqlClient/add-ons/AzureKeyVaultProvider/Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider.csproj

@@ -25,5 +25,6 @@
     <PackageReference Include="Microsoft.SourceLink.GitHub" Version="$(MicrosoftSourceLinkGitHubVersion)" PrivateAssets="All" />
     <PackageReference Include="Azure.Core" Version="$(AzureCoreVersion)" />
     <PackageReference Include="Azure.Security.KeyVault.Keys" Version="$(AzureSecurityKeyVaultKeysVersion)" />
+    <PackageReference Include="Microsoft.Extensions.Caching.Memory" Version="5.0.0" />
   </ItemGroup>
 </Project>

src/Microsoft.Data.SqlClient/add-ons/AzureKeyVaultProvider/SqlColumnEncryptionAzureKeyVaultProvider.cs

@@ -63,6 +63,34 @@ public class SqlColumnEncryptionAzureKeyVaultProvider : SqlColumnEncryptionKeySt
         /// </summary>
         public readonly string[] TrustedEndPoints;
 
+        /// <summary>
+        /// A cache of column encryption keys (once they are decrypted). This is useful for rapidly decrypting multiple data values.
+        /// </summary>
+        private readonly LocalCache<string, byte[]> columnEncryptionKeyCache = 
+            new LocalCache<string, byte[]>() { TimeToLive = TimeSpan.FromHours(2) };
+
+        /// <summary>
+        /// A cache for storing the results of signature verification of column master key metadata.
+        /// </summary>
+        private readonly LocalCache<Tuple<string, bool, string>, bool> columnMasterKeyMetadataSignatureVerificationCache =
+            new LocalCache<Tuple<string, bool, string>, bool>(maxSizeLimit: 2000) { TimeToLive = TimeSpan.FromDays(10) };
+
+        /// <summary>
+        /// Gets or sets the lifespan of the decrypted column encryption key in the cache.
+        /// Once the timespan has elapsed, the decrypted column encryption key is discarded
+        /// and must be revalidated.
+        /// </summary>
+        /// <remarks>
+        /// Internally, there is a cache of column encryption keys (once they are decrypted).
+        /// This is useful for rapidly decrypting multiple data values. The default value is 2 hours.
+        /// Setting the <see cref="ColumnEncryptionKeyCacheTtl"/> to zero disables caching.
+        /// </remarks>
+        public override TimeSpan? ColumnEncryptionKeyCacheTtl
+        {
+            get => columnEncryptionKeyCache.TimeToLive;
+            set => columnEncryptionKeyCache.TimeToLive = value;
+        }
+
         #endregion
 
         #region Constructors
@@ -130,10 +158,16 @@ public override bool VerifyColumnMasterKeyMetadata(string masterKeyPath, bool al
         {
             ValidateNonEmptyAKVPath(masterKeyPath, isSystemOp: true);
 
-            // Also validates key is of RSA type.
-            KeyCryptographer.AddKey(masterKeyPath);
-            byte[] message = CompileMasterKeyMetadata(masterKeyPath, allowEnclaveComputations);
-            return KeyCryptographer.VerifyData(message, signature, masterKeyPath);
+            var key = Tuple.Create(masterKeyPath, allowEnclaveComputations, ToHexString(signature));
+            return GetOrCreateSignatureVerificationResult(key, VerifyColumnMasterKeyMetadata);
+
+            bool VerifyColumnMasterKeyMetadata()
+            {
+                // Also validates key is of RSA type.
+                KeyCryptographer.AddKey(masterKeyPath);
+                byte[] message = CompileMasterKeyMetadata(masterKeyPath, allowEnclaveComputations);
+                return KeyCryptographer.VerifyData(message, signature, masterKeyPath);
+            }
         }
 
         /// <summary>
@@ -153,58 +187,62 @@ public override byte[] DecryptColumnEncryptionKey(string masterKeyPath, string e
             ValidateNotEmpty(encryptedColumnEncryptionKey, nameof(encryptedColumnEncryptionKey));
             ValidateVersionByte(encryptedColumnEncryptionKey[0], s_firstVersion[0]);
 
-            // Also validates whether the key is RSA one or not and then get the key size
-            KeyCryptographer.AddKey(masterKeyPath);
+            return GetOrCreateColumnEncryptionKey(ToHexString(encryptedColumnEncryptionKey), DecryptEncryptionKey);
 
-            int keySizeInBytes = KeyCryptographer.GetKeySize(masterKeyPath);
+            byte[] DecryptEncryptionKey()
+            {
+                // Also validates whether the key is RSA one or not and then get the key size
+                KeyCryptographer.AddKey(masterKeyPath);
 
-            // Get key path length
-            int currentIndex = s_firstVersion.Length;
-            ushort keyPathLength = BitConverter.ToUInt16(encryptedColumnEncryptionKey, currentIndex);
-            currentIndex += sizeof(ushort);
+                int keySizeInBytes = KeyCryptographer.GetKeySize(masterKeyPath);
 
-            // Get ciphertext length
-            ushort cipherTextLength = BitConverter.ToUInt16(encryptedColumnEncryptionKey, currentIndex);
-            currentIndex += sizeof(ushort);
+                // Get key path length
+                int currentIndex = s_firstVersion.Length;
+                ushort keyPathLength = BitConverter.ToUInt16(encryptedColumnEncryptionKey, currentIndex);
+                currentIndex += sizeof(ushort);
 
-            // Skip KeyPath
-            // KeyPath exists only for troubleshooting purposes and doesnt need validation.
-            currentIndex += keyPathLength;
+                // Get ciphertext length
+                ushort cipherTextLength = BitConverter.ToUInt16(encryptedColumnEncryptionKey, currentIndex);
+                currentIndex += sizeof(ushort);
 
-            // validate the ciphertext length
-            if (cipherTextLength != keySizeInBytes)
-            {
-                throw ADP.InvalidCipherTextLength(cipherTextLength, keySizeInBytes, masterKeyPath);
-            }
+                // Skip KeyPath
+                // KeyPath exists only for troubleshooting purposes and doesnt need validation.
+                currentIndex += keyPathLength;
 
-            // Validate the signature length
-            int signatureLength = encryptedColumnEncryptionKey.Length - currentIndex - cipherTextLength;
-            if (signatureLength != keySizeInBytes)
-            {
-                throw ADP.InvalidSignatureLengthTemplate(signatureLength, keySizeInBytes, masterKeyPath);
-            }
+                // validate the ciphertext length
+                if (cipherTextLength != keySizeInBytes)
+                {
+                    throw ADP.InvalidCipherTextLength(cipherTextLength, keySizeInBytes, masterKeyPath);
+                }
 
-            // Get ciphertext
-            byte[] cipherText = encryptedColumnEncryptionKey.Skip(currentIndex).Take(cipherTextLength).ToArray();
-            currentIndex += cipherTextLength;
+                // Validate the signature length
+                int signatureLength = encryptedColumnEncryptionKey.Length - currentIndex - cipherTextLength;
+                if (signatureLength != keySizeInBytes)
+                {
+                    throw ADP.InvalidSignatureLengthTemplate(signatureLength, keySizeInBytes, masterKeyPath);
+                }
 
-            // Get signature
-            byte[] signature = encryptedColumnEncryptionKey.Skip(currentIndex).Take(signatureLength).ToArray();
+                // Get ciphertext
+                byte[] cipherText = encryptedColumnEncryptionKey.Skip(currentIndex).Take(cipherTextLength).ToArray();
+                currentIndex += cipherTextLength;
 
-            // Compute the message to validate the signature
-            byte[] message = encryptedColumnEncryptionKey.Take(encryptedColumnEncryptionKey.Length - signatureLength).ToArray();
+                // Get signature
+                byte[] signature = encryptedColumnEncryptionKey.Skip(currentIndex).Take(signatureLength).ToArray();
 
-            if (null == message)
-            {
-                throw ADP.NullHashFound();
-            }
+                // Compute the message to validate the signature
+                byte[] message = encryptedColumnEncryptionKey.Take(encryptedColumnEncryptionKey.Length - signatureLength).ToArray();
 
-            if (!KeyCryptographer.VerifyData(message, signature, masterKeyPath))
-            {
-                throw ADP.InvalidSignatureTemplate(masterKeyPath);
-            }
+                if (null == message)
+                {
+                    throw ADP.NullHashFound();
+                }
 
-            return KeyCryptographer.UnwrapKey(s_keyWrapAlgorithm, cipherText, masterKeyPath);
+                if (!KeyCryptographer.VerifyData(message, signature, masterKeyPath))
+                {
+                    throw ADP.InvalidSignatureTemplate(masterKeyPath);
+                }
+                return KeyCryptographer.UnwrapKey(s_keyWrapAlgorithm, cipherText, masterKeyPath);
+            }
         }
 
         /// <summary>
@@ -310,6 +348,49 @@ private byte[] CompileMasterKeyMetadata(string masterKeyPath, bool allowEnclaveC
             return Encoding.Unicode.GetBytes(masterkeyMetadata.ToLowerInvariant());
         }
 
+        /// <summary>
+        /// Converts the numeric value of each element of a specified array of bytes to its equivalent hexadecimal string representation.
+        /// </summary>
+        /// <param name="source">An array of bytes to convert.</param>
+        /// <returns>A string of hexadecimal characters</returns>
+        /// <remarks>
+        /// Produces a string of hexadecimal character pairs preceded with "0x", where each pair represents the corresponding element in value; for example, "0x7F2C4A00".
+        /// </remarks>
+        private string ToHexString(byte[] source)
+        {
+            if (source == null)
+            {
+                return null;
+            }
+
+            return "0x" + BitConverter.ToString(source).Replace("-", "");
+        }
+
+        /// <summary>
+        /// Returns the cached decrypted data encryption key, or unwraps the encrypted data encryption if not present.
+        /// </summary>
+        /// <param name="encryptedDataEncryptionKey">Encrypted Data Encryption Key</param>
+        /// <param name="createItem">The delegate function that will decrypt the encrypted column encryption key.</param>
+        /// <returns>The decrypted data encryption key.</returns>
+        /// <remarks>
+        ///
+        /// </remarks>
+        private byte[] GetOrCreateColumnEncryptionKey(string encryptedDataEncryptionKey, Func<byte[]> createItem)
+        {
+            return columnEncryptionKeyCache.GetOrCreate(encryptedDataEncryptionKey, createItem);
+        }
+
+        /// <summary>
+        /// Returns the cached signature verification result, or proceeds to verify if not present.
+        /// </summary>
+        /// <param name="keyInformation">The encryptionKeyId, allowEnclaveComputations and hexadecimal signature.</param>
+        /// <param name="createItem">The delegate function that will perform the verification.</param>
+        /// <returns></returns>
+        private bool GetOrCreateSignatureVerificationResult(Tuple<string, bool, string> keyInformation, Func<bool> createItem)
+        {
+            return columnMasterKeyMetadataSignatureVerificationCache.GetOrCreate(keyInformation, createItem);
+        }
+
         #endregion
     }
 }

src/Microsoft.Data.SqlClient/netcore/ref/Microsoft.Data.SqlClient.cs

@@ -597,6 +597,8 @@ public override void Prepare() { }
         [System.ComponentModel.BrowsableAttribute(false)]
         [System.ComponentModel.DesignerSerializationVisibilityAttribute(0)]
         public Microsoft.Data.Sql.SqlNotificationRequest Notification { get { throw null; } set { } }
+        /// <include file='../../../../doc/snippets/Microsoft.Data.SqlClient/SqlCommand.xml' path='docs/members[@name="SqlCommand"]/RegisterColumnEncryptionKeyStoreProvidersOnCommand/*' />
+        public void RegisterColumnEncryptionKeyStoreProvidersOnCommand(System.Collections.Generic.IDictionary<string, Microsoft.Data.SqlClient.SqlColumnEncryptionKeyStoreProvider> customProviders) { }
         /// <include file='../../../../doc/snippets/Microsoft.Data.SqlClient/SqlCommand.xml' path='docs/members[@name="SqlCommand"]/ResetCommandTimeout/*'/>
         [System.ComponentModel.DesignerSerializationVisibilityAttribute(System.ComponentModel.DesignerSerializationVisibility.Content)]
         public void ResetCommandTimeout() { }