From 7dc5785381902e826e1045890208f71ff0c7264a Mon Sep 17 00:00:00 2001
From: Cheena Malhotra <13396919+cheenamalhotra@users.noreply.github.com>
Date: Tue, 21 Mar 2023 21:26:11 -0700
Subject: [PATCH] Handle Encryption keys changed notification (#1955)

---
 .../MSALEncryptedCacheHelper.cs               |  4 +-
 .../Connection/ConnectionService.cs           | 53 ++++---------------
 .../Contracts/EncryptionKeysChange.cs         | 46 ++++++++++++++++
 3 files changed, 58 insertions(+), 45 deletions(-)
 create mode 100644 src/Microsoft.SqlTools.ServiceLayer/Connection/Contracts/EncryptionKeysChange.cs

diff --git a/src/Microsoft.SqlTools.Authentication/MSALEncryptedCacheHelper.cs b/src/Microsoft.SqlTools.Authentication/MSALEncryptedCacheHelper.cs
index 5371a90b..805b7360 100644
--- a/src/Microsoft.SqlTools.Authentication/MSALEncryptedCacheHelper.cs
+++ b/src/Microsoft.SqlTools.Authentication/MSALEncryptedCacheHelper.cs
@@ -28,7 +28,7 @@ namespace Microsoft.SqlTools.Authentication.Utility
         /// </summary>
         /// <param name="key">(out) Key used for encryption/decryption</param>
         /// <param name="iv">(out) IV used for encryption/decryption</param>
-        public delegate void IvKeyReadCallback(out string key, out string iv);
+        public delegate (string key, string iv) IvKeyReadCallback();
 
         /// <summary>
         /// Lock objects for serialization
@@ -97,7 +97,7 @@ namespace Microsoft.SqlTools.Authentication.Utility
         {
             if (this._key == null || this._iv == null)
             {
-                this._ivKeyReadCallback(out string key, out string iv);
+                (string key, string iv) = this._ivKeyReadCallback();
 
                 if (key != null)
                 {
diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
index e946123e..14802cd7 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
@@ -26,8 +26,6 @@ using Microsoft.SqlTools.Utility;
 using static Microsoft.SqlTools.Shared.Utility.Constants;
 using System.Diagnostics;
 using Microsoft.SqlTools.Authentication.Sql;
-using Microsoft.SqlTools.Credentials;
-using Microsoft.SqlTools.Credentials.Contracts;
 using Microsoft.SqlTools.Authentication;
 using Microsoft.SqlTools.Shared.Utility;
 using System.IO;
@@ -69,6 +67,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// </summary>
         private IAuthenticator authenticator;
 
+        /// <summary>
+        /// IV and Key as received from Encryption Key Notification event.
+        /// </summary>
+        private (string key, string iv) encryptionKeys;
+
         /// <summary>
         /// The SQL connection factory object
         /// </summary>
@@ -1104,6 +1107,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             serviceHost.SetRequestHandler(ChangeDatabaseRequest.Type, HandleChangeDatabaseRequest, true);
             serviceHost.SetRequestHandler(GetConnectionStringRequest.Type, HandleGetConnectionStringRequest, true);
             serviceHost.SetRequestHandler(BuildConnectionInfoRequest.Type, HandleBuildConnectionInfoRequest, true);
+            serviceHost.SetEventHandler(EncryptionKeysChangedNotification.Type, HandleEncryptionKeysNotificationEvent, false);
         }
 
         /// <summary>
@@ -1163,50 +1167,13 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             }
 
             var cachePath = Path.Combine(applicationPath, applicationName, AzureTokenFolder);
-            return new Authenticator(new (ApplicationClientId, applicationName, cachePath, MsalCacheName), ReadCacheIvKey);
+            return new Authenticator(new (ApplicationClientId, applicationName, cachePath, MsalCacheName), () => this.encryptionKeys);
         }
 
-        private void ReadCacheIvKey(out string? key, out string? iv)
+        private Task HandleEncryptionKeysNotificationEvent(EncryptionKeysChangeParams @params, EventContext context)
         {
-            Logger.Verbose("Reading Cached IV and Key from OS credential store.");
-
-            iv = null;
-            key = null;
-            try
-            {
-                // Read Cached Iv for MSAL cache (as Unicode) 
-                Credential ivCred = CredentialService.Instance.ReadCredential(new($"{AzureAccountProviderCredentials}|{MsalCacheName}-iv"));
-                if (!string.IsNullOrEmpty(ivCred.Password))
-                {
-                    iv = ivCred.Password;
-                }
-                else
-                {
-                    throw new Exception($"Could not read credential: {AzureAccountProviderCredentials}|{MsalCacheName}-iv");
-                }
-            }
-            catch (Exception ex)
-            {
-                Logger.Error(ex);
-            }
-
-            try
-            {
-                // Read Cached Key for MSAL cache (as Unicode)
-                Credential keyCred = CredentialService.Instance.ReadCredential(new($"{AzureAccountProviderCredentials}|{MsalCacheName}-key"));
-                if (!string.IsNullOrEmpty(keyCred.Password))
-                {
-                    key = keyCred.Password;
-                }
-                else
-                {
-                    throw new Exception($"Could not read credential: {AzureAccountProviderCredentials}|{MsalCacheName}-key");
-                }
-            }
-            catch (Exception ex)
-            {
-                Logger.Error(ex);
-            }
+            this.encryptionKeys = (@params.Key, @params.Iv);
+            return Task.FromResult(true);
         }
 
         private void RunConnectRequestHandlerTask(ConnectParams connectParams)
diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/Contracts/EncryptionKeysChange.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/Contracts/EncryptionKeysChange.cs
new file mode 100644
index 00000000..de35e8d0
--- /dev/null
+++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/Contracts/EncryptionKeysChange.cs
@@ -0,0 +1,46 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using Microsoft.SqlTools.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
+{
+
+    /// <summary>
+    /// Parameters for the MSAL cache encryption key notification.
+    /// </summary>
+    public class EncryptionKeysChangeParams
+    {
+        /// <summary>
+        /// Buffer encoded IV string for MSAL cache encryption
+        /// </summary>
+        public string Iv { get; set; }
+
+        /// <summary>
+        /// Buffer encoded Key string for MSAL cache encryption
+        /// </summary>
+        public string Key { get; set; }
+
+        /// <summary>
+        /// Default constructor to resolve nullables
+        /// </summary>
+        /// <param name="key">Key for MSAL cache encryption</param>
+        /// <param name="iv">Iv for MSAL cache encryption</param>
+        public EncryptionKeysChangeParams(string key, string iv) {
+            this.Iv = iv;
+            this.Key = key;
+        }
+    }
+
+    /// <summary>
+    /// Defines an event that is sent from the client to notify the encryption keys are changed.
+    /// </summary>
+    public class EncryptionKeysChangedNotification
+    {
+        public static readonly
+            EventType<EncryptionKeysChangeParams> Type =
+            EventType<EncryptionKeysChangeParams>.Create("connection/encryptionKeysChanged");
+    }
+}