Implement async APIs in Always Encrypted Azure Key Vault provider

[cheenamalhotra]

Description

Implements the four asynchronous key store provider APIs on
SqlColumnEncryptionAzureKeyVaultProvider, overriding the base class virtuals
added 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 a KeyedAsyncLock<TKey> helper that gates
  concurrent misses per key so a burst of callers makes one Key Vault request.
• Shared parse and build helpers extracted from the sync methods rather than
  duplicated.

Behavior considerations

• The CEK and signature caches are shared with the sync path, so a key decrypted
  by one is visible to the other.
• Gating is per key and only ever awaited, so no thread blocks and misses for
  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.
• AddKeyAsync deliberately does not share _keyDictionarySemaphore with sync
  AddKey. A sync caller blocking on a gate held across an awaited network call
  would 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.
• Cancellation is checked before argument validation, matching
  SqlColumnEncryptionKeyStoreProvider.
• Validation failures surface through the returned task, not thrown
  synchronously, matching FR-003.
• With caching disabled (ColumnEncryptionKeyCacheTtl of zero) gating is
  bypassed, since there is no entry for a waiter to observe. Callers reach Key
  Vault in parallel where the sync path serialized them.
• VerifyColumnMasterKeyMetadata and VerifyColumnMasterKeyMetadataAsync now
  both reject a null or empty signature with ArgumentNullException /
  ArgumentException. Previously it reached the Azure SDK and failed there. This
  is a deliberate behavior change to the existing sync API, kept in both
  overloads for parity; in-product callers are unaffected because
  SqlSecurityUtility.VerifyColumnMasterKeySignature already rejects it upstream.
  Worth a release note callout.
• Requires Microsoft.Data.SqlClient 7.1 or later at runtime. The NuGet floor
  covers restore, but a runtime downgrade below 7.1 produces a
  TypeLoadException because assembly versions unify at major.0.0.0. Worth a
  release note callout.

Incidental fixes in code the refactor touched

• LocalCache.GetOrCreate compacts on Count >= maxSize rather than ==; the
  equality test could be stepped past under concurrency, permanently disabling
  compaction on the 2000 entry signature cache.
• GetCryptographyClient used TryGetValue then TryAdd, so concurrent
  callers could each use a different CryptographyClient for one key. Now
  GetOrAdd.
• Dropped an unreachable null check on a buffer allocated by new byte[] on the
  preceding line.

No public API removed or changed.

Issues

Addresses #3672 (Step 2)

Testing

AKVUnitTests: async encrypt/decrypt and sign/verify round trips; sync and
async 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 invalid
algorithm 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. LocalCache and KeyedAsyncLock concurrency was
additionally 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
position variable never read after its final update.

[Copilot]

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 SqlColumnEncryptionAzureKeyVaultProvider for encrypt/decrypt and CMK metadata sign/verify, flowing CancellationToken to Azure SDK async APIs.
• Added async-capable LocalCache.GetOrCreateAsync plus a per-key KeyedAsyncLock<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.

[Copilot]

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);

[Copilot]

Pull request overview

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 62.79%. Comparing base (ee529d4) to head (9e0b093).
⚠️ Report is 2 commits behind head on main.

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     

Flag	Coverage Δ
CI-SqlClient	?
PR-SqlClient-Project	62.79% <ø> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[priyankatiwari08]

Overall looks good, just few small comments need to be addressed.

[priyankatiwari08]

This says mutual exclusion, but RemoveIfUnused can evict a gate while another caller is sitting between GetOrAdd and WaitAsync. That caller then acquires the orphaned gate while a later caller creates a fresh one, so two callers can hold the lock for one key at the same time.

That's fine here (worst case is one extra Key Vault fetch), and the RemoveIfUnused doc comment already says as much - but the class summary is the part someone reusing this helper will read. Can we describe it as best-effort deduplication here too?

[cheenamalhotra]

Agreed — the class summary was the misleading part. Reworded in f0e88d8.

