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

[cheenamalhotra]

Description

SqlColumnEncryptionKeyStoreProvider gained async APIs in #4540, but nothing in the driver called them. Every Always Encrypted key store operation still went through the synchronous overloads, so an ExecuteReaderAsync against a key store that performs network I/O (Azure Key Vault, for example) blocked a thread for the duration of the round trip. This PR wires the async APIs into the Always Encrypted execution paths so async callers get a genuine end-to-end async experience.

Async utility layer

Adds async counterparts to the Always Encrypted utility layer, mirroring their synchronous versions:

• SqlSecurityUtility.DecryptSymmetricKeyAsync (both the SqlCipherMetadata and SqlTceCipherInfoEntry overloads), GetKeyFromLocalProvidersAsync, VerifyColumnMasterKeySignatureAsync
• SqlSymmetricKeyCache.GetKeyAsync, which never holds the cache lock across provider I/O. Concurrent decryptions of the same key are allowed to race, and the first writer wins.
• EnclaveDelegate.GetDecryptedKeysToBeSentToEnclaveAsync and GenerateEnclavePackageAsync

Each async method shares its parsing, validation, and error-wrapping logic with the existing synchronous version, so provider failures still surface as the same wrapped SqlException and multi-key fallback still behaves identically. The one deliberate divergence: DecryptSymmetricKeyAsync rethrows OperationCanceledException immediately instead of continuing to try the remaining candidate keys.

Call site integration

The describe-parameter-encryption pipeline was restructured so that sync and async share one parser. ReadDescribeEncryptionParameterResultsCore now records the signature verifications and key decryptions it discovers into a PendingColumnEncryptionKeyOperations list rather than performing them inline; thin sync and async wrappers then execute that work. This avoids duplicating roughly 350 lines of TDS token parsing, and has a useful side effect: no key store call now happens while the describe reader is mid-result-set.

GenerateEnclavePackage and the post-describe execution continuation gained async counterparts built on shared helpers (TryPrepareEnclavePackageGeneration, BuildEnclavePackage, CreateColumnEncryptionKeyInfo), so both modes perform identical validation and throw identical exceptions.

Cancellation

The async provider APIs all accept a CancellationToken, but the Always Encrypted call sites had no way to reach the token supplied to ExecuteReaderAsync / ExecuteNonQueryAsync / ExecuteXmlReaderAsync. The token is now captured in a per-execution field on SqlCommand, set at each async entry point and cleared by the matching cleanup continuation, and passed down into the key operations and enclave package generation.

A field is used rather than signature plumbing because the token would otherwise have to thread through several sync-only overloads of RunExecuteReader and PrepareForTransparentEncryption. This mirrors how existing per-execution mutable state such as _cachedAsyncState is handled.

Query metadata cache

On a metadata cache hit the driver still loaded every column encryption key synchronously on the caller's thread. The lookup is now split into an in-memory probe (TryGetCachedQueryMetadata), which matches cached cipher metadata onto the parameters and reports which keys still need decrypting, and a completion step that loads them. The synchronous lookup is expressed in terms of the same probe, so its behaviour, including stale-key fallback and the hit/miss counters, is unchanged.

Notes for reviewers

A few things worth a careful look:

• usedCache is over-reported on the async cache-hit path. When cached key information turns out to be stale, the returned task falls back to a full describe round trip, but the command has already reported usedCache = true because the caller returned before the fallback was discovered. This is deliberately conservative: over-reporting a cache hit can only cost one extra retry of an already-failing execution, whereas under-reporting it would suppress the TCE_CONVERSION_ERROR_CLIENT_RETRY retry that a genuine cache hit depends on.
• Task.Run is retained on purpose in GetParameterEncryptionDataReader and RunExecuteReaderTdsWithTransparentParameterEncryption. Removing it would be a regression rather than a simplification: PrepareForTransparentEncryption runs synchronously on the caller's thread under ExecuteReaderAsync, and both bodies issue blocking TDS writes before their first await. Task.Run reproduces the thread pool hand-off the previous ContinueWith chain provided and keeps continuations off network callback threads.
• Error semantics were preserved precisely. AsyncHelper.CreateContinuationTaskWithState invokes its failure callback only when the antecedent actually faulted, not on cancellation, hence the explicit catch (OperationCanceledException) { throw; } ahead of the ResetAsyncState() handler.
• Side-effect ordering changed slightly in the describe pipeline: keysToBeSentToEnclave and requiresEnclaveComputations are now populated before a failing signature check rather than after. This is safe because ResetEncryptionState() runs at the start of every execution.
• Async cache hits now cost one extra thread pool dispatch. That is the price of not issuing the command's TDS write from the caller's thread. Sync callers are unaffected, and when the CEK is already in the symmetric key cache the completion runs synchronously anyway.
• A duplicated #if DEBUG failpoint block was replaced with a [Conditional("DEBUG")] method. The call is elided in Release exactly as before, and the static fields that ManualTests reflect on are untouched.

