Snow 1925254 authentication logic refactor

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

@@ -580,7 +580,7 @@ public static boolean isSocksProxyDisabled() {
   }
 
   /**
-   * Executes a HTTP request with the cookie spec set to IGNORE_COOKIES
+   * Executes an HTTP request with the cookie spec set to IGNORE_COOKIES
    *
    * @param httpRequest HttpRequestBase
    * @param retryTimeout retry timeout
@@ -622,7 +622,7 @@ static String executeRequestWithoutCookies(
   }
 
   /**
-   * Executes a HTTP request for Snowflake.
+   * Executes an HTTP request for Snowflake.
    *
    * @param httpRequest HttpRequestBase
    * @param retryTimeout retry timeout
@@ -658,7 +658,7 @@ public static String executeGeneralRequest(
   }
 
   /**
-   * Executes a HTTP request for Snowflake
+   * Executes an HTTP request for Snowflake
    *
    * @param httpRequest HttpRequestBase
    * @param retryTimeout retry timeout
@@ -696,7 +696,7 @@ public static String executeGeneralRequest(
   }
 
   /**
-   * Executes a HTTP request for Snowflake.
+   * Executes an HTTP request for Snowflake.
    *
    * @param httpRequest HttpRequestBase
    * @param retryTimeout retry timeout

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

@@ -7,6 +7,7 @@
 import static net.snowflake.client.core.SFTrustManager.resetOCSPResponseCacherServerURL;
 import static net.snowflake.client.jdbc.SnowflakeUtil.systemGetProperty;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.base.Strings;
@@ -271,7 +272,7 @@ static SFLoginOutput openSession(
       AssertUtil.assertTrue(
           loginInput.getUserName() != null, "missing user name for opening session");
     } else {
-      // OAUTH needs either token or passord
+      // OAUTH needs either token or password
       AssertUtil.assertTrue(
           loginInput.getToken() != null || loginInput.getPassword() != null,
           "missing token or password for opening session");
@@ -707,7 +708,7 @@ private static SFLoginOutput newSession(
               // In RestRequest.execute(), socket timeout is replaced with auth timeout
               // so we can renew the request within auth timeout.
               // auth timeout within socket timeout is thrown without backoff,
-              // and we need to update time remained in socket timeout here to control the
+              // and we need to update time remained in socket timeout here to control
               // the actual socket timeout from customer setting.
               if (loginInput.getSocketTimeoutInMillis() > 0) {
                 if (ex.issocketTimeoutNoBackoff()) {
@@ -737,29 +738,7 @@ private static SFLoginOutput newSession(
         break;
       }
 
-      if (theString == null) {
-        if (lastRestException != null) {
-          logger.error(
-              "Failed to open new session for user: {}, host: {}. Error: {}",
-              loginInput.getUserName(),
-              loginInput.getHostFromServerUrl(),
-              lastRestException);
-          throw lastRestException;
-        } else {
-          SnowflakeSQLException exception =
-              new SnowflakeSQLException(
-                  NO_QUERY_ID,
-                  "empty authentication response",
-                  SqlState.CONNECTION_EXCEPTION,
-                  ErrorCode.CONNECTION_ERROR.getMessageCode());
-          logger.error(
-              "Failed to open new session for user: {}, host: {}. Error: {}",
-              loginInput.getUserName(),
-              loginInput.getHostFromServerUrl(),
-              exception);
-          throw exception;
-        }
-      }
+      handleEmptyAuthResponse(theString, loginInput, lastRestException);
 
       // general method, same as with data binding
       JsonNode jsonNode = mapper.readTree(theString);
@@ -1201,22 +1180,11 @@ static void closeSession(SFLoginInput loginInput) throws SFException, SnowflakeS
   private static String federatedFlowStep4(
       SFLoginInput loginInput, String ssoUrl, String oneTimeToken) throws SnowflakeSQLException {
     String responseHtml = "";
-    try {
 
-      final URL url = new URL(ssoUrl);
-      URI oktaGetUri =
-          new URIBuilder()
-              .setScheme(url.getProtocol())
-              .setHost(url.getHost())
-              .setPath(url.getPath())
-              .setParameter("RelayState", "%2Fsome%2Fdeep%2Flink")
-              .setParameter("onetimetoken", oneTimeToken)
-              .build();
-      HttpGet httpGet = new HttpGet(oktaGetUri);
+    try {
 
-      HeaderGroup headers = new HeaderGroup();
-      headers.addHeader(new BasicHeader(HttpHeaders.ACCEPT, "*/*"));
-      httpGet.setHeaders(headers.getAllHeaders());
+      HttpGet httpGet = new HttpGet();
+      prepareFederatedFlowStep4Request(httpGet, ssoUrl, oneTimeToken);
 
       responseHtml =
           HttpUtil.executeGeneralRequest(
@@ -1276,26 +1244,7 @@ private static String federatedFlowStep3(SFLoginInput loginInput, String tokenUr
       URL url = new URL(tokenUrl);
       URI tokenUri = url.toURI();
       final HttpPost postRequest = new HttpPost(tokenUri);
-
-      String userName;
-      if (Strings.isNullOrEmpty(loginInput.getOKTAUserName())) {
-        userName = loginInput.getUserName();
-      } else {
-        userName = loginInput.getOKTAUserName();
-      }
-      StringEntity params =
-          new StringEntity(
-              "{\"username\":\""
-                  + userName
-                  + "\",\"password\":\""
-                  + loginInput.getPassword()
-                  + "\"}");
-      postRequest.setEntity(params);
-
-      HeaderGroup headers = new HeaderGroup();
-      headers.addHeader(new BasicHeader(HttpHeaders.ACCEPT, "application/json"));
-      headers.addHeader(new BasicHeader(HttpHeaders.CONTENT_TYPE, "application/json"));
-      postRequest.setHeaders(headers.getAllHeaders());
+      setFederatedFlowStep3PostRequestAuthData(postRequest, loginInput);
 
       final String idpResponse =
           HttpUtil.executeRequestWithoutCookies(
@@ -1363,29 +1312,8 @@ private static void federatedFlowStep2(SFLoginInput loginInput, String tokenUrl,
   private static JsonNode federatedFlowStep1(SFLoginInput loginInput) throws SnowflakeSQLException {
     JsonNode dataNode = null;
     try {
-      URIBuilder fedUriBuilder = new URIBuilder(loginInput.getServerUrl());
-      fedUriBuilder.setPath(SF_PATH_AUTHENTICATOR_REQUEST);
-      URI fedUrlUri = fedUriBuilder.build();
-
-      Map<String, Object> data = new HashMap<>();
-      data.put(ClientAuthnParameter.ACCOUNT_NAME.name(), loginInput.getAccountName());
-      data.put(ClientAuthnParameter.AUTHENTICATOR.name(), loginInput.getAuthenticator());
-      data.put(ClientAuthnParameter.CLIENT_APP_ID.name(), loginInput.getAppId());
-      data.put(ClientAuthnParameter.CLIENT_APP_VERSION.name(), loginInput.getAppVersion());
-
-      ClientAuthnDTO authnData = new ClientAuthnDTO(data, null);
-      String json = mapper.writeValueAsString(authnData);
-
-      // attach the login info json body to the post request
-      StringEntity input = new StringEntity(json, StandardCharsets.UTF_8);
-      input.setContentType("application/json");
-      HttpPost postRequest = new HttpPost(fedUrlUri);
-      postRequest.setEntity(input);
-      postRequest.addHeader("accept", "application/json");
-
-      // Add headers for driver name and version
-      postRequest.addHeader(SF_HEADER_CLIENT_APP_ID, loginInput.getAppId());
-      postRequest.addHeader(SF_HEADER_CLIENT_APP_VERSION, loginInput.getAppVersion());
+      StringEntity requestInput = prepareFederatedFlowStep1RequestInput(loginInput);
+      HttpPost postRequest = prepareFederatedFlowStep1PostRequest(loginInput, requestInput);
 
       final String gsResponse =
           HttpUtil.executeGeneralRequest(
@@ -1395,6 +1323,7 @@ private static JsonNode federatedFlowStep1(SFLoginInput loginInput) throws Snowf
               loginInput.getSocketTimeoutInMillis(),
               0,
               loginInput.getHttpClientSettingsKey());
+
       logger.debug("Authenticator-request response: {}", gsResponse);
       JsonNode jsonNode = mapper.readTree(gsResponse);
 
@@ -1790,12 +1719,148 @@ public static boolean isNewRetryStrategyRequest(HttpRequestBase request) {
     URI requestURI = request.getURI();
     String requestPath = requestURI.getPath();
     if (requestPath != null) {
-      if (requestPath.equals(SF_PATH_LOGIN_REQUEST)
+      return requestPath.equals(SF_PATH_LOGIN_REQUEST)
           || requestPath.equals(SF_PATH_AUTHENTICATOR_REQUEST)
-          || requestPath.equals(SF_PATH_TOKEN_REQUEST)) {
-        return true;
-      }
+          || requestPath.equals(SF_PATH_TOKEN_REQUEST);
     }
     return false;
   }
+
+  /**
+   * Prepares an HTTP POST request for the first step of the federated authentication flow.
+   *
+   * @param loginInput The login information for the request.
+   * @param inputData The JSON input data to include in the request.
+   * @return An {@link HttpPost} object ready to execute the federated flow request.
+   * @throws URISyntaxException If the constructed URI is invalid.
+   */
+  private static HttpPost prepareFederatedFlowStep1PostRequest(
+      SFLoginInput loginInput, StringEntity inputData) throws URISyntaxException {
+    URIBuilder fedUriBuilder = new URIBuilder(loginInput.getServerUrl());
+    // TODO: if loginInput.serverUrl contains port or additional segments - it will be ignored and
+    // overwritten here - to be fixed in SNOW-1922872
+    fedUriBuilder.setPath(SF_PATH_AUTHENTICATOR_REQUEST);
+    URI fedUrlUri = fedUriBuilder.build();
+
+    HttpPost postRequest = new HttpPost(fedUrlUri);
+    postRequest.setEntity(inputData);
+    postRequest.addHeader("accept", "application/json");
+
+    postRequest.addHeader(SF_HEADER_CLIENT_APP_ID, loginInput.getAppId());
+    postRequest.addHeader(SF_HEADER_CLIENT_APP_VERSION, loginInput.getAppVersion());
+
+    return postRequest;
+  }
+
+  /**
+   * Prepares the JSON input for the first step of the federated authentication flow.
+   *
+   * @param loginInput The login information for the request.
+   * @return A {@link StringEntity} containing the JSON input for the request.
+   * @throws JsonProcessingException If there is an error generating the JSON input.
+   */
+  private static StringEntity prepareFederatedFlowStep1RequestInput(SFLoginInput loginInput)
+      throws JsonProcessingException {
+    Map<String, Object> data = new HashMap<>();
+    data.put(ClientAuthnParameter.ACCOUNT_NAME.name(), loginInput.getAccountName());
+    data.put(ClientAuthnParameter.AUTHENTICATOR.name(), loginInput.getAuthenticator());
+    data.put(ClientAuthnParameter.CLIENT_APP_ID.name(), loginInput.getAppId());
+    data.put(ClientAuthnParameter.CLIENT_APP_VERSION.name(), loginInput.getAppVersion());
+
+    ClientAuthnDTO authnData = new ClientAuthnDTO(data, null);
+    String json = mapper.writeValueAsString(authnData);
+
+    StringEntity input = new StringEntity(json, StandardCharsets.UTF_8);
+    input.setContentType("application/json");
+    return input;
+  }
+
+  /**
+   * Sets the authentication data for the third step of the federated authentication flow.
+   *
+   * @param postRequest The {@link HttpPost} request to update with authentication data.
+   * @param loginInput The login information for the request.
+   * @throws SnowflakeSQLException If an error occurs while preparing the request.
+   */
+  private static void setFederatedFlowStep3PostRequestAuthData(
+      HttpPost postRequest, SFLoginInput loginInput) throws SnowflakeSQLException {
+    String userName =
+        Strings.isNullOrEmpty(loginInput.getOKTAUserName())
+            ? loginInput.getUserName()
+            : loginInput.getOKTAUserName();
+    try {
+      StringEntity params =
+          new StringEntity(
+              "{\"username\":\""
+                  + userName
+                  + "\",\"password\":\""
+                  + loginInput.getPassword()
+                  + "\"}");
+      postRequest.setEntity(params);
+
+      HeaderGroup headers = new HeaderGroup();
+      headers.addHeader(new BasicHeader(HttpHeaders.ACCEPT, "application/json"));
+      headers.addHeader(new BasicHeader(HttpHeaders.CONTENT_TYPE, "application/json"));
+      postRequest.setHeaders(headers.getAllHeaders());
+    } catch (IOException ex) {
+      handleFederatedFlowError(loginInput, ex);
+    }
+  }
+
+  /**
+   * Prepares an HTTP GET request for the fourth step of the federated authentication flow.
+   *
+   * @param retrieveSamlRequest The {@link HttpRequestBase} to update with the SAML request details.
+   * @param ssoUrl The SSO URL to use for the request.
+   * @param oneTimeToken The one-time token to include in the request.
+   * @throws MalformedURLException If the SSO URL is malformed.
+   * @throws URISyntaxException If the URI for the request cannot be built.
+   */
+  private static void prepareFederatedFlowStep4Request(
+      HttpRequestBase retrieveSamlRequest, String ssoUrl, String oneTimeToken)
+      throws MalformedURLException, URISyntaxException {
+    final URL url = new URL(ssoUrl);
+    URI oktaGetUri =
+        new URIBuilder()
+            .setScheme(url.getProtocol())
+            .setHost(url.getHost())
+            .setPort(url.getPort())
+            .setPath(url.getPath())
+            .setParameter("RelayState", "%2Fsome%2Fdeep%2Flink")
+            .setParameter("onetimetoken", oneTimeToken)
+            .build();
+    retrieveSamlRequest.setURI(oktaGetUri);
+
+    HeaderGroup headers = new HeaderGroup();
+    headers.addHeader(new BasicHeader(HttpHeaders.ACCEPT, "*/*"));
+    retrieveSamlRequest.setHeaders(headers.getAllHeaders());
+  }
+
+  private static void handleEmptyAuthResponse(
+      String theString, SFLoginInput loginInput, Exception lastRestException)
+      throws Exception, SFException {
+    if (theString == null) {
+      if (lastRestException != null) {
+        logger.error(
+            "Failed to open new session for user: {}, host: {}. Error: {}",
+            loginInput.getUserName(),
+            loginInput.getHostFromServerUrl(),
+            lastRestException);
+        throw lastRestException;
+      } else {
+        SnowflakeSQLException exception =
+            new SnowflakeSQLException(
+                NO_QUERY_ID,
+                "empty authentication response",
+                SqlState.CONNECTION_EXCEPTION,
+                ErrorCode.CONNECTION_ERROR.getMessageCode());
+        logger.error(
+            "Failed to open new session for user: {}, host: {}. Error: {}",
+            loginInput.getUserName(),
+            loginInput.getHostFromServerUrl(),
+            exception);
+        throw exception;
+      }
+    }
+  }
 }