SNOW-715524: Add SSO token cache

README.md

@@ -168,6 +168,7 @@ The following table lists all valid connection properties:
 | CLIENT_CONFIG_FILE             | No       | The location of the client configuration json file. In this file you can configure easy logging feature.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                  |
 | ALLOWUNDERSCORESINHOST         | No       | Specifies whether to allow underscores in account names. This impacts PrivateLink customers whose account names contain underscores. In this situation, you must override the default value by setting allowUnderscoresInHost to true.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    |
 | QUERY_TAG                      | No       | Optional string that can be used to tag queries and other SQL statements executed within a connection. The tags are displayed in the output of the QUERY_HISTORY , QUERY_HISTORY_BY_* functions.<br/> To set QUERY_TAG on the statement level you can use SnowflakeDbCommand.QueryTag.                                                                                                                                                                                                                                                                                                                                                                                                                                    |
+| ALLOW_SSO_TOKEN_CACHING        | No       | Specifies whether to cache tokens and use them for SSO authentication.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    |
 
 <br />
 

Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs

@@ -1,4 +1,4 @@
-/*
+/*
  * Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
  */
 
@@ -1020,6 +1020,52 @@ public void TestSSOConnectionTimeoutAfter10s()
             Assert.LessOrEqual(stopwatch.ElapsedMilliseconds, (waitSeconds + 5) * 1000);
         }
 
+        [Test]
+        [Ignore("This test requires manual interaction and therefore cannot be run in CI")]
+        public void TestSSOConnectionWithTokenCaching()
+        {
+            using (IDbConnection conn = new SnowflakeDbConnection())
+            {
+                conn.ConnectionString
+                    = ConnectionStringWithoutAuth
+                        + ";authenticator=externalbrowser;[email protected];allow_sso_token_caching=true;";
+
+                // Authenticate to retrieve and store the token if doesn't exist or invalid
+                conn.Open();
+                Assert.AreEqual(ConnectionState.Open, conn.State);
+
+                conn.Close();
+                Assert.AreEqual(ConnectionState.Closed, conn.State);
+
+                // Authenticate using the token
+                conn.Open();
+                Assert.AreEqual(ConnectionState.Open, conn.State);
+            }
+        }
+
+        [Test]
+        [Ignore("This test requires manual interaction and therefore cannot be run in CI")]
+        public void TestSSOConnectionWithInvalidCachedToken()
+        {
+            using (IDbConnection conn = new SnowflakeDbConnection())
+            {
+                conn.ConnectionString
+                    = ConnectionStringWithoutAuth
+                        + ";authenticator=externalbrowser;[email protected];allow_sso_token_caching=true;";
+
+                var key = SnowflakeCredentialManagerFactory.BuildCredentialKey(testConfig.host, testConfig.user, TokenType.IdToken.ToString());
+                var credentialManager = SnowflakeCredentialManagerInMemoryImpl.Instance;
+                credentialManager.SaveCredentials(key, "wrongToken");
+
+                SnowflakeCredentialManagerFactory.SetCredentialManager(credentialManager);
+
+                conn.Open();
+                Assert.AreEqual(ConnectionState.Open, conn.State);
+
+                SnowflakeCredentialManagerFactory.UseDefaultCredentialManager();
+            }
+        }
+
         [Test]
         [Ignore("This test requires manual interaction and therefore cannot be run in CI")]
         public void TestSSOConnectionWithWrongUser()
@@ -2231,17 +2277,42 @@ public void TestConnectStringWithQueryTag()
             {
                 string expectedQueryTag = "Test QUERY_TAG 12345";
                 conn.ConnectionString = ConnectionString + $";query_tag={expectedQueryTag}";
-                
+
                 conn.Open();
                 var command = conn.CreateCommand();
                 // This query itself will be part of the history and will have the query tag
                 command.CommandText = "SELECT QUERY_TAG FROM table(information_schema.query_history_by_session())";
                 var queryTag = command.ExecuteScalar();
-                
+
                 Assert.AreEqual(expectedQueryTag, queryTag);
             }
         }
-
+
+        [Test]
+        [Ignore("This test requires manual interaction and therefore cannot be run in CI")]
+        public void TestSSOConnectionWithTokenCachingAsync()
+        {
+            using (SnowflakeDbConnection conn = new SnowflakeDbConnection())
+            {
+                conn.ConnectionString
+                    = ConnectionStringWithoutAuth
+                        + ";authenticator=externalbrowser;[email protected];allow_sso_token_caching=true;";
+
+                // 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);
+
+                connectTask = conn.CloseAsync(CancellationToken.None);
+                connectTask.Wait();
+                Assert.AreEqual(ConnectionState.Closed, conn.State);
+
+                // Authenticate using the token
+                connectTask = conn.OpenAsync(CancellationToken.None);
+                connectTask.Wait();
+                Assert.AreEqual(ConnectionState.Open, conn.State);
+            }
+        }
     }
 }
 

Snowflake.Data.Tests/IntegrationTests/SFDbCommandIT.cs

@@ -1624,7 +1624,7 @@ public void TestGetResultsOfUnknownQueryIdWithConfiguredRetry()
                 conn.Close();
             }
         }
