Implement async APIs in Always Encrypted Azure Key Vault provider

BUILDGUIDE.md

@@ -155,6 +155,7 @@ dotnet build -t:<test_target> [optional_parameters]
 |----------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------|
 | `Test`                     | Runs all tests in the repository for all platforms supported by the host OS. _This will take a considerable amount of time and is not recommended_. |
 | `TestAbstractions`         | Runs all tests for Microsoft.Data.SqlClient.Extensions.Abstractions                                                                                 |
+| `TestAkvProvider`          | Runs the unit test project for Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider.                                                      |
 | `TestAzure`                | Runs all tests for Microsoft.Data.SqlClient.Extensions.Azure                                                                                        |
 | `TestSqlClient`            | Runs all tests for Microsoft.Data.SqlClient.                                                                                                        |
 | `TestSqlClientFunctional`  | Runs the "functional" test project for Microsoft.Data.SqlClient. These are a mix of unit and integration tests against live servers.                |

build.proj

@@ -401,7 +401,7 @@
     environments. Please consider running project specific test targets or specific test sets within
     the project.
   -->
-  <Target Name="Test" DependsOnTargets="TestAbstractions;TestAzure;TestSqlClient" />
+  <Target Name="Test" DependsOnTargets="TestAbstractions;TestAzure;TestAkvProvider;TestSqlClient" />
 
   <!-- ================================================================= -->
   <!-- Microsoft.Data.SqlClient Targets -->
@@ -714,7 +714,7 @@
     </PropertyGroup>
     <Message Importance="High" Text=">>> Running manual tests for SqlClient"/>
     <Message Text="    Command: $(DotnetCommand)"/>
     <Exec ConsoleToMsBuild="true" Command="$(DotnetCommand)" />
   </Target>
 
   <!-- TestSqlClientUnit: Runs unit tests for SqlClient -->
@@ -743,7 +743,7 @@
     </PropertyGroup>
     <Message Importance="High" Text=">>> Running unit tests for SqlClient"/>
     <Message Text="    Command: $(DotnetCommand)"/>
     <Exec ConsoleToMsBuild="true" Command="$(DotnetCommand)" />
   </Target>
 
   <!-- ================================================================= -->
@@ -751,6 +751,7 @@
   <PropertyGroup>
     <AkvProviderSrcRoot>$(RepoRoot)src/Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider/</AkvProviderSrcRoot>
     <AkvProviderProjectPath>$(AkvProviderSrcRoot)src/Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider.csproj</AkvProviderProjectPath>
+    <AkvProviderTestProjectPath>$(AkvProviderSrcRoot)test/Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider.Test.csproj</AkvProviderTestProjectPath>
     <AkvProviderArtifactRoot>$(RepoRoot)artifacts/Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider/$(Configuration)/</AkvProviderArtifactRoot>
   </PropertyGroup>
 
@@ -783,7 +784,7 @@
 
     <Message Importance="High" Text=">>> Building AKV Provider"/>
     <Message Text="    Command: $(DotnetCommand)"/>
     <Exec ConsoleToMsBuild="true" Command="$(DotnetCommand)" />
   </Target>
 
   <!-- PackAkvProvider: Packages Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider -->
@@ -827,9 +828,37 @@
       SkipUnchangedFiles="true" />
   </Target>
 
