Always Encrypted Azure Key Vault Support (#934)

Add support for running queries that require a decryption key from Azure Key Vault when using Always Encrypted.

src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs

@@ -9,6 +9,7 @@ using System.Collections.Generic;
 using System.Data;
 using System.Data.Common;
 using Microsoft.Data.SqlClient;
+using Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider;
 using System.Globalization;
 using System.Linq;
 using System.Threading;
@@ -105,6 +106,15 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             }
         }
 
+        static ConnectionService()
+        {
+            SqlColumnEncryptionAzureKeyVaultProvider sqlColumnEncryptionAzureKeyVaultProvider = new SqlColumnEncryptionAzureKeyVaultProvider(AzureActiveDirectoryAuthenticationCallback);
+            SqlConnection.RegisterColumnEncryptionKeyStoreProviders(customProviders: new Dictionary<string, SqlColumnEncryptionKeyStoreProvider>(capacity: 1, comparer: StringComparer.OrdinalIgnoreCase)
+            {
+                { SqlColumnEncryptionAzureKeyVaultProvider.ProviderName, sqlColumnEncryptionAzureKeyVaultProvider }
+            });
+        }
+
 
         /// <summary>
         /// Default constructor should be private since it's a singleton class, but we need a constructor
@@ -117,6 +127,21 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             this.LockedDatabaseManager.ConnectionService = this;
         }
 
+        public static async Task<string> AzureActiveDirectoryAuthenticationCallback(string authority, string resource, string scope)
+        {
+            RequestSecurityTokenParams message = new RequestSecurityTokenParams()
+            {
+                Authority = authority,
+                Provider = "Azure",
+                Resource = resource,
+                Scope = scope
+            };
+
+            RequestSecurityTokenResponse response = await Instance.ServiceHost.SendRequest(SecurityTokenRequest.Type, message, true);
+
+            return response.Token;
+        }
+
         /// <summary>
         /// Returns a connection queue for given type
         /// </summary>

src/Microsoft.SqlTools.ServiceLayer/Connection/Contracts/SecurityTokenRequest.cs

@@ -0,0 +1,50 @@
+using Microsoft.SqlTools.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
+{
+    class RequestSecurityTokenParams 
+    {
+        /// <summary>
+        /// Gets or sets the address of the authority to issue token.
+        /// </summary>
+        public string Authority { get; set; }
+
+        /// <summary>
+        /// Gets or sets the provider that indicates the type of linked account to query.
+        /// </summary>
+        public string Provider { get; set; }
+
+        /// <summary>
+        /// Gets or sets the identifier of the target resource that is the recipient of the requested token.
+        /// </summary>
+        public string Resource { get; set; }
+
+        /// <summary>
+        /// Gets or sets the scope of the authentication request.
+        /// </summary>
+        public string Scope { get; set; }
+    }
+
+    class RequestSecurityTokenResponse
+    {
+        /// <summary>
+        /// Gets or sets the key that uniquely identifies a particular linked account.
+        /// </summary>
+        public string AccountKey { get; set; }
+
+        /// <summary>
+        /// Gets or sets the access token.
+        /// </summary>
+        public string Token { get; set; }
+    }
+
+    /// <summary>
+    /// SecurityToken Request mapping entry 
+    /// </summary>
+    class SecurityTokenRequest
+    {
+        public static readonly
+            RequestType<RequestSecurityTokenParams, RequestSecurityTokenResponse> Type =
+            RequestType<RequestSecurityTokenParams, RequestSecurityTokenResponse>.Create("account/securityTokenRequest");
+    }
+}