The summary now says best-effort deduplication rather than mutual exclusion, and the remarks call out the exact race: a caller sitting between GetOrAdd and WaitAsync can acquire a gate that has just been reclaimed while a later caller creates a fresh one, so two callers can run the guarded work for one key. It also states the precondition that makes that acceptable — the guarded work must be idempotent and duplication merely wasteful, which holds for the key store fetches here — and warns against reuse where exclusion must be absolute.

[priyankatiwari08]

Releaser is a struct, so a copy disposed twice would push CurrentCount to 2 and silently break exclusion for that key. Nothing does that today - every use site is a using on the awaited result - but it's cheap to make impossible. Either guard the release, or document that the releaser must be disposed exactly once.

[cheenamalhotra]

Made it impossible rather than documented, in f0e88d8.

Releaser is now a sealed class instead of a readonly struct, and Dispose does Interlocked.Exchange(ref _gate, null) and returns early if the gate was already taken. A second disposal — or a disposal of a copy, which can no longer exist — is now a no-op, so CurrentCount cannot be pushed above 1. This also brings it in line with the IDisposable contract, which requires disposal to be idempotent.

The cost is one small allocation per acquisition, which is negligible next to the Key Vault round trip it guards. Both call sites (AzureSqlKeyCryptographer and LocalCache) already using the awaited result, so they are unchanged.

[Copilot]

Pull request overview

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

Suppressed comments (2)

src/Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider/src/SqlColumnEncryptionAzureKeyVaultProvider.cs:248

• The cache factory delegate passed to GetOrCreateAsync is a local function that has the same name as the containing method (VerifyColumnMasterKeyMetadataAsync). This name collision makes it easy to misread which method is being invoked, and future refactors could accidentally introduce recursion or bind the wrong method group. Rename the local function to a distinct name (e.g., CoreAsync) and update the delegate reference.

                Tuple<string, bool, string> key = Tuple.Create(masterKeyPath, allowEnclaveComputations, ToHexString(signature));
                return await _columnMasterKeyMetadataSignatureVerificationCache
                    .GetOrCreateAsync(key, VerifyColumnMasterKeyMetadataAsync, cancellationToken)
                    .ConfigureAwait(false);
            }

src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/AKVUnitTests.cs:345

• In CancelledAsyncDecryptionsDoNotAccumulateCreationGates, the CancellationTokenSource is cancelled immediately after scheduling the waiter tasks. Because DecryptColumnEncryptionKeyAsync checks cancellation before reaching LocalCache.GetOrCreateAsync, some/all waiters can cancel before they ever queue on the keyed gate, making the test non-deterministic and potentially vacuous. Add a yield (or small delay) before cancelling to better exercise the “cancel while waiting on the gate” scenario this test is asserting.

                    waiters[j] = Task.Run(() => akvProvider.DecryptColumnEncryptionKeyAsync(
                        _fixture.GeneratedKeyUri, EncryptionAlgorithm, encryptedKey, cts.Token));
                }

                cts.Cancel();

[benrr101]

Blocking on unit tests (preferably arrange/act/assert style) for KeyedAsyncLock class. Consideration of using a dictionary of Lazy objects would be nice, but not a blocker. Otherwise, looks pretty good.

[benrr101]

With a new class like this, I'd really like to see new unit tests for it!

[cheenamalhotra]

Agreed — added in 2802672. The AKV provider had no unit test project at all, so this adds one:

src/Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider/test/, modelled on the Extensions test projects (net462;net8.0;net9.0, xunit, InternalsVisibleTo $(AssemblyName).Test). Wired into build.proj as TestAkvProvider / BuildAkvProviderTests, added to the Test and BuildTests aggregates, the .slnx, the PR test stage matrix, and the target table in BUILDGUIDE.md. It is now the natural home for any future AKV unit tests.

KeyedAsyncLockTest has 10 tests: uncontended acquire, second caller blocked until release, different keys not blocking each other, 32 concurrent callers on one key never overlapping, pre-cancelled token, cancellation while waiting, many cancelled waiters, a cancelled waiter not handing the key to another waiter, double disposal, and per-key single invocation of guarded work under contention. Each asserts GateCount returns to zero, since gate accumulation was the original bug. All handshakes are TaskCompletionSource/state based rather than delay based, with a 30s ceiling so a regression fails instead of hanging.

