Implement async APIs in Always Encrypted Azure Key Vault provider - #4540
Open
cheenamalhotra wants to merge 7 commits into
Open
Implement async APIs in Always Encrypted Azure Key Vault provider#4540cheenamalhotra wants to merge 7 commits into
cheenamalhotra wants to merge 7 commits into
Conversation
Overrides the four async SqlColumnEncryptionKeyStoreProvider APIs in SqlColumnEncryptionAzureKeyVaultProvider with truly asynchronous Azure Key Vault SDK calls, so Always Encrypted async paths no longer block on HTTP I/O (issue #3672, spec phase 2A). - AzureSqlKeyCryptographer: async counterparts for AddKey, SignData, VerifyData, WrapKey and UnwrapKey, all propagating a CancellationToken. AddKeyAsync fetches before locking so no lock is held during network I/O. - LocalCache: GetOrCreateAsync with an async factory, mirroring sync semantics (TTL bypass, compaction, expiration) without holding a lock during I/O. - SqlColumnEncryptionAzureKeyVaultProvider: async overrides for encrypt, decrypt, sign and verify. Blob parse/build logic extracted into shared helpers so sync and async paths stay identical. Sync behavior is unchanged. - AsyncEventScope: reference-type event scope, since SqlClientEventScope is a ref struct and cannot cross an await boundary. - Tests: async round-trip, sync/async interoperability, cache behavior, cancellation and argument validation coverage in the AKV manual tests. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
The reference-type scope wrapper only existed because SqlClientEventScope is a ref struct. Tracking the scope id in a try/finally achieves the same tracing without a new type or an allocation per async call. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Addresses review feedback on sync/async behavioral parity: - LocalCache.GetOrCreateAsync now gates concurrent misses per key, so a burst of concurrent decryptions of the same key issues a single Azure Key Vault request instead of one per caller. Misses for different keys still proceed in parallel, cancellation stays per caller, and a failed or cancelled owner lets the next waiter retry with its own token. - AzureSqlKeyCryptographer.AddKeyAsync double-checks under the semaphore and fetches while holding it, mirroring the deduplication of AddKey. Previously the semaphore only guarded a ConcurrentDictionary write, and a token cancelled mid-flight discarded an already fetched key. - LocalCache.GetOrCreate compacts on Count >= maxSize rather than ==, so a count that overshoots the limit cannot disable compaction permanently. - The async overrides observe the cancellation token before validating arguments, matching SqlColumnEncryptionKeyStoreProvider. - Documented that async argument validation failures surface through the returned task rather than being thrown synchronously. - Tests for concurrent decryption deduplication and for cancellation taking precedence over argument validation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Follow-up to the async provider work, addressing defects found while auditing it for performance and compatibility problems. AddKeyAsync held _keyDictionarySemaphore across the awaited Azure Key Vault fetch. That semaphore is shared with the synchronous AddKey, so a synchronous caller blocked a thread pool thread for the duration of an asynchronous network round trip, which risks thread pool starvation. The semaphore is also global, so fetching one key serialized fetching every other key. The asynchronous path now uses its own per key gate and leaves the synchronous path on its original semaphore. A synchronous and an asynchronous caller may both fetch the same key, which yields an identical result, and this matches the deliberate absence of cross path deduplication in LocalCache. LocalCache.GetOrCreateAsync published its gate before awaiting it, and the try/finally that removed the gate began after the await. A cancelled wait therefore left the gate behind permanently, and the gate dictionary is not bounded by the cache size limit. A loop of a thousand pre-cancelled calls on distinct keys retained a thousand gates. The gate lifetime is now managed by KeyedAsyncLock, which removes the gate when a wait is abandoned and cleans up through a disposable releaser. The per key gating logic now lives in KeyedAsyncLock rather than being repeated, so the release and cleanup ordering is defined in one place. GetCryptographyClient used TryGetValue followed by TryAdd, so concurrent callers could each use a different CryptographyClient instance for the same key. It now uses GetOrAdd, and all callers observe the instance that wins the race. Also documents that gating is bypassed when caching is disabled, and adds a regression test asserting that cancelled asynchronous decryptions do not accumulate creation gates. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
EncryptColumnEncryptionKeyAsync repeated the body of ValidateSignature inline, including both of its trace messages, so a change to one would have silently diverged from the other. The logic now lives in ValidateSignatureAsync next to its synchronous counterpart. ParseEncryptedColumnEncryptionKey carried a null check on a buffer that had just been allocated with new byte[], which no execution can reach. The check moved into the helper when the parsing logic was extracted, and is now dropped. ADP.NullHashFound is left in place because removing it would also strip the associated resource string for no functional gain. Neither change alters behavior. Comparing every statement of the original synchronous encrypt and decrypt paths against the current file confirms all of them are preserved except the unreachable null check and a dead store to a position variable that was never read after its final update. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
9 tasks
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds truly asynchronous Always Encrypted key store operations to the Azure Key Vault provider, wiring the provider’s async overrides to Azure SDK async calls and extending local caching with async entry creation + per-key deduplication gates to avoid bursty duplicate Key Vault requests.
Changes:
- Implemented async overrides in
SqlColumnEncryptionAzureKeyVaultProviderfor encrypt/decrypt and CMK metadata sign/verify, flowingCancellationTokento Azure SDK async APIs. - Added async-capable
LocalCache.GetOrCreateAsyncplus a per-keyKeyedAsyncLock<TKey>to deduplicate concurrent cache misses without blocking threads. - Expanded AKV manual tests to cover async round-trips, cache sharing between sync/async, cancellation, and concurrency semantics.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| src/Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider/src/SqlColumnEncryptionAzureKeyVaultProvider.cs | Adds async overrides for AE AKV provider and extracts shared parsing/message helpers for CEK/signature handling. |
| src/Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider/src/LocalCache.cs | Introduces async cache entry creation with per-key gating and fixes compaction threshold under concurrency. |
| src/Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider/src/KeyedAsyncLock.cs | New helper providing per-key async mutual exclusion with gate cleanup. |
| src/Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider/src/AzureSqlKeyCryptographer.cs | Adds async key fetch/sign/verify/wrap/unwrap APIs and deduplicates concurrent key fetches per key. |
| src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/AKVUnitTests.cs | Adds manual tests for async API behavior, caching/deduplication, and cancellation semantics. |
| src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/ExceptionTestAKVStore.cs | Adds manual tests for async encrypt/decrypt argument validation and decrypt failure modes. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
CancelledAsyncDecryptionsDoNotAccumulateCreationGates cancelled its token before calling DecryptColumnEncryptionKeyAsync. Cancellation is observed before the cache is reached, so no gate was ever created and the assertion held trivially. The test now has one caller take the gate and hold it across the key vault round trip while other callers queue behind it and are cancelled while waiting, which is the path where an abandoned wait could strand a gate. ExceptionTestAKVStore covered argument validation for the asynchronous encrypt and decrypt members only. Adds the same coverage for the asynchronous sign and verify members, mirroring the existing SignInvalidAKVPath cases. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.
Suppressed comments (1)
src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/ExceptionTestAKVStore.cs:84
- These assertions use Assert.Matches with a plain message string that includes regex metacharacters (e.g., the trailing '.' in the null-path case). That makes the check less precise than intended and can produce false positives. Since the goal is substring matching (prefix may vary), Assert.Contains is a better fit here.
string expectedMessage = masterKeyPath == null
? "Azure Key Vault key path cannot be null."
: "Invalid Azure Key Vault key path specified";
Assert.Matches(expectedMessage, signException.Message);
Assert.Matches(expectedMessage, verifyException.Message);
VerifyColumnMasterKeyMetadata and VerifyColumnMasterKeyMetadataAsync now reject a null or empty signature with ArgumentNullException/ArgumentException instead of deferring the failure to the Azure Key Vault SDK. Both overloads validate identically so the sync and async surfaces stay in parity. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: ff71f86b-9c50-45f7-b79f-5aaaf7f98289
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #4540 +/- ##
==========================================
- Coverage 64.78% 62.79% -1.99%
==========================================
Files 288 283 -5
Lines 44418 67410 +22992
==========================================
+ Hits 28774 42332 +13558
- Misses 15644 25078 +9434
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Implements the four asynchronous key store provider APIs on
SqlColumnEncryptionAzureKeyVaultProvider, overriding the base class virtualsadded in #3673. Phase 2A of
specs/002-async-always-encrypted/spec.md.New
EncryptColumnEncryptionKeyAsync,DecryptColumnEncryptionKeyAsync,SignColumnMasterKeyMetadataAsync,VerifyColumnMasterKeyMetadataAsync.These call the Azure SDK's own async methods and flow the cancellation token
to them, rather than completing sync work on a returned task.
LocalCache.GetOrCreateAsync, plus aKeyedAsyncLock<TKey>helper that gatesconcurrent misses per key so a burst of callers makes one Key Vault request.
duplicated.
Behavior considerations
by one is visible to the other.
different keys stay parallel. Cancellation applies to the requesting caller
only; if the gate owner is cancelled or fails, the next waiter retries with
its own token and failures are not cached.
AddKeyAsyncdeliberately does not share_keyDictionarySemaphorewith syncAddKey. A sync caller blocking on a gate held across an awaited network callwould tie up a thread pool thread for that call's duration. Consequence: a
sync and an async caller may both fetch the same key, yielding the same result.
SqlColumnEncryptionKeyStoreProvider.synchronously, matching FR-003.
ColumnEncryptionKeyCacheTtlof zero) gating isbypassed, since there is no entry for a waiter to observe. Callers reach Key
Vault in parallel where the sync path serialized them.
VerifyColumnMasterKeyMetadataandVerifyColumnMasterKeyMetadataAsyncnowboth reject a null or empty
signaturewithArgumentNullException/ArgumentException. Previously it reached the Azure SDK and failed there. Thisis a deliberate behavior change to the existing sync API, kept in both
overloads for parity; in-product callers are unaffected because
SqlSecurityUtility.VerifyColumnMasterKeySignaturealready rejects it upstream.Worth a release note callout.
covers restore, but a runtime downgrade below 7.1 produces a
TypeLoadExceptionbecause assembly versions unify atmajor.0.0.0. Worth arelease note callout.
Incidental fixes in code the refactor touched
LocalCache.GetOrCreatecompacts onCount >= maxSizerather than==; theequality test could be stepped past under concurrency, permanently disabling
compaction on the 2000 entry signature cache.
GetCryptographyClientusedTryGetValuethenTryAdd, so concurrentcallers could each use a different
CryptographyClientfor one key. NowGetOrAdd.new byte[]on thepreceding line.
No public API removed or changed.
Issues
Addresses #3672 (Step 2)
Testing
AKVUnitTests: async encrypt/decrypt and sign/verify round trips; sync andasync keys interchangeable; caching during async decryption and sharing with the
sync path; caching disabled at TTL zero; signature verification caching; 32
concurrent decryptions collapsing to one cache entry; cancelled decryptions not
accumulating gates; cancellation honoured and taking precedence over validation;
master key path validation.
ExceptionTestAKVStore: argument validation for all four members, plus invalidalgorithm version, invalid signature and invalid cipher text length.
These need a live vault and are gated on
DataTestUtility.IsAKVSetupAvailable,so they run in the pipeline.
LocalCacheandKeyedAsyncLockconcurrency wasadditionally verified locally against the built assembly with a standalone
harness covering deduplication, parallelism across keys, cancelled waiters,
owner failure and retry, gate cleanup and compaction.
Sync behavior preservation was checked by comparing every statement of the
original sync encrypt and decrypt methods against the current file. All are
preserved except the unreachable null check above and a dead store to a
positionvariable never read after its final update.