SNOW-715504: MFA token cache support

Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs

@@ -14,6 +14,7 @@
 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;
@@ -2271,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")] // to enroll to mfa authentication edit your user profile
+        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 the default credential manager is a file based one. Uncomment the following line to test in memory implementation.
+            // SnowflakeCredentialManagerFactory.UseInMemoryCredentialManager();
+            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 the default credential manager is a file based one. Uncomment the following line to test in memory implementation.
+            // SnowflakeCredentialManagerFactory.UseInMemoryCredentialManager();
+            // 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;

Snowflake.Data.Tests/UnitTests/ArrowResultSetTest.cs

@@ -33,7 +33,7 @@ public void BeforeTest()
             // by default generate Int32 values from 1 to RowCount
             PrepareTestCase(SFDataType.FIXED, 0, Enumerable.Range(1, RowCount).ToArray());
         }
-        
+
         [Test]
         public void TestResultFormatIsArrow()
         {
@@ -140,7 +140,7 @@ public void TestGetValueReturnsNull()
             var arrowResultSet = new ArrowResultSet(responseData, sfStatement, new CancellationToken());
 
             arrowResultSet.Next();
-            
+
             Assert.AreEqual(true, arrowResultSet.IsDBNull(0));
             Assert.AreEqual(DBNull.Value, arrowResultSet.GetValue(0));
         }
@@ -152,7 +152,7 @@ public void TestGetDecimal()
 
             TestGetNumber(testValues);
         }
-        
+
         [Test]
         public void TestGetNumber64()
         {
@@ -165,7 +165,7 @@ public void TestGetNumber64()
         public void TestGetNumber32()
         {
             var testValues = new int[] { 0, 100, -100, Int32.MaxValue, Int32.MinValue };
-            
+
             TestGetNumber(testValues);
         }
 
@@ -176,7 +176,7 @@ public void TestGetNumber16()
 
             TestGetNumber(testValues);
         }
-        
+
         [Test]
         public void TestGetNumber8()
         {
@@ -200,7 +200,7 @@ private void TestGetNumber(IEnumerable testValues)
                     Assert.AreEqual(expectedValue, _arrowResultSet.GetDecimal(ColumnIndex));
                     Assert.AreEqual(expectedValue, _arrowResultSet.GetDouble(ColumnIndex));
                     Assert.AreEqual(expectedValue, _arrowResultSet.GetFloat(ColumnIndex));
-            
+
                     if (expectedValue >= Int64.MinValue && expectedValue <= Int64.MaxValue)
                     {
                         // get integer value
@@ -230,7 +230,7 @@ public void TestGetBoolean()
             var testValues = new bool[] { true, false };
 
             PrepareTestCase(SFDataType.BOOLEAN, 0, testValues);
-            
+
             foreach (var testValue in testValues)
             {
                 _arrowResultSet.Next();
@@ -245,15 +245,15 @@ public void TestGetReal()
             var testValues = new double[] { 0, Double.MinValue, Double.MaxValue };
 
             PrepareTestCase(SFDataType.REAL, 0, testValues);
-            
+
             foreach (var testValue in testValues)
             {
                 _arrowResultSet.Next();
                 Assert.AreEqual(testValue, _arrowResultSet.GetValue(ColumnIndex));
                 Assert.AreEqual(testValue, _arrowResultSet.GetDouble(ColumnIndex));
             }
         }
-        
+
         [Test]
         public void TestGetText()
         {
@@ -264,15 +264,15 @@ public void TestGetText()
             };
 
             PrepareTestCase(SFDataType.TEXT, 0, testValues);
-            
+
             foreach (var testValue in testValues)
             {
                 _arrowResultSet.Next();
                 Assert.AreEqual(testValue, _arrowResultSet.GetValue(ColumnIndex));
                 Assert.AreEqual(testValue, _arrowResultSet.GetString(ColumnIndex));
             }
         }
-        
+
         [Test]
         public void TestGetTextWithOneChar()
         {
@@ -290,14 +290,14 @@ public void TestGetTextWithOneChar()
 #endif
 
             PrepareTestCase(SFDataType.TEXT, 0, testValues);
-            
+
             foreach (var testValue in testValues)
             {
                 _arrowResultSet.Next();
                 Assert.AreEqual(testValue, _arrowResultSet.GetChar(ColumnIndex));
             }
         }
-        
+
         [Test]
         public void TestGetArray()
         {
@@ -308,7 +308,7 @@ public void TestGetArray()
             };
 
             PrepareTestCase(SFDataType.ARRAY, 0, testValues);
-            
+
             foreach (var testValue in testValues)
             {
                 _arrowResultSet.Next();
@@ -320,7 +320,7 @@ public void TestGetArray()
                 Assert.AreEqual(testValue.Length, str.Length);
             }
         }
-        
+
         [Test]
         public void TestGetBinary()
         {
@@ -342,7 +342,7 @@ public void TestGetBinary()
                     Assert.AreEqual(testValue[j], buffer[j], "position " + j);
             }
         }
-        
+
         [Test]
         public void TestGetDate()
         {
@@ -354,15 +354,15 @@ public void TestGetDate()
             };
 
             PrepareTestCase(SFDataType.DATE, 0, testValues);
-            
+
             foreach (var testValue in testValues)
             {
                 _arrowResultSet.Next();
                 Assert.AreEqual(testValue, _arrowResultSet.GetValue(ColumnIndex));
                 Assert.AreEqual(testValue, _arrowResultSet.GetDateTime(ColumnIndex));
             }
         }
-        
+
         [Test]
         public void TestGetTime()
         {
@@ -384,7 +384,7 @@ public void TestGetTime()
                     Assert.AreEqual(testValue, _arrowResultSet.GetValue(ColumnIndex));
                     Assert.AreEqual(testValue, _arrowResultSet.GetDateTime(ColumnIndex));
                 }
-            }        
+            }
         }
 
         [Test]
@@ -513,10 +513,10 @@ private QueryExecResponseData PrepareResponseData(RecordBatch recordBatch, SFDat
             return new QueryExecResponseData
             {
                 rowType = recordBatch.Schema.FieldsList
-                    .Select(col => 
+                    .Select(col =>
                         new ExecResponseRowType
                         {
-                            name = col.Name, 
+                            name = col.Name,
                             type = sfType.ToString(),
                             scale = scale
                         }).ToList(),
@@ -531,7 +531,7 @@ private string ConvertToBase64String(RecordBatch recordBatch)
         {
             if (recordBatch == null)
                 return "";
-            
+
             using (var stream = new MemoryStream())
             {
                 using (var writer = new ArrowStreamWriter(stream, recordBatch.Schema))
@@ -542,12 +542,12 @@ private string ConvertToBase64String(RecordBatch recordBatch)
                 return Convert.ToBase64String(stream.ToArray());
             }
         }
-        
+
         private SFStatement PrepareStatement()
         {
             SFSession session = new SFSession("user=user;password=password;account=account;", null);
             return new SFStatement(session);
         }
-        
+
     }
 }