Perf | Reduce allocations in Always Encrypted key handling, reorganise

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/AlwaysEncrypted/AeadAes256CbcHmac256EncryptionKey.cs

@@ -0,0 +1,107 @@
+// 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 System.Text;
+
+#nullable enable
+
+namespace Microsoft.Data.SqlClient.AlwaysEncrypted
+{
+    /// <summary>
+    /// Encryption key class containing 4 keys. This class is used by SqlAeadAes256CbcHmac256Algorithm
+    /// 1) root key - Main key that is used to derive the keys used in the encryption algorithm
+    /// 2) encryption key - A derived key that is used to encrypt the plain text and generate cipher text
+    /// 3) mac_key - A derived key that is used to compute HMAC of the cipher text
+    /// 4) iv_key - A derived key that is used to generate a synthetic IV from plain text data.
+    /// </summary>
+    internal sealed class AeadAes256CbcHmac256EncryptionKey : SymmetricKey
+    {
+        /// <summary>
+        /// Key size in bits.
+        /// </summary>
+        public const int KeySizeInBits = 256;
+
+        /// <summary>
+        /// Key size in bytes.
+        /// </summary>
+        public const int KeySizeInBytes = KeySizeInBits / 8;
+
+        /// <summary>
+        /// <see cref="KeySizeInBits"/> as a string, for use in the salt formats below.
+        /// </summary>
+        private const string KeySizeInBitsString = "256";
+
+        /// <summary>
+        /// Encryption Key Salt. This is used to derive the encryption key from the root key.
+        /// </summary>
+        /// <see href="https://learn.microsoft.com/en-us/sql/relational-databases/security/encryption/always-encrypted-cryptography?view=sql-server-ver17#step-2-computing-aes_256_cbc-ciphertext"/>
+        private const string EncryptionKeySaltString = $"Microsoft SQL Server cell encryption key with encryption algorithm:{SqlAeadAes256CbcHmac256Algorithm.AlgorithmName} and key length:{KeySizeInBitsString}";
+
+        /// <summary>
+        /// MAC Key Salt. This is used to derive the MAC key from the root key.
+        /// </summary>
+        /// <see href="https://learn.microsoft.com/en-us/sql/relational-databases/security/encryption/always-encrypted-cryptography?view=sql-server-ver17#step-3-computing-mac"/>
+        private const string MacKeySaltString = $"Microsoft SQL Server cell MAC key with encryption algorithm:{SqlAeadAes256CbcHmac256Algorithm.AlgorithmName} and key length:{KeySizeInBitsString}";
+
+        /// <summary>
+        /// IV Key Salt. This is used to derive the IV key from the root key. This is only used for Deterministic encryption.
+        /// </summary>
+        /// <see href="https://learn.microsoft.com/en-us/sql/relational-databases/security/encryption/always-encrypted-cryptography?view=sql-server-ver17#step-1-generating-the-initialization-vector-iv"/>
+        private const string IvKeySaltString = $"Microsoft SQL Server cell IV key with encryption algorithm:{SqlAeadAes256CbcHmac256Algorithm.AlgorithmName} and key length:{KeySizeInBitsString}";
+
+        private static byte[] EncryptionKeySalt =>
+            field ??= Encoding.Unicode.GetBytes(EncryptionKeySaltString);
+        private static byte[] MacKeySalt =>
+            field ??= Encoding.Unicode.GetBytes(MacKeySaltString);
+        private static byte[] IvKeySalt =>
+            field ??= Encoding.Unicode.GetBytes(IvKeySaltString);
+
+        /// <summary>
+        /// Derives all the required keys from the given root key
+        /// </summary>
+        /// <param name="rootKey">Root key used to derive all the required derived keys</param>
+        public AeadAes256CbcHmac256EncryptionKey(byte[] rootKey) : base(rootKey)
+        {
+            // Key validation
+            if (rootKey.Length != KeySizeInBytes)
+            {
+                throw SQL.InvalidKeySize(SqlAeadAes256CbcHmac256Algorithm.AlgorithmName,
+                                         rootKey.Length,
+                                         KeySizeInBytes);
+            }
+
+            // Derive keys from the root key
+            //
+            // Derive encryption key
+            byte[] buff1 = new byte[KeySizeInBytes];
+            SqlSecurityUtility.GetHMACWithSHA256(EncryptionKeySalt, RootKey, buff1);
+            EncryptionKey = buff1;
+
+            // Derive MAC key
+            byte[] buff2 = new byte[KeySizeInBytes];
+            SqlSecurityUtility.GetHMACWithSHA256(MacKeySalt, RootKey, buff2);
+            MacKey = buff2;
+
+            // Derive IV key
+            byte[] buff3 = new byte[KeySizeInBytes];
+            SqlSecurityUtility.GetHMACWithSHA256(IvKeySalt, RootKey, buff3);
+            IvKey = buff3;
+        }
+
+        /// <summary>
+        /// Encryption key should be used for encryption and decryption
+        /// </summary>
+        public byte[] EncryptionKey { get; }
+
+        /// <summary>
+        /// MAC key should be used to compute and validate HMAC
+        /// </summary>
+        public byte[] MacKey { get; }
+
+        /// <summary>
+        /// IV key should be used to compute synthetic IV from a given plain text
+        /// </summary>
+        public byte[] IvKey { get; }
+    }
+}

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/AlwaysEncrypted/AeadAes256CbcHmac256Factory.cs

