SNOW-715524: Add SSO token cache

Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs

@@ -1,4 +1,4 @@
-/*
+/*
  * Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
  */
 
@@ -21,6 +21,8 @@ namespace Snowflake.Data.Tests.IntegrationTests
     using Snowflake.Data.Tests.Mock;
     using System.Runtime.InteropServices;
     using System.Net.Http;
+    using Snowflake.Data.Core.CredentialManager;
+    using Snowflake.Data.Core.CredentialManager.Infrastructure;
 
     [TestFixture]
     class SFConnectionIT : SFBaseTest
@@ -1046,6 +1048,71 @@ 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()
+        {
+            /*
+             * This test checks that the connector successfully stores an SSO token and uses it for authentication if it exists
+             * 1. Login normally using external browser with allow_sso_token_caching enabled
+             * 2. Login again, this time without a browser, as the connector should be using the SSO token retrieved from step 1
+            */
+
+            using (IDbConnection conn = new SnowflakeDbConnection())
+            {
+                // Set the allow_sso_token_caching property to true to enable token caching
+                // The specified user should be configured for SSO
+                conn.ConnectionString
+                    = ConnectionStringWithoutAuth
+                        + $";authenticator=externalbrowser;user={testConfig.user};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);
+
+                // Authenticate using the SSO token (the connector will automatically use the token and a browser should not pop-up in this step)
+                conn.Open();
+                Assert.AreEqual(ConnectionState.Open, conn.State);
+
+                conn.Close();
+                Assert.AreEqual(ConnectionState.Closed, conn.State);
+            }
+        }
+
+        [Test]
+        [Ignore("This test requires manual interaction and therefore cannot be run in CI")]
+        public void TestSSOConnectionWithInvalidCachedToken()
+        {
+            /*
+             * This test checks that the connector will attempt to re-authenticate using external browser if the token retrieved from the cache is invalid
+             * 1. Create a credential manager and save credentials for the user with a wrong token
+             * 2. Open a connection which initially should try to use the token and then switch to external browser when the token fails
+            */
+
+            using (IDbConnection conn = new SnowflakeDbConnection())
+            {
+                // Set the allow_sso_token_caching property to true to enable token caching
+                conn.ConnectionString
+                    = ConnectionStringWithoutAuth
+                        + $";authenticator=externalbrowser;user={testConfig.user};allow_sso_token_caching=true;";
+
+                // Create a credential manager and save a wrong token for the test user
+                var key = SFCredentialManagerFactory.BuildCredentialKey(testConfig.host, testConfig.user, TokenType.IdToken);
+                var credentialManager = SFCredentialManagerInMemoryImpl.Instance;
+                credentialManager.SaveCredentials(key, "wrongToken");
+
+                // Use the credential manager with the wrong token
+                SFCredentialManagerFactory.SetCredentialManager(credentialManager);
+
+                // Open a connection which should switch to external browser after trying to connect using the wrong token
+                conn.Open();
+                Assert.AreEqual(ConnectionState.Open, conn.State);
+
+                // Switch back to the default credential manager
+                SFCredentialManagerFactory.UseDefaultCredentialManager();
+            }
+        }
+
         [Test]
         [Ignore("This test requires manual interaction and therefore cannot be run in CI")]
         public void TestSSOConnectionWithWrongUser()
@@ -2311,6 +2378,40 @@ public void TestOpenAsyncThrowExceptionWhenOperationIsCancelled()
             }
         }
 
+        [Test]
+        [Ignore("This test requires manual interaction and therefore cannot be run in CI")]
+        public void TestSSOConnectionWithTokenCachingAsync()
+        {
+            /*
+             * This test checks that the connector successfully stores an SSO token and uses it for authentication if it exists
+             * 1. Login normally using external browser with allow_sso_token_caching enabled
+             * 2. Login again, this time without a browser, as the connector should be using the SSO token retrieved from step 1
+            */
+
+            using (SnowflakeDbConnection conn = new SnowflakeDbConnection())
+            {
+                // Set the allow_sso_token_caching property to true to enable token caching
+                // The specified user should be configured for SSO
+                conn.ConnectionString
+                    = ConnectionStringWithoutAuth
+                        + $";authenticator=externalbrowser;user={testConfig.user};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);
+
+                // Authenticate using the SSO token (the connector will automatically use the token and a browser should not pop-up in this step)
+                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);
+            }
+        }
+
         [Test]
         public void TestCloseSessionWhenGarbageCollectorFinalizesConnection()
         {

Snowflake.Data.Tests/Mock/MockExternalBrowser.cs

@@ -0,0 +1,99 @@
+/*
+ * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved.
+ */
+
+using Snowflake.Data.Core;
+using System.Net.Http;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Snowflake.Data.Tests.Mock
+{
+
+    class MockExternalBrowserRestRequester : IMockRestRequester
+    {
+        public string ProofKey { get; set; }
+        public string SSOUrl { get; set; }
+
+        public T Get<T>(IRestRequest request)
+        {
+            throw new System.NotImplementedException();
+        }
+
+        public Task<T> GetAsync<T>(IRestRequest request, CancellationToken cancellationToken)
+        {
+            throw new System.NotImplementedException();
+        }
+
+        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 AuthenticatorRequest)
+            {
+                if (string.IsNullOrEmpty(SSOUrl))
+                {
+                    var body = (AuthenticatorRequest)sfRequest.jsonBody;
+                    var port = body.Data.BrowserModeRedirectPort;
+                    SSOUrl = $"http://localhost:{port}/?token=mockToken";
+                }
+
+                // authenticator
+                var authnResponse = new AuthenticatorResponse
+                {
+                    success = true,
+                    data = new AuthenticatorResponseData
+                    {
+                        proofKey = ProofKey,
+                        ssoUrl = SSOUrl,
+                    }
+                };
+
+                return Task.FromResult<T>((T)(object)authnResponse);
+            }
+            else
+            {
+                // login
+                var loginResponse = new LoginResponse
+                {
+                    success = true,
+                    data = new LoginResponseData
+                    {
+                        sessionId = "",
+                        token = "",
+                        masterToken = "",
+                        masterValidityInSeconds = 0,
+                        authResponseSessionInfo = new SessionInfo
+                        {
+                            databaseName = "",
+                            schemaName = "",
+                            roleName = "",
+                            warehouseName = "",
+                        }
+                    }
+                };
+
+                return Task.FromResult<T>((T)(object)loginResponse);
+            }
+        }
+
+        public HttpResponseMessage Get(IRestRequest request)
+        {
+            throw new System.NotImplementedException();
+        }
+
+        public Task<HttpResponseMessage> GetAsync(IRestRequest request, CancellationToken cancellationToken)
+        {
+            throw new System.NotImplementedException();
+        }
+
+        public void setHttpClient(HttpClient httpClient)
+        {
+            // Nothing to do
+        }
+    }
+}