SNOW-715504: MFA token cache support

Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs

@@ -2,25 +2,25 @@
  * Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
  */
 
+using System;
+using System.Data;
 using System.Data.Common;
+using System.Diagnostics;
 using System.Net;
+using System.Runtime.InteropServices;
+using System.Threading;
+using System.Threading.Tasks;
+using NUnit.Framework;
+using Snowflake.Data.Client;
+using Snowflake.Data.Core;
 using Snowflake.Data.Core.Session;
+using Snowflake.Data.Core.Tools;
+using Snowflake.Data.Log;
+using Snowflake.Data.Tests.Mock;
 using Snowflake.Data.Tests.Util;
 
 namespace Snowflake.Data.Tests.IntegrationTests
 {
-    using NUnit.Framework;
-    using Snowflake.Data.Client;
-    using System.Data;
-    using System;
-    using Snowflake.Data.Core;
-    using System.Threading.Tasks;
-    using System.Threading;
-    using Snowflake.Data.Log;
-    using System.Diagnostics;
-    using Snowflake.Data.Tests.Mock;
-    using System.Runtime.InteropServices;
-    using System.Net.Http;
 
     [TestFixture]
     class SFConnectionIT : SFBaseTest
@@ -2272,6 +2272,52 @@ public void TestUseMultiplePoolsConnectionPoolByDefault()
             Assert.AreEqual(ConnectionPoolType.MultipleConnectionPool, poolVersion);
         }
 
+        [Test]
+        // [Ignore("This test requires manual interaction and therefore cannot be run in CI")]
+        public void TestMFATokenCachingWithPasscodeFromConnectionString()
+        {
+            // Use a connection with MFA enabled and set passcode property for mfa authentication. e.g. ConnectionString + ";authenticator=username_password_mfa;passcode=(set proper passcode)"
+            // ACCOUNT PARAMETER ALLOW_CLIENT_MFA_CACHING should be set to true in the account.
+            // On Mac/Linux OS default credential manager is in memory so please uncomment following line to use file based credential manager
+            // SnowflakeCredentialManagerFactory.UseFileCredentialManager();
+            using (SnowflakeDbConnection conn = new SnowflakeDbConnection())
+            {
+                conn.ConnectionString
+                    = ConnectionString
+                      + ";authenticator=username_password_mfa;application=DuoTest;minPoolSize=0;passcode=(set proper passcode)";
+
+
+                // Authenticate to retrieve and store the token if doesn't exist or invalid
+                Task connectTask = conn.OpenAsync(CancellationToken.None);
+                connectTask.Wait();
+                Assert.AreEqual(ConnectionState.Open, conn.State);
+            }
+        }
+
+        [Test]
+        [Ignore("Requires manual steps and environment with mfa authentication enrolled")] // to enroll to mfa authentication edit your user profile
+        public void TestMfaWithPasswordConnectionUsingPasscodeWithSecureString()
+        {
+            // Use a connection with MFA enabled and Passcode property on connection instance.
+            // ACCOUNT PARAMETER ALLOW_CLIENT_MFA_CACHING should be set to true in the account.
+            // On Mac/Linux OS default credential manager is in memory so please uncomment following line to use file based credential manager
+            // SnowflakeCredentialManagerFactory.UseFileCredentialManager();
+            // arrange
+            using (SnowflakeDbConnection conn = new SnowflakeDbConnection())
+            {
+                conn.Passcode = SecureStringHelper.Encode("$(set proper passcode)");
+                // manual action: stop here in breakpoint to provide proper passcode by: conn.Passcode = SecureStringHelper.Encode("...");
+                conn.ConnectionString = ConnectionString + "minPoolSize=2;application=DuoTest;";
+
+                // act
+                Task connectTask = conn.OpenAsync(CancellationToken.None);
+                connectTask.Wait();
+
+                // assert
+                Assert.AreEqual(ConnectionState.Open, conn.State);
+            }
+        }
+
         [Test]
         [TestCase("connection_timeout=5;")]
         [TestCase("")]

