SNOW-916942: Fix native okta retry logic (#1701)

Author: sfc-gh-ext-simba-jf

src/main/java/net/snowflake/client/core/SessionUtil.java

@@ -651,16 +651,31 @@ 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) {
+                logger.debug("Retrieve new token for Okta authentication.");
+                // 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 +705,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;
@@ -1362,12 +1378,24 @@ private static void handleFederatedFlowError(SFLoginInput loginInput, Exception
    */
   private static String getSamlResponseUsingOkta(SFLoginInput loginInput)
       throws SnowflakeSQLException {
-    JsonNode dataNode = federatedFlowStep1(loginInput);
-    String tokenUrl = dataNode.path("tokenUrl").asText();
-    String ssoUrl = dataNode.path("ssoUrl").asText();
-    federatedFlowStep2(loginInput, tokenUrl, ssoUrl);
-    final String oneTimeToken = federatedFlowStep3(loginInput, tokenUrl);
-    return federatedFlowStep4(loginInput, ssoUrl, oneTimeToken);
+    while (true) {
+      try {
+        JsonNode dataNode = federatedFlowStep1(loginInput);
+        String tokenUrl = dataNode.path("tokenUrl").asText();
+        String ssoUrl = dataNode.path("ssoUrl").asText();
+        federatedFlowStep2(loginInput, tokenUrl, ssoUrl);
+        final String oneTimeToken = federatedFlowStep3(loginInput, tokenUrl);
+        return federatedFlowStep4(loginInput, ssoUrl, oneTimeToken);
+      } catch (SnowflakeSQLException ex) {
+        // This error gets thrown if the okta request encountered a retry-able error that
+        // requires getting a new one-time token.
+        if (ex.getErrorCode() == ErrorCode.AUTHENTICATOR_REQUEST_TIMEOUT.getMessageCode()) {
+          logger.debug("Failed to get Okta SAML response. Retrying without changing retry count.");
+        } else {
+          throw ex;
+        }
+      }
+    }
   }
 
   /**

src/main/java/net/snowflake/client/jdbc/RestRequest.java

@@ -399,6 +399,16 @@ public static CloseableHttpResponse execute(
           break;
         }
 
+        // If this was a request for an Okta one-time token that failed with a retry-able error,
+        // throw exception to renew the token before trying again.
+        if (String.valueOf(httpRequest.getURI()).contains("okta.com/api/v1/authn")) {
+          throw new SnowflakeSQLException(
+              ErrorCode.AUTHENTICATOR_REQUEST_TIMEOUT,
+              retryCount,
+              true,
+              elapsedMilliForTransientIssues / 1000);
+        }
+
         // Make sure that any authenticator specific info that needs to be
         // updated get's updated before the next retry. Ex - JWT token
         // Check to see if customer set socket/connect timeout has been reached,

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.1.
+  @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");
+    }
+  }
 }

src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java

@@ -38,8 +38,10 @@
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Enumeration;
+import java.util.List;
 import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.TimeUnit;
@@ -62,6 +64,7 @@
 import org.apache.http.entity.StringEntity;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
@@ -1150,4 +1153,39 @@ public void testIsAsyncSession() throws SQLException, InterruptedException {
       assertFalse(snowflakeConnection.getSfSession().isAsyncSession());
     }
   }
+
+  // Test for regenerating okta one-time token for versions > 3.15.1
+  @Test
+  @Ignore
+  public void testDataSourceOktaGenerates429StatusCode() throws Exception {
+    // test with username/password authentication
+    // set up DataSource object and ensure connection works
+    Map<String, String> params = getConnectionParameters();
+    SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();
+    ds.setServerName(params.get("host"));
+    ds.setSsl("on".equals(params.get("ssl")));
+    ds.setAccount(params.get("account"));
+    ds.setPortNumber(Integer.parseInt(params.get("port")));
+    ds.setUser(params.get("ssoUser"));
+    ds.setPassword(params.get("ssoPassword"));
+    ds.setAuthenticator("<okta address>");
+    Runnable r =
+        () -> {
+          try {
+            ds.getConnection();
+          } catch (SQLException e) {
+            throw new RuntimeException(e);
+          }
+        };
+    List<Thread> threadList = new ArrayList<>();
+    for (int i = 0;
+        i < 30;
+        ++i) { // https://docs.snowflake.com/en/user-guide/admin-security-fed-auth-use#http-429-errors
+      threadList.add(new Thread(r));
+    }
+    threadList.forEach(Thread::start);
+    for (Thread thread : threadList) {
+      thread.join();
+    }
+  }
 }

src/test/java/net/snowflake/client/jdbc/RestRequestTest.java

@@ -359,6 +359,20 @@ public void testExceptionAuthBasedTimeout() throws IOException {
     }
   }
 
+  @Test
+  public void testExceptionAuthBasedTimeoutFor429ErrorCode() throws IOException {
+    CloseableHttpClient client = mock(CloseableHttpClient.class);
+    when(client.execute(any(HttpUriRequest.class)))
+        .thenAnswer((Answer<CloseableHttpResponse>) invocation -> retryLoginResponse());
+
+    try {
+      execute(client, "login-request.com/?requestId=abcd-1234", 2, 1, 30000, true, false);
+    } catch (SnowflakeSQLException ex) {
+      assertThat(
+          ex.getErrorCode(), equalTo(ErrorCode.AUTHENTICATOR_REQUEST_TIMEOUT.getMessageCode()));
+    }
+  }
+
   @Test
   public void testNoRetry() throws IOException, SnowflakeSQLException {
     boolean telemetryEnabled = TelemetryService.getInstance().isEnabled();