-        
+
         [Test]
         public void TestSetQueryTagOverridesConnectionString()
         {
@@ -1633,14 +1633,14 @@ public void TestSetQueryTagOverridesConnectionString()
                 string expectedQueryTag = "Test QUERY_TAG 12345";
                 string connectQueryTag = "Test 123";
                 conn.ConnectionString = ConnectionString + $";query_tag={connectQueryTag}";
-                
+
                 conn.Open();
                 var command = conn.CreateCommand();
                 ((SnowflakeDbCommand)command).QueryTag =  expectedQueryTag;
                 // This query itself will be part of the history and will have the query tag
                 command.CommandText = "SELECT QUERY_TAG FROM table(information_schema.query_history_by_session())";
                 var queryTag = command.ExecuteScalar();
-                
+
                 Assert.AreEqual(expectedQueryTag, queryTag);
             }
         }

Snowflake.Data.Tests/UnitTests/CredentialManager/SFCredentialManagerTest.cs

@@ -0,0 +1,255 @@
+/*
+ * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved.
+ */
+
+namespace Snowflake.Data.Tests.UnitTests.CredentialManager
+{
+    using Mono.Unix;
+    using Mono.Unix.Native;
+    using Moq;
+    using NUnit.Framework;
+    using Snowflake.Data.Client;
+    using Snowflake.Data.Core.Tools;
+    using System;
+    using System.IO;
+    using System.Runtime.InteropServices;
+
+    public abstract class SFBaseCredentialManagerTest
+    {
+        protected ISnowflakeCredentialManager _credentialManager;
+
+        [Test]
+        public void TestSavingAndRemovingCredentials()
+        {
+            // arrange
+            var key = "mockKey";
+            var expectedToken = "token";
+
+            // act
+            _credentialManager.SaveCredentials(key, expectedToken);
+            var actualToken = _credentialManager.GetCredentials(key);
+
+            // assert
+            Assert.AreEqual(expectedToken, actualToken);
+
+            // act
+            _credentialManager.RemoveCredentials(key);
+            actualToken = _credentialManager.GetCredentials(key);
+
+            // assert
+            Assert.IsTrue(string.IsNullOrEmpty(actualToken));
+        }
+
+        [Test]
+        public void TestSavingCredentialsForAnExistingKey()
+        {
+            // arrange
+            var key = "mockKey";
+            var firstExpectedToken = "mockToken1";
+            var secondExpectedToken = "mockToken2";
+
+            try
+            {
+                // act
+                _credentialManager.SaveCredentials(key, firstExpectedToken);
+
+                // assert
+                Assert.AreEqual(firstExpectedToken, _credentialManager.GetCredentials(key));
+
+                // act
+                _credentialManager.SaveCredentials(key, secondExpectedToken);
+
+                // assert
+                Assert.AreEqual(secondExpectedToken, _credentialManager.GetCredentials(key));
+            }
+            catch (Exception ex)
+            {
+                // assert
+                Assert.Fail("Should not throw an exception: " + ex.Message);
+            }
+        }
+
+        [Test]
+        public void TestRemovingCredentialsForKeyThatDoesNotExist()
+        {
+            // arrange
+            var key = "mockKey";
+
+            try
+            {
+                // act
+                _credentialManager.RemoveCredentials(key);
+
+                // assert
+                Assert.IsTrue(string.IsNullOrEmpty(_credentialManager.GetCredentials(key)));
+            }
+            catch (Exception ex)
+            {
+                // assert
+                Assert.Fail("Should not throw an exception: " + ex.Message);
+            }
+        }
+    }
+
+    [TestFixture]
+    [Platform("Win")]
+    public class SFAdysTechCredentialManagerTest : SFBaseCredentialManagerTest
+    {
+        [SetUp]
+        public void SetUp()
+        {
+            _credentialManager = SnowflakeCredentialManagerAdysTechImpl.Instance;
+        }
+    }
+
+    [TestFixture]
+    public class SFInMemoryCredentialManagerTest : SFBaseCredentialManagerTest
+    {
+        [SetUp]
+        public void SetUp()
+        {
+            _credentialManager = SnowflakeCredentialManagerInMemoryImpl.Instance;
+        }
+    }
+
+    [TestFixture]
+    public class SFFileCredentialManagerTest : SFBaseCredentialManagerTest
+    {
+        [SetUp]
+        public void SetUp()
+        {
+            _credentialManager = SnowflakeCredentialManagerFileImpl.Instance;
+        }
+    }
+
+    [TestFixture]
+    class SFCredentialManagerTest
+    {
+        ISnowflakeCredentialManager _credentialManager;
+
+        [ThreadStatic]
+        private static Mock<FileOperations> t_fileOperations;
+
+        [ThreadStatic]
+        private static Mock<DirectoryOperations> t_directoryOperations;
+
+        [ThreadStatic]
+        private static Mock<UnixOperations> t_unixOperations;
+
+        [ThreadStatic]
+        private static Mock<EnvironmentOperations> t_environmentOperations;
+
+        private static readonly string s_defaultJsonDir = Environment.GetFolderPath(Environment.SpecialFolder.UserProfile);
+
+        private const string CustomJsonDir = "testdirectory";
+
+        private static readonly string s_defaultJsonPath = Path.Combine(s_defaultJsonDir, SnowflakeCredentialManagerFileImpl.CredentialCacheFileName);
+
+        private static readonly string s_customJsonPath = Path.Combine(CustomJsonDir, SnowflakeCredentialManagerFileImpl.CredentialCacheFileName);
+
+        [SetUp] public void SetUp()
+        {
+            t_fileOperations = new Mock<FileOperations>();
+            t_directoryOperations = new Mock<DirectoryOperations>();
+            t_unixOperations = new Mock<UnixOperations>();
+            t_environmentOperations = new Mock<EnvironmentOperations>();
+            SnowflakeCredentialManagerFactory.SetCredentialManager(SnowflakeCredentialManagerInMemoryImpl.Instance);
+        }
+
+        [TearDown] public void TearDown()
+        {
+            SnowflakeCredentialManagerFactory.UseDefaultCredentialManager();
+        }
+
+        [Test]
+        public void TestUsingDefaultCredentialManager()
+        {
+            // arrange
+            SnowflakeCredentialManagerFactory.UseDefaultCredentialManager();
+
+            // act
+            _credentialManager = SnowflakeCredentialManagerFactory.GetCredentialManager();
+
+            // assert
+            if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
+            {
+                Assert.IsInstanceOf<SnowflakeCredentialManagerAdysTechImpl>(_credentialManager);
+            }
+            else
+            {
+                Assert.IsInstanceOf<SnowflakeCredentialManagerInMemoryImpl>(_credentialManager);
+            }
+        }
+
+        [Test]
+        public void TestSettingCustomCredentialManager()
+        {
+            // arrange
+            SnowflakeCredentialManagerFactory.SetCredentialManager(SnowflakeCredentialManagerFileImpl.Instance);
+
+            // act
+            _credentialManager = SnowflakeCredentialManagerFactory.GetCredentialManager();
+
+            // assert
+            Assert.IsInstanceOf<SnowflakeCredentialManagerFileImpl>(_credentialManager);
+        }
+
+        [Test]
+        public void TestThatThrowsErrorWhenCacheFileIsNotCreated()
+        {
+            if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
+            {
+                Assert.Ignore("skip test on Windows");
+            }
+
+            // arrange
+            t_directoryOperations
+                .Setup(d => d.Exists(s_defaultJsonDir))
+                .Returns(false);
+            t_unixOperations
+                .Setup(u => u.CreateFileWithPermissions(s_customJsonPath,
+                    FilePermissions.S_IRUSR | FilePermissions.S_IWUSR | FilePermissions.S_IXUSR))
+                .Returns(-1);
+            t_environmentOperations
+                .Setup(e => e.GetEnvironmentVariable(SnowflakeCredentialManagerFileImpl.CredentialCacheDirectoryEnvironmentName))
+                .Returns(CustomJsonDir);
+            SnowflakeCredentialManagerFactory.SetCredentialManager(new SnowflakeCredentialManagerFileImpl(t_fileOperations.Object, t_directoryOperations.Object, t_unixOperations.Object, t_environmentOperations.Object));
+            _credentialManager = SnowflakeCredentialManagerFactory.GetCredentialManager();
+
+            // act
+            var thrown = Assert.Throws<Exception>(() => _credentialManager.SaveCredentials("key", "token"));
+
+            // assert
+            Assert.That(thrown.Message, Does.Contain("Failed to create the JSON token cache file"));
+        }
+
+        [Test]
+        public void TestThatThrowsErrorWhenCacheFileCanBeAccessedByOthers()
+        {
+            if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
+            {
+                Assert.Ignore("skip test on Windows");
+            }
+
+            // arrange
+            t_unixOperations
+                .Setup(u => u.CreateFileWithPermissions(s_defaultJsonPath,
+                    FilePermissions.S_IRUSR | FilePermissions.S_IWUSR | FilePermissions.S_IXUSR))
+                .Returns(0);
+            t_unixOperations
+                .Setup(u => u.GetFilePermissions(s_defaultJsonPath))
+                .Returns(FileAccessPermissions.AllPermissions);
+            t_environmentOperations
+                .Setup(e => e.GetEnvironmentVariable(SnowflakeCredentialManagerFileImpl.CredentialCacheDirectoryEnvironmentName))
+                .Returns(CustomJsonDir);
+            SnowflakeCredentialManagerFactory.SetCredentialManager(new SnowflakeCredentialManagerFileImpl(t_fileOperations.Object, t_directoryOperations.Object, t_unixOperations.Object, t_environmentOperations.Object));
+            _credentialManager = SnowflakeCredentialManagerFactory.GetCredentialManager();
+
+            // act
+            var thrown = Assert.Throws<Exception>(() => _credentialManager.SaveCredentials("key", "token"));
+
+            // assert
+            Assert.That(thrown.Message, Does.Contain("Permission for the JSON token cache file should contain only the owner access"));
+        }
+    }
+}