No public API changes. No behaviour change for synchronous callers.

Behaviour notes worth a careful look

Exception ordering on malformed describe results. Key store work is now deferred until
after the sp_describe_parameter_encryption result set has been fully read, so that no key
store HTTP happens while the describe reader is mid-result-set. A side effect is that if a
result set is both malformed and contains a bad key, the parsing error now surfaces first
where the key store error used to. Both cases still throw; only which exception wins has
changed.

Task.Run in GetParameterEncryptionDataReaderAsync is deliberate. It looks like a
removable thread-pool hop, but the task being waited on is completed by the SNI network write
callback, so an inline continuation would run blocking TDS reads on a callback thread. There
is a comment on the method saying so.

Column encryption key decryption stays sequential. Distinct cipher metadata entries often
share a key, and GetKeyAsync does not hold its gate across provider I/O, so decrypting in
parallel would let concurrent misses for the same key each issue their own key store call.
Sequential decryption lets the first result warm the cache.

Scope

Async enclave attestation is not in this PR. It adds virtual async members to the public
SqlColumnEncryptionEnclaveProvider hierarchy, so it needs reference assembly updates and API
review, and ships separately. The one remaining blocking call on the async path,
GetEnclaveSession, is marked with a @TODO pointing at that work.

Issues

Part of the async Always Encrypted work tracked by #3672, following on from #4540.

Testing

Unit tests (src/Microsoft.Data.SqlClient/tests/UnitTests/):

• SqlSecurityUtilityAsyncShould (12 tests) covers DecryptSymmetricKeyAsync multi-key fallback, cancellation, provider failure wrapping, VerifyColumnMasterKeySignatureAsync, and SqlSymmetricKeyCache.GetKeyAsync cache behaviour.
• EnclaveDelegateAsyncShould (5 tests) proves the enclave path uses the async provider API, the sync path still uses the sync API, cancellation produces a cancelled task with no provider call, and failures wrap correctly.
• SqlQueryMetadataCacheAsyncShould (4 tests) covers async API usage on cache completion, stale-key degradation to a miss, cancellation, and a regression guard that the sync lookup still uses the sync provider API.

Manual test (AsyncKeyStoreProviderTests): registers a recording key store provider and asserts that sync execution hits the sync overload, that ExecuteReaderAsync / ExecuteNonQueryAsync hit the async overload and never the sync one, and that cancelling the caller's token cancels the key store operation. This test requires a live SQL Server plus key store, so it is not part of local verification.

Verified locally: Debug and Release builds of the driver on net8.0 with 0 warnings, ManualTests and FunctionalTests compile, and 160 targeted unit tests pass. Three failures in NativeColumnEncryptionKeyBaseline are pre-existing on macOS and unrelated (they fail in Interop.AppleCrypto.X509StoreAddCertificate against the system keychain).

Not verified locally: end-to-end Always Encrypted behaviour against a real server and Azure Key Vault. That validation is being done separately, which is why this is a draft.

Guidelines

Please review the contribution guidelines before submitting a pull request:

• Contributing
• Code of Conduct
• Best Practices
• Coding Style
• Review Process

[Copilot]

Pull request overview

This PR wires the async SqlColumnEncryptionKeyStoreProvider APIs into the Always Encrypted async execution paths so ExecuteReaderAsync / ExecuteNonQueryAsync / ExecuteXmlReaderAsync no longer block threads on key store network I/O (e.g., Azure Key Vault), while preserving sync behavior and error semantics.

Changes:

• Adds async counterparts in the Always Encrypted utility/enclave layers (e.g., SqlSecurityUtility.DecryptSymmetricKeyAsync, SqlSymmetricKeyCache.GetKeyAsync, EnclaveDelegate.GenerateEnclavePackageAsync).
• Refactors describe-parameter-encryption processing to share parsing between sync/async and to defer key-store operations until after result-set parsing.
• Flows the caller’s CancellationToken into AE key-store operations via per-execution state on SqlCommand, and updates query metadata cache to split probe vs. async completion.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSecurityUtility.cs	Adds async key decryption + CMK signature verification paths with cancellation and consistent wrapping.
src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSymmetricKeyCache.cs	Adds GetKeyAsync with lock-release around provider I/O to avoid blocking sync callers.
src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlQueryMetadataCache.cs	Splits cache lookup into in-memory probe + awaitable completion that loads CEKs asynchronously.
src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.cs	Introduces per-execution _asyncExecutionCancellationToken for async AE flows.
src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs	Refactors describe-parameter-encryption processing to defer key-store ops and add async completion.
src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs	Adds async enclave package generation and replaces continuation chains with async/await + Task.Run dispatch.
src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.NonQuery.cs	Captures/clears per-execution cancellation token for async non-query execution.
src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Xml.cs	Captures/clears per-execution cancellation token for async XML reader execution.
src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.cs	Adds async key resolution for enclave package generation and shares key validation/projection.
src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.Crypto.cs	Adds GenerateEnclavePackageAsync and shared helper for building the enclave package.
src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlSecurityUtilityAsyncShould.cs	Adds unit tests covering async decryption/signature verification/cancellation/cache behavior.
src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/EnclaveDelegateAsyncShould.cs	Adds unit tests asserting enclave async path uses async provider APIs and honors cancellation.
src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs	Adds unit tests for async metadata cache completion and stale-key degradation.
src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/AsyncKeyStoreProviderTests.cs	Adds manual integration test verifying sync vs async provider API usage and token flow.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 14 out of 14 changed files in this pull request and generated no new comments.

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.cs:160