+  <!-- TestAkvProvider: Runs Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider unit tests -->
+  <Target Name="TestAkvProvider">
+    <PropertyGroup>
+      <!--
+        Note: These tests exclusively use project references, so neither ReferenceType nor any
+        package version arguments are specified in this command.
+      -->
+      <LogFilePrefix>AkvProviderTests-$(OS)</LogFilePrefix>
+      <LogFilePrefix Condition="'$(TestFramework)' != ''">$(LogFilePrefix)-$(TestFramework)</LogFilePrefix>
+
+      <DotnetCommand>
+        "$(DotnetPath)dotnet" test "$(AkvProviderTestProjectPath)"
+        -p:Configuration=$(Configuration)
+        $(TestBlameArgument)
+        $(TestCodeCoverageArgument)
+        $(TestFiltersArgument)
+        $(TestFrameworkArgument)
+        --results-directory "$(TestResultsFolderPath)"
+        --logger:"trx;LogFilePrefix=$(LogFilePrefix)"
+      </DotnetCommand>
+      <!-- Convert more than one whitespace character into one space -->
+      <DotnetCommand>$([System.Text.RegularExpressions.Regex]::Replace($(DotnetCommand), "\s+", " "))</DotnetCommand>
+    </PropertyGroup>
+
+    <Message Importance="High" Text=">>> Running tests for AKV Provider"/>
+    <Message Text="    Command: $(DotnetCommand)"/>
+    <Exec ConsoleToMsBuild="true" Command="$(DotnetCommand)" />
+  </Target>
+
   <!-- ================================================================= -->
   <!-- Microsoft.Data.SqlClient.Extensions.Abstractions Targets -->
-
   <PropertyGroup>
     <AbstractionsSrcRoot>$(RepoRoot)src/Microsoft.Data.SqlClient.Extensions/Abstractions/</AbstractionsSrcRoot>
     <AbstractionsProjectPath>$(AbstractionsSrcRoot)src/Abstractions.csproj</AbstractionsProjectPath>
@@ -1179,7 +1208,21 @@
   </PropertyGroup>
 
   <!-- BuildTests: Builds all test projects for all supported OS combinations. -->
-  <Target Name="BuildTests" DependsOnTargets="BuildAbstractionsTests;BuildSqlClientTestsUnit;BuildSqlClientTestsFunctional;BuildSqlClientTestsManual;BuildSqlClientTestsPerformance;BuildSqlClientTestsStress;BuildAzureTests" />
+  <Target Name="BuildTests" DependsOnTargets="BuildAbstractionsTests;BuildAkvProviderTests;BuildSqlClientTestsUnit;BuildSqlClientTestsFunctional;BuildSqlClientTestsManual;BuildSqlClientTestsPerformance;BuildSqlClientTestsStress;BuildAzureTests" />
+
+  <!-- BuildAkvProviderTests: Builds the AKV Provider test project (OS-agnostic). -->
+  <Target Name="BuildAkvProviderTests">
+    <PropertyGroup>
+      <DotnetCommand>
+        "$(DotnetPath)dotnet" build "$(AkvProviderTestProjectPath)"
+        -p:Configuration=$(Configuration)
+      </DotnetCommand>
+      <DotnetCommand>$([System.Text.RegularExpressions.Regex]::Replace($(DotnetCommand), "\s+", " "))</DotnetCommand>
+    </PropertyGroup>
+    <Message Importance="High" Text=">>> Building AKV Provider Tests"/>
+    <Message Text="    Command: $(DotnetCommand)"/>
+    <Exec ConsoleToMsBuild="true" Command="$(DotnetCommand)" />
+  </Target>
 
   <!-- BuildAbstractionsTests: Builds the Abstractions test project (OS-agnostic). -->
   <Target Name="BuildAbstractionsTests">

eng/pipelines/pr/stages/test-stages.yml

@@ -155,6 +155,21 @@ stages:
               packageShortName: "Azure"
               testDisplayName: "azure"
 
+          # TestAkvProvider
+          - template: /eng/pipelines/pr/jobs/test-buildproj-job.yml@self
+            parameters:
+              buildConfiguration: ${{ parameters.buildConfiguration }}
+              buildSuffix: ${{ parameters.buildSuffix }}
+              dotnetVerbosity: ${{ parameters.dotnetVerbosity }}
+              platformDisplayName: ${{ platform.displayName }}
+              platformDotnet: ${{ platform.dotnet }}
+              platformImage: ${{ platform.image }}
+              poolName: ${{ parameters.poolName }}
+              testResultsArtifactBaseName: ${{ parameters.testResultsArtifactBaseName }}
+
+              packageShortName: "AkvProvider"
+              testDisplayName: "akvprovider"
+
           # TestSqlClientFunctional
           - template: /eng/pipelines/pr/jobs/test-buildproj-job.yml@self
             parameters:

