Fix Okta authentication Deadlock issue (#139)

snowflakedb · Jun 17, 2019 · 50c4517 · 50c4517
1 parent ae1c36c
commit 50c4517
Show file tree

Hide file tree

Showing 11 changed files with 100 additions and 7 deletions.
diff --git a/Snowflake.Data.Tests/Mock/MockCloseSessionGone.cs b/Snowflake.Data.Tests/Mock/MockCloseSessionGone.cs
@@ -29,6 +29,11 @@ public Task<HttpResponseMessage> GetAsync(IRestRequest request, CancellationToke
             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)).Result;

diff --git a/Snowflake.Data.Tests/Mock/MockOkta.cs b/Snowflake.Data.Tests/Mock/MockOkta.cs
@@ -36,7 +36,7 @@ public Task<HttpResponseMessage> GetAsync(IRestRequest request, CancellationToke
 
         public T Post<T>(IRestRequest postRequest)
         {
-            throw new System.NotImplementedException();
+            return Task.Run(async () => await PostAsync<T>(postRequest, CancellationToken.None)).Result;
         }
 
         public Task<T> PostAsync<T>(IRestRequest postRequest, CancellationToken cancellationToken)
@@ -66,5 +66,10 @@ public Task<T> PostAsync<T>(IRestRequest postRequest, CancellationToken cancella
                 return Task.FromResult<T>((T)(object)tokenResponse);
             }
         }
+
+        public HttpResponseMessage Get(IRestRequest request)
+        {
+            return Task.Run(async () => await GetAsync(request, CancellationToken.None)).Result;
+        }
     }
 }
diff --git a/Snowflake.Data.Tests/Mock/MockRestSessionExpired.cs b/Snowflake.Data.Tests/Mock/MockRestSessionExpired.cs
@@ -121,5 +121,10 @@ public Task<HttpResponseMessage> GetAsync(IRestRequest request, CancellationToke
         {
             return Task.FromResult<HttpResponseMessage>(null);
         }
+
+        public HttpResponseMessage Get(IRestRequest request)
+        {
+            return null;
+        }
     }
 }
diff --git a/Snowflake.Data.Tests/Mock/MockRestSessionExpiredInQueryExec.cs b/Snowflake.Data.Tests/Mock/MockRestSessionExpiredInQueryExec.cs
@@ -139,6 +139,11 @@ public Task<HttpResponseMessage> GetAsync(IRestRequest request, CancellationToke
         {
             return Task.FromResult<HttpResponseMessage>(null);
         }
+
+        public HttpResponseMessage Get(IRestRequest request)
+        {
+            return null;
+        }
     }
 }
 
diff --git a/Snowflake.Data.Tests/Mock/MockServiceName.cs b/Snowflake.Data.Tests/Mock/MockServiceName.cs
@@ -98,5 +98,10 @@ public Task<HttpResponseMessage> GetAsync(IRestRequest request, CancellationToke
         {
             return Task.FromResult<HttpResponseMessage>(null);
         }
+
+        public HttpResponseMessage Get(IRestRequest request)
+        {
+            return null;
+        }
     }
 }
diff --git a/Snowflake.Data.Tests/SFConnectionIT.cs b/Snowflake.Data.Tests/SFConnectionIT.cs
@@ -92,9 +92,10 @@ public void TestLoginTimeout()
                     conn.Open();
                     Assert.Fail();
                 }
-                catch(SnowflakeDbException e)
+                catch (AggregateException e)
                 {
-                    Assert.AreEqual(SFError.REQUEST_TIMEOUT.GetAttribute<SFErrorAttr>().errorCode,e.ErrorCode);
+                    Assert.AreEqual(SFError.REQUEST_TIMEOUT.GetAttribute<SFErrorAttr>().errorCode, 
+                        ((SnowflakeDbException)e.InnerException).ErrorCode);
                 }
                 Assert.AreEqual(5, conn.ConnectionTimeout);
             }
@@ -301,9 +302,14 @@ public void TestOktaConnection()
         {
             using (IDbConnection conn = new SnowflakeDbConnection())
             {
-                conn.ConnectionString = "scheme=http;host=10.211.55.3;port=8080;user=qa@snowflakecomputing.com;password=Test123!;" +
+                conn.ConnectionString = "scheme=http;host=testaccount.reg.snowflakecomputing.com;port=8082;user=qa@snowflakecomputing.com;password=Test123!;" +
                     "account=testaccount;role=sysadmin;db=testdb;schema=public;warehouse=regress;authenticator=https://snowflakecomputing.okta.com";
                 conn.Open();
+                using (IDbCommand command = conn.CreateCommand())
+                {
+                    command.CommandText = "SELECT 1";
+                    Assert.AreEqual("1", command.ExecuteScalar().ToString());
+                }
             }
         }
     }

diff --git a/Snowflake.Data/Core/Authenticator/BasicAuthenticator.cs b/Snowflake.Data/Core/Authenticator/BasicAuthenticator.cs
@@ -27,6 +27,15 @@ async Task IAuthenticator.AuthenticateAsync(CancellationToken cancellationToken)
             session.ProcessLoginResponse(response);
         }
 