Two checks that they are not vacuous: reverting Dispose to the non-idempotent form fails Releaser_DisposedTwice_DoesNotHandKeyToTwoCallers with SemaphoreFullException, and reverting the gate-cleanup makes the contention tests fail. Removing the cleanup in the AcquireAsync catch alone does not fail anything — that path only covers the window where a caller creates a gate and is then cancelled before WaitAsync observes the token, which cannot be driven deterministically from outside the class. It stays as defence in depth; every externally reachable path is covered.

10/10 pass on net8.0 and net9.0 via dotnet msbuild build.proj -t:TestAkvProvider; all three TFMs build clean.

[benrr101]

One alternative to consider rather than building our own lock primitive would be to utilize a ConcurrentDictionary<TKey, Lazy<Task<T>> - that'd handle the dedup of work automatically, but would require a bit more consciousness around faults/cancellation in the tasks. Just something to consider.

[cheenamalhotra]

Considered it seriously — it is the more elegant shape in general, but it does not come out ahead here, for three reasons specific to this code.

Cancellation. A shared Task cannot carry a per-caller token: whoever starts the work fixes its cancellation, and one caller giving up must not cancel the work others are awaiting. The usual fix is to run the shared task with CancellationToken.None and let each caller observe its own token via Task.WaitAsync(token) — but this project targets netstandard2.0, where WaitAsync does not exist, so that shim would have to be hand-rolled. The gate gives this for free: SemaphoreSlim.WaitAsync(token) cancels only the caller that asked, which is exactly the semantics documented on AcquireAsync.

Faults. A cached Lazy<Task<T>> that faults is cached faulted forever, so every later caller replays a stale Key Vault failure. Avoiding that needs explicit eviction on fault, which is the same lifetime bookkeeping RemoveIfUnused already does — moved, not removed.

It would not replace the existing cache. Neither call site wants a task cache as its store. LocalCache is authoritatively a MemoryCache with TTL and size-based compaction, and AzureSqlKeyCryptographer keeps a ConcurrentDictionary<string, KeyVaultKey> shared with the synchronous path. A task dictionary would be a second store to keep coherent with the first, whereas the gate leaves exactly one store and just serialises the misses.

So the trade is: drop one small, now unit-tested primitive, and gain a second cache plus fault eviction plus a hand-rolled WaitAsync. Happy to revisit if we later move off netstandard2.0 or grow a call site that genuinely wants the shared task itself.

[Copilot]

Pull request overview

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

Suppressed comments (1)

src/Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider/test/KeyedAsyncLockTest.cs:80

• This test asserts mutual exclusion (maxActive == 1), but KeyedAsyncLock’s own remarks explicitly document that exclusion is best-effort and that two callers can run concurrently for the same key if a gate is reclaimed between GetOrAdd and WaitAsync. As written, the test is either inconsistent with the intended semantics or risks becoming flaky if the documented race is hit. Align the implementation/docs and the test’s expectation (either make exclusion guaranteed or relax/replace this assertion).

        [Fact]
        public async Task AcquireAsync_ConcurrentCallersOnOneKey_NeverOverlap()
        {
            const int callerCount = 32;

[cheenamalhotra]

Good catch, and correct — fixed in 7a6ac36.

The grant is now gated on '$(SigningKeyPath)' == '', matching Logging/Abstractions/Azure exactly, with a comment explaining why. Verified both ways locally:

• Signed (-p:SigningKeyPath=<key>): builds clean, and the generated AssemblyInfo.cs contains zero InternalsVisibleTo attributes.
• Unsigned: all 10 KeyedAsyncLockTest tests still pass on net8.0.

This does mean the unit tests only build against an unsigned provider, which is the same trade the sibling packages already make — build.proj runs TestAkvProvider without a signing key, so CI is unaffected.

[Copilot]

Pull request overview

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