Support updating access token when found expired for OE (#1772)

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

@@ -18,6 +18,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
     /// </summary>
     public class ConnectionInfo
     {
+        private static readonly DateTime UnixEpochUtc = new DateTime(1970, 1, 1, 0, 0, 0, 0, DateTimeKind.Utc);
+
         /// <summary>
         /// Constructor
         /// </summary>
@@ -27,7 +29,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             OwnerUri = ownerUri;
             ConnectionDetails = details;
             ConnectionId = Guid.NewGuid();
-            IntellisenseMetrics = new InteractionMetrics<double>(new int[] {50, 100, 200, 500, 1000, 2000});
+            IntellisenseMetrics = new InteractionMetrics<double>(new int[] { 50, 100, 200, 500, 1000, 2000 });
         }
 
         /// <summary>
@@ -73,10 +75,16 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// </summary>
         public bool IsSqlDb { get; set; }
 
+        /// <summary>
         /// Returns true if the sql connection is to a DW instance
         /// </summary>
         public bool IsSqlDW { get; set; }
 
+        /// <summary>
+        /// Returns true if Authentication mode is AzureMFA, determines if Access token is in use.
+        /// </summary>
+        public bool IsAzureAuth { get; set; }
+
         /// <summary>
         /// Returns the connection Engine Edition
         /// </summary>
@@ -158,8 +166,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         public void RemoveConnection(string connectionType)
         {
             Validate.IsNotNullOrEmptyString("Connection Type", connectionType);
-            DbConnection connection;
+            ConnectionTypeToConnectionMap.TryRemove(connectionType, out _);
-            ConnectionTypeToConnectionMap.TryRemove(connectionType, out connection);
         }
 
         /// <summary>
@@ -169,18 +176,44 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         {
             foreach (var type in AllConnectionTypes)
             {
-                DbConnection connection;
+                ConnectionTypeToConnectionMap.TryRemove(type, out _);
-                ConnectionTypeToConnectionMap.TryRemove(type, out connection);
             }
         }
 
         /// <summary>
         /// Updates the Auth Token and Expires On fields
         /// </summary>
-        public void UpdateAuthToken(string token, int expiresOn)
+        public bool TryUpdateAccessToken(SecurityToken? securityToken)
         {
-            ConnectionDetails.AzureAccountToken = token;
+            if (securityToken != null && !string.IsNullOrEmpty(securityToken.Token) && IsAzureAuth && IsAccessTokenExpired)
-            ConnectionDetails.ExpiresOn = expiresOn;
+            {
+                ConnectionDetails.AzureAccountToken = securityToken.Token;
+                ConnectionDetails.ExpiresOn = securityToken.ExpiresOn;
+                return true;
+            }
+            return false;
+        }
+
+        /// <summary>
+        /// Returns true if Access token saved in connection details is expired or about to expire in 2 minutes.
+        /// </summary>
+        private bool IsAccessTokenExpired
+        {
+            get
+            {
+                if (IsAzureAuth && ConnectionDetails.ExpiresOn != null && double.TryParse(ConnectionDetails.ExpiresOn.ToString(), out var expiresOn))
+                {
+                    DateTime dateTime = UnixEpochUtc.AddSeconds(expiresOn);
+
+                    // Check if access token is already expired or shall expire in 2 minutes.
+                    if (dateTime <= DateTime.UtcNow.AddMinutes(2))
+                    {
+                        Logger.Verbose("Access token found expired.");
+                        return true;
+                    }
+                }
+                return false;
+            }
         }
     }
 }

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

@@ -304,7 +304,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
                 return;
             }
             this.TokenUpdateUris.Remove(tokenRefreshedParams.Uri, out var result);
-            connection.UpdateAuthToken(tokenRefreshedParams.Token, tokenRefreshedParams.ExpiresOn);
+            connection.TryUpdateAccessToken(new SecurityToken() { Token = tokenRefreshedParams.Token, ExpiresOn = tokenRefreshedParams.ExpiresOn });
         }
 
         /// <summary>