src/Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider/src/AzureSqlKeyCryptographer.cs

@@ -8,6 +8,7 @@
 using System;
 using System.Collections.Concurrent;
 using System.Threading;
+using System.Threading.Tasks;
 using static Azure.Security.KeyVault.Keys.Cryptography.SignatureAlgorithm;
 using Microsoft.Data.SqlClient.Internal;
 
@@ -32,9 +33,17 @@ internal sealed class AzureSqlKeyCryptographer : IDisposable
 
         /// <summary>
         /// SemaphoreSlim to ensure thread safety when accessing the key dictionary or making network calls to Azure Key Vault to fetch keys.
+        /// Used by the synchronous path only. The asynchronous path uses <see cref="_keyFetchLock"/> so that a synchronous caller
+        /// never blocks a thread waiting on an asynchronous, network bound fetch.
         /// </summary>
         private SemaphoreSlim _keyDictionarySemaphore = new(1, 1);
 
+        /// <summary>
+        /// Gates concurrent asynchronous fetches of the same key, so that only one caller retrieves a given key from
+        /// Azure Key Vault. Fetches of different keys are not serialized against one another.
+        /// </summary>
+        private readonly KeyedAsyncLock<string> _keyFetchLock = new();
+
         /// <summary>
         /// Holds references to the Azure Key Vault CryptographyClient objects and maps them to their corresponding Azure Key Vault Key Identifier (URI).
         /// </summary>
@@ -89,6 +98,47 @@ internal void AddKey(string keyIdentifierUri)
             }
         }
 
+        /// <summary>
+        /// Asynchronously adds the key, specified by the Key Identifier URI, to the cache.
+        /// Validates the key type and fetches the key from Azure Key Vault if it is not already cached.
+        /// </summary>
+        /// <param name="keyIdentifierUri">The key vault key identifier URI</param>
+        /// <param name="cancellationToken">Token used to request cancellation of the operation</param>
+        /// <remarks>
+        /// Mirrors the deduplication of the synchronous path: only one caller fetches a given key from Azure Key Vault.
+        /// The gate is per key and is only ever awaited, so no thread is blocked while a fetch is in flight and fetches of
+        /// different keys proceed in parallel. The gate is deliberately not shared with the synchronous path, because a
+        /// synchronous caller blocking on a gate held across an awaited network call would consume a thread pool thread for
+        /// the duration of that call. A synchronous and an asynchronous caller may therefore both fetch the same key; the
+        /// result is identical and the last write wins.
+        /// </remarks>
+        internal async Task AddKeyAsync(string keyIdentifierUri, CancellationToken cancellationToken = default)
+        {
+            if (_keyDictionary.ContainsKey(keyIdentifierUri))
+            {
+                return;
+            }
+
+            using (await _keyFetchLock.AcquireAsync(keyIdentifierUri, cancellationToken).ConfigureAwait(false))
+            {
+                // Another caller may have fetched the key while this one waited for the gate.
+                if (_keyDictionary.ContainsKey(keyIdentifierUri))
+                {
+                    return;
+                }
+
+                ParseAKVPath(keyIdentifierUri, out Uri vaultUri, out string keyName, out string keyVersion);
+
+                // Fetch the KeyClient for the Key vault URI.
+                KeyClient keyClient = GetOrCreateKeyClient(vaultUri);
+
+                // Fetch the key from Azure Key Vault.
+                KeyVaultKey key = await FetchKeyFromKeyVaultAsync(keyClient, keyName, keyVersion, cancellationToken).ConfigureAwait(false);
+
+                _keyDictionary.AddOrUpdate(keyIdentifierUri, key, (k, v) => key);
+            }
+        }
+
         /// <summary>
         /// Returns the key specified by the Key Identifier URI
         /// </summary>
