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>
@@ -72,11 +74,17 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// Returns true if the db connection is to a SQL db instance
         /// </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 connection);
+            ConnectionTypeToConnectionMap.TryRemove(connectionType, out _);
         }
 
         /// <summary>
@@ -169,18 +176,44 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         {
             foreach (var type in AllConnectionTypes)
             {
-                DbConnection connection;
-                ConnectionTypeToConnectionMap.TryRemove(type, out connection);
+                ConnectionTypeToConnectionMap.TryRemove(type, out _);
             }
-        } 
+        }
 
         /// <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;
-            ConnectionDetails.ExpiresOn = expiresOn;
+            if (securityToken != null && !string.IsNullOrEmpty(securityToken.Token) && IsAzureAuth && IsAccessTokenExpired)
+            {
+                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/Connection/Contracts/SecurityTokenRequest.cs

@@ -7,7 +7,7 @@ using Microsoft.SqlTools.Hosting.Protocol.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
 {
-    class RequestSecurityTokenParams 
+    class RequestSecurityTokenParams
     {
         /// <summary>
         /// Gets or sets the address of the authority to issue token.