Feature | AKV Provider Upgrade to use new Azure key Vault libraries

BUILDGUIDE.md

@@ -112,6 +112,7 @@ Manual Tests require the below setup to run:
 |AADSecurePrincipalId | (Optional) The Application Id of a registered application which has been granted permission to the database defined in the AADPasswordConnectionString. | {Application ID} |
 |AADSecurePrincipalSecret | (Optional) A Secret defined for a registered application which has been granted permission to the database defined in the AADPasswordConnectionString. | {Secret} |
 |AzureKeyVaultURL | (Optional) Azure Key Vault Identifier URL | `https://{keyvaultname}.vault.azure.net/` |
+|AzureKeyVaultTenantId | (Optional) The Azure Active Directory tenant (directory) Id of the service principal. | _{Tenant ID of Active Directory}_ |
 |AzureKeyVaultClientId | (Optional) "Application (client) ID" of an Active Directory registered application, granted access to the Azure Key Vault specified in `AZURE_KEY_VAULT_URL`. Requires the key permissions Get, List, Import, Decrypt, Encrypt, Unwrap, Wrap, Verify, and Sign. | _{Client Application ID}_ |
 |AzureKeyVaultClientSecret | (Optional) "Client Secret" of the Active Directory registered application, granted access to the Azure Key Vault specified in `AZURE_KEY_VAULT_URL` | _{Client Application Secret}_ |
 |SupportsLocalDb | (Optional) Whether or not a LocalDb instance of SQL Server is installed on the machine running the tests. |`true` OR `false`|

src/Microsoft.Data.SqlClient/add-ons/AzureKeyVaultProvider/AzureKeyVaultProviderTokenCredential.cs

