Use the async Always Encrypted key store APIs on async execution paths

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.Crypto.cs

@@ -5,6 +5,8 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace Microsoft.Data.SqlClient
 {
@@ -155,11 +157,92 @@ internal EnclavePackage GenerateEnclavePackage(SqlConnectionAttestationProtocol
                 throw new RetryableEnclaveQueryExecutionException(e.Message, e);
             }
 
-            List<ColumnEncryptionKeyInfo> decryptedKeysToBeSentToEnclave = GetDecryptedKeysToBeSentToEnclave(keysToBeSentToEnclave, enclaveSessionParameters.ServerName, connection, command);
+            List<ColumnEncryptionKeyInfo> decryptedKeysToBeSentToEnclave = GetDecryptedKeysToBeSentToEnclave(keysToBeSentToEnclave, connection, command);
+            return BuildEnclavePackage(decryptedKeysToBeSentToEnclave, queryText, counter, sqlEnclaveSession, enclaveSessionParameters.ServerName);
+        }
+
+        /// <summary>
+        /// Asynchronously encrypts the byte package containing keys with the session key.
+        /// </summary>
+        /// <remarks>
+        /// Async counterpart of <see cref="GenerateEnclavePackage"/>. Only the column encryption key decryption
+        /// step performs I/O, so only that step is awaited; the remaining work is local cryptography and is
+        /// performed inline.
+        /// </remarks>
+        /// <param name="attestationProtocol">attestation protocol</param>
+        /// <param name="keysToBeSentToEnclave">Keys to be sent to enclave</param>
+        /// <param name="queryText">Text of the query being executed</param>
+        /// <param name="enclaveType">enclave type</param>
+        /// <param name="enclaveSessionParameters">The set of parameters required for enclave session.</param>
+        /// <param name="connection">connection executing the query</param>
+        /// <param name="command">command executing the query</param>
+        /// <param name="cancellationToken">Token used to request cancellation of the operation</param>
+        internal async Task<EnclavePackage> GenerateEnclavePackageAsync(
+            SqlConnectionAttestationProtocol attestationProtocol,
+            ConcurrentDictionary<int, SqlTceCipherInfoEntry> keysToBeSentToEnclave,
+            string queryText,
+            string enclaveType,
+            EnclaveSessionParameters enclaveSessionParameters,
+            SqlConnection connection,
+            SqlCommand command,
+            CancellationToken cancellationToken)
+        {
+            SqlEnclaveSession sqlEnclaveSession;
+            long counter;
+
+            try
+            {
+                // @TODO: GetEnclaveSession is still synchronous. On a session cache hit it is pure in-memory
+                // work, but on a miss it performs blocking attestation HTTP. Making it awaitable requires the
+                // async enclave provider hierarchy (Phase 3 of the async Always Encrypted spec), which adds
+                // public virtual API and therefore ships separately. Until then this is the one remaining
+                // blocking call on the asynchronous Always Encrypted path.
+                GetEnclaveSession(
+                    attestationProtocol,
+                    enclaveType,
+                    enclaveSessionParameters,
+                    generateCustomData: false,
+                    isRetry: false,
+                    sqlEnclaveSession: out sqlEnclaveSession,
+                    counter: out counter,
+                    customData: out _,
+                    customDataLength: out _,
+                    throwIfNull: true
+                );
+            }
+            catch (Exception e)
+            {
+                throw new RetryableEnclaveQueryExecutionException(e.Message, e);
+            }
+
+            List<ColumnEncryptionKeyInfo> decryptedKeysToBeSentToEnclave = await GetDecryptedKeysToBeSentToEnclaveAsync(
+                    keysToBeSentToEnclave,
+                    connection,
+                    command,
+                    cancellationToken)
+                .ConfigureAwait(false);
+
+            return BuildEnclavePackage(decryptedKeysToBeSentToEnclave, queryText, counter, sqlEnclaveSession, enclaveSessionParameters.ServerName);
+        }
+
+        /// <summary>
+        /// Assembles the enclave package from already-decrypted column encryption keys.
+        /// </summary>
+        /// <remarks>
+        /// Shared by the synchronous and asynchronous package generation paths. All work here is local
+        /// cryptography, so there is no asynchronous counterpart.
+        /// </remarks>
+        private EnclavePackage BuildEnclavePackage(
+            List<ColumnEncryptionKeyInfo> decryptedKeysToBeSentToEnclave,
+            string queryText,
+            long counter,
+            SqlEnclaveSession sqlEnclaveSession,
+            string serverName)
+        {
             byte[] queryStringHashBytes = ComputeQueryStringHash(queryText);
             byte[] keyBytePackage = GenerateBytePackageForKeys(counter, queryStringHashBytes, decryptedKeysToBeSentToEnclave);
             byte[] sessionKey = sqlEnclaveSession.GetSessionKey();
-            byte[] encryptedBytePackage = EncryptBytePackage(keyBytePackage, sessionKey, enclaveSessionParameters.ServerName);
+            byte[] encryptedBytePackage = EncryptBytePackage(keyBytePackage, sessionKey, serverName);
             byte[] enclaveSessionHandle = BitConverter.GetBytes(sqlEnclaveSession.SessionId);
             byte[] byteArrayToBeSentToEnclave = CombineByteArrays(enclaveSessionHandle, encryptedBytePackage);
             return new EnclavePackage(byteArrayToBeSentToEnclave, sqlEnclaveSession);

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.cs

@@ -8,6 +8,8 @@
 using System.Collections.Generic;
 using System.Security.Cryptography;
 using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace Microsoft.Data.SqlClient
 {
@@ -46,46 +48,90 @@ private byte[] GetUintBytes(string enclaveType, int intValue, string variableNam
         /// Decrypt the keys that need to be sent to the enclave
         /// </summary>
         /// <param name="keysTobeSentToEnclave">Keys that need to sent to the enclave</param>
-        /// <param name="serverName"></param>
-        /// <param name="connection"></param>
-        /// <param name="command"></param>
+        /// <param name="connection">Connection executing the query</param>
+        /// <param name="command">Command executing the query</param>
         /// <returns></returns>
-        private List<ColumnEncryptionKeyInfo> GetDecryptedKeysToBeSentToEnclave(ConcurrentDictionary<int, SqlTceCipherInfoEntry> keysTobeSentToEnclave, string serverName, SqlConnection connection, SqlCommand command)
+        internal List<ColumnEncryptionKeyInfo> GetDecryptedKeysToBeSentToEnclave(ConcurrentDictionary<int, SqlTceCipherInfoEntry> keysTobeSentToEnclave, SqlConnection connection, SqlCommand command)
         {
-            List<ColumnEncryptionKeyInfo> decryptedKeysToBeSentToEnclave = new List<ColumnEncryptionKeyInfo>();
+            List<ColumnEncryptionKeyInfo> decryptedKeysToBeSentToEnclave = new List<ColumnEncryptionKeyInfo>(keysTobeSentToEnclave.Count);
 
             foreach (SqlTceCipherInfoEntry cipherInfo in keysTobeSentToEnclave.Values)
             {
                 SqlSecurityUtility.DecryptSymmetricKey(cipherInfo, out SqlClientSymmetricKey sqlClientSymmetricKey, out SqlEncryptionKeyInfo encryptionkeyInfoChosen, connection, command);
 
-                if (sqlClientSymmetricKey == null)
-                {
-                    throw SQL.NullArgumentInternal(nameof(sqlClientSymmetricKey), nameof(EnclaveDelegate), nameof(GetDecryptedKeysToBeSentToEnclave));
-                }
-                if (cipherInfo.ColumnEncryptionKeyValues == null)
-                {
-                    throw SQL.NullArgumentInternal(nameof(cipherInfo.ColumnEncryptionKeyValues), nameof(EnclaveDelegate), nameof(GetDecryptedKeysToBeSentToEnclave));
-                }
-                if (!(cipherInfo.ColumnEncryptionKeyValues.Count > 0))
-                {
-                    throw SQL.ColumnEncryptionKeysNotFound();
-                }
+                decryptedKeysToBeSentToEnclave.Add(CreateColumnEncryptionKeyInfo(cipherInfo, sqlClientSymmetricKey));
+            }
+            return decryptedKeysToBeSentToEnclave;
+        }
 
-                //cipherInfo.CekId is always 0, hence used cipherInfo.ColumnEncryptionKeyValues[0].cekId. Even when cek has multiple ColumnEncryptionKeyValues
-                //the cekid and the plaintext value will remain the same, what varies is the encrypted cek value, since the cek can be encrypted by 
-                //multiple CMKs
-                decryptedKeysToBeSentToEnclave.Add(
-                    new ColumnEncryptionKeyInfo(
-                        sqlClientSymmetricKey.RootKey,
-                        cipherInfo.ColumnEncryptionKeyValues[0].databaseId,
-                        cipherInfo.ColumnEncryptionKeyValues[0].cekMdVersion,
-                        cipherInfo.ColumnEncryptionKeyValues[0].cekId
-                    )
-                );
+        /// <summary>
+        /// Asynchronously decrypts the keys that need to be sent to the enclave.
+        /// </summary>
+        /// <remarks>
+        /// Async counterpart of <see cref="GetDecryptedKeysToBeSentToEnclave"/>. Each column encryption key is
+        /// resolved through <see cref="SqlSecurityUtility.DecryptSymmetricKeyAsync(SqlTceCipherInfoEntry, SqlConnection, SqlCommand, CancellationToken)"/>
+        /// so that key store providers performing network I/O (for example Azure Key Vault) do not block a
+        /// thread while the enclave package is being assembled.
+        /// </remarks>
+        /// <param name="keysTobeSentToEnclave">Keys that need to sent to the enclave</param>
+        /// <param name="connection">Connection executing the query</param>
+        /// <param name="command">Command executing the query</param>
+        /// <param name="cancellationToken">Token used to request cancellation of the operation</param>
+        internal async Task<List<ColumnEncryptionKeyInfo>> GetDecryptedKeysToBeSentToEnclaveAsync(
+            ConcurrentDictionary<int, SqlTceCipherInfoEntry> keysTobeSentToEnclave,
+            SqlConnection connection,
+            SqlCommand command,
+            CancellationToken cancellationToken)
+        {
+            List<ColumnEncryptionKeyInfo> decryptedKeysToBeSentToEnclave = new List<ColumnEncryptionKeyInfo>(keysTobeSentToEnclave.Count);
+
+            foreach (SqlTceCipherInfoEntry cipherInfo in keysTobeSentToEnclave.Values)
+            {
+                (SqlClientSymmetricKey sqlClientSymmetricKey, SqlEncryptionKeyInfo _) =
+                    await SqlSecurityUtility.DecryptSymmetricKeyAsync(cipherInfo, connection, command, cancellationToken)
+                        .ConfigureAwait(false);
+
+                decryptedKeysToBeSentToEnclave.Add(CreateColumnEncryptionKeyInfo(cipherInfo, sqlClientSymmetricKey));
             }
+
             return decryptedKeysToBeSentToEnclave;
         }
 
+        /// <summary>
+        /// Validates a decrypted column encryption key and projects it into the shape the enclave expects.
+        /// </summary>
+        /// <remarks>
+        /// Shared by the synchronous and asynchronous decryption paths so that both apply identical validation
+        /// and produce identical <see cref="ColumnEncryptionKeyInfo"/> instances.
+        /// </remarks>
+        private static ColumnEncryptionKeyInfo CreateColumnEncryptionKeyInfo(
+            SqlTceCipherInfoEntry cipherInfo,
+            SqlClientSymmetricKey sqlClientSymmetricKey)
+        {
+            if (sqlClientSymmetricKey == null)
+            {
+                throw SQL.NullArgumentInternal(nameof(sqlClientSymmetricKey), nameof(EnclaveDelegate), nameof(CreateColumnEncryptionKeyInfo));
+            }
+            if (cipherInfo.ColumnEncryptionKeyValues == null)
+            {
+                throw SQL.NullArgumentInternal(nameof(cipherInfo.ColumnEncryptionKeyValues), nameof(EnclaveDelegate), nameof(CreateColumnEncryptionKeyInfo));
+            }
+            if (!(cipherInfo.ColumnEncryptionKeyValues.Count > 0))
+            {
+                throw SQL.ColumnEncryptionKeysNotFound();
+            }
+
+            //cipherInfo.CekId is always 0, hence used cipherInfo.ColumnEncryptionKeyValues[0].cekId. Even when cek has multiple ColumnEncryptionKeyValues
+            //the cekid and the plaintext value will remain the same, what varies is the encrypted cek value, since the cek can be encrypted by 
+            //multiple CMKs
+            return new ColumnEncryptionKeyInfo(
+                sqlClientSymmetricKey.RootKey,
+                cipherInfo.ColumnEncryptionKeyValues[0].databaseId,
+                cipherInfo.ColumnEncryptionKeyValues[0].cekMdVersion,
+                cipherInfo.ColumnEncryptionKeyValues[0].cekId
+            );
+        }
+
         /// <summary>
         /// Generate a byte package consisting of decrypted keys and some headers expected by the enclave
         /// </summary>