SNOW-916949: Fix Okta retry for SSO/SAML endpoints

Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs

@@ -17,6 +17,7 @@ namespace Snowflake.Data.Tests.IntegrationTests
     using System.Diagnostics;
     using Snowflake.Data.Tests.Mock;
     using System.Runtime.InteropServices;
+    using System.Net.Http;
 
     [TestFixture]
     class SFConnectionIT : SFBaseTest
@@ -817,6 +818,46 @@ public void TestOktaConnection()
             }
         }
 
+        [Test]
+        public void TestOktaConnectionUntilMaxTimeout()
+        {
+            var expectedMaxRetryCount = 15;
+            var expectedMaxConnectionTimeout = 450;
+            var oktaUrl = "https://test.okta.com";
+            var mockRestRequester = new MockOktaRestRequester()
+            {
+                TokenUrl = $"{oktaUrl}/api/v1/sessions?additionalFields=cookieToken",
+                SSOUrl = $"{oktaUrl}/app/testaccount/sso/saml",
+                ResponseContent = new StringContent("<form=error}"),
+                MaxRetryCount = expectedMaxRetryCount,
+                MaxRetryTimeout = expectedMaxConnectionTimeout
+            };
+            using (DbConnection conn = new MockSnowflakeDbConnection(mockRestRequester))
+            {
+                try
+                {
+                    conn.ConnectionString
+                        = ConnectionStringWithoutAuth
+                        + String.Format(
+                            ";authenticator={0};user=test;password=test;MAXHTTPRETRIES={1};RETRY_TIMEOUT={2};",
+                            oktaUrl,
+                            expectedMaxRetryCount,
+                            expectedMaxConnectionTimeout);
+                    conn.Open();
+                    Assert.Fail();
+                }
+                catch (Exception e)
+                {
+                    Assert.IsInstanceOf<SnowflakeDbException>(e);
+                    Assert.AreEqual(SFError.INTERNAL_ERROR.GetAttribute<SFErrorAttr>().errorCode, ((SnowflakeDbException)e).ErrorCode);
+                    Assert.IsTrue(e.Message.Contains(
+                        $"The retry count has reached its limit of {expectedMaxRetryCount} and " +
+                        $"the timeout elapsed has reached its limit of {expectedMaxConnectionTimeout} " +
+                        "while trying to authenticate through Okta"));
+                }
+            }
+        }
+
         [Test]
         [Ignore("This test requires manual setup and therefore cannot be run in CI")]
         public void TestOkta2ConnectionsFollowingEachOther()
@@ -2056,6 +2097,54 @@ public void TestExplicitTransactionOperationsTracked()
                 Assert.AreEqual(false, conn.HasActiveExplicitTransaction());
             }
         }
+
+
+        [Test]
+        public void TestAsyncOktaConnectionUntilMaxTimeout()
+        {
+            var expectedMaxRetryCount = 15;
+            var expectedMaxConnectionTimeout = 450;
+            var oktaUrl = "https://test.okta.com";
+            var mockRestRequester = new MockOktaRestRequester()
+            {
+                TokenUrl = $"{oktaUrl}/api/v1/sessions?additionalFields=cookieToken",
+                SSOUrl = $"{oktaUrl}/app/testaccount/sso/saml",
+                ResponseContent = new StringContent("<form=error}"),
+                MaxRetryCount = expectedMaxRetryCount,
+                MaxRetryTimeout = expectedMaxConnectionTimeout
+            };
+            using (DbConnection conn = new MockSnowflakeDbConnection(mockRestRequester))
+            {
+                try
+                {
+                    conn.ConnectionString
+                        = ConnectionStringWithoutAuth
+                        + String.Format(
+                            ";authenticator={0};user=test;password=test;MAXHTTPRETRIES={1};RETRY_TIMEOUT={2};",
+                            oktaUrl,
+                            expectedMaxRetryCount,
+                            expectedMaxConnectionTimeout);
+                    Task connectTask = conn.OpenAsync(CancellationToken.None);
+                    connectTask.Wait();
+                    Assert.Fail();
+                }
+                catch (Exception e)
+                {
+                    Assert.IsInstanceOf<SnowflakeDbException>(e.InnerException);
+                    Assert.AreEqual(SFError.INTERNAL_ERROR.GetAttribute<SFErrorAttr>().errorCode, ((SnowflakeDbException)e.InnerException).ErrorCode);
+                    Exception oktaException;
+#if NETFRAMEWORK
+                    oktaException = e.InnerException.InnerException.InnerException;
+#else
+                    oktaException = e.InnerException.InnerException;
+#endif
+                    Assert.IsTrue(oktaException.Message.Contains(
+                        $"The retry count has reached its limit of {expectedMaxRetryCount} and " +
+                        $"the timeout elapsed has reached its limit of {expectedMaxConnectionTimeout} " +
+                        "while trying to authenticate through Okta"));
+                }
+            }
+        }
     }
 }
 