Snowflake.Data.Tests/Mock/MockLoginMFATokenCacheRestRequester.cs

@@ -0,0 +1,89 @@
+using System;
+using System.Collections.Generic;
+using System.Net.Http;
+using System.Threading;
+using System.Threading.Tasks;
+using Snowflake.Data.Core;
+
+namespace Snowflake.Data.Tests.Mock
+{
+    using Microsoft.IdentityModel.Tokens;
+
+    class MockLoginMFATokenCacheRestRequester: IMockRestRequester
+    {
+        internal Queue<LoginRequest> LoginRequests { get; } = new();
+
+        internal Queue<LoginResponseData> LoginResponses { get; } = new();
+
+        public T Get<T>(IRestRequest request)
+        {
+            return Task.Run(async () => await (GetAsync<T>(request, CancellationToken.None)).ConfigureAwait(false)).Result;
+        }
+
+        public Task<T> GetAsync<T>(IRestRequest request, CancellationToken cancellationToken)
+        {
+            return Task.FromResult<T>((T)(object)null);
+        }
+
+        public Task<HttpResponseMessage> GetAsync(IRestRequest request, CancellationToken cancellationToken)
+        {
+            return Task.FromResult<HttpResponseMessage>(null);
+        }
+
+        public HttpResponseMessage Get(IRestRequest request)
+        {
+            return null;
+        }
+
+        public T Post<T>(IRestRequest postRequest)
+        {
+            return Task.Run(async () => await (PostAsync<T>(postRequest, CancellationToken.None)).ConfigureAwait(false)).Result;
+        }
+
+        public Task<T> PostAsync<T>(IRestRequest postRequest, CancellationToken cancellationToken)
+        {
+            SFRestRequest sfRequest = (SFRestRequest)postRequest;
+            if (sfRequest.jsonBody is LoginRequest)
+            {
+                LoginRequests.Enqueue((LoginRequest) sfRequest.jsonBody);
+                var responseData = this.LoginResponses.IsNullOrEmpty() ? new LoginResponseData()
+                {
+                    token = "session_token",
+                    masterToken = "master_token",
+                    authResponseSessionInfo = new SessionInfo(),
+                    nameValueParameter = new List<NameValueParameter>()
+                } : this.LoginResponses.Dequeue();
+                var authnResponse = new LoginResponse
+                {
+                    data = responseData,
+                    success = true
+                };
+
+                // login request return success
+                return Task.FromResult<T>((T)(object)authnResponse);
+            }
+            else if (sfRequest.jsonBody is CloseResponse)
+            {
+                var authnResponse = new CloseResponse()
+                {
+                    success = true
+                };
+
+                // login request return success
+                return Task.FromResult<T>((T)(object)authnResponse);
+            }
+            throw new NotImplementedException();
+        }
+
+        public void setHttpClient(HttpClient httpClient)
+        {
+            // Nothing to do
+        }
+
+        public void Reset()
+        {
+            LoginRequests.Clear();
+            LoginResponses.Clear();
+        }
+    }
+}

Snowflake.Data.Tests/Mock/MockSnowflakeDbConnection.cs

@@ -78,10 +78,10 @@ public override Task OpenAsync(CancellationToken cancellationToken)
                 cancellationToken);
 
         }
-        
+
         private void SetMockSession()
         {
-            SfSession = new SFSession(ConnectionString, Password, _restRequester);
+            SfSession = new SFSession(ConnectionString, Password, Passcode, EasyLoggingStarter.Instance, _restRequester);
 
             _connectionTimeout = (int)SfSession.connectionTimeout.TotalSeconds;
 
@@ -92,7 +92,7 @@ private void OnSessionEstablished()
         {
             _connectionState = ConnectionState.Open;
         }
-        
+
         protected override bool CanReuseSession(TransactionRollbackStatus transactionRollbackStatus)
         {
             return false;