@@ -21,8 +21,10 @@ namespace Microsoft.Data.SqlClient.AlwaysEncrypted
     /// </remarks>
     internal sealed class AeadAes256CbcHmac256Factory : EncryptionAlgorithmFactory
     {
+        private static readonly AeadAes256CbcHmac256Factory s_singletonInstance = new();
+
         /// <summary>
-        /// Factory classes cache the <see cref="SqlAeadAes256CbcHmac256EncryptionKey" /> objects to avoid recomputation of the derived keys.
+        /// Factory classes cache the <see cref="AeadAes256CbcHmac256EncryptionKey" /> objects to avoid recomputation of the derived keys.
         /// </summary>
         private readonly ConcurrentDictionary<string, SqlAeadAes256CbcHmac256Algorithm> _encryptionAlgorithms =
             new(concurrencyLevel: 4 * Environment.ProcessorCount /* default value in ConcurrentDictionary */, capacity: 2);
@@ -32,7 +34,7 @@ private AeadAes256CbcHmac256Factory() { }
         /// <summary>
         /// Access the instance of the factory class for the AEAD_AES_256_CBC_HMAC_SHA256 encryption algorithm.
         /// </summary>
-        public static AeadAes256CbcHmac256Factory Instance => field ??= new();
+        public static AeadAes256CbcHmac256Factory Instance => s_singletonInstance;
 
         /// <summary>
         /// Creates an instance of the <see cref="SqlAeadAes256CbcHmac256Algorithm" /> class with a given root key.
@@ -41,18 +43,18 @@ private AeadAes256CbcHmac256Factory() { }
         /// <param name="encryptionType">Encryption type. Expected values are either Deterministic or Randomized.</param>
         /// <param name="encryptionAlgorithm">Cryptographic algorithm.</param>
         /// <returns>An implementation of the AEAD_AES_256_CBC_HMAC_SHA256 cryptographic algorithm.</returns>
-        internal override SqlClientEncryptionAlgorithm Create(SqlClientSymmetricKey encryptionKey, SqlClientEncryptionType encryptionType, string encryptionAlgorithm)
+        internal override SqlClientEncryptionAlgorithm Create(SymmetricKey encryptionKey, EncryptionType encryptionType, string encryptionAlgorithm)
         {
             // Callers should have validated the encryption algorithm and the encryption key
             Debug.Assert(string.Equals(encryptionAlgorithm, SqlAeadAes256CbcHmac256Algorithm.AlgorithmName, StringComparison.OrdinalIgnoreCase));
 
             // Validate encryption type
-            if (encryptionType is not SqlClientEncryptionType.Deterministic and not SqlClientEncryptionType.Randomized)
+            if (encryptionType is not EncryptionType.Deterministic and not EncryptionType.Randomized)
             {
                 throw SQL.InvalidEncryptionType(SqlAeadAes256CbcHmac256Algorithm.AlgorithmName,
                                                 encryptionType,
-                                                SqlClientEncryptionType.Deterministic,
-                                                SqlClientEncryptionType.Randomized);
+                                                EncryptionType.Deterministic,
+                                                EncryptionType.Randomized);
             }
 
             // Get the cached cryptographic algorithm if one exists or create a new one, add it to cache and use it
@@ -70,7 +72,7 @@ internal override SqlClientEncryptionAlgorithm Create(SqlClientSymmetricKey encr
 
             if (!_encryptionAlgorithms.TryGetValue(algorithmKey, out SqlAeadAes256CbcHmac256Algorithm? aesAlgorithm))
             {
-                SqlAeadAes256CbcHmac256EncryptionKey encryptedKey = new(encryptionKey.RootKey, SqlAeadAes256CbcHmac256Algorithm.AlgorithmName);
+                AeadAes256CbcHmac256EncryptionKey encryptedKey = new(encryptionKey.RootKey);
                 aesAlgorithm = new SqlAeadAes256CbcHmac256Algorithm(encryptedKey, encryptionType, SqlAeadAes256CbcHmac256Algorithm.CurrentVersion);
 
                 // In case multiple threads reach here at the same time, the first one adds the value

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/AlwaysEncrypted/EncryptionAlgorithmFactory.cs

@@ -20,5 +20,5 @@ internal abstract class EncryptionAlgorithmFactory
     /// <param name="encryptionType">Encryption Type, some algorithms will need this</param>
     /// <param name="encryptionAlgorithm">Cryptographic algorithm name. Needed for extracting version bits</param>
     /// <returns>Return a newly created SqlClientEncryptionAlgorithm instance</returns>
-    internal abstract SqlClientEncryptionAlgorithm Create(SqlClientSymmetricKey encryptionKey, SqlClientEncryptionType encryptionType, string encryptionAlgorithm);
+    internal abstract SqlClientEncryptionAlgorithm Create(SymmetricKey encryptionKey, EncryptionType encryptionType, string encryptionAlgorithm);
 }

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/AlwaysEncrypted/EncryptionAlgorithmFactoryList.cs

@@ -24,14 +24,14 @@ internal static class EncryptionAlgorithmFactoryList
     /// <param name="type">Encryption type (read from SQL Server.)</param>
     /// <param name="algorithmName">Name of the cryptographic algorithm.</param>
     /// <param name="encryptionAlgorithm">Specified cryptographic algorithm's implementation.</param>
-    public static void GetAlgorithm(SqlClientSymmetricKey key, byte type, string algorithmName, out SqlClientEncryptionAlgorithm encryptionAlgorithm)
+    public static void GetAlgorithm(SymmetricKey key, byte type, string algorithmName, out SqlClientEncryptionAlgorithm encryptionAlgorithm)
     {
         EncryptionAlgorithmFactory factory = algorithmName switch
         {
             SqlAeadAes256CbcHmac256Algorithm.AlgorithmName => AeadAes256CbcHmac256Factory.Instance,
             _ => throw SQL.UnknownColumnEncryptionAlgorithm(algorithmName, RegisteredCipherAlgorithmNames)
         };
 
-        encryptionAlgorithm = factory.Create(key, (SqlClientEncryptionType)type, algorithmName);
+        encryptionAlgorithm = factory.Create(key, (EncryptionType)type, algorithmName);
     }
 }

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/AlwaysEncrypted/EncryptionType.cs

@@ -0,0 +1,17 @@
+// 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.
+
+namespace Microsoft.Data.SqlClient.AlwaysEncrypted
+{
+    /// <summary>
+    /// Encryption types supported in TCE. Corresponds to EncryptionAlgoType in MS-TDS.
+    /// </summary>
+    /// <see href="https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-tds/7091f6f6-b83d-4ed2-afeb-ba5013dfb18f"/>
+    internal enum EncryptionType
+    {
+        PlainText = 0x00,
+        Deterministic = 0x01,
+        Randomized = 0x02
+    }
+}

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlClientSymmetricKey.cs → src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/AlwaysEncrypted/SymmetricKey.cs

@@ -1,45 +1,36 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// 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.
 
-namespace Microsoft.Data.SqlClient
+#nullable enable
+
+namespace Microsoft.Data.SqlClient.AlwaysEncrypted
 {
     /// <summary>
     /// Base class containing raw key bytes for symmetric key algorithms. Some encryption algorithms can use the key directly while others derive sub keys from this.
     /// If an algorithm needs to derive more keys, have a derived class from this and use it in the corresponding encryption algorithm.
     /// </summary>
-    internal class SqlClientSymmetricKey
+    internal class SymmetricKey
     {
-        /// <summary>
-        /// The underlying key material
-        /// </summary>
-        protected readonly byte[] _rootKey;
-
         /// <summary>
         /// Constructor that initializes the root key.
         /// </summary>
-        /// <param name="rootKey">root key</param>
-        internal SqlClientSymmetricKey(byte[] rootKey)
+        /// <param name="rootKey">Root key</param>
+        public SymmetricKey(byte[]? rootKey)
         {
             // Key validation
-            if (rootKey == null || rootKey.Length == 0)
+            if (rootKey is null || rootKey.Length == 0)
             {
                 throw SQL.NullColumnEncryptionKeySysErr();
             }
 
-            _rootKey = rootKey;
+            RootKey = rootKey;
         }
 
         /// <summary>
-        /// Returns a copy of the plain text key
+        /// Returns the plain text key.
         /// This is needed for actual encryption/decryption.
         /// </summary>
-        internal virtual byte[] RootKey
-        {
-            get
-            {
-                return _rootKey;
-            }
-        }
+        public byte[] RootKey { get; }
     }
 }

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/AlwaysEncrypted/SymmetricKeyCache.cs

@@ -0,0 +1,117 @@
+// 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 System;
+using System.Diagnostics;
+using System.Text;
+using System.Threading;
+using Microsoft.Extensions.Caching.Memory;
+
+#nullable enable
+
+namespace Microsoft.Data.SqlClient.AlwaysEncrypted
+{
+    /// <summary>
+    /// Implements a cache of Symmetric Keys (once they are decrypted). Useful for rapidly decrypting multiple data values.
+    /// </summary>
+    internal sealed class SymmetricKeyCache
+    {
+        private static readonly SymmetricKeyCache s_singletonInstance = new();
+        private static readonly SemaphoreSlim s_cacheLock = new(1, 1);
+
+        private readonly MemoryCache _cache;
+
+        private SymmetricKeyCache()
+        {
+            _cache = new MemoryCache(new MemoryCacheOptions());
+        }
+
+        public static SymmetricKeyCache Instance => s_singletonInstance;
+
+        /// <summary>
+        /// Retrieves Symmetric Key (in plaintext) given the encryption material.
+        /// </summary>
+        public SymmetricKey GetKey(SqlEncryptionKeyInfo keyInfo, SqlConnection connection, SqlCommand? command)
+        {
+            string serverName = connection.DataSource;
+            Debug.Assert(serverName is not null, @"serverName should not be null.");
+            int capacity = serverName!.Length + SqlSecurityUtility.GetBase64LengthFromByteLength(keyInfo.encryptedKey.Length) + keyInfo.keyStoreName.Length + 2 /* separators */;
+            StringBuilder cacheLookupKeyBuilder = new(serverName, capacity);
+
+            cacheLookupKeyBuilder.Append(':');
+            cacheLookupKeyBuilder.Append(Convert.ToBase64String(keyInfo.encryptedKey));
+            cacheLookupKeyBuilder.Append(':');
+            cacheLookupKeyBuilder.Append(keyInfo.keyStoreName);
+
+            string cacheLookupKey = cacheLookupKeyBuilder.ToString();
+
+            Debug.Assert(cacheLookupKey.Length <= capacity, "We needed to allocate a larger array");
+
+            // Lookup the key in cache
+            if (!(_cache.TryGetValue(cacheLookupKey, out SymmetricKey? encryptionKey))
+                // A null cryptographic key is never added to the cache, but this null check satisfies the nullability warning.
+                || encryptionKey is null)
+            {
+                // Acquire the lock to ensure thread safety when modifying the cache, and to guarantee that only one thread calls
+                // DecryptColumnEncryptionKey on a user-provided SqlColumnEncryptionKeyStoreProvider at a time.
+                s_cacheLock.Wait();
+
+                try
+                {
+                    // Perform a second check to see if the key was added to the cache while waiting for the lock, to avoid redundant work.
+                    if (!(_cache.TryGetValue(cacheLookupKey, out encryptionKey))
+                        || encryptionKey is null)
+                    {
+                        Debug.Assert(SqlConnection.ColumnEncryptionTrustedMasterKeyPaths is not null, @"SqlConnection.ColumnEncryptionTrustedMasterKeyPaths should not be null");
+
+                        SqlSecurityUtility.ThrowIfKeyPathIsNotTrustedForServer(serverName, keyInfo.keyPath);
+
+                        // Key Not found, attempt to look up the provider and decrypt CEK
+                        if (!SqlSecurityUtility.TryGetColumnEncryptionKeyStoreProvider(keyInfo.keyStoreName, out SqlColumnEncryptionKeyStoreProvider provider, connection, command))
+                        {
+                            throw SQL.UnrecognizedKeyStoreProviderName(keyInfo.keyStoreName,
+                                    SqlConnection.GetColumnEncryptionSystemKeyStoreProvidersNames(),
+                                    SqlSecurityUtility.GetListOfProviderNamesThatWereSearched(connection, command));
+                        }
+
+                        // Decrypt the CEK
+                        // We will simply bubble up the exception from the DecryptColumnEncryptionKey function.
+                        byte[] plaintextKey;
+                        try
+                        {
+                            // AKV provider registration supports multi-user scenarios, so it is not safe to cache the CEK in the global provider.
+                            // The CEK cache is a global cache, and is shared across all connections.
+                            // To prevent conflicts between CEK caches, global providers should not use their own CEK caches
+                            provider.ColumnEncryptionKeyCacheTtl = TimeSpan.Zero;
+                            plaintextKey = provider.DecryptColumnEncryptionKey(keyInfo.keyPath, keyInfo.algorithmName, keyInfo.encryptedKey);
+                        }
+                        catch (Exception e)
+                        {
+                            // Generate a new exception and throw.
+                            string keyHex = SqlSecurityUtility.GetBytesAsString(keyInfo.encryptedKey, fLast: true, countOfBytes: 10);
+                            throw SQL.KeyDecryptionFailed(keyInfo.keyStoreName, keyHex, e);
+                        }
+
+                        encryptionKey = new SymmetricKey(plaintextKey);
+
+                        // If the cache TTL is zero, don't even bother inserting to the cache.
+                        if (SqlConnection.ColumnEncryptionKeyCacheTtl != TimeSpan.Zero)
+                        {
+                            // In case multiple threads reach here at the same time, the first one wins.
+                            // The allocated memory will be reclaimed by Garbage Collector.
+                            _cache.Set(cacheLookupKey, encryptionKey, absoluteExpirationRelativeToNow: SqlConnection.ColumnEncryptionKeyCacheTtl);
+                        }
+                    }
+                }
+                finally
+                {
+                    // Release the lock to allow other threads to access the cache
+                    s_cacheLock.Release();
+                }
+            }
+
+            return encryptionKey;
+        }
+    }
+}