Snowflake.Data.Tests/Mock/MockOkta.cs

@@ -16,6 +16,8 @@ class MockOktaRestRequester : IMockRestRequester
         public string TokenUrl { get; set; }
         public string SSOUrl { get; set; }
         public StringContent ResponseContent { get; set; }
+        public int MaxRetryCount { get; set; }
+        public int MaxRetryTimeout { get; set; }
 
         public T Get<T>(IRestRequest request)
         {
@@ -31,6 +33,8 @@ public Task<HttpResponseMessage> GetAsync(IRestRequest request, CancellationToke
         {
             var response = new HttpResponseMessage(System.Net.HttpStatusCode.OK);
             response.Content = ResponseContent;
+            response.Content.Headers.Add(OktaAuthenticator.RetryCountHeader, MaxRetryCount.ToString());
+            response.Content.Headers.Add(OktaAuthenticator.TimeoutElapsedHeader, MaxRetryTimeout.ToString());
             return Task.FromResult(response);
         }
 
@@ -43,18 +47,45 @@ public Task<T> PostAsync<T>(IRestRequest postRequest, CancellationToken cancella
         {
             if (postRequest is SFRestRequest)
             {
-                // authenticator
-                var authnResponse = new AuthenticatorResponse 
+                if (((SFRestRequest)postRequest).jsonBody is AuthenticatorRequest)
                 {
-                    success = true,
-                    data = new AuthenticatorResponseData 
+                    // authenticator
+                    var authnResponse = new AuthenticatorResponse
                     {
-                        tokenUrl = TokenUrl,
-                        ssoUrl = SSOUrl,
-                    }
-                };
+                        success = true,
+                        data = new AuthenticatorResponseData
+                        {
+                            tokenUrl = TokenUrl,
+                            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)authnResponse);
+                    return Task.FromResult<T>((T)(object)loginResponse);
+                }
             }
             else
             {

Snowflake.Data.Tests/UnitTests/HttpUtilTest.cs

@@ -58,6 +58,24 @@ public void TestIsLoginUrl(string requestUrl, bool expectedIsLoginEndpoint)
             Assert.AreEqual(expectedIsLoginEndpoint, isLoginEndpoint);
         }
 
+        // Parameters: request url, expected value
+        [TestCase("https://dev.okta.com/sso/saml", true)]
+        [TestCase("https://test.snowflakecomputing.com/session/v1/login-request", false)]
+        [TestCase("https://test.snowflakecomputing.com/session/authenticator-request", false)]
+        [TestCase("https://test.snowflakecomputing.com/session/token-request", false)]
+        [Test]
+        public void TestIsOktaSSORequest(string requestUrl, bool expectedIsOktaSSORequest)
+        {
+            // given
+            var uri = new Uri(requestUrl);
+
+            // when
+            bool isOktaSSORequest = HttpUtil.IsOktaSSORequest(uri.Host, uri.AbsolutePath);
+
+            // then
+            Assert.AreEqual(expectedIsOktaSSORequest, isOktaSSORequest);
+        }
+
         // Parameters: time in seconds
         [TestCase(4)]
         [TestCase(8)]