@@ -0,0 +1,90 @@
+// 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.Linq;
+using System.Net.Http;
+using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
+using Azure.Core;
+
+namespace Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider
+{
+    internal class AzureKeyVaultProviderTokenCredential : TokenCredential
+    {
+        private const string Bearer = "Bearer ";
+
+        private AuthenticationCallback Callback { get; set; }
+
+        private string Authority { get; set; }
+
+        private string Resource { get; set; }
+
+        internal AzureKeyVaultProviderTokenCredential(AuthenticationCallback authenticationCallback, string masterKeyPath)
+        {
+            Callback = authenticationCallback;
+            using (HttpClient httpClient = new HttpClient())
+            {
+                HttpResponseMessage response = httpClient.GetAsync(masterKeyPath).GetAwaiter().GetResult();
+                string challenge = response?.Headers.WwwAuthenticate.FirstOrDefault()?.ToString();
+                string trimmedChallenge = ValidateChallenge(challenge);
+                string[] pairs = trimmedChallenge.Split(new string[] { "," }, StringSplitOptions.RemoveEmptyEntries);
+
+                if (pairs != null && pairs.Length > 0)
+                {
+                    for (int i = 0; i < pairs.Length; i++)
+                    {
+                        string[] pair = pairs[i]?.Split('=');
+
+                        if (pair.Length == 2)
+                        {
+                            string key = pair[0]?.Trim().Trim(new char[] { '\"' });
+                            string value = pair[1]?.Trim().Trim(new char[] { '\"' });
+
+                            if (!string.IsNullOrEmpty(key))
+                            {
+                                if (key.Equals("authorization", StringComparison.InvariantCultureIgnoreCase))
+                                {
+                                    Authority = value;
+                                }
+                                else if (key.Equals("resource", StringComparison.InvariantCultureIgnoreCase))
+                                {
+                                    Resource = value;
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        public override AccessToken GetToken(TokenRequestContext requestContext, CancellationToken cancellationToken)
+        {
+            string token = Callback.Invoke(Authority, Resource, string.Empty).GetAwaiter().GetResult();
+            return new AccessToken(token, DateTimeOffset.Now.AddHours(1));
+        }
+
+        public override ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken)
+        {
+            string token = Callback.Invoke(Authority, Resource, string.Empty).GetAwaiter().GetResult();
+            Task<AccessToken> task = Task.FromResult(new AccessToken(token, DateTimeOffset.Now.AddHours(1)));
+            return new ValueTask<AccessToken>(task);
+        }
+
+        private static string ValidateChallenge(string challenge)
+        {
+            if (string.IsNullOrEmpty(challenge))
+                throw new ArgumentNullException(nameof(challenge));
+
+            string trimmedChallenge = challenge.Trim();
+
+            if (!trimmedChallenge.StartsWith(Bearer))
+                throw new ArgumentException("Challenge is not Bearer", nameof(challenge));
+
+            return trimmedChallenge.Substring(Bearer.Length);
+        }
+    }
+}

src/Microsoft.Data.SqlClient/add-ons/AzureKeyVaultProvider/AzureSqlKeyCryptographer.cs

@@ -0,0 +1,233 @@
+// 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 Azure.Core;
+using Azure.Security.KeyVault.Keys;
+using Azure.Security.KeyVault.Keys.Cryptography;
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Globalization;
+using System.Threading.Tasks;
+using static Azure.Security.KeyVault.Keys.Cryptography.SignatureAlgorithm;
+
+
+namespace Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider
+{
+    internal class AzureSqlKeyCryptographer
+    {
+        /// <summary>
+        /// TokenCredential to be used with the KeyClient
+        /// </summary>
+        private TokenCredential TokenCredential { get; set; }
+
+        /// <summary>
+        /// AuthenticationCallback to be used with the KeyClient for legacy support.
+        /// </summary>
+        private AuthenticationCallback AuthenticationCallback { get; set; }
+
+        /// <summary>
+        /// A flag to determine whether to use AuthenticationCallback with the KeyClient for legacy support.
+        /// </summary>
+        private readonly bool isUsingLegacyAuthentication = false;
+
+        /// <summary>
+        /// A mapping of the KeyClient objects to the corresponding Azure Key Vault URI
+        /// </summary>
+        private readonly ConcurrentDictionary<Uri, KeyClient> keyClientDictionary = new ConcurrentDictionary<Uri, KeyClient>();
+
+        /// <summary>
+        /// Holds references to the fetch key tasks and maps them to their corresponding Azure Key Vault Key Identifier (URI).
+        /// These tasks will be used for returning the key in the event that the fetch task has not finished depositing the 
+        /// key into the key dictionary.
+        /// </summary>
+        private readonly ConcurrentDictionary<string, Task<Azure.Response<KeyVaultKey>>> _keyFetchTaskDictionary = new ConcurrentDictionary<string, Task<Azure.Response<KeyVaultKey>>>();
+
+        /// <summary>
+        /// Holds references to the Azure Key Vault keys and maps them to their corresponding Azure Key Vault Key Identifier (URI).
+        /// </summary>
+        private readonly ConcurrentDictionary<string, KeyVaultKey> _keyDictionary = new ConcurrentDictionary<string, KeyVaultKey>();
+
+        /// <summary>
+        /// Holds references to the Azure Key Vault CryptographyClient objects and maps them to their corresponding Azure Key Vault Key Identifier (URI).
+        /// </summary>
+        private readonly ConcurrentDictionary<string, CryptographyClient> cryptoClientDictionary = new ConcurrentDictionary<string, CryptographyClient>();
+
+        /// <summary>
+        /// Constructs a new KeyCryptographer
+        /// </summary>
+        /// <param name="tokenCredential"></param>
+        internal AzureSqlKeyCryptographer(TokenCredential tokenCredential)
+        {
+            TokenCredential = tokenCredential;
+        }
+
+        internal AzureSqlKeyCryptographer(AuthenticationCallback authenticationCallback)
+        {
+            AuthenticationCallback = authenticationCallback;
+            isUsingLegacyAuthentication = true;
+        }
+
+        /// <summary>
+        /// Adds the key, specified by the Key Identifier URI, to the cache.
+        /// </summary>
+        /// <param name="keyIdentifierUri"></param>
+        internal void AddKey(string keyIdentifierUri)
+        {
+            if (TheKeyHasNotBeenCached(keyIdentifierUri))
+            {
+                if (isUsingLegacyAuthentication)
+                {
+                    TokenCredential = new AzureKeyVaultProviderTokenCredential(AuthenticationCallback, keyIdentifierUri);
+                }
+
+                ParseAKVPath(keyIdentifierUri, out Uri vaultUri, out string keyName);
+                CreateKeyClient(vaultUri);
+                FetchKey(vaultUri, keyName, keyIdentifierUri);
+            }
+
+            bool TheKeyHasNotBeenCached(string k) => !_keyDictionary.ContainsKey(k) && !_keyFetchTaskDictionary.ContainsKey(k);
+        }
+
+        /// <summary>
+        /// Returns the key specified by the Key Identifier URI
+        /// </summary>
+        /// <param name="keyIdentifierUri"></param>
+        /// <returns></returns>
+        internal KeyVaultKey GetKey(string keyIdentifierUri)
+        {
+            if (_keyDictionary.ContainsKey(keyIdentifierUri))
+            {
+                return _keyDictionary[keyIdentifierUri];
+            }
+
+            if (_keyFetchTaskDictionary.ContainsKey(keyIdentifierUri))
+            {
+                return Task.Run(() => _keyFetchTaskDictionary[keyIdentifierUri]).Result;
+            }
+
+            // Not a public exception - not likely to occur.
+            throw new KeyNotFoundException($"The key with identifier {keyIdentifierUri} was not found.");
+        }
+
+        /// <summary>
+        /// Gets the public Key size in bytes.
+        /// </summary>
+        /// <param name="keyIdentifierUri">The key vault key identifier URI</param>
+        /// <returns></returns>
+        internal int GetKeySize(string keyIdentifierUri)
+        {
+            return GetKey(keyIdentifierUri).Key.N.Length;
+        }
+
+        /// <summary>
+        /// Generates signature based on 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>
+        /// <returns></returns>
+        internal byte[] SignData(byte[] message, string keyIdentifierUri)
+        {
+            CryptographyClient cryptographyClient = GetCryptographyClient(keyIdentifierUri);
+            return cryptographyClient.SignData(RS256, message).Signature;
+        }
+
+        internal bool VerifyData(byte[] message, byte[] signature, string keyIdentifierUri)
+        {
+            CryptographyClient cryptographyClient = GetCryptographyClient(keyIdentifierUri);
+            return cryptographyClient.VerifyData(RS256, message, signature).IsValid;
+        }
+
+        internal byte[] UnwrapKey(KeyWrapAlgorithm keyWrapAlgorithm, byte[] encryptedKey, string keyIdentifierUri)
+        {
+            CryptographyClient cryptographyClient = GetCryptographyClient(keyIdentifierUri);
+            return cryptographyClient.UnwrapKey(keyWrapAlgorithm, encryptedKey).Key;
+        }
+
+        internal byte[] WrapKey(KeyWrapAlgorithm keyWrapAlgorithm, byte[] key, string keyIdentifierUri)
+        {
+            CryptographyClient cryptographyClient = GetCryptographyClient(keyIdentifierUri);
+            return cryptographyClient.WrapKey(keyWrapAlgorithm, key).EncryptedKey;
+        }
+
+        private CryptographyClient GetCryptographyClient(string keyIdentifierUri)
+        {
+            if (cryptoClientDictionary.ContainsKey(keyIdentifierUri))
+            {
+                return cryptoClientDictionary[keyIdentifierUri];
+            }
+
+            CryptographyClient cryptographyClient = new CryptographyClient(GetKey(keyIdentifierUri).Id, TokenCredential);
+            cryptoClientDictionary[keyIdentifierUri] = cryptographyClient;
+
+            return cryptographyClient;
+        }
+
+        /// <summary>
+        /// 
+        /// </summary>
+        /// <param name="vaultUri">The Azure Key Vault URI</param>
+        /// <param name="keyName">The name of the Azure Key Vault key</param>
+        /// <param name="keyResourceUri">The Azure Key Vault key identifier</param>
+        private void FetchKey(Uri vaultUri, string keyName, string keyResourceUri)
+        {
+            var fetchKeyTask = FetchKeyFromKeyVault(vaultUri, keyName);
+            _keyFetchTaskDictionary[keyResourceUri] = fetchKeyTask;
+
+            fetchKeyTask
+                .ContinueWith(k => ValidateRsaKey(k.Result))
+                .ContinueWith(k => _keyDictionary[keyResourceUri] = k.Result);
+
+            Task.Run(() => fetchKeyTask);
+        }
+
+        /// <summary>
+        /// Looks up the KeyClient object by it's URI and then fethces the key by name.
+        /// </summary>
+        /// <param name="vaultUri">The Azure Key Vault URI</param>
+        /// <param name="keyName">Then name of the key</param>
+        /// <returns></returns>
+        private Task<Azure.Response<KeyVaultKey>> FetchKeyFromKeyVault(Uri vaultUri, string keyName) => keyClientDictionary[vaultUri].GetKeyAsync(keyName);
+
+        /// <summary>
+        /// Validates that a key is of type RSA
+        /// </summary>
+        /// <param name="key"></param>
+        /// <returns></returns>
+        private KeyVaultKey ValidateRsaKey(KeyVaultKey key)
+        {
+            if (key.KeyType != KeyType.Rsa && key.KeyType != KeyType.RsaHsm)
+            {
+                throw new FormatException(string.Format(CultureInfo.InvariantCulture, Strings.NonRsaKeyTemplate, key.KeyType));
+            }
+
+            return key;
+        }
+
+        /// <summary>
+        /// Instantiates and adds a KeyClient to the KeyClient dictionary
+        /// </summary>
+        /// <param name="vaultUri">The Azure Key Vault URI</param>
+        private void CreateKeyClient(Uri vaultUri)
+        {
+            if (!keyClientDictionary.ContainsKey(vaultUri))
+            {
+                keyClientDictionary[vaultUri] = new KeyClient(vaultUri, TokenCredential);
+            }
+        }
+
+        /// <summary>
+        /// Validates and parses the Azure Key Vault URI and key name.
+        /// </summary>
+        /// <param name="masterKeyPath">The Azure Key Vault key identifier</param>
+        /// <param name="vaultUri">The Azure Key Vault URI</param>
+        /// <param name="masterKeyName">The name of the key</param>
+        private void ParseAKVPath(string masterKeyPath, out Uri vaultUri, out string masterKeyName)
+        {
+            Uri masterKeyPathUri = new Uri(masterKeyPath);
+            vaultUri = new Uri(masterKeyPathUri.GetLeftPart(UriPartial.Authority));
+            masterKeyName = masterKeyPathUri.Segments[2];
+        }
+    }
+}

src/Microsoft.Data.SqlClient/add-ons/AzureKeyVaultProvider/Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider.csproj

@@ -22,9 +22,7 @@
     <PackageReference Condition="$(ReferenceType)=='Package'" Include="Microsoft.Data.SqlClient" Version="$(TestMicrosoftDataSqlClientVersion)" />
   </ItemGroup>
   <ItemGroup>
-    <PackageReference Include="Microsoft.Azure.KeyVault" Version="$(MicrosoftAzureKeyVaultVersion)" />
-    <PackageReference Include="Microsoft.Azure.KeyVault.WebKey" Version="$(MicrosoftAzureKeyVaultWebKeyVersion)" />
-    <PackageReference Include="Microsoft.Rest.ClientRuntime" Version="$(MicrosoftRestClientRuntimeVersion)" />
-    <PackageReference Include="Microsoft.Rest.ClientRuntime.Azure" Version="$(MicrosoftRestClientRuntimeAzureVersion)" />
+    <PackageReference Include="Azure.Core" Version="$(AzureCoreVersion)" />
+    <PackageReference Include="Azure.Security.KeyVault.Keys" Version="$(AzureSecurityKeyVaultKeysVersion)" />
   </ItemGroup>
 </Project>