+        void IAuthenticator.Authenticate()
+        {
+            var loginRequest = BuildLoginRequest();
+
+            var response = session.restRequester.Post<AuthnResponse>(loginRequest);
+
+            session.ProcessLoginResponse(response);
+        }
+
         private SFRestRequest BuildLoginRequest()
         {
             // build uri

diff --git a/Snowflake.Data/Core/Authenticator/IAuthenticator.cs b/Snowflake.Data/Core/Authenticator/IAuthenticator.cs
@@ -22,6 +22,12 @@ internal interface IAuthenticator
         /// <returns></returns>
         /// <exception cref="SnowflakeDbException"></exception>
         Task AuthenticateAsync(CancellationToken cancellationToken);
+
+        /// <summary>
+        /// Process the authentication synchronously
+        /// </summary>
+        /// <exception cref="SnowflakeDbException"></exception>
+        void Authenticate();
     }
 
     /// <summary>

diff --git a/Snowflake.Data/Core/Authenticator/OktaAuthenticator.cs b/Snowflake.Data/Core/Authenticator/OktaAuthenticator.cs
@@ -78,6 +78,42 @@ async Task IAuthenticator.AuthenticateAsync(CancellationToken cancellationToken)
             session.ProcessLoginResponse(authnResponse);   
         }
 
+        void IAuthenticator.Authenticate()
+        {
+            logger.Info("Okta Authentication");
+
+            logger.Debug("step 1: get sso and token url");
+            var authenticatorRestRequest = BuildAuthenticatorRestRequest();
+            var authenticatorResponse = session.restRequester.Post<AuthnResponse>(authenticatorRestRequest);
+            FilterFailedResponse(authenticatorResponse);
+            Uri ssoUrl = new Uri(authenticatorResponse.data.ssoUrl);
+            Uri tokenUrl = new Uri(authenticatorResponse.data.tokenUrl);
+
+            logger.Debug("step 2: verify urls fetched from step 1");
+            logger.Debug("Checking sso url");
+            VerifyUrls(ssoUrl, oktaUrl);
+            logger.Debug("Checking token url");
+            VerifyUrls(tokenUrl, oktaUrl);
+
+            logger.Debug("step 3: get idp onetime token");
+            IdpTokenRestRequest idpTokenRestRequest = BuildIdpTokenRestRequest(tokenUrl);
+            var idpResponse =  session.restRequester.Post<IdpTokenResponse>(idpTokenRestRequest);
+            string onetimeToken = idpResponse.CookieToken;
+
+            logger.Debug("step 4: get SAML reponse from sso");
+            var samlRestRequest = BuildSAMLRestRequest(ssoUrl, onetimeToken);
+            var samlRawResponse = session.restRequester.Get(samlRestRequest);
+            var samlRawHtmlString = Task.Run(async () => await samlRawResponse.Content.ReadAsStringAsync()).Result;
+
+            logger.Debug("step 5: verify postback url in SAML reponse");
+            VerifyPostbackUrl(samlRawHtmlString);
+
+            logger.Debug("step 6: send SAML reponse to snowflake to login");
+            var loginRestRequest = BuildOktaLoginRestRequest(samlRawHtmlString);
+            var authnResponse = session.restRequester.Post<AuthnResponse>(loginRestRequest);
+            session.ProcessLoginResponse(authnResponse);   
+        }
+
         private SFRestRequest BuildAuthenticatorRestRequest()
         {
             var fedUrl = session.BuildUri(RestPath.SF_AUTHENTICATOR_REQUEST_PATH);

diff --git a/Snowflake.Data/Core/RestRequester.cs b/Snowflake.Data/Core/RestRequester.cs
@@ -21,11 +21,13 @@ internal interface IRestRequester
 
         T Post<T>(IRestRequest postRequest);
 
-        T Get<T>(IRestRequest request);
-
         Task<T> GetAsync<T>(IRestRequest request, CancellationToken cancellationToken);
 
+        T Get<T>(IRestRequest request);
+
         Task<HttpResponseMessage> GetAsync(IRestRequest request, CancellationToken cancellationToken);
+
+        HttpResponseMessage Get(IRestRequest request);
     }
 
     internal class RestRequester : IRestRequester
@@ -80,6 +82,15 @@ public Task<HttpResponseMessage> GetAsync(IRestRequest request, CancellationToke
 
             return SendAsync(message, request.GetRestTimeout(), cancellationToken);
         }
+
+        public HttpResponseMessage Get(IRestRequest request)
+        {
+            HttpRequestMessage message = request.ToRequestMessage(HttpMethod.Get);
+            logger.Debug($"Http method: {message.ToString()}, http request message: {message.ToString()}");
+
+            //Run synchronous in a new thread-pool task.
+            return Task.Run(async () => await GetAsync(request, CancellationToken.None)).Result;
+        }
 
         private async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, 
                                                           TimeSpan timeoutPerRestRequest,  

diff --git a/Snowflake.Data/Core/SFSession.cs b/Snowflake.Data/Core/SFSession.cs
@@ -133,7 +133,7 @@ internal void Open()
                 authenticator = AuthenticatorFactory.GetAuthenticator(this);
             }
 
-            authenticator.AuthenticateAsync(CancellationToken.None).GetAwaiter().GetResult();
+            authenticator.Authenticate();
         }
 
         internal async Task OpenAsync(CancellationToken cancellationToken)