From 50c7217a94b78df92c6f3c7f7357164b1ae9e6c2 Mon Sep 17 00:00:00 2001 From: Cheena Malhotra Date: Fri, 21 Aug 2026 17:04:18 -0700 Subject: [PATCH 1/7] Use the async Always Encrypted key store APIs on async execution paths Async command execution now resolves Always Encrypted key material through the asynchronous key store provider APIs, so a provider that performs network I/O (for example Azure Key Vault) no longer blocks a thread pool thread for the duration of every column encryption key decryption, column master key signature verification, and enclave package generation. Utility layer (Phase 4): * SqlSecurityUtility gains DecryptSymmetricKeyAsync, GetKeyFromLocalProvidersAsync and VerifyColumnMasterKeySignatureAsync. * SqlSymmetricKeyCache gains GetKeyAsync, which never holds the cache lock across provider I/O. Call sites (Phase 5): * sp_describe_parameter_encryption result parsing is now shared by the sync and async paths. Signature verifications and key decryptions are recorded while parsing and executed afterwards, so no key store call is made while the reader is still positioned mid-stream. * The two callback-based describe-parameter-encryption continuations collapse into a single async method, preserving the existing thread pool hand-off and error/cancellation semantics. * EnclaveDelegate and SqlCommand gain async enclave package generation, used by the async execution path. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../Data/SqlClient/EnclaveDelegate.Crypto.cs | 81 ++- .../Data/SqlClient/EnclaveDelegate.cs | 97 +++- .../Data/SqlClient/SqlCommand.Encryption.cs | 391 +++++++++----- .../Data/SqlClient/SqlCommand.Reader.cs | 268 +++++++--- .../Data/SqlClient/SqlSecurityUtility.cs | 230 +++++++++ .../Data/SqlClient/SqlSymmetricKeyCache.cs | 154 +++++- .../EnclaveDelegateAsyncShould.cs | 239 +++++++++ .../SqlSecurityUtilityAsyncShould.cs | 486 ++++++++++++++++++ 8 files changed, 1685 insertions(+), 261 deletions(-) create mode 100644 src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/EnclaveDelegateAsyncShould.cs create mode 100644 src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlSecurityUtilityAsyncShould.cs diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.Crypto.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.Crypto.cs index 8b461087bd..610ae10fe0 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.Crypto.cs +++ b/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 { @@ -156,10 +158,87 @@ internal EnclavePackage GenerateEnclavePackage(SqlConnectionAttestationProtocol } List decryptedKeysToBeSentToEnclave = GetDecryptedKeysToBeSentToEnclave(keysToBeSentToEnclave, enclaveSessionParameters.ServerName, connection, command); + return BuildEnclavePackage(decryptedKeysToBeSentToEnclave, queryText, counter, sqlEnclaveSession, enclaveSessionParameters.ServerName); + } + + /// + /// Asynchronously encrypts the byte package containing keys with the session key. + /// + /// + /// Async counterpart of . 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. + /// + /// attestation protocol + /// Keys to be sent to enclave + /// Text of the query being executed + /// enclave type + /// The set of parameters required for enclave session. + /// connection executing the query + /// command executing the query + /// Token used to request cancellation of the operation + internal async Task GenerateEnclavePackageAsync( + SqlConnectionAttestationProtocol attestationProtocol, + ConcurrentDictionary keysToBeSentToEnclave, + string queryText, + string enclaveType, + EnclaveSessionParameters enclaveSessionParameters, + SqlConnection connection, + SqlCommand command, + CancellationToken cancellationToken) + { + SqlEnclaveSession sqlEnclaveSession; + long counter; + + try + { + 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 decryptedKeysToBeSentToEnclave = await GetDecryptedKeysToBeSentToEnclaveAsync( + keysToBeSentToEnclave, + enclaveSessionParameters.ServerName, + connection, + command, + cancellationToken) + .ConfigureAwait(false); + + return BuildEnclavePackage(decryptedKeysToBeSentToEnclave, queryText, counter, sqlEnclaveSession, enclaveSessionParameters.ServerName); + } + + /// + /// Assembles the enclave package from already-decrypted column encryption keys. + /// + /// + /// Shared by the synchronous and asynchronous package generation paths. All work here is local + /// cryptography, so there is no asynchronous counterpart. + /// + private EnclavePackage BuildEnclavePackage( + List 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); diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.cs index dcd5fd200a..1c6ef2fd35 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.cs +++ b/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 { @@ -50,7 +52,7 @@ private byte[] GetUintBytes(string enclaveType, int intValue, string variableNam /// /// /// - private List GetDecryptedKeysToBeSentToEnclave(ConcurrentDictionary keysTobeSentToEnclave, string serverName, SqlConnection connection, SqlCommand command) + internal List GetDecryptedKeysToBeSentToEnclave(ConcurrentDictionary keysTobeSentToEnclave, string serverName, SqlConnection connection, SqlCommand command) { List decryptedKeysToBeSentToEnclave = new List(); @@ -58,34 +60,81 @@ private List GetDecryptedKeysToBeSentToEnclave(Concurre { 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 - ) - ); + /// + /// Asynchronously decrypts the keys that need to be sent to the enclave. + /// + /// + /// Async counterpart of . Each column encryption key is + /// resolved through + /// 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. + /// + /// Keys that need to sent to the enclave + /// Name of the server the keys are being resolved for + /// Connection executing the query + /// Command executing the query + /// Token used to request cancellation of the operation + internal async Task> GetDecryptedKeysToBeSentToEnclaveAsync( + ConcurrentDictionary keysTobeSentToEnclave, + string serverName, + SqlConnection connection, + SqlCommand command, + CancellationToken cancellationToken) + { + List decryptedKeysToBeSentToEnclave = new List(); + + 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; } + /// + /// Validates a decrypted column encryption key and projects it into the shape the enclave expects. + /// + /// + /// Shared by the synchronous and asynchronous decryption paths so that both apply identical validation + /// and produce identical instances. + /// + private static ColumnEncryptionKeyInfo CreateColumnEncryptionKeyInfo( + SqlTceCipherInfoEntry cipherInfo, + SqlClientSymmetricKey sqlClientSymmetricKey) + { + 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(); + } + + //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 + ); + } + /// /// Generate a byte package consisting of decrypted keys and some headers expected by the enclave /// diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs index 3b0fc6e93f..fc6a7b67b0 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs @@ -14,7 +14,6 @@ using System.Threading.Tasks; using Microsoft.Data.Common; using Microsoft.Data.SqlClient.Connection; -using Microsoft.Data.SqlClient.Utilities; namespace Microsoft.Data.SqlClient { @@ -243,8 +242,32 @@ private EnclaveSessionParameters GetEnclaveSessionParameters() => _activeConnection.EnclaveAttestationUrl, _activeConnection.Database); - // @TODO: Isn't this doing things asynchronously? We should just have a purely asynchronous and a purely synchronous pathway instead of this mix of check this check that and flags. - private SqlDataReader GetParameterEncryptionDataReader( + /// + /// Schedules asynchronous consumption of the sp_describe_parameter_encryption results. + /// + /// + /// + /// The returned task completes once the describe-parameter-encryption results have been read and every + /// column encryption key has been decrypted through the asynchronous key store provider APIs. + /// + /// + /// The work is dispatched with because the body issues blocking TDS + /// reads before reaching its first suspension point, and PrepareForTransparentEncryption is + /// invoked synchronously on the caller's thread by the asynchronous execution entry points. Dispatching + /// keeps that thread free, matching the behaviour of the continuation chain this replaced, while every + /// subsequent await releases the pooled thread for the duration of key store provider I/O. + /// + /// + /// Receives the task representing the pending work + /// + /// Task representing the pending network write of the describe-parameter-encryption request, or + /// null when that write completed synchronously. + /// + /// Reader over the describe-parameter-encryption results + /// Map of encryption RPC requests to their original RPC requests + /// Whether describe parameter encryption was required + /// Indicates if this is a retry from a failed call + private void GetParameterEncryptionDataReader( out Task returnTask, Task fetchInputParameterEncryptionInfoTask, SqlDataReader describeParameterEncryptionDataReader, @@ -252,143 +275,102 @@ private SqlDataReader GetParameterEncryptionDataReader( bool describeParameterEncryptionNeeded, bool isRetry) { - returnTask = AsyncHelper.CreateContinuationTaskWithState( - taskToContinue: fetchInputParameterEncryptionInfoTask, - state: this, - onSuccess: sqlCommand => - { - bool processFinallyBlockAsync = true; - bool decrementAsyncCountInFinallyBlockAsync = true; - - try - { - // Check for any exceptions on network write, before reading. - sqlCommand.CheckThrowSNIException(); - - // If it is async, then TryFetchInputParameterEncryptionInfo -> - // RunExecuteReaderTds would have incremented the async count. Decrement it - // when we are about to complete async execute reader. - SqlConnectionInternal internalConnectionTds = sqlCommand._activeConnection.GetOpenTdsConnection(); - if (internalConnectionTds is not null) - { - internalConnectionTds.DecrementAsyncCount(); - decrementAsyncCountInFinallyBlockAsync = false; - } - - // Complete executereader. - // @TODO: If we can remove this reference, this could be a static lambda - describeParameterEncryptionDataReader = sqlCommand.CompleteAsyncExecuteReader( - isInternal: false, - forDescribeParameterEncryption: true); - Debug.Assert(sqlCommand._stateObj is null, "non-null state object in PrepareForTransparentEncryption."); - - // Read the results of describe parameter encryption. - sqlCommand.ReadDescribeEncryptionParameterResults( - describeParameterEncryptionDataReader, - describeParameterEncryptionRpcOriginalRpcMap, - isRetry); - - #if DEBUG - // Failpoint to force the thread to halt to simulate cancellation of SqlCommand. - if (_sleepAfterReadDescribeEncryptionParameterResults) - { - Thread.Sleep(TimeSpan.FromSeconds(10)); - } - #endif - } - catch (Exception e) - { - processFinallyBlockAsync = ADP.IsCatchableExceptionType(e); - throw; - } - finally - { - sqlCommand.PrepareTransparentEncryptionFinallyBlock( - closeDataReader: processFinallyBlockAsync, - decrementAsyncCount: decrementAsyncCountInFinallyBlockAsync, - clearDataStructures: processFinallyBlockAsync, - wasDescribeParameterEncryptionNeeded: describeParameterEncryptionNeeded, - describeParameterEncryptionRpcOriginalRpcMap: describeParameterEncryptionRpcOriginalRpcMap, - describeParameterEncryptionDataReader: describeParameterEncryptionDataReader); - } - }, - onFailure: static (sqlCommand, exception) => - { - sqlCommand.CachedAsyncState?.ResetAsyncState(); - if (exception is not null) - { - throw exception; - } - }); - - return describeParameterEncryptionDataReader; + returnTask = Task.Run(() => GetParameterEncryptionDataReaderAsync( + fetchInputParameterEncryptionInfoTask, + describeParameterEncryptionDataReader, + describeParameterEncryptionRpcOriginalRpcMap, + describeParameterEncryptionNeeded, + isRetry)); } - private SqlDataReader GetParameterEncryptionDataReaderAsync( - out Task returnTask, + /// + /// Awaits the describe-parameter-encryption request and consumes its results asynchronously. + /// + /// + /// Failure of resets the cached async state and + /// skips the transparent encryption finally block; cancellation of it does neither. Failures raised + /// while consuming the results run the finally block but leave the cached async state alone. These + /// semantics are inherited from the callback-based continuation this method replaced. + /// + private async Task GetParameterEncryptionDataReaderAsync( + Task fetchInputParameterEncryptionInfoTask, SqlDataReader describeParameterEncryptionDataReader, ReadOnlyDictionary<_SqlRPC, _SqlRPC> describeParameterEncryptionRpcOriginalRpcMap, bool describeParameterEncryptionNeeded, bool isRetry) { - returnTask = Task.Run(() => + if (fetchInputParameterEncryptionInfoTask is not null) { - bool processFinallyBlockAsync = true; - bool decrementAsyncCountInFinallyBlockAsync = true; - try { - // Check for any exception on network write before reading. - CheckThrowSNIException(); - - // If it is async, then TryFetchInputParameterEncryptionInfo -> - // RunExecuteReaderTds would have incremented the async count. Decrement it - // when we are about to complete async execute reader. - SqlConnectionInternal internalConnectionTds = _activeConnection.GetOpenTdsConnection(); - if (internalConnectionTds is not null) - { - internalConnectionTds.DecrementAsyncCount(); - decrementAsyncCountInFinallyBlockAsync = false; - } - - // Complete executereader. - describeParameterEncryptionDataReader = CompleteAsyncExecuteReader( - isInternal: false, - forDescribeParameterEncryption: true); - Debug.Assert(_stateObj is null, "non-null state object in PrepareForTransparentEncryption."); - - // Read the results of describe parameter encryption. - ReadDescribeEncryptionParameterResults( - describeParameterEncryptionDataReader, - describeParameterEncryptionRpcOriginalRpcMap, - isRetry); - - #if DEBUG - // Failpoint to force the thread to halt to simulate cancellation of SqlCommand. - if (_sleepAfterReadDescribeEncryptionParameterResults) - { - Thread.Sleep(TimeSpan.FromSeconds(10)); - } - #endif + await fetchInputParameterEncryptionInfoTask.ConfigureAwait(false); } - catch (Exception e) + catch (OperationCanceledException) { - processFinallyBlockAsync = ADP.IsCatchableExceptionType(e); throw; } - finally + catch { - PrepareTransparentEncryptionFinallyBlock( - closeDataReader: processFinallyBlockAsync, - decrementAsyncCount: decrementAsyncCountInFinallyBlockAsync, - clearDataStructures: processFinallyBlockAsync, - wasDescribeParameterEncryptionNeeded: describeParameterEncryptionNeeded, - describeParameterEncryptionRpcOriginalRpcMap: describeParameterEncryptionRpcOriginalRpcMap, - describeParameterEncryptionDataReader: describeParameterEncryptionDataReader); + CachedAsyncState?.ResetAsyncState(); + throw; } - }); + } + + bool processFinallyBlockAsync = true; + bool decrementAsyncCountInFinallyBlockAsync = true; - return describeParameterEncryptionDataReader; + try + { + // Check for any exceptions on network write, before reading. + CheckThrowSNIException(); + + // If it is async, then TryFetchInputParameterEncryptionInfo -> + // RunExecuteReaderTds would have incremented the async count. Decrement it + // when we are about to complete async execute reader. + SqlConnectionInternal internalConnectionTds = _activeConnection.GetOpenTdsConnection(); + if (internalConnectionTds is not null) + { + internalConnectionTds.DecrementAsyncCount(); + decrementAsyncCountInFinallyBlockAsync = false; + } + + // Complete executereader. + describeParameterEncryptionDataReader = CompleteAsyncExecuteReader( + isInternal: false, + forDescribeParameterEncryption: true); + Debug.Assert(_stateObj is null, "non-null state object in PrepareForTransparentEncryption."); + + // Read the results of describe parameter encryption. + await ReadDescribeEncryptionParameterResultsAsync( + describeParameterEncryptionDataReader, + describeParameterEncryptionRpcOriginalRpcMap, + isRetry, + CancellationToken.None) + .ConfigureAwait(false); + + #if DEBUG + // Failpoint to force the thread to halt to simulate cancellation of SqlCommand. + if (_sleepAfterReadDescribeEncryptionParameterResults) + { + Thread.Sleep(TimeSpan.FromSeconds(10)); + } + #endif + } + catch (Exception e) + { + processFinallyBlockAsync = ADP.IsCatchableExceptionType(e); + throw; + } + finally + { + PrepareTransparentEncryptionFinallyBlock( + closeDataReader: processFinallyBlockAsync, + decrementAsyncCount: decrementAsyncCountInFinallyBlockAsync, + clearDataStructures: processFinallyBlockAsync, + wasDescribeParameterEncryptionNeeded: describeParameterEncryptionNeeded, + describeParameterEncryptionRpcOriginalRpcMap: describeParameterEncryptionRpcOriginalRpcMap, + describeParameterEncryptionDataReader: describeParameterEncryptionDataReader); + } } private void InvalidateEnclaveSession() @@ -662,7 +644,7 @@ private void PrepareForTransparentEncryption( // execution pending. Note that this should be done outside the task's // continuation delegate. processFinallyBlock = false; - describeParameterEncryptionDataReader = GetParameterEncryptionDataReader( + GetParameterEncryptionDataReader( out returnTask, fetchInputParameterEncryptionInfoTask, describeParameterEncryptionDataReader, @@ -682,8 +664,9 @@ private void PrepareForTransparentEncryption( // execution pending. Note that this should be done outside the task's // continuation delegate. processFinallyBlock = false; - describeParameterEncryptionDataReader = GetParameterEncryptionDataReaderAsync( + GetParameterEncryptionDataReader( out returnTask, + fetchInputParameterEncryptionInfoTask: null, describeParameterEncryptionDataReader, describeParameterEncryptionRpcOriginalRpcMap, describeParameterEncryptionNeeded, @@ -786,9 +769,106 @@ private void PrepareTransparentEncryptionFinallyBlock( /// Readonly dictionary with the map of parameter encryption rpc requests with the corresponding original rpc requests. /// Indicates if this is a retry from a failed call. private void ReadDescribeEncryptionParameterResults( - SqlDataReader ds, // @TODO: Rename something more obvious + SqlDataReader ds, ReadOnlyDictionary<_SqlRPC, _SqlRPC> describeParameterEncryptionRpcOriginalRpcMap, bool isRetry) + { + PendingColumnEncryptionKeyOperations pending = new PendingColumnEncryptionKeyOperations(); + ReadDescribeEncryptionParameterResultsCore(ds, describeParameterEncryptionRpcOriginalRpcMap, isRetry, pending); + + foreach (ColumnMasterKeySignatureVerification verification in pending.SignatureVerifications) + { + SqlSecurityUtility.VerifyColumnMasterKeySignature( + verification.KeyStoreName, + verification.KeyPath, + isEnclaveEnabled: true, + verification.Signature, + _activeConnection, + this); + } + + foreach (SqlCipherMetadata cipherMetadata in pending.KeyDecryptions) + { + SqlSecurityUtility.DecryptSymmetricKey(cipherMetadata, _activeConnection, this); + } + + CacheQueryMetadataIfNeeded(); + } + + /// + /// Asynchronously reads the output of sp_describe_parameter_encryption. + /// + /// + /// Async counterpart of . Result set parsing is + /// shared with the synchronous path; only the column master key signature verifications and column + /// encryption key decryptions differ, and those are the operations that may reach out to a key store + /// over the network. + /// + /// Resultset from calling to sp_describe_parameter_encryption + /// Readonly dictionary with the map of parameter encryption rpc requests with the corresponding original rpc requests. + /// Indicates if this is a retry from a failed call. + /// Token used to request cancellation of the operation + private async Task ReadDescribeEncryptionParameterResultsAsync( + SqlDataReader ds, + ReadOnlyDictionary<_SqlRPC, _SqlRPC> describeParameterEncryptionRpcOriginalRpcMap, + bool isRetry, + CancellationToken cancellationToken) + { + PendingColumnEncryptionKeyOperations pending = new PendingColumnEncryptionKeyOperations(); + ReadDescribeEncryptionParameterResultsCore(ds, describeParameterEncryptionRpcOriginalRpcMap, isRetry, pending); + + foreach (ColumnMasterKeySignatureVerification verification in pending.SignatureVerifications) + { + await SqlSecurityUtility.VerifyColumnMasterKeySignatureAsync( + verification.KeyStoreName, + verification.KeyPath, + isEnclaveEnabled: true, + verification.Signature, + _activeConnection, + this, + cancellationToken) + .ConfigureAwait(false); + } + + foreach (SqlCipherMetadata cipherMetadata in pending.KeyDecryptions) + { + await SqlSecurityUtility.DecryptSymmetricKeyAsync(cipherMetadata, _activeConnection, this, cancellationToken) + .ConfigureAwait(false); + } + + CacheQueryMetadataIfNeeded(); + } + + /// + /// Adds the encryption metadata for the current query to the query metadata cache when applicable. + /// + private void CacheQueryMetadataIfNeeded() + { + // If we are not in Batch RPC mode, update the query cache with the encryption MD. + if (!_batchRPCMode && ShouldCacheEncryptionMetadata && _parameters?.Count > 0) + { + SqlQueryMetadataCache.GetInstance().AddQueryMetadata(this, ignoreQueriesWithReturnValueParams: true); + } + } + + /// + /// Parses the result sets returned by sp_describe_parameter_encryption. + /// + /// + /// Column master key signature verification and column encryption key decryption are not performed + /// here; they are recorded in so that the caller can execute them either + /// synchronously or asynchronously. Deferring them also means no key store network call is made while + /// the describe-parameter-encryption reader is still positioned mid-stream. + /// + /// Resultset from calling to sp_describe_parameter_encryption + /// Readonly dictionary with the map of parameter encryption rpc requests with the corresponding original rpc requests. + /// Indicates if this is a retry from a failed call. + /// Collects the key store operations that the caller must complete + private void ReadDescribeEncryptionParameterResultsCore( + SqlDataReader ds, // @TODO: Rename something more obvious + ReadOnlyDictionary<_SqlRPC, _SqlRPC> describeParameterEncryptionRpcOriginalRpcMap, + bool isRetry, + PendingColumnEncryptionKeyOperations pending) { // @TODO: This should be SqlTceCipherInfoTable Dictionary columnEncryptionKeyTable = new Dictionary(); @@ -820,7 +900,7 @@ private void ReadDescribeEncryptionParameterResults( } // 1) Read the first result set that contains the column encryption key list - bool enclaveMetadataExists = ReadDescribeEncryptionParameterResultsKeys(ds, columnEncryptionKeyTable); + bool enclaveMetadataExists = ReadDescribeEncryptionParameterResultsKeys(ds, columnEncryptionKeyTable, pending); if (!enclaveMetadataExists && !ds.NextResult()) { throw SQL.UnexpectedDescribeParamFormatParameterMetadata(); @@ -854,7 +934,7 @@ private void ReadDescribeEncryptionParameterResults( int receivedMetadataCount = 0; if (!enclaveMetadataExists || ds.NextResult()) { - receivedMetadataCount = ReadDescribeEncryptionParameterResultsMetadata(ds, rpc, columnEncryptionKeyTable); + receivedMetadataCount = ReadDescribeEncryptionParameterResultsMetadata(ds, rpc, columnEncryptionKeyTable, pending); } // When the RPC object gets reused, the parameter array has more parameters that the valid params for the command. @@ -902,12 +982,38 @@ private void ReadDescribeEncryptionParameterResults( } } } + } - // If we are not in Batch RPC mode, update the query cache with the encryption MD. - if (!_batchRPCMode && ShouldCacheEncryptionMetadata && _parameters?.Count > 0) + /// + /// Key store operations discovered while parsing sp_describe_parameter_encryption results, deferred so + /// that they can be executed either synchronously or asynchronously. + /// + private sealed class PendingColumnEncryptionKeyOperations + { + /// Column master key signatures that must be verified. + internal List SignatureVerifications { get; } = new(); + + /// Column encryption keys that must be decrypted. + internal List KeyDecryptions { get; } = new(); + } + + /// + /// Describes a pending column master key signature verification. + /// + private readonly struct ColumnMasterKeySignatureVerification + { + internal ColumnMasterKeySignatureVerification(string keyStoreName, string keyPath, byte[] signature) { - SqlQueryMetadataCache.GetInstance().AddQueryMetadata(this, ignoreQueriesWithReturnValueParams: true); + KeyStoreName = keyStoreName; + KeyPath = keyPath; + Signature = signature; } + + internal string KeyStoreName { get; } + + internal string KeyPath { get; } + + internal byte[] Signature { get; } } private void ReadDescribeEncryptionParameterResultsAttestation(SqlDataReader ds, bool isRetry) @@ -960,7 +1066,8 @@ private void ReadDescribeEncryptionParameterResultsAttestation(SqlDataReader ds, private bool ReadDescribeEncryptionParameterResultsKeys( SqlDataReader ds, - Dictionary columnEncryptionKeyTable) + Dictionary columnEncryptionKeyTable, + PendingColumnEncryptionKeyOperations pending) { bool enclaveMetadataExists = true; while (ds.Read()) @@ -1060,13 +1167,10 @@ private bool ReadDescribeEncryptionParameterResultsKeys( length: keySignatureLength); } - SqlSecurityUtility.VerifyColumnMasterKeySignature( - providerName, - keyPath, - isEnclaveEnabled: isRequestedByEnclave, - keySignature, - _activeConnection, - this); + // Defer signature verification: it may reach a key store over the network and must not + // run while this reader is still positioned mid-result-set. + pending.SignatureVerifications.Add( + new ColumnMasterKeySignatureVerification(providerName, keyPath, keySignature)); // Lookup the key, failing which throw an exception // @TODO: Seriously, we *just* did this, why are we looking it up again?? @@ -1101,7 +1205,8 @@ private bool ReadDescribeEncryptionParameterResultsKeys( private int ReadDescribeEncryptionParameterResultsMetadata( SqlDataReader ds, _SqlRPC rpc, - Dictionary columnEncryptionKeyTable) + Dictionary columnEncryptionKeyTable, + PendingColumnEncryptionKeyOperations pending) { Debug.Assert(rpc is not null, "Describe Parameter Encryption requested for non-TCE spec proc"); @@ -1156,8 +1261,10 @@ private int ReadDescribeEncryptionParameterResultsMetadata( encryptionType: columnEncryptionType, normalizationRuleVersion: columnNormalizationRuleVersion); - // Decrypt the symmetric key. This will also validate and throw if needed. - SqlSecurityUtility.DecryptSymmetricKey(sqlParameter.CipherMetadata, _activeConnection, this); + // Defer decryption of the symmetric key: it may reach a key store over the network + // and must not run while this reader is still positioned mid-result-set. Decryption + // also validates the metadata and will throw if it is invalid. + pending.KeyDecryptions.Add(sqlParameter.CipherMetadata); // This is effective only for _batchRPCMode even though we set it for // non-_batchRPCMode also, since for non-_batchRPCMode, param options diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs index ec4d7709da..0fbb7cdf65 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs @@ -893,10 +893,111 @@ private void FinishExecuteReader( private void GenerateEnclavePackage() { + if (!TryPrepareEnclavePackageGeneration( + out string enclaveType, + out SqlConnectionAttestationProtocol attestationProtocol)) + { + return; + } + + // Generate the enclave package + try + { + ThrowIfForcedRetryableEnclaveQueryExecutionException(); + + enclavePackage = EnclaveDelegate.Instance.GenerateEnclavePackage( + attestationProtocol, + keysToBeSentToEnclave, + CommandText, + enclaveType, + GetEnclaveSessionParameters(), + _activeConnection, + command: this); + } + catch (EnclaveDelegate.RetryableEnclaveQueryExecutionException) + { + throw; + } + catch (Exception e) + { + throw SQL.ExceptionWhenGeneratingEnclavePackage(e); + } + } + + /// + /// Asynchronously generates the enclave package for the current command. + /// + /// + /// Async counterpart of . Column encryption keys destined for the + /// enclave are decrypted through the asynchronous key store provider APIs, so a key store that performs + /// network I/O does not block a thread while the package is generated. + /// + /// Token used to request cancellation of the operation + private async Task GenerateEnclavePackageAsync(CancellationToken cancellationToken) + { + if (!TryPrepareEnclavePackageGeneration( + out string enclaveType, + out SqlConnectionAttestationProtocol attestationProtocol)) + { + return; + } + + // Generate the enclave package + try + { + ThrowIfForcedRetryableEnclaveQueryExecutionException(); + + enclavePackage = await EnclaveDelegate.Instance.GenerateEnclavePackageAsync( + attestationProtocol, + keysToBeSentToEnclave, + CommandText, + enclaveType, + GetEnclaveSessionParameters(), + _activeConnection, + command: this, + cancellationToken) + .ConfigureAwait(false); + } + catch (EnclaveDelegate.RetryableEnclaveQueryExecutionException) + { + throw; + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + // Propagate cancellation unwrapped so the returned Task is cancelled rather than faulted + // with an enclave-specific exception that would suggest an attestation problem. + throw; + } + catch (Exception e) + { + throw SQL.ExceptionWhenGeneratingEnclavePackage(e); + } + } + + /// + /// Validates the connection state required to generate an enclave package. + /// + /// + /// Shared by the synchronous and asynchronous enclave package generation paths so that both perform + /// identical validation and throw identical exceptions. + /// + /// The validated enclave type reported by the server + /// The validated attestation protocol configured on the connection + /// + /// false when there are no keys to send to the enclave and package generation must be skipped; + /// otherwise true. + /// + private bool TryPrepareEnclavePackageGeneration( + out string enclaveType, + out SqlConnectionAttestationProtocol attestationProtocol) + { + enclaveType = null; + attestationProtocol = SqlConnectionAttestationProtocol.NotSpecified; + // Skip processing if there are no keys to send to enclave if (keysToBeSentToEnclave is null || keysToBeSentToEnclave.IsEmpty) { - return; + return false; } // Validate attestation url is provided when necessary @@ -908,49 +1009,37 @@ private void GenerateEnclavePackage() } // Validate enclave type - string enclaveType = _activeConnection.Parser.EnclaveType; + enclaveType = _activeConnection.Parser.EnclaveType; if (string.IsNullOrWhiteSpace(enclaveType)) { throw SQL.EnclaveTypeNullForEnclaveBasedQuery(); } // Validate protocol type - SqlConnectionAttestationProtocol attestationProtocol = _activeConnection.AttestationProtocol; + attestationProtocol = _activeConnection.AttestationProtocol; if (attestationProtocol is SqlConnectionAttestationProtocol.NotSpecified) { throw SQL.AttestationProtocolNotSpecifiedForGeneratingEnclavePackage(); } - // Generate the enclave package - try - { - #if DEBUG - // @TODO: These should be wrapped with something other than DEBUG since we don't even run tests in debug mode - // Test-only code for forcing a retryable exception to occur - if (_forceRetryableEnclaveQueryExecutionExceptionDuringGenerateEnclavePackage) - { - _forceRetryableEnclaveQueryExecutionExceptionDuringGenerateEnclavePackage = false; - throw new EnclaveDelegate.RetryableEnclaveQueryExecutionException("testing", null); - } - #endif + return true; + } - enclavePackage = EnclaveDelegate.Instance.GenerateEnclavePackage( - attestationProtocol, - keysToBeSentToEnclave, - CommandText, - enclaveType, - GetEnclaveSessionParameters(), - _activeConnection, - command: this); - } - catch (EnclaveDelegate.RetryableEnclaveQueryExecutionException) + /// + /// Test-only failpoint that forces a retryable enclave failure during package generation. + /// + [Conditional("DEBUG")] + private void ThrowIfForcedRetryableEnclaveQueryExecutionException() + { + #if DEBUG + // @TODO: These should be wrapped with something other than DEBUG since we don't even run tests in debug mode + // Test-only code for forcing a retryable exception to occur + if (_forceRetryableEnclaveQueryExecutionExceptionDuringGenerateEnclavePackage) { - throw; - } - catch (Exception e) - { - throw SQL.ExceptionWhenGeneratingEnclavePackage(e); + _forceRetryableEnclaveQueryExecutionExceptionDuringGenerateEnclavePackage = false; + throw new EnclaveDelegate.RetryableEnclaveQueryExecutionException("testing", null); } + #endif } private Task InternalExecuteReaderAsync( @@ -1759,55 +1848,21 @@ private SqlDataReader RunExecuteReaderTdsWithTransparentParameterEncryption( // @TODO: I guess this means async execution? Using tasks as the primary means of determining async vs sync is clunky. It would be better to have separate async vs sync pathways. long parameterEncryptionStart = ADP.TimerCurrent(); - // @TODO: This can totally be a non-generic TCS - // @TODO: This is a prime candidate for proper async-await execution - TaskCompletionSource completion = new TaskCompletionSource(); - AsyncHelper.ContinueTaskWithState( - taskToContinue: describeParameterEncryptionTask, - taskCompletionSource: completion, - state: this, - onSuccess: sqlCommand => - { - sqlCommand.GenerateEnclavePackage(); - sqlCommand.RunExecuteReaderTds( - cmdBehavior, - runBehavior, - returnStream, - isAsync, - TdsParserStaticMethods.GetRemainingTimeout(timeout, parameterEncryptionStart), - out Task subTask, - asyncWrite, - isRetry, - ds); - - if (subTask is null) - { - // @TODO: Why would this ever be the case? We should structure this so that it doesn't need to be checked. - completion.SetResult(null); - } - else - { - AsyncHelper.ContinueTaskWithState( - taskToContinue: subTask, - taskCompletionSource: completion, - state: completion, - onSuccess: static state => state.SetResult(null)); - } - }, - onFailure: static (sqlCommand, exception) => - { - sqlCommand.CachedAsyncState?.ResetAsyncState(); - if (exception is not null) - { - throw exception; - } - }, - onCancellation: static sqlCommand => - { - sqlCommand.CachedAsyncState?.ResetAsyncState(); - }); + // The remainder of the execution issues blocking TDS writes, so it must never run inline on + // the caller's thread (nor on a network callback thread). Task.Run reproduces the thread pool + // hand-off that the previous ContinueWith-based continuation provided. + task = Task.Run(() => ContinueRunExecuteReaderTdsAsync( + describeParameterEncryptionTask, + cmdBehavior, + runBehavior, + returnStream, + isAsync, + timeout, + parameterEncryptionStart, + asyncWrite, + isRetry, + ds)); - task = completion.Task; return ds; } else @@ -1827,6 +1882,63 @@ private SqlDataReader RunExecuteReaderTdsWithTransparentParameterEncryption( } } + /// + /// Completes transparent parameter encryption processing and then executes the command. + /// + /// + /// + /// Awaits the describe-parameter-encryption round trip, generates the enclave package using the + /// asynchronous key store provider APIs, and then issues the actual command execution. This replaces a + /// callback-based continuation chain so that key store network I/O yields the thread instead of + /// blocking it. + /// + /// + /// Error semantics match the previous continuation chain: a failure or cancellation of the + /// describe-parameter-encryption round trip resets the cached async state, while a failure that occurs + /// after that point is surfaced without resetting it. + /// + /// + private async Task ContinueRunExecuteReaderTdsAsync( + Task describeParameterEncryptionTask, + CommandBehavior cmdBehavior, + RunBehavior runBehavior, + bool returnStream, + bool isAsync, + int timeout, + long parameterEncryptionStart, + bool asyncWrite, + bool isRetry, + SqlDataReader ds) + { + try + { + await describeParameterEncryptionTask.ConfigureAwait(false); + } + catch + { + CachedAsyncState?.ResetAsyncState(); + throw; + } + + await GenerateEnclavePackageAsync(CancellationToken.None).ConfigureAwait(false); + + RunExecuteReaderTds( + cmdBehavior, + runBehavior, + returnStream, + isAsync, + TdsParserStaticMethods.GetRemainingTimeout(timeout, parameterEncryptionStart), + out Task subTask, + asyncWrite, + isRetry, + ds); + + if (subTask is not null) + { + await subTask.ConfigureAwait(false); + } + } + private SqlDataReader RunExecuteReaderWithRetry( CommandBehavior cmdBehavior, RunBehavior runBehavior, diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSecurityUtility.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSecurityUtility.cs index 79ebc60d27..b5e92f4c09 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSecurityUtility.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSecurityUtility.cs @@ -8,6 +8,8 @@ using System.Reflection; using System.Security.Cryptography; using System.Text; +using System.Threading; +using System.Threading.Tasks; using Microsoft.Data.Common; using Microsoft.Data.SqlClient.AlwaysEncrypted; @@ -214,6 +216,96 @@ internal static void DecryptSymmetricKey(SqlCipherMetadata md, SqlConnection con return; } + /// + /// Asynchronously decrypts the symmetric key and saves it in metadata. In addition, initializes + /// the SqlClientEncryptionAlgorithm for rapid decryption. + /// + /// + /// This is the async counterpart of + /// and mutates in exactly the same way. + /// + internal static async Task DecryptSymmetricKeyAsync(SqlCipherMetadata md, SqlConnection connection, SqlCommand command, CancellationToken cancellationToken) + { + Debug.Assert(md is not null, "md should not be null in DecryptSymmetricKeyAsync."); + + (SqlClientSymmetricKey symKey, SqlEncryptionKeyInfo encryptionkeyInfoChosen) = + await DecryptSymmetricKeyAsync(md.EncryptionInfo, connection, command, cancellationToken).ConfigureAwait(false); + + // Given the symmetric key instantiate a SqlClientEncryptionAlgorithm object and cache it in metadata + md.CipherAlgorithm = null; + SqlClientEncryptionAlgorithm cipherAlgorithm = null; + string algorithmName = ValidateAndGetEncryptionAlgorithmName(md.CipherAlgorithmId, md.CipherAlgorithmName); // may throw + EncryptionAlgorithmFactoryList.GetAlgorithm(symKey, md.EncryptionType, algorithmName, out cipherAlgorithm); // will validate algorithm name and type + Debug.Assert(cipherAlgorithm is not null); + md.CipherAlgorithm = cipherAlgorithm; + md.EncryptionKeyInfo = encryptionkeyInfoChosen; + } + + /// + /// Asynchronously decrypts the symmetric key. + /// + /// + /// + /// Async methods cannot declare out parameters, so the two values reported by + /// + /// are returned as a tuple instead (spec Design Decision 4). + /// + /// + /// Like the sync overload, each candidate key is tried in turn and failures are remembered so that the + /// last one can be rethrown if every candidate fails. Unlike the sync overload, a cancellation of + /// is never swallowed: it abandons the loop immediately instead of + /// causing the remaining candidates to be attempted. + /// + /// + internal static async Task<(SqlClientSymmetricKey Key, SqlEncryptionKeyInfo KeyInfoChosen)> DecryptSymmetricKeyAsync( + SqlTceCipherInfoEntry sqlTceCipherInfoEntry, + SqlConnection connection, + SqlCommand command, + CancellationToken cancellationToken) + { + Debug.Assert(connection is not null, "Connection should not be null."); + Debug.Assert(sqlTceCipherInfoEntry is not null, "sqlTceCipherInfoEntry should not be null in DecryptSymmetricKeyAsync."); + Debug.Assert(sqlTceCipherInfoEntry.ColumnEncryptionKeyValues is not null, + "sqlTceCipherInfoEntry.ColumnEncryptionKeyValues should not be null in DecryptSymmetricKeyAsync."); + + SqlClientSymmetricKey sqlClientSymmetricKey = null; + SqlEncryptionKeyInfo encryptionkeyInfoChosen = null; + Exception lastException = null; + SqlSymmetricKeyCache globalCekCache = SqlSymmetricKeyCache.GetInstance(); + + foreach (SqlEncryptionKeyInfo keyInfo in sqlTceCipherInfoEntry.ColumnEncryptionKeyValues) + { + try + { + sqlClientSymmetricKey = ShouldUseInstanceLevelProviderFlow(keyInfo.keyStoreName, connection, command) ? + await GetKeyFromLocalProvidersAsync(keyInfo, connection, command, cancellationToken).ConfigureAwait(false) : + await globalCekCache.GetKeyAsync(keyInfo, connection, command, cancellationToken).ConfigureAwait(false); + encryptionkeyInfoChosen = keyInfo; + break; + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + // The caller cancelled. Do not treat this as a failure of this particular key and do not + // attempt the remaining keys; propagate so that the returned Task is cancelled. + throw; + } + catch (Exception e) + { + lastException = e; + } + } + + if (sqlClientSymmetricKey is null) + { + Debug.Assert(lastException is not null, "CEK decryption failed without raising exceptions"); + throw lastException; + } + + Debug.Assert(encryptionkeyInfoChosen is not null, "encryptionkeyInfoChosen must have a value."); + + return (sqlClientSymmetricKey, encryptionkeyInfoChosen); + } + /// /// Decrypts the symmetric key and saves it in metadata. /// @@ -286,6 +378,58 @@ private static SqlClientSymmetricKey GetKeyFromLocalProviders(SqlEncryptionKeyIn return new SqlClientSymmetricKey(plaintextKey); } + /// + /// Asynchronous counterpart of . + /// + /// + /// Instance-level providers are never backed by the global CEK cache, so this method performs the + /// provider call directly. A cancellation of is propagated + /// unwrapped rather than being reported as SQL.KeyDecryptionFailed. + /// + private static async Task GetKeyFromLocalProvidersAsync( + SqlEncryptionKeyInfo keyInfo, + SqlConnection connection, + SqlCommand command, + CancellationToken cancellationToken) + { + string serverName = connection.DataSource; + Debug.Assert(serverName is not null, @"serverName should not be null."); + + Debug.Assert(SqlConnection.ColumnEncryptionTrustedMasterKeyPaths is not null, @"SqlConnection.ColumnEncryptionTrustedMasterKeyPaths should not be null"); + + ThrowIfKeyPathIsNotTrustedForServer(serverName, keyInfo.keyPath); + if (!TryGetColumnEncryptionKeyStoreProvider(keyInfo.keyStoreName, out SqlColumnEncryptionKeyStoreProvider provider, connection, command)) + { + throw SQL.UnrecognizedKeyStoreProviderName(keyInfo.keyStoreName, + SqlConnection.GetColumnEncryptionSystemKeyStoreProvidersNames(), + GetListOfProviderNamesThatWereSearched(connection, command)); + } + + cancellationToken.ThrowIfCancellationRequested(); + + // Decrypt the CEK + // We will simply bubble up the exception from the DecryptColumnEncryptionKeyAsync function. + byte[] plaintextKey; + try + { + plaintextKey = await provider + .DecryptColumnEncryptionKeyAsync(keyInfo.keyPath, keyInfo.algorithmName, keyInfo.encryptedKey, cancellationToken) + .ConfigureAwait(false); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + throw; + } + catch (Exception e) + { + // Generate a new exception and throw. + string keyHex = GetBytesAsString(keyInfo.encryptedKey, fLast: true, countOfBytes: 10); + throw SQL.KeyDecryptionFailed(keyInfo.keyStoreName, keyHex, e); + } + + return new SqlClientSymmetricKey(plaintextKey); + } + /// /// Calculates the length of the Base64 string used to represent a byte[] with the specified length. /// @@ -360,6 +504,92 @@ internal static void VerifyColumnMasterKeySignature(string keyStoreName, string } } + /// + /// Asynchronously verifies Column Master Key Signature. + /// + /// + /// Mirrors , including wrapping failures in + /// SQL.UnableToVerifyColumnMasterKeySignature. The single exception is a cancellation of + /// , which is propagated unwrapped so the returned Task is + /// cancelled rather than faulted with an . + /// + internal static async Task VerifyColumnMasterKeySignatureAsync( + string keyStoreName, + string keyPath, + bool isEnclaveEnabled, + byte[] CMKSignature, + SqlConnection connection, + SqlCommand command, + CancellationToken cancellationToken) + { + bool isValidSignature = false; + + try + { + Debug.Assert(SqlConnection.ColumnEncryptionTrustedMasterKeyPaths is not null, + @"SqlConnection.ColumnEncryptionTrustedMasterKeyPaths should not be null"); + + if (CMKSignature is null || CMKSignature.Length == 0) + { + throw SQL.ColumnMasterKeySignatureNotFound(keyPath); + } + + ThrowIfKeyPathIsNotTrustedForServer(connection.DataSource, keyPath); + + // Attempt to look up the provider and verify CMK Signature + if (!TryGetColumnEncryptionKeyStoreProvider(keyStoreName, out SqlColumnEncryptionKeyStoreProvider provider, connection, command)) + { + throw SQL.InvalidKeyStoreProviderName(keyStoreName, + SqlConnection.GetColumnEncryptionSystemKeyStoreProvidersNames(), + GetListOfProviderNamesThatWereSearched(connection, command)); + } + + if (ShouldUseInstanceLevelProviderFlow(keyStoreName, connection, command)) + { + cancellationToken.ThrowIfCancellationRequested(); + + isValidSignature = await provider + .VerifyColumnMasterKeyMetadataAsync(keyPath, isEnclaveEnabled, CMKSignature, cancellationToken) + .ConfigureAwait(false); + } + else + { + SignatureVerificationResult cachedResult = ColumnMasterKeyMetadataSignatureVerificationCache.Instance + .GetSignatureVerificationResult(keyStoreName, keyPath, isEnclaveEnabled, CMKSignature); + + if (cachedResult == SignatureVerificationResult.NotFound) + { + cancellationToken.ThrowIfCancellationRequested(); + + // Cache miss: verify with the provider and cache the result. + // Exceptions from VerifyColumnMasterKeyMetadataAsync bubble up to the outer catch. + isValidSignature = await provider + .VerifyColumnMasterKeyMetadataAsync(keyPath, isEnclaveEnabled, CMKSignature, cancellationToken) + .ConfigureAwait(false); + ColumnMasterKeyMetadataSignatureVerificationCache.Instance + .AddSignatureVerificationResult(keyStoreName, keyPath, isEnclaveEnabled, CMKSignature, isValidSignature); + } + else + { + isValidSignature = cachedResult == SignatureVerificationResult.True; + } + } + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + throw; + } + catch (Exception e) + { + throw SQL.UnableToVerifyColumnMasterKeySignature(e); + } + + if (!isValidSignature) + { + throw SQL.ColumnMasterKeySignatureVerificationFailed(keyPath); + } + } + // Instance-level providers will be used if at least one is registered on a connection or command and // the required provider is not a system provider. System providers are pre-registered globally and // must use the global provider flow diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSymmetricKeyCache.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSymmetricKeyCache.cs index 6618288b32..cce481b0a1 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSymmetricKeyCache.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSymmetricKeyCache.cs @@ -6,6 +6,7 @@ using System.Diagnostics; using System.Text; using System.Threading; +using System.Threading.Tasks; using Microsoft.Extensions.Caching.Memory; namespace Microsoft.Data.SqlClient @@ -36,22 +37,7 @@ internal SqlClientSymmetricKey GetKey(SqlEncryptionKeyInfo keyInfo, SqlConnectio { string serverName = connection.DataSource; Debug.Assert(serverName is not null, @"serverName should not be null."); - StringBuilder cacheLookupKeyBuilder = new(serverName, capacity: serverName.Length + SqlSecurityUtility.GetBase64LengthFromByteLength(keyInfo.encryptedKey.Length) + keyInfo.keyStoreName.Length + 2/*separators*/); - -#if DEBUG - int capacity = cacheLookupKeyBuilder.Capacity; -#endif //DEBUG - - cacheLookupKeyBuilder.Append(":"); - cacheLookupKeyBuilder.Append(Convert.ToBase64String(keyInfo.encryptedKey)); - cacheLookupKeyBuilder.Append(":"); - cacheLookupKeyBuilder.Append(keyInfo.keyStoreName); - - string cacheLookupKey = cacheLookupKeyBuilder.ToString(); - -#if DEBUG - Debug.Assert(cacheLookupKey.Length <= capacity, "We needed to allocate a larger array"); -#endif //DEBUG + string cacheLookupKey = CreateCacheLookupKey(keyInfo, serverName); // Acquire the lock to ensure thread safety when accessing the cache _cacheLock.Wait(); @@ -110,5 +96,141 @@ internal SqlClientSymmetricKey GetKey(SqlEncryptionKeyInfo keyInfo, SqlConnectio _cacheLock.Release(); } } + + /// + /// Asynchronously retrieves Symmetric Key (in plaintext) given the encryption material. + /// + /// + /// + /// _cacheLock is a process-wide gate that the synchronous path blocks on + /// with . Holding it across the awaited provider call would therefore + /// stall every synchronous caller's thread for the duration of that (potentially remote) call. To avoid + /// that, this method follows the check-release-fetch-relock pattern required by FR-012: the gate is held + /// only for the cache lookup and for the cache insertion, never across I/O. + /// + /// + /// The consequence, accepted by FR-014, is that concurrent async callers that miss the cache for the + /// same key may each decrypt it. That is weaker than the synchronous path, which serializes the fetch. + /// The results are equivalent, and the first caller to publish its result wins: subsequent callers + /// return the already-cached instance so that all callers observe the same key object. + /// + /// + internal async Task GetKeyAsync(SqlEncryptionKeyInfo keyInfo, SqlConnection connection, SqlCommand command, CancellationToken cancellationToken) + { + string serverName = connection.DataSource; + Debug.Assert(serverName is not null, @"serverName should not be null."); + string cacheLookupKey = CreateCacheLookupKey(keyInfo, serverName); + + // Acquire the lock to ensure thread safety when accessing the cache, and release it before + // performing any I/O so that synchronous callers are never blocked behind a remote call. + await _cacheLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + if (_cache.TryGetValue(cacheLookupKey, out SqlClientSymmetricKey cachedKey)) + { + return cachedKey; + } + } + finally + { + _cacheLock.Release(); + } + + 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)); + } + + cancellationToken.ThrowIfCancellationRequested(); + + // Decrypt the CEK + // We will simply bubble up the exception from the DecryptColumnEncryptionKeyAsync 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 + // + // Unlike the sync path this assignment happens outside the gate, so two async callers may write + // it concurrently. The write is idempotent (always TimeSpan.Zero), so the outcome is the same. + provider.ColumnEncryptionKeyCacheTtl = new TimeSpan(0); + plaintextKey = await provider + .DecryptColumnEncryptionKeyAsync(keyInfo.keyPath, keyInfo.algorithmName, keyInfo.encryptedKey, cancellationToken) + .ConfigureAwait(false); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + // Cancellation is propagated unwrapped so the returned Task is cancelled rather than faulted. + throw; + } + 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); + } + + SqlClientSymmetricKey decryptedKey = new SqlClientSymmetricKey(plaintextKey); + + // If the cache TTL is zero, don't even bother inserting to the cache. + TimeSpan cacheTtl = SqlConnection.ColumnEncryptionKeyCacheTtl; + if (cacheTtl == TimeSpan.Zero) + { + return decryptedKey; + } + + await _cacheLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + // Another caller may have populated the entry while this one was decrypting. In that case the + // first one wins, exactly as in the sync path, so that all callers observe the same key object. + if (_cache.TryGetValue(cacheLookupKey, out SqlClientSymmetricKey concurrentlyCachedKey)) + { + return concurrentlyCachedKey; + } + + _cache.Set(cacheLookupKey, decryptedKey, absoluteExpirationRelativeToNow: cacheTtl); + } + finally + { + _cacheLock.Release(); + } + + return decryptedKey; + } + + /// + /// Builds the cache lookup key for the given encryption material. Pure function shared by the sync + /// and async lookup paths. + /// + private static string CreateCacheLookupKey(SqlEncryptionKeyInfo keyInfo, string serverName) + { + StringBuilder cacheLookupKeyBuilder = new(serverName, capacity: serverName.Length + SqlSecurityUtility.GetBase64LengthFromByteLength(keyInfo.encryptedKey.Length) + keyInfo.keyStoreName.Length + 2/*separators*/); + +#if DEBUG + int capacity = cacheLookupKeyBuilder.Capacity; +#endif //DEBUG + + cacheLookupKeyBuilder.Append(":"); + cacheLookupKeyBuilder.Append(Convert.ToBase64String(keyInfo.encryptedKey)); + cacheLookupKeyBuilder.Append(":"); + cacheLookupKeyBuilder.Append(keyInfo.keyStoreName); + + string cacheLookupKey = cacheLookupKeyBuilder.ToString(); + +#if DEBUG + Debug.Assert(cacheLookupKey.Length <= capacity, "We needed to allocate a larger array"); +#endif //DEBUG + + return cacheLookupKey; + } } } diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/EnclaveDelegateAsyncShould.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/EnclaveDelegateAsyncShould.cs new file mode 100644 index 0000000000..6302c5164f --- /dev/null +++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/EnclaveDelegateAsyncShould.cs @@ -0,0 +1,239 @@ +// 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. + +// The Always Encrypted enclave layer is internal and nullable-oblivious, and several of its members +// legitimately accept or produce nulls. Nullable analysis is disabled for this file so the tests can +// mirror those signatures exactly. +#nullable disable + +using System; +using System.Collections.Concurrent; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using Xunit; + +namespace Microsoft.Data.SqlClient.UnitTests.AlwaysEncrypted +{ + /// + /// Tests for the asynchronous column encryption key resolution used when building an enclave package. + /// + /// + /// The enclave package itself cannot be produced without an attested enclave session, so these tests + /// target the one step of that flow which performs key store I/O. + /// + public class EnclaveDelegateAsyncShould + { + private const string ProviderName = "TEST_ENCLAVE_ASYNC_PROVIDER"; + private const string KeyPath = "test-enclave-key-path"; + private const string EncryptionAlgorithm = "RSA_OAEP"; + + private const int DatabaseId = 42; + private const int CekId = 7; + private const ulong CekMdVersion = 3; + + // The AEAD_AES_256_CBC_HMAC_SHA256 algorithm requires a 256 bit root key. + private static byte[] NewPlaintextKey(byte seed) + { + byte[] key = new byte[32]; + for (int i = 0; i < key.Length; i++) + { + key[i] = (byte)(seed + i); + } + return key; + } + + private static SqlConnection NewConnection(SqlColumnEncryptionKeyStoreProvider provider) + { + SqlConnection connection = new SqlConnection("Data Source=async-ae-enclave-unit-test;Column Encryption Setting=Enabled"); + connection.RegisterColumnEncryptionKeyStoreProvidersOnConnection( + new Dictionary { [ProviderName] = provider }); + return connection; + } + + /// + /// Produces a cipher info entry whose cache lookup key cannot collide with any other test. The CEK + /// cache is a process-wide singleton shared by every test in this assembly. + /// + private static SqlTceCipherInfoEntry NewCipherInfoEntry(int ordinal = 0) + { + SqlTceCipherInfoEntry entry = new SqlTceCipherInfoEntry(ordinal); + entry.Add( + encryptedKey: Guid.NewGuid().ToByteArray(), + databaseId: DatabaseId, + cekId: CekId, + cekVersion: 1, + cekMdVersion: CekMdVersion, + keyPath: KeyPath, + keyStoreName: ProviderName, + algorithmName: EncryptionAlgorithm); + return entry; + } + + private static ConcurrentDictionary NewKeyTable(params SqlTceCipherInfoEntry[] entries) + { + ConcurrentDictionary table = new(); + for (int i = 0; i < entries.Length; i++) + { + table[i] = entries[i]; + } + return table; + } + + /// + /// The whole point of the change: the enclave key resolution must go through the provider's async + /// API so a key store performing network I/O never blocks the pooled thread. + /// + [Fact] + public async Task GetDecryptedKeysToBeSentToEnclaveAsync_UsesAsyncProviderApi() + { + byte[] expectedKey = NewPlaintextKey(seed: 11); + TestEnclaveKeyStoreProvider provider = new TestEnclaveKeyStoreProvider { PlaintextKey = expectedKey }; + using SqlConnection connection = NewConnection(provider); + + List keys = await EnclaveDelegate.Instance.GetDecryptedKeysToBeSentToEnclaveAsync( + NewKeyTable(NewCipherInfoEntry()), + serverName: "async-ae-enclave-unit-test", + connection, + command: null, + CancellationToken.None); + + ColumnEncryptionKeyInfo key = Assert.Single(keys); + Assert.Equal(expectedKey, key.DecryptedKeyBytes); + Assert.Equal(DatabaseId, key.DatabaseId); + Assert.Equal(CekId, key.KeyId); + Assert.Equal(CekMdVersion, key.KeyMetadataVersion); + + Assert.Equal(1, provider.DecryptAsyncCallCount); + Assert.Equal(0, provider.DecryptCallCount); + } + + /// + /// Every requested key must be resolved, and each one must use the async provider API. + /// + [Fact] + public async Task GetDecryptedKeysToBeSentToEnclaveAsync_ResolvesEveryRequestedKey() + { + TestEnclaveKeyStoreProvider provider = new TestEnclaveKeyStoreProvider { PlaintextKey = NewPlaintextKey(seed: 23) }; + using SqlConnection connection = NewConnection(provider); + + List keys = await EnclaveDelegate.Instance.GetDecryptedKeysToBeSentToEnclaveAsync( + NewKeyTable(NewCipherInfoEntry(ordinal: 0), NewCipherInfoEntry(ordinal: 1), NewCipherInfoEntry(ordinal: 2)), + serverName: "async-ae-enclave-unit-test", + connection, + command: null, + CancellationToken.None); + + Assert.Equal(3, keys.Count); + Assert.Equal(3, provider.DecryptAsyncCallCount); + Assert.Equal(0, provider.DecryptCallCount); + } + + /// + /// The synchronous path must keep using the synchronous provider API so that custom providers which + /// only override the synchronous members continue to work unchanged. + /// + [Fact] + public void GetDecryptedKeysToBeSentToEnclave_UsesSyncProviderApi() + { + TestEnclaveKeyStoreProvider provider = new TestEnclaveKeyStoreProvider { PlaintextKey = NewPlaintextKey(seed: 37) }; + using SqlConnection connection = NewConnection(provider); + + List keys = EnclaveDelegate.Instance.GetDecryptedKeysToBeSentToEnclave( + NewKeyTable(NewCipherInfoEntry()), + serverName: "async-ae-enclave-unit-test", + connection, + command: null); + + Assert.Single(keys); + Assert.Equal(1, provider.DecryptCallCount); + Assert.Equal(0, provider.DecryptAsyncCallCount); + } + + /// + /// Cancellation must surface as a cancelled task rather than being wrapped in an + /// Always Encrypted specific exception, and no provider call must be made. + /// + [Fact] + public async Task GetDecryptedKeysToBeSentToEnclaveAsync_WhenCancelled_ProducesCancelledTask() + { + TestEnclaveKeyStoreProvider provider = new TestEnclaveKeyStoreProvider { PlaintextKey = NewPlaintextKey(seed: 53) }; + using SqlConnection connection = NewConnection(provider); + + using CancellationTokenSource cts = new(); + cts.Cancel(); + + await Assert.ThrowsAnyAsync( + () => EnclaveDelegate.Instance.GetDecryptedKeysToBeSentToEnclaveAsync( + NewKeyTable(NewCipherInfoEntry()), + serverName: "async-ae-enclave-unit-test", + connection, + command: null, + cts.Token)); + + Assert.Equal(0, provider.DecryptAsyncCallCount); + } + + /// + /// A provider failure must propagate to the caller — wrapped exactly as the synchronous path wraps + /// it — rather than producing a partially populated enclave key list. + /// + [Fact] + public async Task GetDecryptedKeysToBeSentToEnclaveAsync_WhenProviderFails_Propagates() + { + TestEnclaveKeyStoreProvider provider = new TestEnclaveKeyStoreProvider + { + DecryptAsyncCallback = _ => throw new InvalidOperationException("key store unavailable"), + }; + using SqlConnection connection = NewConnection(provider); + + SqlException exception = await Assert.ThrowsAsync( + () => EnclaveDelegate.Instance.GetDecryptedKeysToBeSentToEnclaveAsync( + NewKeyTable(NewCipherInfoEntry()), + serverName: "async-ae-enclave-unit-test", + connection, + command: null, + CancellationToken.None)); + + Assert.IsType(exception.InnerException); + Assert.Equal("key store unavailable", exception.InnerException.Message); + } + + private sealed class TestEnclaveKeyStoreProvider : SqlColumnEncryptionKeyStoreProvider + { + private int _decryptCallCount; + private int _decryptAsyncCallCount; + + internal byte[] PlaintextKey { get; set; } = new byte[32]; + + internal Func> DecryptAsyncCallback { get; set; } + + internal int DecryptCallCount => _decryptCallCount; + + internal int DecryptAsyncCallCount => _decryptAsyncCallCount; + + public override byte[] DecryptColumnEncryptionKey(string masterKeyPath, string encryptionAlgorithm, byte[] encryptedColumnEncryptionKey) + { + Interlocked.Increment(ref _decryptCallCount); + return PlaintextKey; + } + + public override Task DecryptColumnEncryptionKeyAsync(string masterKeyPath, string encryptionAlgorithm, byte[] encryptedColumnEncryptionKey, CancellationToken cancellationToken = default) + { + Interlocked.Increment(ref _decryptAsyncCallCount); + Func> callback = DecryptAsyncCallback; + return callback is not null ? callback(cancellationToken) : Task.FromResult(PlaintextKey); + } + + public override byte[] EncryptColumnEncryptionKey(string masterKeyPath, string encryptionAlgorithm, byte[] columnEncryptionKey) + => throw new NotSupportedException(); + + public override byte[] SignColumnMasterKeyMetadata(string masterKeyPath, bool allowEnclaveComputations) + => throw new NotSupportedException(); + + public override bool VerifyColumnMasterKeyMetadata(string masterKeyPath, bool allowEnclaveComputations, byte[] signature) + => throw new NotSupportedException(); + } + } +} diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlSecurityUtilityAsyncShould.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlSecurityUtilityAsyncShould.cs new file mode 100644 index 0000000000..83710c614e --- /dev/null +++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlSecurityUtilityAsyncShould.cs @@ -0,0 +1,486 @@ +// 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. + +// The Always Encrypted utility layer is internal and nullable-oblivious, and several of its members +// legitimately accept or produce nulls. Nullable analysis is disabled for this file so the tests can +// mirror those signatures exactly. +#nullable disable + +using System; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using Xunit; + +namespace Microsoft.Data.SqlClient.UnitTests.AlwaysEncrypted +{ + /// + /// Tests for the async counterparts of the Always Encrypted utility layer: + /// .DecryptSymmetricKeyAsync / VerifyColumnMasterKeySignatureAsync + /// and .GetKeyAsync. + /// + public class SqlSecurityUtilityAsyncShould + { + private const string ProviderName = "TEST_ASYNC_PROVIDER"; + private const string SecondProviderName = "TEST_ASYNC_PROVIDER_2"; + private const string KeyPath = "test-key-path"; + private const string SecondKeyPath = "test-key-path-2"; + private const string EncryptionAlgorithm = "RSA_OAEP"; + + // The AEAD_AES_256_CBC_HMAC_SHA256 algorithm requires a 256 bit root key. + private static byte[] NewPlaintextKey(byte seed) + { + byte[] key = new byte[32]; + for (int i = 0; i < key.Length; i++) + { + key[i] = (byte)(seed + i); + } + return key; + } + + /// + /// Produces key material whose cache lookup key cannot collide with any other test. The CEK + /// cache is a process-wide singleton shared by every test in this assembly. + /// + private static SqlEncryptionKeyInfo NewKeyInfo(string providerName = ProviderName, string keyPath = KeyPath) + => new SqlEncryptionKeyInfo + { + encryptedKey = Guid.NewGuid().ToByteArray(), + databaseId = 1, + cekId = 1, + cekVersion = 1, + cekMdVersion = 1, + keyPath = keyPath, + keyStoreName = providerName, + algorithmName = EncryptionAlgorithm, + }; + + private static SqlConnection NewConnection(params (string Name, SqlColumnEncryptionKeyStoreProvider Provider)[] providers) + { + SqlConnection connection = new SqlConnection("Data Source=async-ae-unit-test;Column Encryption Setting=Enabled"); + Dictionary map = new(); + foreach ((string name, SqlColumnEncryptionKeyStoreProvider provider) in providers) + { + map[name] = provider; + } + connection.RegisterColumnEncryptionKeyStoreProvidersOnConnection(map); + return connection; + } + + private static SqlTceCipherInfoEntry NewCipherInfoEntry(params SqlEncryptionKeyInfo[] keyInfos) + { + SqlTceCipherInfoEntry entry = new SqlTceCipherInfoEntry(ordinal: 0); + foreach (SqlEncryptionKeyInfo keyInfo in keyInfos) + { + entry.Add(keyInfo.encryptedKey, keyInfo.databaseId, keyInfo.cekId, keyInfo.cekVersion, + keyInfo.cekMdVersion, keyInfo.keyPath, keyInfo.keyStoreName, keyInfo.algorithmName); + } + return entry; + } + + #region DecryptSymmetricKeyAsync - multi key fallback + + /// + /// Verifies that the per-key fallback loop of the sync overload is preserved: when the first + /// candidate key fails, the next one is attempted and its key info is the one reported. + /// + [Fact] + public async Task DecryptSymmetricKeyAsync_WhenFirstKeyFails_FallsBackToSecondKey() + { + byte[] expectedKey = NewPlaintextKey(seed: 7); + TestKeyStoreProvider failing = new TestKeyStoreProvider { DecryptAsyncCallback = _ => throw new InvalidOperationException("first key unavailable") }; + TestKeyStoreProvider succeeding = new TestKeyStoreProvider { PlaintextKey = expectedKey }; + + using SqlConnection connection = NewConnection((ProviderName, failing), (SecondProviderName, succeeding)); + + SqlTceCipherInfoEntry entry = NewCipherInfoEntry( + NewKeyInfo(ProviderName, KeyPath), + NewKeyInfo(SecondProviderName, SecondKeyPath)); + + (SqlClientSymmetricKey key, SqlEncryptionKeyInfo keyInfoChosen) = + await SqlSecurityUtility.DecryptSymmetricKeyAsync(entry, connection, command: null, CancellationToken.None); + + Assert.NotNull(key); + Assert.Equal(expectedKey, key.RootKey); + Assert.Equal(SecondProviderName, keyInfoChosen.keyStoreName); + Assert.Equal(1, failing.DecryptAsyncCallCount); + Assert.Equal(1, succeeding.DecryptAsyncCallCount); + } + + /// + /// Verifies that when every candidate key fails, the last exception is rethrown, matching the + /// sync overload. + /// + [Fact] + public async Task DecryptSymmetricKeyAsync_WhenAllKeysFail_ThrowsLastException() + { + TestKeyStoreProvider failing = new TestKeyStoreProvider { DecryptAsyncCallback = _ => throw new InvalidOperationException("boom") }; + using SqlConnection connection = NewConnection((ProviderName, failing)); + + SqlTceCipherInfoEntry entry = NewCipherInfoEntry(NewKeyInfo(), NewKeyInfo()); + + // Failures from the provider are reported as SQL.KeyDecryptionFailed. + SqlException exception = await Assert.ThrowsAsync( + () => SqlSecurityUtility.DecryptSymmetricKeyAsync(entry, connection, command: null, CancellationToken.None)); + + Assert.Contains("boom", exception.ToString()); + Assert.Equal(2, failing.DecryptAsyncCallCount); + } + + /// + /// Verifies that the SqlCipherMetadata overload mutates the metadata exactly as the sync overload does. + /// + [Fact] + public async Task DecryptSymmetricKeyAsync_WithCipherMetadata_PopulatesAlgorithmAndKeyInfo() + { + TestKeyStoreProvider provider = new TestKeyStoreProvider { PlaintextKey = NewPlaintextKey(seed: 3) }; + using SqlConnection connection = NewConnection((ProviderName, provider)); + + SqlEncryptionKeyInfo keyInfo = NewKeyInfo(); + SqlCipherMetadata metadata = new SqlCipherMetadata( + NewCipherInfoEntry(keyInfo), + ordinal: 0, + cipherAlgorithmId: (byte)TdsEnums.AEAD_AES_256_CBC_HMAC_SHA256, + cipherAlgorithmName: null, + encryptionType: (byte)SqlClientEncryptionType.Deterministic, + normalizationRuleVersion: 0x01); + + await SqlSecurityUtility.DecryptSymmetricKeyAsync(metadata, connection, command: null, CancellationToken.None); + + Assert.True(metadata.IsAlgorithmInitialized()); + Assert.NotNull(metadata.EncryptionKeyInfo); + Assert.Equal(keyInfo.keyStoreName, metadata.EncryptionKeyInfo.keyStoreName); + } + + #endregion + + #region DecryptSymmetricKeyAsync - cancellation + + /// + /// Verifies that a token cancelled before the call is observed without invoking the provider. + /// + [Fact] + public async Task DecryptSymmetricKeyAsync_WithCancelledToken_CancelsWithoutInvokingProvider() + { + TestKeyStoreProvider provider = new TestKeyStoreProvider { PlaintextKey = NewPlaintextKey(seed: 1) }; + using SqlConnection connection = NewConnection((ProviderName, provider)); + SqlTceCipherInfoEntry entry = NewCipherInfoEntry(NewKeyInfo()); + + using CancellationTokenSource cts = new CancellationTokenSource(); + cts.Cancel(); + + await Assert.ThrowsAnyAsync( + () => SqlSecurityUtility.DecryptSymmetricKeyAsync(entry, connection, command: null, cts.Token)); + + Assert.Equal(0, provider.DecryptAsyncCallCount); + } + + /// + /// Verifies that a cancellation observed mid-flight is neither swallowed by the per-key fallback + /// loop nor wrapped in SQL.KeyDecryptionFailed, and that the remaining candidate keys are not tried. + /// + [Fact] + public async Task DecryptSymmetricKeyAsync_WhenCancelledMidFlight_DoesNotAttemptRemainingKeys() + { + using CancellationTokenSource cts = new CancellationTokenSource(); + TaskCompletionSource providerEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + TestKeyStoreProvider provider = new TestKeyStoreProvider + { + DecryptAsyncCallback = async token => + { + providerEntered.TrySetResult(true); + await Task.Delay(Timeout.Infinite, token); + return NewPlaintextKey(seed: 2); + } + }; + + using SqlConnection connection = NewConnection((ProviderName, provider)); + SqlTceCipherInfoEntry entry = NewCipherInfoEntry(NewKeyInfo(), NewKeyInfo()); + + Task decryptTask = SqlSecurityUtility.DecryptSymmetricKeyAsync(entry, connection, command: null, cts.Token); + await providerEntered.Task; + cts.Cancel(); + + await Assert.ThrowsAnyAsync(() => decryptTask); + Assert.True(decryptTask.IsCanceled); + + // The second key was never attempted: cancellation abandons the loop. + Assert.Equal(1, provider.DecryptAsyncCallCount); + } + + #endregion + + #region VerifyColumnMasterKeySignatureAsync + + /// + /// Verifies that a valid signature completes and an invalid one is reported the same way the + /// sync path reports it. + /// + [Fact] + public async Task VerifyColumnMasterKeySignatureAsync_ReportsProviderResult() + { + TestKeyStoreProvider provider = new TestKeyStoreProvider { VerifyResult = true }; + using SqlConnection connection = NewConnection((ProviderName, provider)); + + await SqlSecurityUtility.VerifyColumnMasterKeySignatureAsync( + ProviderName, KeyPath, isEnclaveEnabled: false, CMKSignature: new byte[] { 1, 2, 3 }, + connection, command: null, CancellationToken.None); + + Assert.Equal(1, provider.VerifyAsyncCallCount); + + provider.VerifyResult = false; + await Assert.ThrowsAsync(() => SqlSecurityUtility.VerifyColumnMasterKeySignatureAsync( + ProviderName, KeyPath, isEnclaveEnabled: false, CMKSignature: new byte[] { 1, 2, 3 }, + connection, command: null, CancellationToken.None)); + } + + /// + /// Verifies that provider failures are still wrapped in SQL.UnableToVerifyColumnMasterKeySignature. + /// + [Fact] + public async Task VerifyColumnMasterKeySignatureAsync_WhenProviderThrows_WrapsException() + { + TestKeyStoreProvider provider = new TestKeyStoreProvider + { + VerifyAsyncCallback = _ => throw new InvalidOperationException("verification exploded") + }; + using SqlConnection connection = NewConnection((ProviderName, provider)); + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => SqlSecurityUtility.VerifyColumnMasterKeySignatureAsync( + ProviderName, KeyPath, isEnclaveEnabled: false, CMKSignature: new byte[] { 1, 2, 3 }, + connection, command: null, CancellationToken.None)); + + Assert.Contains("verification exploded", exception.ToString()); + } + + /// + /// Verifies that cancellation surfaces as a cancelled Task rather than being wrapped in + /// SQL.UnableToVerifyColumnMasterKeySignature, both before and during the provider call. + /// + [Fact] + public async Task VerifyColumnMasterKeySignatureAsync_WhenCancelled_ProducesCancelledTask() + { + TestKeyStoreProvider provider = new TestKeyStoreProvider { VerifyResult = true }; + using SqlConnection connection = NewConnection((ProviderName, provider)); + + using CancellationTokenSource preCancelled = new CancellationTokenSource(); + preCancelled.Cancel(); + + Task preCancelledTask = SqlSecurityUtility.VerifyColumnMasterKeySignatureAsync( + ProviderName, KeyPath, isEnclaveEnabled: false, CMKSignature: new byte[] { 1, 2, 3 }, + connection, command: null, preCancelled.Token); + + await Assert.ThrowsAnyAsync(() => preCancelledTask); + Assert.True(preCancelledTask.IsCanceled); + Assert.Equal(0, provider.VerifyAsyncCallCount); + + using CancellationTokenSource cts = new CancellationTokenSource(); + TaskCompletionSource providerEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + provider.VerifyAsyncCallback = async token => + { + providerEntered.TrySetResult(true); + await Task.Delay(Timeout.Infinite, token); + return true; + }; + + Task verifyTask = SqlSecurityUtility.VerifyColumnMasterKeySignatureAsync( + ProviderName, "unique-" + Guid.NewGuid(), isEnclaveEnabled: false, CMKSignature: new byte[] { 4, 5, 6 }, + connection, command: null, cts.Token); + + await providerEntered.Task; + cts.Cancel(); + + await Assert.ThrowsAnyAsync(() => verifyTask); + Assert.True(verifyTask.IsCanceled); + } + + #endregion + + #region SqlSymmetricKeyCache.GetKeyAsync + + /// + /// Verifies that a key cached by the sync path is reused by the async path without a second + /// provider call, and vice versa. + /// + [Fact] + public async Task GetKeyAsync_ReusesEntriesCachedBySyncPath_AndViceVersa() + { + SqlSymmetricKeyCache cache = SqlSymmetricKeyCache.GetInstance(); + + TestKeyStoreProvider provider = new TestKeyStoreProvider { PlaintextKey = NewPlaintextKey(seed: 11) }; + using SqlConnection connection = NewConnection((ProviderName, provider)); + + // Sync first, then async. + SqlEncryptionKeyInfo syncFirst = NewKeyInfo(); + SqlClientSymmetricKey fromSync = cache.GetKey(syncFirst, connection, command: null); + SqlClientSymmetricKey fromAsync = await cache.GetKeyAsync(syncFirst, connection, command: null, CancellationToken.None); + + Assert.Same(fromSync, fromAsync); + Assert.Equal(1, provider.DecryptCallCount); + Assert.Equal(0, provider.DecryptAsyncCallCount); + + // Async first, then sync. + SqlEncryptionKeyInfo asyncFirst = NewKeyInfo(); + SqlClientSymmetricKey asyncKey = await cache.GetKeyAsync(asyncFirst, connection, command: null, CancellationToken.None); + SqlClientSymmetricKey syncKey = cache.GetKey(asyncFirst, connection, command: null); + + Assert.Same(asyncKey, syncKey); + Assert.Equal(1, provider.DecryptCallCount); + Assert.Equal(1, provider.DecryptAsyncCallCount); + } + + /// + /// Verifies the documented consequence of check-release-fetch-relock (FR-012/FR-014): concurrent + /// async cache misses for the same key may each decrypt it, but all callers end up observing the + /// same cached key instance, and neither caller deadlocks behind the process-wide cache gate. + /// + [Fact] + public async Task GetKeyAsync_WithConcurrentCacheMisses_FetchesTwiceButPublishesOneKey() + { + SqlSymmetricKeyCache cache = SqlSymmetricKeyCache.GetInstance(); + + TaskCompletionSource release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + int entered = 0; + TaskCompletionSource bothEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + TestKeyStoreProvider provider = new TestKeyStoreProvider + { + DecryptAsyncCallback = async _ => + { + if (Interlocked.Increment(ref entered) == 2) + { + bothEntered.TrySetResult(true); + } + await release.Task; + return NewPlaintextKey(seed: 21); + } + }; + + using SqlConnection connection = NewConnection((ProviderName, provider)); + SqlEncryptionKeyInfo keyInfo = NewKeyInfo(); + + Task first = cache.GetKeyAsync(keyInfo, connection, command: null, CancellationToken.None); + Task second = cache.GetKeyAsync(keyInfo, connection, command: null, CancellationToken.None); + + // Both callers reach the provider concurrently, which proves the cache gate is not held + // across the awaited decryption. + await bothEntered.Task; + release.TrySetResult(true); + + SqlClientSymmetricKey firstKey = await first; + SqlClientSymmetricKey secondKey = await second; + + Assert.Equal(2, provider.DecryptAsyncCallCount); + Assert.Same(firstKey, secondKey); + + // Subsequent callers, sync or async, see the published instance without another fetch. + Assert.Same(firstKey, cache.GetKey(keyInfo, connection, command: null)); + Assert.Same(firstKey, await cache.GetKeyAsync(keyInfo, connection, command: null, CancellationToken.None)); + Assert.Equal(2, provider.DecryptAsyncCallCount); + Assert.Equal(0, provider.DecryptCallCount); + } + + /// + /// Verifies that cancellation of GetKeyAsync is observed before the provider call and while it + /// is in flight, and that nothing is published to the cache in either case. + /// + [Fact] + public async Task GetKeyAsync_WhenCancelled_ProducesCancelledTaskAndDoesNotCache() + { + SqlSymmetricKeyCache cache = SqlSymmetricKeyCache.GetInstance(); + + TestKeyStoreProvider provider = new TestKeyStoreProvider { PlaintextKey = NewPlaintextKey(seed: 31) }; + using SqlConnection connection = NewConnection((ProviderName, provider)); + SqlEncryptionKeyInfo keyInfo = NewKeyInfo(); + + using CancellationTokenSource preCancelled = new CancellationTokenSource(); + preCancelled.Cancel(); + + Task preCancelledTask = cache.GetKeyAsync(keyInfo, connection, command: null, preCancelled.Token); + await Assert.ThrowsAnyAsync(() => preCancelledTask); + Assert.True(preCancelledTask.IsCanceled); + Assert.Equal(0, provider.DecryptAsyncCallCount); + + using CancellationTokenSource cts = new CancellationTokenSource(); + TaskCompletionSource providerEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + provider.DecryptAsyncCallback = async token => + { + providerEntered.TrySetResult(true); + await Task.Delay(Timeout.Infinite, token); + return NewPlaintextKey(seed: 32); + }; + + Task inFlight = cache.GetKeyAsync(keyInfo, connection, command: null, cts.Token); + await providerEntered.Task; + cts.Cancel(); + + // Cancellation is not wrapped in SQL.KeyDecryptionFailed. + await Assert.ThrowsAnyAsync(() => inFlight); + Assert.True(inFlight.IsCanceled); + + // Nothing was cached, so a subsequent successful call still reaches the provider. + provider.DecryptAsyncCallback = null; + SqlClientSymmetricKey key = await cache.GetKeyAsync(keyInfo, connection, command: null, CancellationToken.None); + Assert.Equal(provider.PlaintextKey, key.RootKey); + } + + #endregion + + /// + /// A key store provider whose sync and async paths are separately observable, so that tests can + /// assert which path was taken and how many times. + /// + private sealed class TestKeyStoreProvider : SqlColumnEncryptionKeyStoreProvider + { + private int _decryptCallCount; + private int _decryptAsyncCallCount; + private int _verifyAsyncCallCount; + + internal byte[] PlaintextKey { get; set; } = new byte[32]; + + internal bool VerifyResult { get; set; } + + internal Func> DecryptAsyncCallback { get; set; } + + internal Func> VerifyAsyncCallback { get; set; } + + internal int DecryptCallCount => _decryptCallCount; + + internal int DecryptAsyncCallCount => _decryptAsyncCallCount; + + internal int VerifyAsyncCallCount => _verifyAsyncCallCount; + + public override byte[] DecryptColumnEncryptionKey(string masterKeyPath, string encryptionAlgorithm, byte[] encryptedColumnEncryptionKey) + { + Interlocked.Increment(ref _decryptCallCount); + return PlaintextKey; + } + + public override Task DecryptColumnEncryptionKeyAsync(string masterKeyPath, string encryptionAlgorithm, byte[] encryptedColumnEncryptionKey, CancellationToken cancellationToken = default) + { + Interlocked.Increment(ref _decryptAsyncCallCount); + Func> callback = DecryptAsyncCallback; + return callback is not null ? callback(cancellationToken) : Task.FromResult(PlaintextKey); + } + + public override byte[] EncryptColumnEncryptionKey(string masterKeyPath, string encryptionAlgorithm, byte[] columnEncryptionKey) + => throw new NotSupportedException(); + + public override byte[] SignColumnMasterKeyMetadata(string masterKeyPath, bool allowEnclaveComputations) + => throw new NotSupportedException(); + + public override bool VerifyColumnMasterKeyMetadata(string masterKeyPath, bool allowEnclaveComputations, byte[] signature) + => VerifyResult; + + public override Task VerifyColumnMasterKeyMetadataAsync(string masterKeyPath, bool allowEnclaveComputations, byte[] signature, CancellationToken cancellationToken = default) + { + Interlocked.Increment(ref _verifyAsyncCallCount); + Func> callback = VerifyAsyncCallback; + return callback is not null ? callback(cancellationToken) : Task.FromResult(VerifyResult); + } + } + } +} From 9f01623ac2d3ae5d5d227b9dbe64c02e13073782 Mon Sep 17 00:00:00 2001 From: Cheena Malhotra Date: Fri, 21 Aug 2026 17:23:58 -0700 Subject: [PATCH 2/7] Thread the caller's CancellationToken into async Always Encrypted key store calls The async key store provider APIs all accept a CancellationToken, but the Always Encrypted call sites had no way to reach the token supplied to ExecuteReaderAsync / ExecuteNonQueryAsync / ExecuteXmlReaderAsync, so key store I/O (for example an Azure Key Vault round trip) could not be cancelled. Capture the token in a per-execution field on SqlCommand, set at each async entry point and cleared by the matching cleanup continuation, and pass it to the describe-parameter-encryption key operations and to enclave package generation. A field is used rather than signature plumbing because the token would otherwise have to be threaded through several sync-only overloads of RunExecuteReader and PrepareForTransparentEncryption. Also add a manual test that registers a recording key store provider and asserts that the synchronous execution paths use the synchronous provider API, the asynchronous paths use the asynchronous provider API, and that cancelling the caller's token cancels the key store operation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 82aca8ca-2078-4a47-860c-6f6ee78cbb02 --- .../Data/SqlClient/SqlCommand.Encryption.cs | 2 +- .../Data/SqlClient/SqlCommand.NonQuery.cs | 4 + .../Data/SqlClient/SqlCommand.Reader.cs | 6 +- .../Data/SqlClient/SqlCommand.Xml.cs | 4 + .../Microsoft/Data/SqlClient/SqlCommand.cs | 14 ++ .../AsyncKeyStoreProviderTests.cs | 156 ++++++++++++++++++ 6 files changed, 184 insertions(+), 2 deletions(-) create mode 100644 src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/AsyncKeyStoreProviderTests.cs diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs index fc6a7b67b0..e3f7b60315 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs @@ -345,7 +345,7 @@ await ReadDescribeEncryptionParameterResultsAsync( describeParameterEncryptionDataReader, describeParameterEncryptionRpcOriginalRpcMap, isRetry, - CancellationToken.None) + _asyncExecutionCancellationToken) .ConfigureAwait(false); #if DEBUG diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.NonQuery.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.NonQuery.cs index ef3e8afe25..1b653a290d 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.NonQuery.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.NonQuery.cs @@ -331,6 +331,8 @@ private void BeginExecuteNonQueryInternalReadStage(TaskCompletionSource private void CleanupAfterExecuteNonQueryAsync(Task task, TaskCompletionSource source, Guid operationId) { + _asyncExecutionCancellationToken = default; + if (task.IsFaulted) { Exception e = task.Exception?.InnerException; @@ -697,6 +699,8 @@ private Task InternalExecuteNonQueryAsync(CancellationToken cancellationTok registration = cancellationToken.Register(callback: s_cancelIgnoreFailure, state: this); } + _asyncExecutionCancellationToken = cancellationToken; + Task returnedTask = source.Task; returnedTask = RegisterForConnectionCloseNotification(returnedTask); diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs index 0fbb7cdf65..c496676ce8 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs @@ -597,6 +597,8 @@ private void CleanupExecuteReaderAsync( TaskCompletionSource source, Guid operationId) { + _asyncExecutionCancellationToken = default; + if (task.IsFaulted) { Exception e = task.Exception.InnerException; @@ -1087,6 +1089,8 @@ private Task InternalExecuteReaderAsync( registration = cancellationToken.Register(s_cancelIgnoreFailure, state: this); } + _asyncExecutionCancellationToken = cancellationToken; + Task returnedTask = source.Task; ExecuteReaderAsyncCallContext context = null; try @@ -1920,7 +1924,7 @@ private async Task ContinueRunExecuteReaderTdsAsync( throw; } - await GenerateEnclavePackageAsync(CancellationToken.None).ConfigureAwait(false); + await GenerateEnclavePackageAsync(_asyncExecutionCancellationToken).ConfigureAwait(false); RunExecuteReaderTds( cmdBehavior, diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Xml.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Xml.cs index d74bda54d6..c237132660 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Xml.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Xml.cs @@ -366,6 +366,8 @@ private void CleanupAfterExecuteXmlReaderAsync( TaskCompletionSource source, Guid operationId) { + _asyncExecutionCancellationToken = default; + if (task.IsFaulted) { Exception e = task.Exception?.InnerException; @@ -501,6 +503,8 @@ private Task InternalExecuteXmlReaderAsync(CancellationToken cancella registration = cancellationToken.Register(callback: s_cancelIgnoreFailure, state: this); } + _asyncExecutionCancellationToken = cancellationToken; + // @TODO: This can be cleaned up to lines if InnerConnection is always SqlInternalConnection ExecuteXmlReaderAsyncCallContext context = null; if (_activeConnection?.InnerConnection is SqlConnectionInternal sqlInternalConnection) diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.cs index 294310a08f..2c218705fe 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.cs @@ -153,6 +153,20 @@ public sealed partial class SqlCommand : DbCommand, ICloneable /// private AsyncState _cachedAsyncState = null; + /// + /// Cancellation token supplied by the caller of the asynchronous execution that is currently in + /// flight, or when the command is executing synchronously. + /// + /// + /// The asynchronous execution entry points reach the Always Encrypted machinery through several + /// layers of synchronous, callback-driven code that carry no cancellation token. Rather than + /// threading a token through every one of those signatures — including the synchronous paths that + /// would only ever pass — the token is recorded here for the + /// duration of the operation. It is currently consumed only by the asynchronous Always Encrypted + /// key store calls, which are the sole cancellable I/O in that machinery. + /// + private CancellationToken _asyncExecutionCancellationToken; + private int _currentlyExecutingBatch; /// diff --git a/src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/AsyncKeyStoreProviderTests.cs b/src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/AsyncKeyStoreProviderTests.cs new file mode 100644 index 0000000000..77218fe080 --- /dev/null +++ b/src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/AsyncKeyStoreProviderTests.cs @@ -0,0 +1,156 @@ +// 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.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Data.SqlClient.ManualTesting.Tests.AlwaysEncrypted.Setup; +using Xunit; + +namespace Microsoft.Data.SqlClient.ManualTesting.Tests.AlwaysEncrypted +{ + /// + /// Verifies that the asynchronous command execution paths dispatch column encryption key + /// operations through the asynchronous key store provider APIs, while the synchronous + /// paths continue to use the synchronous APIs. + /// + [Trait("Set", "AE")] + public sealed class AsyncKeyStoreProviderTests : IClassFixture + { + private readonly SQLSetupStrategy _fixture; + + public AsyncKeyStoreProviderTests(SQLSetupStrategyCertStoreProvider fixture) + { + _fixture = fixture; + } + + [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.IsTargetReadyForAeWithKeyStore))] + [ClassData(typeof(AEConnectionStringProvider))] + public void SyncExecutionUsesSyncKeyStoreProviderApi(string connectionString) + { + RecordingKeyStoreProvider provider = new(); + + using SqlConnection connection = new(connectionString); + connection.Open(); + connection.RegisterColumnEncryptionKeyStoreProvidersOnConnection(Wrap(provider)); + + using SqlCommand command = CreateEncryptedParameterCommand(connection); + Assert.Throws(() => command.ExecuteReader()); + + Assert.True(provider.SyncCalls > 0, "Expected the synchronous provider API to be used by ExecuteReader."); + Assert.Equal(0, provider.AsyncCalls); + } + + [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.IsTargetReadyForAeWithKeyStore))] + [ClassData(typeof(AEConnectionStringProvider))] + public async Task ExecuteReaderAsyncUsesAsyncKeyStoreProviderApi(string connectionString) + { + await AssertAsyncApiUsed(connectionString, command => command.ExecuteReaderAsync()); + } + + [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.IsTargetReadyForAeWithKeyStore))] + [ClassData(typeof(AEConnectionStringProvider))] + public async Task ExecuteNonQueryAsyncUsesAsyncKeyStoreProviderApi(string connectionString) + { + await AssertAsyncApiUsed(connectionString, command => command.ExecuteNonQueryAsync()); + } + + [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.IsTargetReadyForAeWithKeyStore))] + [ClassData(typeof(AEConnectionStringProvider))] + public async Task AsyncExecutionHonoursCancellationToken(string connectionString) + { + RecordingKeyStoreProvider provider = new(); + + using SqlConnection connection = new(connectionString); + await connection.OpenAsync(); + connection.RegisterColumnEncryptionKeyStoreProvidersOnConnection(Wrap(provider)); + + using CancellationTokenSource cts = new(); + using SqlCommand command = CreateEncryptedParameterCommand(connection); + + // The provider cancels the token from inside the async decrypt call, which proves the + // caller's token was threaded all the way down into the key store operation. + provider.OnAsyncCall = () => cts.Cancel(); + + await Assert.ThrowsAnyAsync( + () => command.ExecuteReaderAsync(cts.Token)); + + Assert.True(provider.AsyncCalls > 0, "Expected the asynchronous provider API to be used."); + } + + private async Task AssertAsyncApiUsed(string connectionString, Func execute) + { + RecordingKeyStoreProvider provider = new(); + + using SqlConnection connection = new(connectionString); + await connection.OpenAsync(); + connection.RegisterColumnEncryptionKeyStoreProvidersOnConnection(Wrap(provider)); + + using SqlCommand command = CreateEncryptedParameterCommand(connection); + await Assert.ThrowsAsync(() => execute(command)); + + Assert.True(provider.AsyncCalls > 0, "Expected the asynchronous provider API to be used."); + Assert.Equal(0, provider.SyncCalls); + } + + private static Dictionary Wrap( + RecordingKeyStoreProvider provider) => + new() { { DummyKeyStoreProvider.Name, provider } }; + + private SqlCommand CreateEncryptedParameterCommand(SqlConnection connection) + { + SqlCommand command = new( + $"SELECT * FROM [{_fixture.CustomKeyStoreProviderTestTable.Name}] WHERE CustomerID = @id", + connection, + transaction: null, + SqlCommandColumnEncryptionSetting.Enabled); + command.Parameters.AddWithValue("id", 9); + return command; + } + + /// + /// Records which key store provider overload the driver invoked. Both overloads throw so + /// that the test never depends on real key material; the assertions are made on the counters. + /// + private sealed class RecordingKeyStoreProvider : SqlColumnEncryptionKeyStoreProvider + { + private int _syncCalls; + private int _asyncCalls; + + public int SyncCalls => Volatile.Read(ref _syncCalls); + + public int AsyncCalls => Volatile.Read(ref _asyncCalls); + + public Action OnAsyncCall { get; set; } + + public override byte[] DecryptColumnEncryptionKey( + string masterKeyPath, + string encryptionAlgorithm, + byte[] encryptedColumnEncryptionKey) + { + Interlocked.Increment(ref _syncCalls); + throw new NotImplementedException(); + } + + public override Task DecryptColumnEncryptionKeyAsync( + string masterKeyPath, + string encryptionAlgorithm, + byte[] encryptedColumnEncryptionKey, + CancellationToken cancellationToken = default) + { + Interlocked.Increment(ref _asyncCalls); + OnAsyncCall?.Invoke(); + cancellationToken.ThrowIfCancellationRequested(); + throw new NotImplementedException(); + } + + public override byte[] EncryptColumnEncryptionKey( + string masterKeyPath, + string encryptionAlgorithm, + byte[] columnEncryptionKey) => + throw new NotImplementedException(); + } + } +} From 53e0253c79245c9ee12372528fdd8ab4c866159c Mon Sep 17 00:00:00 2001 From: Cheena Malhotra Date: Mon, 24 Aug 2026 09:56:54 -0700 Subject: [PATCH 3/7] Await column encryption key loads on the query metadata cache hit path On a query metadata cache hit the driver still loaded every column encryption key synchronously, on the caller's thread, and then generated the enclave package with the synchronous key store provider API. Both can perform key store network I/O, so an ExecuteReaderAsync that hit the metadata cache could still block the caller. Split the cache lookup into an in-memory probe (TryGetCachedQueryMetadata), which matches the cached cipher metadata onto the parameters and reports which keys still need decrypting, and a completion step that loads those keys. The synchronous lookup is now expressed in terms of the same probe, so its behaviour, including the stale key fallback and the cache hit/miss counters, is unchanged. Asynchronous execution now takes the probe and returns a task for the key loads, which means the cache hit path produces a describe-parameter-encryption task exactly like the cache miss path does. As a result it also picks up the asynchronous enclave package generation and no longer issues the command's TDS write from the caller's thread. When the cached key information turns out to be stale, that task falls back to a full describe parameter encryption round trip. The command still reports usedCache = true in that case because the caller has already returned; this is deliberately conservative, since over-reporting a cache hit can only cost one extra retry of an already failing execution, whereas under-reporting it would suppress the retry a genuine cache hit needs. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 82aca8ca-2078-4a47-860c-6f6ee78cbb02 --- .../Data/SqlClient/SqlCommand.Encryption.cs | 111 +++++++++- .../Data/SqlClient/SqlQueryMetadataCache.cs | 167 +++++++++++--- .../SqlQueryMetadataCacheAsyncShould.cs | 205 ++++++++++++++++++ 3 files changed, 446 insertions(+), 37 deletions(-) create mode 100644 src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs index e3f7b60315..ec342b9000 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs @@ -556,6 +556,105 @@ private void PrepareForTransparentEncryption( bool asyncWrite, out bool usedCache, bool isRetry) + { + returnTask = null; + usedCache = false; + + // If we are not in _batchRPC and not already retrying, attempt to fetch the cipher MD for + // each parameter from the cache. If this succeeds then return immediately, otherwise just + // fall back to the full crypto MD discovery. + if (!_batchRPCMode && !isRetry && _parameters?.Count > 0) + { + SqlQueryMetadataCache cache = SqlQueryMetadataCache.GetInstance(); + + if (!isAsync) + { + if (cache.GetQueryMetadataIfExists(this)) + { + usedCache = true; + return; + } + } + else if (cache.TryGetCachedQueryMetadata(this, out SqlQueryMetadataCache.CachedQueryMetadata metadata)) + { + // The cached metadata matched, but the column encryption keys still have to be + // loaded, which may require key store network I/O. Hand the rest of the work to a + // task so that the caller's thread is not blocked on it. + usedCache = true; + returnTask = CompleteCachedQueryMetadataAsync( + metadata, + timeout, + completion, + asyncWrite, + isRetry); + return; + } + } + + PrepareForTransparentEncryptionCore( + isAsync, + timeout, + completion, + out returnTask, + asyncWrite, + isRetry); + } + + /// + /// Finishes a query metadata cache lookup whose column encryption keys still had to be loaded, + /// falling back to a full describe parameter encryption round trip if the cached key information + /// turned out to be stale. + /// + /// + /// The command reports usedCache = true even when this method falls back, because the + /// caller has already returned by the time the fallback is discovered. That is deliberately + /// conservative: over-reporting a cache hit can only cause one additional retry attempt of an + /// already failing execution, whereas under-reporting it would suppress the + /// retry that a genuine cache hit needs. + /// + private async Task CompleteCachedQueryMetadataAsync( + SqlQueryMetadataCache.CachedQueryMetadata metadata, + int timeout, + TaskCompletionSource completion, + bool asyncWrite, + bool isRetry) + { + bool cacheHit = await SqlQueryMetadataCache.GetInstance() + .CompleteCachedQueryMetadataAsync(this, metadata, _asyncExecutionCancellationToken) + .ConfigureAwait(false); + + if (cacheHit) + { + return; + } + + // The cached key information was stale, so run the full describe parameter encryption + // round trip. It issues blocking TDS writes, but we are already off the caller's thread. + PrepareForTransparentEncryptionCore( + isAsync: true, + timeout, + completion, + out Task describeTask, + asyncWrite, + isRetry); + + if (describeTask is not null) + { + await describeTask.ConfigureAwait(false); + } + } + + /// + /// Performs transparent parameter encryption preparation by issuing a full + /// sp_describe_parameter_encryption round trip, bypassing the query metadata cache. + /// + private void PrepareForTransparentEncryptionCore( + bool isAsync, + int timeout, + TaskCompletionSource completion, // @TODO: Only used for debug checks + out Task returnTask, + bool asyncWrite, + bool isRetry) { Debug.Assert(_activeConnection != null, "_activeConnection should not be null in PrepareForTransparentEncryption."); @@ -574,18 +673,6 @@ private void PrepareForTransparentEncryption( bool describeParameterEncryptionNeeded = false; SqlDataReader describeParameterEncryptionDataReader = null; returnTask = null; - usedCache = false; - - // If we are not in _batchRPC and not already retrying, attempt to fetch the cipher MD for each parameter from the cache. - // If this succeeds then return immediately, otherwise just fall back to the full crypto MD discovery. - if (!_batchRPCMode && - !isRetry && - _parameters?.Count > 0 && - SqlQueryMetadataCache.GetInstance().GetQueryMetadataIfExists(this)) - { - usedCache = true; - return; - } // A flag to indicate if finallyblock needs to execute. bool processFinallyBlock = true; diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlQueryMetadataCache.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlQueryMetadataCache.cs index 40618faf03..24dbe48812 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlQueryMetadataCache.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlQueryMetadataCache.cs @@ -9,6 +9,7 @@ using System.Diagnostics; using System.Text; using System.Threading; +using System.Threading.Tasks; using Microsoft.Extensions.Caching.Memory; namespace Microsoft.Data.SqlClient @@ -47,8 +48,54 @@ internal static SqlQueryMetadataCache GetInstance() /// /// Retrieves the query metadata for a specific query from the cache. /// + /// + /// Column encryption keys are loaded synchronously, which may block the calling thread on key + /// store I/O. Asynchronous execution paths should use + /// followed by instead. + /// internal bool GetQueryMetadataIfExists(SqlCommand sqlCommand) { + if (!TryGetCachedQueryMetadata(sqlCommand, out CachedQueryMetadata metadata)) + { + return false; + } + + foreach (SqlCipherMetadata cipherMetadata in metadata.KeysToLoad) + { + try + { + SqlSecurityUtility.DecryptSymmetricKey(cipherMetadata, sqlCommand.Connection, sqlCommand); + } + catch (Exception ex) when (ex is SqlException or ArgumentException) + { + OnKeyLoadFailed(sqlCommand, clearParameterMetadata: true); + return false; + } + catch (Exception) + { + OnKeyLoadFailed(sqlCommand, clearParameterMetadata: false); + throw; + } + } + + CompleteCachedQueryMetadata(sqlCommand, metadata); + return true; + } + + /// + /// Performs the purely in-memory portion of a query metadata cache lookup. + /// + /// + /// On success every parameter has been assigned a private copy of its cached cipher metadata and + /// describes the column encryption keys that still have to be loaded. + /// The caller must then either load those keys and call , + /// or call which does both. Splitting the lookup this + /// way lets asynchronous execution paths await key store I/O instead of blocking on it. + /// + internal bool TryGetCachedQueryMetadata(SqlCommand sqlCommand, out CachedQueryMetadata metadata) + { + metadata = default; + // Return immediately if caching is disabled. if (!SqlConnection.ColumnEncryptionQueryMetadataCacheEnabled) { @@ -98,6 +145,7 @@ internal bool GetQueryMetadataIfExists(SqlCommand sqlCommand) // Create a copy of the cipherMD in order to load the key. // The key shouldn't be loaded in the cached version for security reasons. + List keysToLoad = new(); foreach (SqlParameter param in sqlCommand.Parameters) { SqlCipherMetadata cipherMdCopy = null; @@ -117,43 +165,112 @@ internal bool GetQueryMetadataIfExists(SqlCommand sqlCommand) if (cipherMdCopy is not null) { - // Try to get the encryption key. If the key information is stale, this might fail. - // In this case, just fail the cache lookup. - try - { - SqlSecurityUtility.DecryptSymmetricKey(cipherMdCopy, sqlCommand.Connection, sqlCommand); - } - catch (Exception ex) when (ex is SqlException or ArgumentException) - { - // Invalidate the cache entry. - InvalidateCacheEntry(sqlCommand); + keysToLoad.Add(cipherMdCopy); + } + } - foreach (SqlParameter paramToCleanup in sqlCommand.Parameters) - { - paramToCleanup.CipherMetadata = null; - } + metadata = new CachedQueryMetadata(enclaveLookupKey, keysToLoad); + return true; + } - IncrementCacheMisses(); - return false; - } - catch (Exception) - { - // Invalidate the cache entry. - InvalidateCacheEntry(sqlCommand); - throw; - } + /// + /// Loads the outstanding column encryption keys for a cache lookup using the asynchronous key + /// store provider APIs and then finishes the lookup. + /// + /// + /// if the lookup completed as a cache hit; if the + /// cached key information was stale, in which case the entry has been invalidated and the caller + /// must fall back to a full describe parameter encryption round trip. + /// + internal async Task CompleteCachedQueryMetadataAsync( + SqlCommand sqlCommand, + CachedQueryMetadata metadata, + CancellationToken cancellationToken) + { + foreach (SqlCipherMetadata cipherMetadata in metadata.KeysToLoad) + { + try + { + await SqlSecurityUtility.DecryptSymmetricKeyAsync( + cipherMetadata, + sqlCommand.Connection, + sqlCommand, + cancellationToken).ConfigureAwait(false); + } + catch (Exception ex) when (ex is SqlException or ArgumentException) + { + // The key information is stale, so fail the cache lookup. + OnKeyLoadFailed(sqlCommand, clearParameterMetadata: true); + return false; + } + catch (Exception) + { + OnKeyLoadFailed(sqlCommand, clearParameterMetadata: false); + throw; } } + CompleteCachedQueryMetadata(sqlCommand, metadata); + return true; + } + + /// + /// Finishes a successful cache lookup by publishing the enclave keys and recording the hit. + /// + private void CompleteCachedQueryMetadata(SqlCommand sqlCommand, CachedQueryMetadata metadata) + { ConcurrentDictionary enclaveKeys = - _cache.Get>(enclaveLookupKey); + _cache.Get>(metadata.EnclaveLookupKey); if (enclaveKeys is not null) { sqlCommand.keysToBeSentToEnclave = CreateCopyOfEnclaveKeys(enclaveKeys); } IncrementCacheHits(); - return true; + } + + /// + /// Rolls back a cache lookup whose column encryption key load failed. + /// + private void OnKeyLoadFailed(SqlCommand sqlCommand, bool clearParameterMetadata) + { + // Invalidate the cache entry. + InvalidateCacheEntry(sqlCommand); + + if (!clearParameterMetadata) + { + return; + } + + foreach (SqlParameter paramToCleanup in sqlCommand.Parameters) + { + paramToCleanup.CipherMetadata = null; + } + + IncrementCacheMisses(); + } + + /// + /// The in-memory result of a query metadata cache probe, before its column encryption keys + /// have been loaded. + /// + internal readonly struct CachedQueryMetadata + { + internal CachedQueryMetadata(string enclaveLookupKey, List keysToLoad) + { + EnclaveLookupKey = enclaveLookupKey; + KeysToLoad = keysToLoad; + } + + /// + /// The cache key under which this query's enclave keys are stored. + /// + internal string EnclaveLookupKey { get; } + + /// + /// The parameter cipher metadata whose column encryption keys still have to be decrypted. + /// + internal List KeysToLoad { get; } } /// diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs new file mode 100644 index 0000000000..16446cf85a --- /dev/null +++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs @@ -0,0 +1,205 @@ +// 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. + +// The Always Encrypted utility layer is internal and nullable-oblivious, and several of its members +// legitimately accept or produce nulls. Nullable analysis is disabled for this file so the tests can +// mirror those signatures exactly. +#nullable disable + +using System; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using Xunit; + +namespace Microsoft.Data.SqlClient.UnitTests.AlwaysEncrypted +{ + /// + /// Tests for the split of lookups into an in-memory probe + /// (TryGetCachedQueryMetadata) and an awaitable column encryption key load + /// (CompleteCachedQueryMetadataAsync), which is what lets asynchronous execution paths + /// avoid blocking the caller's thread on key store I/O. + /// + public class SqlQueryMetadataCacheAsyncShould + { + private const string ProviderName = "TEST_METADATA_CACHE_PROVIDER"; + private const string KeyPath = "metadata-cache-key-path"; + private const string EncryptionAlgorithm = "RSA_OAEP"; + + // The AEAD_AES_256_CBC_HMAC_SHA256 algorithm requires a 256 bit root key. + private static byte[] NewPlaintextKey(byte seed) + { + byte[] key = new byte[32]; + for (int i = 0; i < key.Length; i++) + { + key[i] = (byte)(seed + i); + } + return key; + } + + /// + /// Builds a command whose single parameter carries cipher metadata backed by unique key + /// material, so that neither the process-wide query metadata cache nor the process-wide + /// column encryption key cache can collide with another test. + /// + private static SqlCommand NewCommandWithCipherMetadata(SqlColumnEncryptionKeyStoreProvider provider) + { + SqlConnection connection = new SqlConnection( + "Data Source=async-ae-metadata-cache-test;Initial Catalog=async-ae-db;Column Encryption Setting=Enabled"); + connection.RegisterColumnEncryptionKeyStoreProvidersOnConnection( + new Dictionary { { ProviderName, provider } }); + + SqlCommand command = new SqlCommand($"SELECT * FROM T WHERE C = @p -- {Guid.NewGuid()}", connection); + SqlParameter parameter = command.Parameters.AddWithValue("@p", 1); + + SqlTceCipherInfoEntry entry = new SqlTceCipherInfoEntry(ordinal: 0); + entry.Add( + encryptedKey: Guid.NewGuid().ToByteArray(), + databaseId: 1, + cekId: 1, + cekVersion: 1, + cekMdVersion: 1, + keyPath: KeyPath, + keyStoreName: ProviderName, + algorithmName: EncryptionAlgorithm); + + parameter.CipherMetadata = new SqlCipherMetadata( + entry, + ordinal: 0, + cipherAlgorithmId: 2, // AEAD_AES_256_CBC_HMAC_SHA256 + cipherAlgorithmName: null, + encryptionType: 1, // Deterministic + normalizationRuleVersion: 1); + + return command; + } + + private static SqlCommand NewCachedCommand(SqlColumnEncryptionKeyStoreProvider provider) + { + SqlCommand command = NewCommandWithCipherMetadata(provider); + SqlQueryMetadataCache.GetInstance().AddQueryMetadata(command, ignoreQueriesWithReturnValueParams: true); + return command; + } + + [Fact] + public async Task CompleteCachedQueryMetadataAsync_UsesAsyncKeyStoreProviderApi() + { + TestKeyStoreProvider provider = new TestKeyStoreProvider { PlaintextKey = NewPlaintextKey(seed: 41) }; + SqlCommand command = NewCachedCommand(provider); + + Assert.True(SqlQueryMetadataCache.GetInstance() + .TryGetCachedQueryMetadata(command, out SqlQueryMetadataCache.CachedQueryMetadata metadata)); + + // The probe must not have touched the key store: that is the whole point of the split. + Assert.Equal(0, provider.DecryptCallCount); + Assert.Equal(0, provider.DecryptAsyncCallCount); + Assert.NotEmpty(metadata.KeysToLoad); + + Assert.True(await SqlQueryMetadataCache.GetInstance() + .CompleteCachedQueryMetadataAsync(command, metadata, CancellationToken.None)); + + Assert.Equal(1, provider.DecryptAsyncCallCount); + Assert.Equal(0, provider.DecryptCallCount); + Assert.True(command.Parameters[0].CipherMetadata.IsAlgorithmInitialized()); + } + + [Fact] + public async Task CompleteCachedQueryMetadataAsync_WhenKeyIsStale_ReportsMissAndClearsMetadata() + { + TestKeyStoreProvider provider = new TestKeyStoreProvider + { + DecryptAsyncCallback = _ => throw new ArgumentException("stale key"), + }; + SqlCommand command = NewCachedCommand(provider); + + Assert.True(SqlQueryMetadataCache.GetInstance() + .TryGetCachedQueryMetadata(command, out SqlQueryMetadataCache.CachedQueryMetadata metadata)); + + // A stale key must degrade to a cache miss rather than surface as an execution failure, so + // that the caller falls back to a full describe parameter encryption round trip. + Assert.False(await SqlQueryMetadataCache.GetInstance() + .CompleteCachedQueryMetadataAsync(command, metadata, CancellationToken.None)); + + Assert.Null(command.Parameters[0].CipherMetadata); + + // The entry must have been invalidated so the next lookup does not repeat the failure. + Assert.False(SqlQueryMetadataCache.GetInstance().TryGetCachedQueryMetadata(command, out _)); + } + + /// + /// The synchronous lookup is now expressed in terms of the same probe, so guard that it still + /// uses the synchronous key store provider API and still reports a hit. + /// + [Fact] + public void GetQueryMetadataIfExists_StillUsesSyncKeyStoreProviderApi() + { + TestKeyStoreProvider provider = new TestKeyStoreProvider { PlaintextKey = NewPlaintextKey(seed: 47) }; + SqlCommand command = NewCachedCommand(provider); + + Assert.True(SqlQueryMetadataCache.GetInstance().GetQueryMetadataIfExists(command)); + + Assert.Equal(1, provider.DecryptCallCount); + Assert.Equal(0, provider.DecryptAsyncCallCount); + Assert.True(command.Parameters[0].CipherMetadata.IsAlgorithmInitialized()); + } + + [Fact] + public async Task CompleteCachedQueryMetadataAsync_WhenCancelled_DoesNotReportHit() + { + TestKeyStoreProvider provider = new TestKeyStoreProvider { PlaintextKey = NewPlaintextKey(seed: 43) }; + SqlCommand command = NewCachedCommand(provider); + + Assert.True(SqlQueryMetadataCache.GetInstance() + .TryGetCachedQueryMetadata(command, out SqlQueryMetadataCache.CachedQueryMetadata metadata)); + + using CancellationTokenSource cts = new(); + cts.Cancel(); + + await Assert.ThrowsAnyAsync( + () => SqlQueryMetadataCache.GetInstance() + .CompleteCachedQueryMetadataAsync(command, metadata, cts.Token)); + } + + /// + /// A key store provider whose sync and async paths are separately observable, so that tests can + /// assert which path was taken and how many times. + /// + private sealed class TestKeyStoreProvider : SqlColumnEncryptionKeyStoreProvider + { + private int _decryptCallCount; + private int _decryptAsyncCallCount; + + internal byte[] PlaintextKey { get; set; } = new byte[32]; + + internal Func> DecryptAsyncCallback { get; set; } + + internal int DecryptCallCount => _decryptCallCount; + + internal int DecryptAsyncCallCount => _decryptAsyncCallCount; + + public override byte[] DecryptColumnEncryptionKey(string masterKeyPath, string encryptionAlgorithm, byte[] encryptedColumnEncryptionKey) + { + Interlocked.Increment(ref _decryptCallCount); + return PlaintextKey; + } + + public override Task DecryptColumnEncryptionKeyAsync(string masterKeyPath, string encryptionAlgorithm, byte[] encryptedColumnEncryptionKey, CancellationToken cancellationToken = default) + { + Interlocked.Increment(ref _decryptAsyncCallCount); + cancellationToken.ThrowIfCancellationRequested(); + Func> callback = DecryptAsyncCallback; + return callback is not null ? callback(cancellationToken) : Task.FromResult(PlaintextKey); + } + + public override byte[] EncryptColumnEncryptionKey(string masterKeyPath, string encryptionAlgorithm, byte[] columnEncryptionKey) + => throw new NotSupportedException(); + + public override byte[] SignColumnMasterKeyMetadata(string masterKeyPath, bool allowEnclaveComputations) + => throw new NotSupportedException(); + + public override bool VerifyColumnMasterKeyMetadata(string masterKeyPath, bool allowEnclaveComputations, byte[] signature) + => throw new NotSupportedException(); + } + } +} From fbf775efe3c911c9e5e6ed2c92ebe866ebd0b609 Mon Sep 17 00:00:00 2001 From: Cheena Malhotra Date: Mon, 24 Aug 2026 10:37:58 -0700 Subject: [PATCH 4/7] Name the describe-parameter-encryption task methods for what they do GetParameterEncryptionDataReader returned void with an out Task and did nothing but dispatch asynchronous work, so its name gave no hint that the caller was receiving a pending operation. The name was inherited from main, where it was actively misleading: the overload without the Async suffix was the one that continued from a pending network write, while the overload with the suffix was the one used when that write had already completed. Rename the dispatcher to GetParameterEncryptionDataReaderAsync and have it return the task directly instead of through an out parameter, and rename the async body to ConsumeDescribeParameterEncryptionResultsAsync, which describes what it actually does. Both names now carry the Async suffix and both return a task. Also make the Task.Run rationale precise. It is not merely a convenience: the body issues blocking TDS reads before it reaches a suspension point, and PrepareForTransparentEncryption runs synchronously on the caller's thread under the async entry points, so awaiting inline would put those reads on the caller's thread whenever the network write is null or already complete. Neither continuation this replaced could run inline, since one used Task.Run and the other used ContinueWith without ExecuteSynchronously. No behaviour change. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 82aca8ca-2078-4a47-860c-6f6ee78cbb02 --- .../Data/SqlClient/SqlCommand.Encryption.cs | 34 +++++++++---------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs index ec342b9000..49cc356e76 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs @@ -251,14 +251,18 @@ private EnclaveSessionParameters GetEnclaveSessionParameters() => /// column encryption key has been decrypted through the asynchronous key store provider APIs. /// /// - /// The work is dispatched with because the body issues blocking TDS - /// reads before reaching its first suspension point, and PrepareForTransparentEncryption is - /// invoked synchronously on the caller's thread by the asynchronous execution entry points. Dispatching - /// keeps that thread free, matching the behaviour of the continuation chain this replaced, while every - /// subsequent await releases the pooled thread for the duration of key store provider I/O. + /// The work is dispatched with rather than awaited inline because + /// the body issues blocking TDS reads before reaching a suspension point, and + /// PrepareForTransparentEncryption runs synchronously on the caller's thread under the + /// asynchronous execution entry points. Awaiting inline would run those reads on the caller's thread + /// whenever is null or already + /// completed. This preserves the behaviour of the two continuations it replaced, which reached the + /// thread pool via Task.Run and via ContinueWith without + /// respectively, so neither could run + /// inline either. Once dispatched, every await releases the pooled thread for the duration of + /// key store provider I/O. /// /// - /// Receives the task representing the pending work /// /// Task representing the pending network write of the describe-parameter-encryption request, or /// null when that write completed synchronously. @@ -267,21 +271,19 @@ private EnclaveSessionParameters GetEnclaveSessionParameters() => /// Map of encryption RPC requests to their original RPC requests /// Whether describe parameter encryption was required /// Indicates if this is a retry from a failed call - private void GetParameterEncryptionDataReader( - out Task returnTask, + /// A task that completes once the describe-parameter-encryption results have been consumed + private Task GetParameterEncryptionDataReaderAsync( Task fetchInputParameterEncryptionInfoTask, SqlDataReader describeParameterEncryptionDataReader, ReadOnlyDictionary<_SqlRPC, _SqlRPC> describeParameterEncryptionRpcOriginalRpcMap, bool describeParameterEncryptionNeeded, - bool isRetry) - { - returnTask = Task.Run(() => GetParameterEncryptionDataReaderAsync( + bool isRetry) => + Task.Run(() => ConsumeDescribeParameterEncryptionResultsAsync( fetchInputParameterEncryptionInfoTask, describeParameterEncryptionDataReader, describeParameterEncryptionRpcOriginalRpcMap, describeParameterEncryptionNeeded, isRetry)); - } /// /// Awaits the describe-parameter-encryption request and consumes its results asynchronously. @@ -292,7 +294,7 @@ private void GetParameterEncryptionDataReader( /// while consuming the results run the finally block but leave the cached async state alone. These /// semantics are inherited from the callback-based continuation this method replaced. /// - private async Task GetParameterEncryptionDataReaderAsync( + private async Task ConsumeDescribeParameterEncryptionResultsAsync( Task fetchInputParameterEncryptionInfoTask, SqlDataReader describeParameterEncryptionDataReader, ReadOnlyDictionary<_SqlRPC, _SqlRPC> describeParameterEncryptionRpcOriginalRpcMap, @@ -731,8 +733,7 @@ private void PrepareForTransparentEncryptionCore( // execution pending. Note that this should be done outside the task's // continuation delegate. processFinallyBlock = false; - GetParameterEncryptionDataReader( - out returnTask, + returnTask = GetParameterEncryptionDataReaderAsync( fetchInputParameterEncryptionInfoTask, describeParameterEncryptionDataReader, describeParameterEncryptionRpcOriginalRpcMap, @@ -751,8 +752,7 @@ private void PrepareForTransparentEncryptionCore( // execution pending. Note that this should be done outside the task's // continuation delegate. processFinallyBlock = false; - GetParameterEncryptionDataReader( - out returnTask, + returnTask = GetParameterEncryptionDataReaderAsync( fetchInputParameterEncryptionInfoTask: null, describeParameterEncryptionDataReader, describeParameterEncryptionRpcOriginalRpcMap, From 7daff149fad12c7b37a32f5d953b84095743cdb9 Mon Sep 17 00:00:00 2001 From: Cheena Malhotra Date: Mon, 24 Aug 2026 10:46:02 -0700 Subject: [PATCH 5/7] Address PR review feedback on async Always Encrypted changes Report CreateColumnEncryptionKeyInfo as the throwing member in the internal null-argument diagnostics. The throws moved into that shared helper and are now reachable from both the sync and async decrypt paths, so naming GetDecryptedKeysToBeSentToEnclave was misleading. This also matches the convention used by EncryptBytePackage and ComputeQueryStringHash in the same type. Use TdsEnums.AEAD_AES_256_CBC_HMAC_SHA256 and SqlClientEncryptionType.Deterministic in the query metadata cache tests instead of bare literals with explanatory comments. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 82aca8ca-2078-4a47-860c-6f6ee78cbb02 --- .../src/Microsoft/Data/SqlClient/EnclaveDelegate.cs | 4 ++-- .../AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.cs index 1c6ef2fd35..6a8264fe3b 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.cs @@ -113,11 +113,11 @@ private static ColumnEncryptionKeyInfo CreateColumnEncryptionKeyInfo( { if (sqlClientSymmetricKey == null) { - throw SQL.NullArgumentInternal(nameof(sqlClientSymmetricKey), nameof(EnclaveDelegate), nameof(GetDecryptedKeysToBeSentToEnclave)); + throw SQL.NullArgumentInternal(nameof(sqlClientSymmetricKey), nameof(EnclaveDelegate), nameof(CreateColumnEncryptionKeyInfo)); } if (cipherInfo.ColumnEncryptionKeyValues == null) { - throw SQL.NullArgumentInternal(nameof(cipherInfo.ColumnEncryptionKeyValues), nameof(EnclaveDelegate), nameof(GetDecryptedKeysToBeSentToEnclave)); + throw SQL.NullArgumentInternal(nameof(cipherInfo.ColumnEncryptionKeyValues), nameof(EnclaveDelegate), nameof(CreateColumnEncryptionKeyInfo)); } if (!(cipherInfo.ColumnEncryptionKeyValues.Count > 0)) { diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs index 16446cf85a..3b830f2da3 100644 --- a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs +++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs @@ -67,9 +67,9 @@ private static SqlCommand NewCommandWithCipherMetadata(SqlColumnEncryptionKeySto parameter.CipherMetadata = new SqlCipherMetadata( entry, ordinal: 0, - cipherAlgorithmId: 2, // AEAD_AES_256_CBC_HMAC_SHA256 + cipherAlgorithmId: TdsEnums.AEAD_AES_256_CBC_HMAC_SHA256, cipherAlgorithmName: null, - encryptionType: 1, // Deterministic + encryptionType: (byte)SqlClientEncryptionType.Deterministic, normalizationRuleVersion: 1); return command; From a71f1b6abf1bf71e473a512e8ba261fedb3f24a8 Mon Sep 17 00:00:00 2001 From: Cheena Malhotra Date: Mon, 24 Aug 2026 11:11:24 -0700 Subject: [PATCH 6/7] Address self-review findings on the async Always Encrypted path A pass over the async Always Encrypted work looking for correctness, performance, and consistency problems. The behavioural changes: Capture the cancellation token on the originating thread. The token lives in a field so it does not have to be threaded through layers of callback-driven code that have no use for it, but the Always Encrypted work is handed to the thread pool and can outlive the cleanup continuation that resets that field. Reading it late could silently observe CancellationToken.None. It is now read once at the two synchronous capture points and passed on as an explicit parameter. Make the symmetric key cache insert non-cancellable. Cancelling between a successful decrypt and the cache insert threw away a key that had already cost a key store round trip, so the next caller paid for it again. Drop a dead serverName parameter from both GetDecryptedKeysToBeSentToEnclave overloads, and give the result list a capacity hint. The rest is allocation and clarity work: lazily allocate the pending key operation lists and hand them out as IReadOnlyList, walk them by index so the interface enumerator is not boxed, and make IsEnclaveEnabled an explicit field on ColumnMasterKeySignatureVerification rather than a hard-coded true that only happens to be correct because of where the call sits. Also documents three decisions that look like missed optimisations but are not. The Task.Run in GetParameterEncryptionDataReaderAsync cannot be awaited inline because the task it waits on is completed by the SNI network callback, so an inline continuation would run blocking TDS reads on a callback thread. Column encryption key decryption stays sequential because distinct cipher metadata entries often share a key, and GetKeyAsync does not hold its gate across provider I/O, so parallel misses for the same key would each issue their own key store call. Sequential decryption lets the first result warm the cache. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 82aca8ca-2078-4a47-860c-6f6ee78cbb02 --- .../Data/SqlClient/EnclaveDelegate.Crypto.cs | 8 +- .../Data/SqlClient/EnclaveDelegate.cs | 13 +- .../Data/SqlClient/SqlCommand.Encryption.cs | 125 ++++++++++++++---- .../Data/SqlClient/SqlCommand.Reader.cs | 24 +++- .../Microsoft/Data/SqlClient/SqlCommand.cs | 8 ++ .../Data/SqlClient/SqlQueryMetadataCache.cs | 20 +-- .../Data/SqlClient/SqlSymmetricKeyCache.cs | 5 +- .../EnclaveDelegateAsyncShould.cs | 5 - .../SqlSecurityUtilityAsyncShould.cs | 35 +++++ 9 files changed, 189 insertions(+), 54 deletions(-) diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.Crypto.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.Crypto.cs index 610ae10fe0..3bc150f87e 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.Crypto.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.Crypto.cs @@ -157,7 +157,7 @@ internal EnclavePackage GenerateEnclavePackage(SqlConnectionAttestationProtocol throw new RetryableEnclaveQueryExecutionException(e.Message, e); } - List decryptedKeysToBeSentToEnclave = GetDecryptedKeysToBeSentToEnclave(keysToBeSentToEnclave, enclaveSessionParameters.ServerName, connection, command); + List decryptedKeysToBeSentToEnclave = GetDecryptedKeysToBeSentToEnclave(keysToBeSentToEnclave, connection, command); return BuildEnclavePackage(decryptedKeysToBeSentToEnclave, queryText, counter, sqlEnclaveSession, enclaveSessionParameters.ServerName); } @@ -192,6 +192,11 @@ internal async Task GenerateEnclavePackageAsync( 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, @@ -212,7 +217,6 @@ internal async Task GenerateEnclavePackageAsync( List decryptedKeysToBeSentToEnclave = await GetDecryptedKeysToBeSentToEnclaveAsync( keysToBeSentToEnclave, - enclaveSessionParameters.ServerName, connection, command, cancellationToken) diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.cs index 6a8264fe3b..6e2f7d398f 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.cs @@ -48,13 +48,12 @@ private byte[] GetUintBytes(string enclaveType, int intValue, string variableNam /// Decrypt the keys that need to be sent to the enclave /// /// Keys that need to sent to the enclave - /// - /// - /// + /// Connection executing the query + /// Command executing the query /// - internal List GetDecryptedKeysToBeSentToEnclave(ConcurrentDictionary keysTobeSentToEnclave, string serverName, SqlConnection connection, SqlCommand command) + internal List GetDecryptedKeysToBeSentToEnclave(ConcurrentDictionary keysTobeSentToEnclave, SqlConnection connection, SqlCommand command) { - List decryptedKeysToBeSentToEnclave = new List(); + List decryptedKeysToBeSentToEnclave = new List(keysTobeSentToEnclave.Count); foreach (SqlTceCipherInfoEntry cipherInfo in keysTobeSentToEnclave.Values) { @@ -75,18 +74,16 @@ internal List GetDecryptedKeysToBeSentToEnclave(Concurr /// thread while the enclave package is being assembled. /// /// Keys that need to sent to the enclave - /// Name of the server the keys are being resolved for /// Connection executing the query /// Command executing the query /// Token used to request cancellation of the operation internal async Task> GetDecryptedKeysToBeSentToEnclaveAsync( ConcurrentDictionary keysTobeSentToEnclave, - string serverName, SqlConnection connection, SqlCommand command, CancellationToken cancellationToken) { - List decryptedKeysToBeSentToEnclave = new List(); + List decryptedKeysToBeSentToEnclave = new List(keysTobeSentToEnclave.Count); foreach (SqlTceCipherInfoEntry cipherInfo in keysTobeSentToEnclave.Values) { diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs index 49cc356e76..6a2ae6c32d 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs @@ -262,6 +262,14 @@ private EnclaveSessionParameters GetEnclaveSessionParameters() => /// inline either. Once dispatched, every await releases the pooled thread for the duration of /// key store provider I/O. /// + /// + /// Do not "optimise" this into an inline await of + /// in order to save a thread pool dispatch. + /// That task is completed by the network write callback, so an inline continuation would run the + /// blocking TDS reads below on an SNI callback thread. The replaced ContinueWith deliberately + /// omitted for exactly this reason. The + /// dispatch this costs is negligible next to the round trip it accompanies. + /// /// /// /// Task representing the pending network write of the describe-parameter-encryption request, or @@ -271,19 +279,22 @@ private EnclaveSessionParameters GetEnclaveSessionParameters() => /// Map of encryption RPC requests to their original RPC requests /// Whether describe parameter encryption was required /// Indicates if this is a retry from a failed call + /// Token used to request cancellation of the key store operations /// A task that completes once the describe-parameter-encryption results have been consumed private Task GetParameterEncryptionDataReaderAsync( Task fetchInputParameterEncryptionInfoTask, SqlDataReader describeParameterEncryptionDataReader, ReadOnlyDictionary<_SqlRPC, _SqlRPC> describeParameterEncryptionRpcOriginalRpcMap, bool describeParameterEncryptionNeeded, - bool isRetry) => + bool isRetry, + CancellationToken cancellationToken) => Task.Run(() => ConsumeDescribeParameterEncryptionResultsAsync( fetchInputParameterEncryptionInfoTask, describeParameterEncryptionDataReader, describeParameterEncryptionRpcOriginalRpcMap, describeParameterEncryptionNeeded, - isRetry)); + isRetry, + cancellationToken)); /// /// Awaits the describe-parameter-encryption request and consumes its results asynchronously. @@ -299,7 +310,8 @@ private async Task ConsumeDescribeParameterEncryptionResultsAsync( SqlDataReader describeParameterEncryptionDataReader, ReadOnlyDictionary<_SqlRPC, _SqlRPC> describeParameterEncryptionRpcOriginalRpcMap, bool describeParameterEncryptionNeeded, - bool isRetry) + bool isRetry, + CancellationToken cancellationToken) { if (fetchInputParameterEncryptionInfoTask is not null) { @@ -347,7 +359,7 @@ await ReadDescribeEncryptionParameterResultsAsync( describeParameterEncryptionDataReader, describeParameterEncryptionRpcOriginalRpcMap, isRetry, - _asyncExecutionCancellationToken) + cancellationToken) .ConfigureAwait(false); #if DEBUG @@ -562,6 +574,11 @@ private void PrepareForTransparentEncryption( returnTask = null; usedCache = false; + // Capture the caller's cancellation token once, here, while we are still on the thread that + // started the execution. Everything below may run after the operation has completed and + // cleared the field, so reading it later could silently observe CancellationToken.None. + CancellationToken cancellationToken = isAsync ? _asyncExecutionCancellationToken : CancellationToken.None; + // If we are not in _batchRPC and not already retrying, attempt to fetch the cipher MD for // each parameter from the cache. If this succeeds then return immediately, otherwise just // fall back to the full crypto MD discovery. @@ -588,7 +605,8 @@ private void PrepareForTransparentEncryption( timeout, completion, asyncWrite, - isRetry); + isRetry, + cancellationToken); return; } } @@ -599,7 +617,8 @@ private void PrepareForTransparentEncryption( completion, out returnTask, asyncWrite, - isRetry); + isRetry, + cancellationToken); } /// @@ -619,10 +638,11 @@ private async Task CompleteCachedQueryMetadataAsync( int timeout, TaskCompletionSource completion, bool asyncWrite, - bool isRetry) + bool isRetry, + CancellationToken cancellationToken) { bool cacheHit = await SqlQueryMetadataCache.GetInstance() - .CompleteCachedQueryMetadataAsync(this, metadata, _asyncExecutionCancellationToken) + .CompleteCachedQueryMetadataAsync(this, metadata, cancellationToken) .ConfigureAwait(false); if (cacheHit) @@ -638,7 +658,8 @@ private async Task CompleteCachedQueryMetadataAsync( completion, out Task describeTask, asyncWrite, - isRetry); + isRetry, + cancellationToken); if (describeTask is not null) { @@ -650,13 +671,18 @@ private async Task CompleteCachedQueryMetadataAsync( /// Performs transparent parameter encryption preparation by issuing a full /// sp_describe_parameter_encryption round trip, bypassing the query metadata cache. /// + /// + /// is captured by the caller while it is still on the thread + /// that started the execution, and is used only by the asynchronous key store provider calls. + /// private void PrepareForTransparentEncryptionCore( bool isAsync, int timeout, TaskCompletionSource completion, // @TODO: Only used for debug checks out Task returnTask, bool asyncWrite, - bool isRetry) + bool isRetry, + CancellationToken cancellationToken) { Debug.Assert(_activeConnection != null, "_activeConnection should not be null in PrepareForTransparentEncryption."); @@ -738,7 +764,8 @@ private void PrepareForTransparentEncryptionCore( describeParameterEncryptionDataReader, describeParameterEncryptionRpcOriginalRpcMap, describeParameterEncryptionNeeded, - isRetry); + isRetry, + cancellationToken); decrementAsyncCountInFinallyBlock = false; } @@ -757,7 +784,8 @@ private void PrepareForTransparentEncryptionCore( describeParameterEncryptionDataReader, describeParameterEncryptionRpcOriginalRpcMap, describeParameterEncryptionNeeded, - isRetry); + isRetry, + cancellationToken); decrementAsyncCountInFinallyBlock = false; } @@ -863,20 +891,23 @@ private void ReadDescribeEncryptionParameterResults( PendingColumnEncryptionKeyOperations pending = new PendingColumnEncryptionKeyOperations(); ReadDescribeEncryptionParameterResultsCore(ds, describeParameterEncryptionRpcOriginalRpcMap, isRetry, pending); - foreach (ColumnMasterKeySignatureVerification verification in pending.SignatureVerifications) + IReadOnlyList verifications = pending.SignatureVerifications; + for (int i = 0; i < verifications.Count; i++) { + ColumnMasterKeySignatureVerification verification = verifications[i]; SqlSecurityUtility.VerifyColumnMasterKeySignature( verification.KeyStoreName, verification.KeyPath, - isEnclaveEnabled: true, + verification.IsEnclaveEnabled, verification.Signature, _activeConnection, this); } - foreach (SqlCipherMetadata cipherMetadata in pending.KeyDecryptions) + IReadOnlyList keyDecryptions = pending.KeyDecryptions; + for (int i = 0; i < keyDecryptions.Count; i++) { - SqlSecurityUtility.DecryptSymmetricKey(cipherMetadata, _activeConnection, this); + SqlSecurityUtility.DecryptSymmetricKey(keyDecryptions[i], _activeConnection, this); } CacheQueryMetadataIfNeeded(); @@ -904,12 +935,14 @@ private async Task ReadDescribeEncryptionParameterResultsAsync( PendingColumnEncryptionKeyOperations pending = new PendingColumnEncryptionKeyOperations(); ReadDescribeEncryptionParameterResultsCore(ds, describeParameterEncryptionRpcOriginalRpcMap, isRetry, pending); - foreach (ColumnMasterKeySignatureVerification verification in pending.SignatureVerifications) + IReadOnlyList verifications = pending.SignatureVerifications; + for (int i = 0; i < verifications.Count; i++) { + ColumnMasterKeySignatureVerification verification = verifications[i]; await SqlSecurityUtility.VerifyColumnMasterKeySignatureAsync( verification.KeyStoreName, verification.KeyPath, - isEnclaveEnabled: true, + verification.IsEnclaveEnabled, verification.Signature, _activeConnection, this, @@ -917,9 +950,15 @@ await SqlSecurityUtility.VerifyColumnMasterKeySignatureAsync( .ConfigureAwait(false); } - foreach (SqlCipherMetadata cipherMetadata in pending.KeyDecryptions) + // Decrypted sequentially rather than with Task.WhenAll. Distinct SqlCipherMetadata entries + // frequently share a column encryption key, and SqlSymmetricKeyCache.GetKeyAsync deliberately + // does not hold its gate across provider I/O, so concurrent misses for the same key would each + // issue their own key store call. Sequential decryption lets the first result populate the cache + // for the rest, which matters more than overlapping the few distinct keys a query uses. + IReadOnlyList keyDecryptions = pending.KeyDecryptions; + for (int i = 0; i < keyDecryptions.Count; i++) { - await SqlSecurityUtility.DecryptSymmetricKeyAsync(cipherMetadata, _activeConnection, this, cancellationToken) + await SqlSecurityUtility.DecryptSymmetricKeyAsync(keyDecryptions[i], _activeConnection, this, cancellationToken) .ConfigureAwait(false); } @@ -942,10 +981,19 @@ private void CacheQueryMetadataIfNeeded() /// Parses the result sets returned by sp_describe_parameter_encryption. /// /// + /// /// Column master key signature verification and column encryption key decryption are not performed /// here; they are recorded in so that the caller can execute them either /// synchronously or asynchronously. Deferring them also means no key store network call is made while /// the describe-parameter-encryption reader is still positioned mid-stream. + /// + /// + /// Deferring these operations changes which exception surfaces when a query hits more than one + /// problem. Previously verification and decryption were interleaved with parsing, so a bad signature + /// on the first key was raised before the parameter metadata result set was parsed at all. Now parsing + /// always completes first, so a malformed result set is reported in preference to a key store failure. + /// The set of exceptions that can be thrown, and the state left behind on failure, are unchanged. + /// /// /// Resultset from calling to sp_describe_parameter_encryption /// Readonly dictionary with the map of parameter encryption rpc requests with the corresponding original rpc requests. @@ -1077,11 +1125,30 @@ private void ReadDescribeEncryptionParameterResultsCore( /// private sealed class PendingColumnEncryptionKeyOperations { + private List _signatureVerifications; + private List _keyDecryptions; + /// Column master key signatures that must be verified. - internal List SignatureVerifications { get; } = new(); + internal IReadOnlyList SignatureVerifications => + (IReadOnlyList)_signatureVerifications ?? + Array.Empty(); /// Column encryption keys that must be decrypted. - internal List KeyDecryptions { get; } = new(); + internal IReadOnlyList KeyDecryptions => + (IReadOnlyList)_keyDecryptions ?? Array.Empty(); + + /// + /// Records a column master key signature that must be verified before the query runs. + /// + internal void AddSignatureVerification(string keyStoreName, string keyPath, bool isEnclaveEnabled, byte[] signature) => + (_signatureVerifications ??= new List()) + .Add(new ColumnMasterKeySignatureVerification(keyStoreName, keyPath, isEnclaveEnabled, signature)); + + /// + /// Records a column encryption key that must be decrypted before the query runs. + /// + internal void AddKeyDecryption(SqlCipherMetadata cipherMetadata) => + (_keyDecryptions ??= new List()).Add(cipherMetadata); } /// @@ -1089,10 +1156,11 @@ private sealed class PendingColumnEncryptionKeyOperations /// private readonly struct ColumnMasterKeySignatureVerification { - internal ColumnMasterKeySignatureVerification(string keyStoreName, string keyPath, byte[] signature) + internal ColumnMasterKeySignatureVerification(string keyStoreName, string keyPath, bool isEnclaveEnabled, byte[] signature) { KeyStoreName = keyStoreName; KeyPath = keyPath; + IsEnclaveEnabled = isEnclaveEnabled; Signature = signature; } @@ -1100,6 +1168,12 @@ internal ColumnMasterKeySignatureVerification(string keyStoreName, string keyPat internal string KeyPath { get; } + /// + /// Whether the server reported that this key is required by the enclave. Carried explicitly + /// rather than assumed, because it selects which signature the key store validates. + /// + internal bool IsEnclaveEnabled { get; } + internal byte[] Signature { get; } } @@ -1256,8 +1330,7 @@ private bool ReadDescribeEncryptionParameterResultsKeys( // Defer signature verification: it may reach a key store over the network and must not // run while this reader is still positioned mid-result-set. - pending.SignatureVerifications.Add( - new ColumnMasterKeySignatureVerification(providerName, keyPath, keySignature)); + pending.AddSignatureVerification(providerName, keyPath, isRequestedByEnclave, keySignature); // Lookup the key, failing which throw an exception // @TODO: Seriously, we *just* did this, why are we looking it up again?? @@ -1351,7 +1424,7 @@ private int ReadDescribeEncryptionParameterResultsMetadata( // Defer decryption of the symmetric key: it may reach a key store over the network // and must not run while this reader is still positioned mid-result-set. Decryption // also validates the metadata and will throw if it is invalid. - pending.KeyDecryptions.Add(sqlParameter.CipherMetadata); + pending.AddKeyDecryption(sqlParameter.CipherMetadata); // This is effective only for _batchRPCMode even though we set it for // non-_batchRPCMode also, since for non-_batchRPCMode, param options diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs index c496676ce8..be1a311d12 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs @@ -1030,11 +1030,16 @@ private bool TryPrepareEnclavePackageGeneration( /// /// Test-only failpoint that forces a retryable enclave failure during package generation. /// + /// + /// Both guards are required and are not redundant: elides the + /// call sites in non-DEBUG builds, while the #if DEBUG is needed because the body still + /// compiles there and the field it reads is itself DEBUG-only. + /// + // @TODO: These should be wrapped with something other than DEBUG since we don't even run tests in debug mode [Conditional("DEBUG")] private void ThrowIfForcedRetryableEnclaveQueryExecutionException() { #if DEBUG - // @TODO: These should be wrapped with something other than DEBUG since we don't even run tests in debug mode // Test-only code for forcing a retryable exception to occur if (_forceRetryableEnclaveQueryExecutionExceptionDuringGenerateEnclavePackage) { @@ -1855,6 +1860,11 @@ private SqlDataReader RunExecuteReaderTdsWithTransparentParameterEncryption( // The remainder of the execution issues blocking TDS writes, so it must never run inline on // the caller's thread (nor on a network callback thread). Task.Run reproduces the thread pool // hand-off that the previous ContinueWith-based continuation provided. + // Capture the caller's cancellation token here, on the thread that started the + // execution. The body below runs after a thread pool hand-off and may observe the field + // already cleared by the execution's cleanup continuation. + CancellationToken cancellationToken = isAsync ? _asyncExecutionCancellationToken : CancellationToken.None; + task = Task.Run(() => ContinueRunExecuteReaderTdsAsync( describeParameterEncryptionTask, cmdBehavior, @@ -1865,7 +1875,8 @@ private SqlDataReader RunExecuteReaderTdsWithTransparentParameterEncryption( parameterEncryptionStart, asyncWrite, isRetry, - ds)); + ds, + cancellationToken)); return ds; } @@ -1912,7 +1923,8 @@ private async Task ContinueRunExecuteReaderTdsAsync( long parameterEncryptionStart, bool asyncWrite, bool isRetry, - SqlDataReader ds) + SqlDataReader ds, + CancellationToken cancellationToken) { try { @@ -1920,11 +1932,15 @@ private async Task ContinueRunExecuteReaderTdsAsync( } catch { + // Unlike ConsumeDescribeParameterEncryptionResultsAsync, cancellation is reset here as + // well as failure. That asymmetry is deliberate: this method replaced + // AsyncHelper.ContinueTaskWithState, which supplied an onCancellation callback, whereas + // that one replaced CreateContinuationTaskWithState, which did not. CachedAsyncState?.ResetAsyncState(); throw; } - await GenerateEnclavePackageAsync(_asyncExecutionCancellationToken).ConfigureAwait(false); + await GenerateEnclavePackageAsync(cancellationToken).ConfigureAwait(false); RunExecuteReaderTds( cmdBehavior, diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.cs index 2c218705fe..80ff971ed8 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.cs @@ -158,12 +158,20 @@ public sealed partial class SqlCommand : DbCommand, ICloneable /// flight, or when the command is executing synchronously. /// /// + /// /// The asynchronous execution entry points reach the Always Encrypted machinery through several /// layers of synchronous, callback-driven code that carry no cancellation token. Rather than /// threading a token through every one of those signatures — including the synchronous paths that /// would only ever pass — the token is recorded here for the /// duration of the operation. It is currently consumed only by the asynchronous Always Encrypted /// key store calls, which are the sole cancellable I/O in that machinery. + /// + /// + /// Read this field only on the thread that started the execution, and pass the captured value on as + /// an explicit parameter. The Always Encrypted work is handed to the thread pool and can still be + /// running when the execution's cleanup continuation resets this field, so a read taken after that + /// hand-off may silently observe . + /// /// private CancellationToken _asyncExecutionCancellationToken; diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlQueryMetadataCache.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlQueryMetadataCache.cs index 24dbe48812..a087018289 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlQueryMetadataCache.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlQueryMetadataCache.cs @@ -60,11 +60,12 @@ internal bool GetQueryMetadataIfExists(SqlCommand sqlCommand) return false; } - foreach (SqlCipherMetadata cipherMetadata in metadata.KeysToLoad) + IReadOnlyList keysToLoad = metadata.KeysToLoad; + for (int i = 0; i < keysToLoad.Count; i++) { try { - SqlSecurityUtility.DecryptSymmetricKey(cipherMetadata, sqlCommand.Connection, sqlCommand); + SqlSecurityUtility.DecryptSymmetricKey(keysToLoad[i], sqlCommand.Connection, sqlCommand); } catch (Exception ex) when (ex is SqlException or ArgumentException) { @@ -145,7 +146,9 @@ internal bool TryGetCachedQueryMetadata(SqlCommand sqlCommand, out CachedQueryMe // Create a copy of the cipherMD in order to load the key. // The key shouldn't be loaded in the cached version for security reasons. - List keysToLoad = new(); + // Allocated lazily: a cached query whose parameters are all plaintext needs no list at all, + // and this runs on every metadata cache hit. + List keysToLoad = null; foreach (SqlParameter param in sqlCommand.Parameters) { SqlCipherMetadata cipherMdCopy = null; @@ -165,7 +168,7 @@ internal bool TryGetCachedQueryMetadata(SqlCommand sqlCommand, out CachedQueryMe if (cipherMdCopy is not null) { - keysToLoad.Add(cipherMdCopy); + (keysToLoad ??= new List()).Add(cipherMdCopy); } } @@ -187,12 +190,13 @@ internal async Task CompleteCachedQueryMetadataAsync( CachedQueryMetadata metadata, CancellationToken cancellationToken) { - foreach (SqlCipherMetadata cipherMetadata in metadata.KeysToLoad) + IReadOnlyList keysToLoad = metadata.KeysToLoad; + for (int i = 0; i < keysToLoad.Count; i++) { try { await SqlSecurityUtility.DecryptSymmetricKeyAsync( - cipherMetadata, + keysToLoad[i], sqlCommand.Connection, sqlCommand, cancellationToken).ConfigureAwait(false); @@ -259,7 +263,7 @@ internal readonly struct CachedQueryMetadata internal CachedQueryMetadata(string enclaveLookupKey, List keysToLoad) { EnclaveLookupKey = enclaveLookupKey; - KeysToLoad = keysToLoad; + KeysToLoad = (IReadOnlyList)keysToLoad ?? Array.Empty(); } /// @@ -270,7 +274,7 @@ internal CachedQueryMetadata(string enclaveLookupKey, List ke /// /// The parameter cipher metadata whose column encryption keys still have to be decrypted. /// - internal List KeysToLoad { get; } + internal IReadOnlyList KeysToLoad { get; } } /// diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSymmetricKeyCache.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSymmetricKeyCache.cs index cce481b0a1..de7febfbe8 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSymmetricKeyCache.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSymmetricKeyCache.cs @@ -187,7 +187,10 @@ internal async Task GetKeyAsync(SqlEncryptionKeyInfo keyI return decryptedKey; } - await _cacheLock.WaitAsync(cancellationToken).ConfigureAwait(false); + // Deliberately not cancellable. The expensive, remote part of this method is already done; the + // section below is in-memory only. Honouring cancellation here would throw away a completed key + // decryption and leave the next caller to repeat it. + await _cacheLock.WaitAsync().ConfigureAwait(false); try { // Another caller may have populated the entry while this one was decrypting. In that case the diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/EnclaveDelegateAsyncShould.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/EnclaveDelegateAsyncShould.cs index 6302c5164f..7650bcee2c 100644 --- a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/EnclaveDelegateAsyncShould.cs +++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/EnclaveDelegateAsyncShould.cs @@ -94,7 +94,6 @@ public async Task GetDecryptedKeysToBeSentToEnclaveAsync_UsesAsyncProviderApi() List keys = await EnclaveDelegate.Instance.GetDecryptedKeysToBeSentToEnclaveAsync( NewKeyTable(NewCipherInfoEntry()), - serverName: "async-ae-enclave-unit-test", connection, command: null, CancellationToken.None); @@ -120,7 +119,6 @@ public async Task GetDecryptedKeysToBeSentToEnclaveAsync_ResolvesEveryRequestedK List keys = await EnclaveDelegate.Instance.GetDecryptedKeysToBeSentToEnclaveAsync( NewKeyTable(NewCipherInfoEntry(ordinal: 0), NewCipherInfoEntry(ordinal: 1), NewCipherInfoEntry(ordinal: 2)), - serverName: "async-ae-enclave-unit-test", connection, command: null, CancellationToken.None); @@ -142,7 +140,6 @@ public void GetDecryptedKeysToBeSentToEnclave_UsesSyncProviderApi() List keys = EnclaveDelegate.Instance.GetDecryptedKeysToBeSentToEnclave( NewKeyTable(NewCipherInfoEntry()), - serverName: "async-ae-enclave-unit-test", connection, command: null); @@ -167,7 +164,6 @@ public async Task GetDecryptedKeysToBeSentToEnclaveAsync_WhenCancelled_ProducesC await Assert.ThrowsAnyAsync( () => EnclaveDelegate.Instance.GetDecryptedKeysToBeSentToEnclaveAsync( NewKeyTable(NewCipherInfoEntry()), - serverName: "async-ae-enclave-unit-test", connection, command: null, cts.Token)); @@ -191,7 +187,6 @@ public async Task GetDecryptedKeysToBeSentToEnclaveAsync_WhenProviderFails_Propa SqlException exception = await Assert.ThrowsAsync( () => EnclaveDelegate.Instance.GetDecryptedKeysToBeSentToEnclaveAsync( NewKeyTable(NewCipherInfoEntry()), - serverName: "async-ae-enclave-unit-test", connection, command: null, CancellationToken.None)); diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlSecurityUtilityAsyncShould.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlSecurityUtilityAsyncShould.cs index 83710c614e..e025ca229a 100644 --- a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlSecurityUtilityAsyncShould.cs +++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlSecurityUtilityAsyncShould.cs @@ -383,6 +383,41 @@ public async Task GetKeyAsync_WithConcurrentCacheMisses_FetchesTwiceButPublishes Assert.Equal(0, provider.DecryptCallCount); } + /// + /// A decryption that has already succeeded must not be thrown away just because the caller + /// cancelled while the result was being published. The remote work is done at that point, so + /// discarding it would only force the next caller to repeat it. + /// + [Fact] + public async Task GetKeyAsync_WhenCancelledAfterDecryptSucceeds_StillPublishesTheKey() + { + SqlSymmetricKeyCache cache = SqlSymmetricKeyCache.GetInstance(); + + using CancellationTokenSource cts = new CancellationTokenSource(); + TestKeyStoreProvider provider = new TestKeyStoreProvider + { + DecryptAsyncCallback = _ => + { + // Cancel between a successful provider call and the cache insertion. + cts.Cancel(); + return Task.FromResult(NewPlaintextKey(seed: 57)); + } + }; + + using SqlConnection connection = NewConnection((ProviderName, provider)); + SqlEncryptionKeyInfo keyInfo = NewKeyInfo(); + + SqlClientSymmetricKey key = await cache.GetKeyAsync(keyInfo, connection, command: null, cts.Token); + + Assert.NotNull(key); + Assert.Equal(1, provider.DecryptAsyncCallCount); + + // The key was published despite the cancellation, so no second decryption is needed. + Assert.Same(key, cache.GetKey(keyInfo, connection, command: null)); + Assert.Equal(1, provider.DecryptAsyncCallCount); + Assert.Equal(0, provider.DecryptCallCount); + } + /// /// Verifies that cancellation of GetKeyAsync is observed before the provider call and while it /// is in flight, and that nothing is published to the cache in either case. From a9a072e1f7d265cec10fa2655ac12c040d87794f Mon Sep 17 00:00:00 2001 From: Cheena Malhotra Date: Mon, 24 Aug 2026 12:22:05 -0700 Subject: [PATCH 7/7] Do not evict cached query metadata when a key load is cancelled CompleteCachedQueryMetadataAsync treated every exception from the column encryption key load as evidence that the cached metadata was stale, and invalidated the entry. Cancellation is not evidence of anything of the sort, so cancelling an execution threw away a perfectly good cache entry and made the next caller pay for another describe parameter encryption round trip. Cancellation now propagates without touching the cache. This failure mode cannot arise on the synchronous path, so only the asynchronous overload handles it and the two stay otherwise identical. The existing cancellation test only asserted that the operation threw, which is why this went unnoticed. It now also asserts that the entry survives and that a subsequent uncancelled load still succeeds, and it fails without the fix. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 82aca8ca-2078-4a47-860c-6f6ee78cbb02 --- .../Data/SqlClient/SqlQueryMetadataCache.cs | 8 ++++++++ .../SqlQueryMetadataCacheAsyncShould.cs | 16 +++++++++++++++- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlQueryMetadataCache.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlQueryMetadataCache.cs index a087018289..5eb5652115 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlQueryMetadataCache.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlQueryMetadataCache.cs @@ -207,6 +207,14 @@ await SqlSecurityUtility.DecryptSymmetricKeyAsync( OnKeyLoadFailed(sqlCommand, clearParameterMetadata: true); return false; } + catch (OperationCanceledException) + { + // Cancellation says nothing about whether the cached metadata is still valid, so + // leave the entry in place for the next caller instead of forcing it to pay for + // another describe parameter encryption round trip. This case cannot arise on the + // synchronous path, which is why only this overload handles it. + throw; + } catch (Exception) { OnKeyLoadFailed(sqlCommand, clearParameterMetadata: false); diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs index 3b830f2da3..3ada42de24 100644 --- a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs +++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs @@ -144,8 +144,13 @@ public void GetQueryMetadataIfExists_StillUsesSyncKeyStoreProviderApi() Assert.True(command.Parameters[0].CipherMetadata.IsAlgorithmInitialized()); } + /// + /// Cancellation says nothing about whether the cached metadata is still valid, so a cancelled + /// key load must not report a hit and must not evict the entry. Evicting it would make the next + /// caller pay for a full describe parameter encryption round trip. + /// [Fact] - public async Task CompleteCachedQueryMetadataAsync_WhenCancelled_DoesNotReportHit() + public async Task CompleteCachedQueryMetadataAsync_WhenCancelled_DoesNotReportHitOrEvictTheEntry() { TestKeyStoreProvider provider = new TestKeyStoreProvider { PlaintextKey = NewPlaintextKey(seed: 43) }; SqlCommand command = NewCachedCommand(provider); @@ -159,6 +164,15 @@ public async Task CompleteCachedQueryMetadataAsync_WhenCancelled_DoesNotReportHi await Assert.ThrowsAnyAsync( () => SqlQueryMetadataCache.GetInstance() .CompleteCachedQueryMetadataAsync(command, metadata, cts.Token)); + + // The entry must still be there, and a subsequent uncancelled load must succeed. + Assert.True(SqlQueryMetadataCache.GetInstance() + .TryGetCachedQueryMetadata(command, out SqlQueryMetadataCache.CachedQueryMetadata retryMetadata)); + + Assert.True(await SqlQueryMetadataCache.GetInstance() + .CompleteCachedQueryMetadataAsync(command, retryMetadata, CancellationToken.None)); + + Assert.True(command.Parameters[0].CipherMetadata.IsAlgorithmInitialized()); } ///