Added retry logic for authentication with Okta

snowflakedb · Mar 7, 2024 · 8891c79 · 8891c79
1 parent d4ca9f0
commit 8891c79
Show file tree

Hide file tree

Showing 2 changed files with 101 additions and 11 deletions.
diff --git a/src/main/java/net/snowflake/client/core/SessionUtil.java b/src/main/java/net/snowflake/client/core/SessionUtil.java
@@ -651,16 +651,30 @@ private static SFLoginOutput newSession(
                   loginInput.getHttpClientSettingsKey());
         } catch (SnowflakeSQLException ex) {
           if (ex.getErrorCode() == ErrorCode.AUTHENTICATOR_REQUEST_TIMEOUT.getMessageCode()) {
-            if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT) {
-              SessionUtilKeyPair s =
-                  new SessionUtilKeyPair(
-                      loginInput.getPrivateKey(),
-                      loginInput.getPrivateKeyFile(),
-                      loginInput.getPrivateKeyFilePwd(),
-                      loginInput.getAccountName(),
-                      loginInput.getUserName());
-
-              data.put(ClientAuthnParameter.TOKEN.name(), s.issueJwtToken());
+            if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT
+                || authenticatorType == ClientAuthnDTO.AuthenticatorType.OKTA) {
+
+              if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT) {
+                SessionUtilKeyPair s =
+                    new SessionUtilKeyPair(
+                        loginInput.getPrivateKey(),
+                        loginInput.getPrivateKeyFile(),
+                        loginInput.getPrivateKeyFilePwd(),
+                        loginInput.getAccountName(),
+                        loginInput.getUserName());
+
+                data.put(ClientAuthnParameter.TOKEN.name(), s.issueJwtToken());
+              } else if (authenticatorType == ClientAuthnDTO.AuthenticatorType.OKTA) {
+                // If we need to retry, we need to get a new Okta token
+                tokenOrSamlResponse = getSamlResponseUsingOkta(loginInput);
+                data.put(ClientAuthnParameter.RAW_SAML_RESPONSE.name(), tokenOrSamlResponse);
+                authnData.setData(data);
+                String updatedJson = mapper.writeValueAsString(authnData);
+
+                StringEntity updatedInput = new StringEntity(updatedJson, StandardCharsets.UTF_8);
+                updatedInput.setContentType("application/json");
+                postRequest.setEntity(updatedInput);
+              }
 
               long elapsedSeconds = ex.getElapsedSeconds();
 
@@ -690,7 +704,8 @@ private static SFLoginOutput newSession(
                 }
               }
 
-              // JWT renew should not count as a retry, so we pass back the current retry count.
+              // JWT or Okta renew should not count as a retry, so we pass back the current retry
+              // count.
               retryCount = ex.getRetryCount();
 
               continue;

diff --git a/src/test/java/net/snowflake/client/core/SessionUtilLatestIT.java b/src/test/java/net/snowflake/client/core/SessionUtilLatestIT.java
@@ -390,4 +390,79 @@ public void testOktaAuthGetFail() throws Throwable {
       assertEquals(SqlState.IO_ERROR, e.getSQLState());
     }
   }
+
+  private SFLoginInput createOktaLoginInput() {
+    SFLoginInput input = new SFLoginInput();
+    input.setServerUrl("https://testauth.okta.com");
+    input.setUserName("MOCK_USERNAME");
+    input.setPassword("MOCK_PASSWORD");
+    input.setAccountName("MOCK_ACCOUNT_NAME");
+    input.setAppId("MOCK_APP_ID");
+    input.setOCSPMode(OCSPMode.FAIL_OPEN);
+    input.setHttpClientSettingsKey(new HttpClientSettingsKey(OCSPMode.FAIL_OPEN));
+    input.setLoginTimeout(1000);
+    input.setSessionParameters(new HashMap<>());
+    input.setAuthenticator("https://testauth.okta.com");
+    return input;
+  }
+
+  // Testing retry with Okta calls the service to get a new unique token. This is valid after
+  // version 3.15.0.
+  @Test
+  public void testOktaAuthRetry() throws Throwable {
+    SFLoginInput loginInput = createOktaLoginInput();
+    Map<SFSessionProperty, Object> connectionPropertiesMap = initConnectionPropertiesMap();
+    SnowflakeSQLException ex =
+        new SnowflakeSQLException(ErrorCode.AUTHENTICATOR_REQUEST_TIMEOUT, 0, true, 0);
+    try (MockedStatic<HttpUtil> mockedHttpUtil = mockStatic(HttpUtil.class)) {
+      mockedHttpUtil
+          .when(
+              () ->
+                  HttpUtil.executeGeneralRequest(
+                      Mockito.any(HttpPost.class),
+                      Mockito.anyInt(),
+                      Mockito.anyInt(),
+                      Mockito.anyInt(),
+                      Mockito.anyInt(),
+                      Mockito.nullable(HttpClientSettingsKey.class)))
+          .thenReturn(
+              "{\"data\":{\"tokenUrl\":\"https://testauth.okta.com/api/v1/authn\","
+                  + "\"ssoUrl\":\"https://testauth.okta.com/app/snowflake/abcdefghijklmnopqrstuvwxyz/sso/saml\","
+                  + "\"proofKey\":null},\"code\":null,\"message\":null,\"success\":true}")
+          .thenThrow(ex)
+          .thenReturn(
+              "{\"data\":{\"tokenUrl\":\"https://testauth.okta.com/api/v1/authn\","
+                  + "\"ssoUrl\":\"https://testauth.okta.com/app/snowflake/abcdefghijklmnopqrstuvwxyz/sso/saml\","
+                  + "\"proofKey\":null},\"code\":null,\"message\":null,\"success\":true}");
+
+      mockedHttpUtil
+          .when(
+              () ->
+                  HttpUtil.executeRequestWithoutCookies(
+                      Mockito.any(HttpRequestBase.class),
+                      Mockito.anyInt(),
+                      Mockito.anyInt(),
+                      Mockito.anyInt(),
+                      Mockito.anyInt(),
+                      Mockito.anyInt(),
+                      Mockito.nullable(AtomicBoolean.class),
+                      Mockito.nullable(HttpClientSettingsKey.class)))
+          .thenReturn(
+              "{\"expiresAt\":\"2023-10-13T19:18:09.000Z\",\"status\":\"SUCCESS\",\"sessionToken\":\"testsessiontoken\"}");
+
+      mockedHttpUtil
+          .when(
+              () ->
+                  HttpUtil.executeGeneralRequest(
+                      Mockito.any(HttpGet.class),
+                      Mockito.anyInt(),
+                      Mockito.anyInt(),
+                      Mockito.anyInt(),
+                      Mockito.anyInt(),
+                      Mockito.nullable(HttpClientSettingsKey.class)))
+          .thenReturn("<body><form action=\"https://testauth.okta.com\"></form></body>");
+
+      SessionUtil.openSession(loginInput, connectionPropertiesMap, "ALL");
+    }
+  }
 }