@@ -133,6 +183,34 @@ internal bool VerifyData(byte[] message, byte[] signature, string keyIdentifierU
             return cryptographyClient.VerifyData(RS256, message, signature).IsValid;
         }
 
+        /// <summary>
+        /// Asynchronously generates a signature based on the RSA PKCS#v1.5 scheme using a specified Azure Key Vault Key URL.
+        /// </summary>
+        /// <param name="message">The data to sign</param>
+        /// <param name="keyIdentifierUri">The key vault key identifier URI</param>
+        /// <param name="cancellationToken">Token used to request cancellation of the operation</param>
+        internal async Task<byte[]> SignDataAsync(byte[] message, string keyIdentifierUri, CancellationToken cancellationToken = default)
+        {
+            CryptographyClient cryptographyClient = GetCryptographyClient(keyIdentifierUri);
+            SignResult result = await cryptographyClient.SignDataAsync(RS256, message, cancellationToken).ConfigureAwait(false);
+            return result.Signature;
+        }
+
+        /// <summary>
+        /// Asynchronously verifies a signature based on the RSA PKCS#v1.5 scheme using a specified Azure Key Vault Key URL.
+        /// </summary>
+        /// <param name="message">The signed data</param>
+        /// <param name="signature">The signature to verify</param>
+        /// <param name="keyIdentifierUri">The key vault key identifier URI</param>
+        /// <param name="cancellationToken">Token used to request cancellation of the operation</param>
+        internal async Task<bool> VerifyDataAsync(byte[] message, byte[] signature, string keyIdentifierUri, CancellationToken cancellationToken = default)
+        {
+            CryptographyClient cryptographyClient = GetCryptographyClient(keyIdentifierUri);
+            SqlClientEventSource.Log.TryTraceEvent("Sending request to verify data");
+            VerifyResult result = await cryptographyClient.VerifyDataAsync(RS256, message, signature, cancellationToken).ConfigureAwait(false);
+            return result.IsValid;
+        }
+
         internal byte[] UnwrapKey(KeyWrapAlgorithm keyWrapAlgorithm, byte[] encryptedKey, string keyIdentifierUri)
         {
             CryptographyClient cryptographyClient = GetCryptographyClient(keyIdentifierUri);
@@ -147,16 +225,53 @@ internal byte[] WrapKey(KeyWrapAlgorithm keyWrapAlgorithm, byte[] key, string ke
             return cryptographyClient.WrapKey(keyWrapAlgorithm, key).EncryptedKey;
         }
 
+        /// <summary>
+        /// Asynchronously unwraps the given encrypted key using the specified Azure Key Vault key.
+        /// </summary>
+        /// <param name="keyWrapAlgorithm">The key wrap algorithm</param>
+        /// <param name="encryptedKey">The encrypted key to unwrap</param>
+        /// <param name="keyIdentifierUri">The key vault key identifier URI</param>
+        /// <param name="cancellationToken">Token used to request cancellation of the operation</param>
+        internal async Task<byte[]> UnwrapKeyAsync(KeyWrapAlgorithm keyWrapAlgorithm, byte[] encryptedKey, string keyIdentifierUri, CancellationToken cancellationToken = default)
+        {
+            CryptographyClient cryptographyClient = GetCryptographyClient(keyIdentifierUri);
+            SqlClientEventSource.Log.TryTraceEvent("Sending request to unwrap key.");
+            UnwrapResult result = await cryptographyClient.UnwrapKeyAsync(keyWrapAlgorithm, encryptedKey, cancellationToken).ConfigureAwait(false);
+            return result.Key;
+        }
+
+        /// <summary>
+        /// Asynchronously wraps the given key using the specified Azure Key Vault key.
+        /// </summary>
+        /// <param name="keyWrapAlgorithm">The key wrap algorithm</param>
+        /// <param name="key">The key to wrap</param>
+        /// <param name="keyIdentifierUri">The key vault key identifier URI</param>
+        /// <param name="cancellationToken">Token used to request cancellation of the operation</param>
+        internal async Task<byte[]> WrapKeyAsync(KeyWrapAlgorithm keyWrapAlgorithm, byte[] key, string keyIdentifierUri, CancellationToken cancellationToken = default)
+        {
+            CryptographyClient cryptographyClient = GetCryptographyClient(keyIdentifierUri);
+            SqlClientEventSource.Log.TryTraceEvent("Sending request to wrap key.");
+            WrapResult result = await cryptographyClient.WrapKeyAsync(keyWrapAlgorithm, key, cancellationToken).ConfigureAwait(false);
+            return result.EncryptedKey;
+        }
+
+        /// <summary>
+        /// Returns the CryptographyClient for the given key identifier URI, creating it if necessary.
+        /// </summary>
+        /// <remarks>
+        /// Concurrent callers may each construct a client, but all of them observe the single instance that wins the
+        /// race, so a key is never used through two different clients.
+        /// </remarks>
         private CryptographyClient GetCryptographyClient(string keyIdentifierUri)
         {
             if (_cryptoClientDictionary.TryGetValue(keyIdentifierUri, out CryptographyClient client))
             {
                 return client;
             }
 
-            CryptographyClient cryptographyClient = new(GetKey(keyIdentifierUri).Id, TokenCredential);
-            _cryptoClientDictionary.TryAdd(keyIdentifierUri, cryptographyClient);
-            return cryptographyClient;
+            return _cryptoClientDictionary.GetOrAdd(
+                keyIdentifierUri,
+                uri => new CryptographyClient(GetKey(uri).Id, TokenCredential));
         }
 
         /// <summary>
@@ -171,6 +286,35 @@ private KeyVaultKey FetchKeyFromKeyVault(KeyClient keyClient, string keyName, st
 
             Azure.Response<KeyVaultKey> keyResponse = keyClient?.GetKey(keyName, keyVersion);
 
+            return ValidateKeyResponse(keyResponse, keyName, keyVersion);
+        }
+
+        /// <summary>
+        /// Asynchronously fetches the column master key from the Azure Key Vault.
+        /// </summary>
+        /// <param name="keyClient">The KeyClient instance</param>
+        /// <param name="keyName">The name of the Azure Key Vault key</param>
+        /// <param name="keyVersion">The version of the Azure Key Vault key</param>
+        /// <param name="cancellationToken">Token used to request cancellation of the operation</param>
+        private async Task<KeyVaultKey> FetchKeyFromKeyVaultAsync(KeyClient keyClient, string keyName, string keyVersion, CancellationToken cancellationToken)
+        {
+            SqlClientEventSource.Log.TryTraceEvent("Fetching key name={0}", keyName);
+
+            Azure.Response<KeyVaultKey> keyResponse = keyClient is null
+                ? null
+                : await keyClient.GetKeyAsync(keyName, keyVersion, cancellationToken).ConfigureAwait(false);
+
+            return ValidateKeyResponse(keyResponse, keyName, keyVersion);
+        }
+
+        /// <summary>
+        /// Validates the response received from Azure Key Vault and ensures the returned key is an RSA key.
+        /// </summary>
+        /// <param name="keyResponse">The response returned by Azure Key Vault</param>
+        /// <param name="keyName">The name of the Azure Key Vault key</param>
+        /// <param name="keyVersion">The version of the Azure Key Vault key</param>
+        private static KeyVaultKey ValidateKeyResponse(Azure.Response<KeyVaultKey> keyResponse, string keyName, string keyVersion)
+        {
             // Handle the case where the key response is null or contains an error
             // This can happen if the key does not exist or if there is an issue with the KeyClient.
             // In such cases, we log the error and throw an exception.