@@ -585,6 +585,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
                 connectionInfo.MajorVersion = serverInfo.ServerMajorVersion;
                 connectionInfo.IsSqlDb = serverInfo.EngineEditionId == (int)DatabaseEngineEdition.SqlDatabase;
                 connectionInfo.IsSqlDW = (serverInfo.EngineEditionId == (int)DatabaseEngineEdition.SqlDataWarehouse);
+                // Determines that access token is used for creating connection.
+                connectionInfo.IsAzureAuth = connectionInfo.ConnectionDetails.AuthenticationType == "AzureMFA";
                 connectionInfo.EngineEdition = (DatabaseEngineEdition)serverInfo.EngineEditionId;
                 // Azure Data Studio supports SQL Server 2014 and later releases.
                 response.IsSupportedVersion = serverInfo.IsCloud || serverInfo.ServerMajorVersion >= 12;

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

@@ -7,7 +7,7 @@ using Microsoft.SqlTools.Hosting.Protocol.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
 {
-    class RefreshTokenParams 
+    public class RefreshTokenParams
     {
         /// <summary>
         /// ID of the tenant
@@ -39,14 +39,14 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
     /// <summary>
     /// Refresh token request mapping entry 
     /// </summary>
-    class RefreshTokenNotification
+    public class RefreshTokenNotification
     {
         public static readonly
             EventType<RefreshTokenParams> Type =
             EventType<RefreshTokenParams>.Create("account/refreshToken");
     }
 
-    class TokenRefreshedParams
+    public class TokenRefreshedParams
     {
         /// <summary>
         /// Gets or sets the refresh token.
@@ -64,7 +64,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
         public string Uri { get; set; }
     }
 
-    class TokenRefreshedNotification
+    public class TokenRefreshedNotification
     {
         public static readonly
         EventType<TokenRefreshedParams> Type =

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

@@ -0,0 +1,25 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
+{
+    public class SecurityToken
+    {
+        /// <summary>
+        /// Gets or sets the refresh token.
+        /// </summary>
+        public string Token { get; set; }
+
+        /// <summmary>
+        /// Gets or sets the token expiration, a Unix epoch 
+        /// </summary>
+        public int? ExpiresOn { get; set; }
+
+        /// <summmary>
+        /// Gets or sets the token type, e.g. 'Bearer'
+        /// </summary>
+        public string? TokenType { get; set; }
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Contracts/ExpandRequest.cs

@@ -4,6 +4,7 @@
 //
 
 using Microsoft.SqlTools.Hosting.Protocol.Contracts;
+using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
 {
@@ -49,6 +50,11 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
         /// Path identifying the node to expand. See <see cref="NodeInfo.NodePath"/> for details
         /// </summary>
         public string NodePath { get; set; }
+
+        /// <summary>
+        /// Security token for AzureMFA authentication for refresing access token on connection.
+        /// </summary>
+        public SecurityToken? SecurityToken { get; set; }
     }
 
     /// <summary>

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Contracts/RefreshRequest.cs

@@ -10,7 +10,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
     /// <summary>
     /// Parameters to the <see cref="ExpandRequest"/>.
     /// </summary>
-    public class RefreshParams: ExpandParams
+    public class RefreshParams : ExpandParams
     {
     }
 

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/TreeNode.cs

@@ -240,14 +240,14 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
         /// Expands this node and returns its children
         /// </summary>
         /// <returns>Children as an IList. This is the raw children collection, not a copy</returns>
-        public IList<TreeNode> Expand(string name, CancellationToken cancellationToken)
+        public IList<TreeNode> Expand(string name, CancellationToken cancellationToken, string? accessToken = null)
         {
             // TODO consider why solution explorer has separate Children and Items options
             if (children.IsInitialized)
             {
                 return children;
             }
-            PopulateChildren(false, name, cancellationToken);
+            PopulateChildren(false, name, cancellationToken, accessToken);
             return children;
         }
 
@@ -255,19 +255,19 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
         /// Expands this node and returns its children
         /// </summary>
         /// <returns>Children as an IList. This is the raw children collection, not a copy</returns>
-        public IList<TreeNode> Expand(CancellationToken cancellationToken)
+        public IList<TreeNode> Expand(CancellationToken cancellationToken, string? accessToken = null)
         {
-            return Expand(null, cancellationToken);
+            return Expand(null, cancellationToken, accessToken);
         }
 
         /// <summary>
         /// Refresh this node and returns its children
         /// </summary>
         /// <returns>Children as an IList. This is the raw children collection, not a copy</returns>
-        public virtual IList<TreeNode> Refresh(CancellationToken cancellationToken)
+        public virtual IList<TreeNode> Refresh(CancellationToken cancellationToken, string? accessToken = null)
         {
             // TODO consider why solution explorer has separate Children and Items options
-            PopulateChildren(true, null, cancellationToken);
+            PopulateChildren(true, null, cancellationToken, accessToken);
             return children;
         }
 
@@ -319,7 +319,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
             return Parent as T;
         }
 
-        protected virtual void PopulateChildren(bool refresh, string name, CancellationToken cancellationToken)
+        protected virtual void PopulateChildren(bool refresh, string name, CancellationToken cancellationToken, string? accessToken = null)
         {
             Logger.Write(TraceEventType.Verbose, string.Format(CultureInfo.InvariantCulture, "Populating oe node :{0}", this.GetNodePath()));
             Debug.Assert(IsAlwaysLeaf == false);
@@ -335,6 +335,9 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
             children.Clear();
             BeginChildrenInit();
 
+            // Update access token for future queries
+            context.UpdateAccessToken(accessToken);
+
             try
             {
                 ErrorMessage = null;

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs

@@ -212,7 +212,8 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
                 }
                 else
                 {
-                    RunExpandTask(session, expandParams);
+                    bool refreshNeeded = session.ConnectionInfo.TryUpdateAccessToken(expandParams.SecurityToken);
+                    RunExpandTask(session, expandParams, refreshNeeded);
                     return true;
                 }
             };
@@ -239,6 +240,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             }
             else
             {
+                session.ConnectionInfo.TryUpdateAccessToken(refreshParams.SecurityToken);
                 RunExpandTask(session, refreshParams, true);
             }
             await context.SendResult(true);
@@ -373,12 +375,12 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
 
         }
 
-        internal Task<ExpandResponse> ExpandNode(ObjectExplorerSession session, string nodePath, bool forceRefresh = false)
+        internal Task<ExpandResponse> ExpandNode(ObjectExplorerSession session, string nodePath, bool forceRefresh = false, SecurityToken? securityToken = null)
         {
-            return Task.Run(() => QueueExpandNodeRequest(session, nodePath, forceRefresh));
+            return Task.Run(() => QueueExpandNodeRequest(session, nodePath, forceRefresh, securityToken));
         }
 
-        internal ExpandResponse QueueExpandNodeRequest(ObjectExplorerSession session, string nodePath, bool forceRefresh = false)
+        internal ExpandResponse QueueExpandNodeRequest(ObjectExplorerSession session, string nodePath, bool forceRefresh = false, SecurityToken? securityToken = null)
         {
             NodeInfo[] nodes = null;
             TreeNode node = session.Root.FindNodeByPath(nodePath);
@@ -432,15 +434,21 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
                            waitForLockTimeout: timeout,
                            bindOperation: (bindingContext, cancelToken) =>
                            {
+                               if (!session.ConnectionInfo.IsAzureAuth)
+                               {
+                                   // explicitly set null here to prevent setting access token for non-Azure auth modes.
+                                   securityToken = null;
+                               }
+
                                if (forceRefresh)
                                {
                                    Logger.Verbose($"Forcing refresh for {nodePath}");
-                                   nodes = node.Refresh(cancelToken).Select(x => x.ToNodeInfo()).ToArray();
+                                   nodes = node.Refresh(cancelToken, securityToken?.Token).Select(x => x.ToNodeInfo()).ToArray();
                                }
                                else
                                {
                                    Logger.Verbose($"Expanding {nodePath}");
-                                   nodes = node.Expand(cancelToken).Select(x => x.ToNodeInfo()).ToArray();
+                                   nodes = node.Expand(cancelToken, securityToken?.Token).Select(x => x.ToNodeInfo()).ToArray();
                                }
                                response.Nodes = nodes;
                                response.ErrorMessage = node.ErrorMessage;
@@ -635,7 +643,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
         private async Task ExpandNodeAsync(ObjectExplorerSession session, ExpandParams expandParams, CancellationToken cancellationToken, bool forceRefresh = false)
         {
             ExpandResponse response = null;
-            response = await ExpandNode(session, expandParams.NodePath, forceRefresh);
+            response = await ExpandNode(session, expandParams.NodePath, forceRefresh, expandParams.SecurityToken);
             if (cancellationToken.IsCancellationRequested)
             {
                 Logger.Write(TraceEventType.Verbose, "OE expand canceled");

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/DatabaseTreeNode.cs

@@ -16,7 +16,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
 {
     internal partial class DatabaseTreeNode
     {
-        public DatabaseTreeNode(ServerNode serverNode, string databaseName): this()
+        public DatabaseTreeNode(ServerNode serverNode, string databaseName) : this()
         {
             Parent = serverNode;
             NodeValue = databaseName;
@@ -52,12 +52,12 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
             }
         }
 
-        protected override void PopulateChildren(bool refresh, string name, CancellationToken cancellationToken)
+        protected override void PopulateChildren(bool refresh, string name, CancellationToken cancellationToken, string? accessToken = null)
         {
             var smoQueryContext = this.GetContextAs<SmoQueryContext>();
             if (IsAccessible(smoQueryContext))
             {
-                base.PopulateChildren(refresh, name, cancellationToken);
+                base.PopulateChildren(refresh, name, cancellationToken, accessToken);
             }
             else
             {

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoExtensions.cs

@@ -0,0 +1,27 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+using Microsoft.SqlServer.Management.Smo;
+using Microsoft.SqlTools.ServiceLayer.Connection;
+
+namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
+{
+    internal static class SmoExtensions
+    {
+        /// <summary>
+        /// Updates access token on the connection context of <paramref name="sqlObj"/> instance.
+        /// </summary>
+        /// <param name="sqlObj">(this) SMO SQL Object containing connection context.</param>
+        /// <param name="accessToken">Access token</param>
+        public static void UpdateAccessToken(this SqlSmoObject sqlObj, string accessToken)
+        {
+            if (sqlObj != null && !string.IsNullOrEmpty(accessToken)
+                && sqlObj.ExecutionManager != null
+                && sqlObj.ExecutionManager.ConnectionContext != null)
+            {
+                sqlObj.ExecutionManager.ConnectionContext.AccessToken = new AzureAccessToken(accessToken);
+            }
+        }
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoQueryContext.cs

@@ -46,7 +46,8 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
         /// <summary>
         /// The server SMO will query against
         /// </summary>
-        public Server Server { 
+        public Server Server
+        {
             get
             {
                 return GetObjectWithOpenedConnection(server);
@@ -56,7 +57,8 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
         /// <summary>
         /// Optional Database context object to query against
         /// </summary>
-        public Database Database { 
+        public Database Database
+        {
             get
             {
                 return GetObjectWithOpenedConnection(database);
@@ -70,7 +72,8 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
         /// <summary>
         /// Parent of a give node to use for queries
         /// </summary>
-        public SmoObjectBase Parent { 
+        public SmoObjectBase Parent
+        {
             get
             {
                 return GetObjectWithOpenedConnection(parent);
@@ -147,7 +150,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
         {
             get
             {
-                if(validFor == 0)
+                if (validFor == 0)
                 {
                     validFor = ServerVersionHelper.GetValidForFlag(SqlServerType, Database);
                 }
@@ -169,6 +172,31 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
             return smoObj;
         }
 
+        /// <summary>
+        /// Updates access token on parent connection context.
+        /// </summary>
+        /// <param name="accessToken">Acquired access token</param>
+        public void UpdateAccessToken(string? accessToken)
+        {
+            if (!string.IsNullOrEmpty(accessToken))
+            {
+                // Update all applicable nodes that could contain access token
+                // to prevent stale token from being re-used.
+                if (server != null)
+                {
+                    (server as SqlSmoObject).UpdateAccessToken(accessToken);
+                }
+                if (database != null)
+                {
+                    (database as SqlSmoObject).UpdateAccessToken(accessToken);
+                }
+                if (parent != null)
+                {
+                    (parent as SqlSmoObject).UpdateAccessToken(accessToken);
+                }
+            }
+        }
+
         /// <summary>
         /// Ensures the server objects connection context is open. This is used by all child objects, and 
         /// the only way to easily access is via the server object. This should be called during access of

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoWrapper.cs

@@ -14,11 +14,21 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
     /// </summary>
     internal class SmoWrapper
     {
+        /// <summary>
+        /// Creates instance of <see cref="Server"/> from provided <paramref name="serverConn"/> instance.
+        /// </summary>
+        /// <param name="serverConn">Server connection instance.</param>
+        /// <returns>Server instance.</returns>
         public virtual Server CreateServer(ServerConnection serverConn)
         {
             return serverConn == null ? null : new Server(serverConn);
         }
 
+        /// <summary>
+        /// Checks if connection is open on the <paramref name="smoObj"/> instance.
+        /// </summary>
+        /// <param name="smoObj">SMO Object containing connection context.</param>
+        /// <returns>True if connection is open, otherwise false.</returns>
         public virtual bool IsConnectionOpen(SmoObjectBase smoObj)
         {
             SqlSmoObject sqlObj = smoObj as SqlSmoObject;
@@ -28,6 +38,10 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
                 && sqlObj.ExecutionManager.ConnectionContext.IsOpen;
         }
 
+        /// <summary>
+        /// Opens connection on the connection context of <paramref name="smoObj"/> instance.
+        /// </summary>
+        /// <param name="smoObj">SMO Object containing connection context.</param>
         public virtual void OpenConnection(SmoObjectBase smoObj)
         {
             SqlSmoObject sqlObj = smoObj as SqlSmoObject;

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/Utility/LiveConnectionHelper.cs

@@ -18,7 +18,8 @@ using System.Threading;
 
 namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.Utility
 {
-    public class LiveConnectionException : Exception {       
+    public class LiveConnectionException : Exception
+    {
         public LiveConnectionException(string message)
             : base(message) { }
     }

test/Microsoft.SqlTools.ServiceLayer.UnitTests/ObjectExplorer/ObjectExplorerServiceTests.cs

@@ -295,7 +295,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
         public void FindNodeCanExpandParentNodes()
         {
             var mockTreeNode = new Mock<TreeNode>();
-            object[] populateChildrenArguments = { ItExpr.Is<bool>(x => x == false), ItExpr.IsNull<string>(), new CancellationToken() };
+            object[] populateChildrenArguments = { ItExpr.Is<bool>(x => x == false), ItExpr.IsNull<string>(), new CancellationToken(), ItExpr.IsNull<string>() };
             mockTreeNode.Protected().Setup("PopulateChildren", populateChildrenArguments);
             mockTreeNode.Object.IsAlwaysLeaf = false;