• _asyncExecutionCancellationToken is documented as representing the token “currently in flight”, but it looks like it’s only cleared in the async cleanup continuations. If an async entry point throws during setup (e.g., Task.Factory.FromAsync throws before the continuation runs), the token can remain set on the SqlCommand instance, unnecessarily retaining the caller’s CancellationTokenSource and violating the stated invariant. Consider clearing this field in the synchronous exception paths of InternalExecuteReaderAsync / InternalExecuteNonQueryAsync / InternalExecuteXmlReaderAsync as well (before returning a faulted Task).

        /// <summary>
        /// Cancellation token supplied by the caller of the asynchronous execution that is currently in
        /// flight, or <see cref="CancellationToken.None"/> when the command is executing synchronously.
        /// </summary>
        /// <remarks>

[Copilot]

Pull request overview

Copilot reviewed 14 out of 14 changed files in this pull request and generated no new comments.

Suppressed comments (4)

Previously missed (1) — in code that hasn't changed since the last review.

src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs:90

• SqlCommand (and its underlying SqlConnection) should be disposed in this test to avoid leaking resources across the unit test process; use a using declaration for the cached command instance.

This issue also appears in the following locations of the same file:

• line 113
• line 137
• line 150

            TestKeyStoreProvider provider = new TestKeyStoreProvider { PlaintextKey = NewPlaintextKey(seed: 41) };
            SqlCommand command = NewCachedCommand(provider);

src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs:115

• This test creates a SqlCommand with an owned SqlConnection but never disposes it; prefer a using declaration so the connection/command are cleaned up even if the test fails.

            };
            SqlCommand command = NewCachedCommand(provider);

src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs:152

• Use a using declaration for the SqlCommand created by NewCachedCommand so the command/connection are disposed deterministically.

            TestKeyStoreProvider provider = new TestKeyStoreProvider { PlaintextKey = NewPlaintextKey(seed: 43) };
            SqlCommand command = NewCachedCommand(provider);

src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlQueryMetadataCacheAsyncShould.cs:139

• Dispose the SqlCommand created for this test (it owns a SqlConnection) to avoid resource leakage across the unit test run.

            TestKeyStoreProvider provider = new TestKeyStoreProvider { PlaintextKey = NewPlaintextKey(seed: 47) };
            SqlCommand command = NewCachedCommand(provider);

[Copilot]

Pull request overview

Copilot reviewed 14 out of 14 changed files in this pull request and generated 1 comment.

[cheenamalhotra]

Good catch on the eviction, fixed in a9a072e.

CompleteCachedQueryMetadataAsync now has an OperationCanceledException catch ahead of the general one that rethrows without touching the cache. Cancellation is not evidence that the cached metadata went stale, so evicting made the next caller pay for another sp_describe_parameter_encryption round trip for no reason. This failure mode cannot arise on the synchronous path, so only the async overload handles it and the two stay otherwise identical.

One correction: misses were not being incremented. The general catch calls OnKeyLoadFailed(clearParameterMetadata: false), and IncrementCacheMisses sits behind the early return that flag guards, so only the invalidation was happening.

The reason this slipped through is that the existing cancellation test only asserted that the call threw. It now also asserts the entry survives and that a subsequent uncancelled load succeeds, and I verified it fails against the old code at that assertion.

[Copilot]

Pull request overview

Copilot reviewed 14 out of 14 changed files in this pull request and generated no new comments.

Suppressed comments (2)

Previously missed (2) — in code that hasn't changed since the last review.

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSymmetricKeyCache.cs:155

• The comment says the provider exception will be “simply bubbled up”, but this method wraps non-cancellation failures in SQL.KeyDecryptionFailed. Please update the comment so it matches the actual error behavior (only OperationCanceledException is propagated unwrapped).

            // Decrypt the CEK
            // We will simply bubble up the exception from the DecryptColumnEncryptionKeyAsync function.
            byte[] plaintextKey;

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlSecurityUtility.cs:412

• The comment says the provider exception will be “simply bubbled up”, but this helper wraps non-cancellation failures in SQL.KeyDecryptionFailed. Please update the comment so it matches the actual behavior (only OperationCanceledException is propagated unwrapped).

            // Decrypt the CEK
            // We will simply bubble up the exception from the DecryptColumnEncryptionKeyAsync function.
            byte[] plaintextKey;