SNOW 500881 JWT expiration issue (#715)

*SNOW-500881 JWT expiration fix

Co-authored-by: sfc-gh-hchaturvedi <[email protected]>
Co-authored-by: Lorna Barber <[email protected]>

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012-2019 Snowflake Computing Inc. All rights reserved.
+ * Copyright (c) 2012-2022 Snowflake Computing Inc. All rights reserved.
  */
 
 package net.snowflake.client.core;
@@ -471,6 +471,24 @@ public static RequestConfig getDefaultRequestConfigWithSocketTimeout(
         .build();
   }
 
+  /**
+   * Return a request configuration inheriting from the default request configuration of the shared
+   * HttpClient with a different socket and connect timeout.
+   *
+   * @param requestSocketAndConnectTimeout - custom socket and connect timeout in milli-seconds
+   * @param withoutCookies - whether this request should ignore cookies or not
+   * @return RequestConfig object
+   */
+  public static RequestConfig getDefaultRequestConfigWithSocketAndConnectTimeout(
+      int requestSocketAndConnectTimeout, boolean withoutCookies) {
+    final String cookieSpec = withoutCookies ? IGNORE_COOKIES : DEFAULT;
+    return RequestConfig.copy(DefaultRequestConfig)
+        .setSocketTimeout(requestSocketAndConnectTimeout)
+        .setConnectTimeout(requestSocketAndConnectTimeout)
+        .setCookieSpec(cookieSpec)
+        .build();
+  }
+
   /**
    * Return a request configuration inheriting from the default request configuration of the shared
    * HttpClient with the coopkie spec set to ignore.
@@ -517,6 +535,9 @@ public static boolean isSocksProxyDisabled() {
    *
    * @param httpRequest HttpRequestBase
    * @param retryTimeout retry timeout
+   * @param authTimeout authenticator specific timeout
+   * @param socketTimeout socket timeout (in ms)
+   * @param retryCount retry count for the request
    * @param injectSocketTimeout injecting socket timeout
    * @param canceling canceling?
    * @param ocspAndProxyKey OCSP mode and proxy settings for httpclient
@@ -527,13 +548,19 @@ public static boolean isSocksProxyDisabled() {
   static String executeRequestWithoutCookies(
       HttpRequestBase httpRequest,
       int retryTimeout,
+      int authTimeout,
+      int socketTimeout,
+      int retryCount,
       int injectSocketTimeout,
       AtomicBoolean canceling,
       HttpClientSettingsKey ocspAndProxyKey)
       throws SnowflakeSQLException, IOException {
     return executeRequestInternal(
         httpRequest,
         retryTimeout,
+        authTimeout,
+        socketTimeout,
+        retryCount,
         injectSocketTimeout,
         canceling,
         true, // no cookie
@@ -548,17 +575,28 @@ static String executeRequestWithoutCookies(
    *
    * @param httpRequest HttpRequestBase
    * @param retryTimeout retry timeout
+   * @param authTimeout authenticator specific timeout
+   * @param socketTimeout socket timeout (in ms)
+   * @param retryCount retry count for the request
    * @param ocspAndProxyKey OCSP mode and proxy settings for httpclient
    * @return response
    * @throws SnowflakeSQLException if Snowflake error occurs
    * @throws IOException raises if a general IO error occurs
    */
   public static String executeGeneralRequest(
-      HttpRequestBase httpRequest, int retryTimeout, HttpClientSettingsKey ocspAndProxyKey)
+      HttpRequestBase httpRequest,
+      int retryTimeout,
+      int authTimeout,
+      int socketTimeout,
+      int retryCount,
+      HttpClientSettingsKey ocspAndProxyKey)
       throws SnowflakeSQLException, IOException {
     return executeRequest(
         httpRequest,
         retryTimeout,
+        authTimeout,
+        socketTimeout,
+        retryCount,
         0, // no inject socket timeout
         null, // no canceling
         false, // no retry parameter
@@ -571,17 +609,28 @@ public static String executeGeneralRequest(
    *
    * @param httpRequest HttpRequestBase
    * @param retryTimeout retry timeout
+   * @param authTimeout authenticator specific timeout
+   * @param socketTimeout socket timeout (in ms)
+   * @param retryCount retry count for the request
    * @param httpClient client object used to communicate with other machine
    * @return response
    * @throws SnowflakeSQLException if Snowflake error occurs
    * @throws IOException raises if a general IO error occurs
    */
   public static String executeGeneralRequest(
-      HttpRequestBase httpRequest, int retryTimeout, CloseableHttpClient httpClient)
+      HttpRequestBase httpRequest,
+      int retryTimeout,
+      int authTimeout,
+      int socketTimeout,
+      int retryCount,
+      CloseableHttpClient httpClient)
       throws SnowflakeSQLException, IOException {
     return executeRequestInternal(
         httpRequest,
         retryTimeout,
+        authTimeout,
+        socketTimeout,
+        retryCount,
         0, // no inject socket timeout
         null, // no canceling
         false, // with cookie
@@ -596,6 +645,9 @@ public static String executeGeneralRequest(
    *
    * @param httpRequest HttpRequestBase
    * @param retryTimeout retry timeout
+   * @param authTimeout authenticator timeout
+   * @param socketTimeout socket timeout (in ms)
+   * @param retryCount retry count for the request
    * @param injectSocketTimeout injecting socket timeout
    * @param canceling canceling?
    * @param includeRetryParameters whether to include retry parameters in retried requests
@@ -608,6 +660,9 @@ public static String executeGeneralRequest(
   public static String executeRequest(
       HttpRequestBase httpRequest,
       int retryTimeout,
+      int authTimeout,
+      int socketTimeout,
+      int retryCount,
       int injectSocketTimeout,
       AtomicBoolean canceling,
       boolean includeRetryParameters,
@@ -617,6 +672,9 @@ public static String executeRequest(
     return executeRequestInternal(
         httpRequest,
         retryTimeout,
+        authTimeout,
+        socketTimeout,
+        retryCount,
         injectSocketTimeout,
         canceling,
         false, // with cookie (do we need cookie?)
@@ -635,6 +693,9 @@ public static String executeRequest(
    *
    * @param httpRequest request object contains all the information
    * @param retryTimeout retry timeout (in seconds)
+   * @param authTimeout authenticator specific timeout (in seconds)
+   * @param socketTimeout socket timeout (in ms)
+   * @param retryCount retry count for the request
    * @param injectSocketTimeout simulate socket timeout
    * @param canceling canceling flag
    * @param withoutCookies whether this request should ignore cookies
@@ -649,6 +710,9 @@ public static String executeRequest(
   private static String executeRequestInternal(
       HttpRequestBase httpRequest,
       int retryTimeout,
+      int authTimeout,
+      int socketTimeout,
+      int retryCount,
       int injectSocketTimeout,
       AtomicBoolean canceling,
       boolean withoutCookies,
@@ -667,12 +731,16 @@ private static String executeRequestInternal(
     String theString;
     StringWriter writer = null;
     CloseableHttpResponse response = null;
+
     try {
       response =
           RestRequest.execute(
               httpClient,
               httpRequest,
               retryTimeout,
+              authTimeout,
+              socketTimeout,
+              retryCount,
               injectSocketTimeout,
               canceling,
               withoutCookies,

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019 Snowflake Computing Inc. All rights reserved.
+ * Copyright (c) 2019-2022 Snowflake Computing Inc. All rights reserved.
  */
 package net.snowflake.client.core;
 
@@ -232,6 +232,9 @@ public void flush() {
           HttpUtil.executeGeneralRequest(
               postRequest,
               1000,
+              0,
+              0,
+              0,
               ocspAndProxyKey != null
                   ? ocspAndProxyKey
                   : new HttpClientSettingsKey(OCSPMode.FAIL_OPEN));

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012-2021 Snowflake Computing Inc. All rights reserved.
+ * Copyright (c) 2012-2022 Snowflake Computing Inc. All rights reserved.
  */
 
 package net.snowflake.client.core;
@@ -765,5 +765,11 @@ public SFConnectionHandler getSfConnectionHandler() {
 
   public abstract int getNetworkTimeoutInMilli();
 
+  public abstract int getAuthTimeout();
+
   public abstract SnowflakeConnectString getSnowflakeConnectionString();
+
+  public abstract int getHttpClientConnectionTimeout();
+
+  public abstract int getHttpClientSocketTimeout();
 }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012-2019 Snowflake Computing Inc. All rights reserved.
+ * Copyright (c) 2012-2022 Snowflake Computing Inc. All rights reserved.
  */
 
 package net.snowflake.client.core;
@@ -25,6 +25,7 @@ public class SFLoginInput {
   private String oktaUserName;
   private String accountName;
   private int loginTimeout = -1; // default is invalid
+  private int authTimeout = 0;
   private String userName;
   private String password;
   private boolean passcodeInPassword;
@@ -139,6 +140,15 @@ SFLoginInput setLoginTimeout(int loginTimeout) {
     return this;
   }
 
+  int getAuthTimeout() {
+    return authTimeout;
+  }
+
+  SFLoginInput setAuthTimeout(int authTimeout) {
+    this.authTimeout = authTimeout;
+    return this;
+  }
+
   public String getUserName() {
     return userName;
   }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012-2019 Snowflake Computing Inc. All rights reserved.
+ * Copyright (c) 2012-2022 Snowflake Computing Inc. All rights reserved.
  */
 
 package net.snowflake.client.core;
@@ -80,6 +80,7 @@ public class SFSession extends SFBaseSession {
    */
   private int networkTimeoutInMilli = 0; // in milliseconds
 
+  private int authTimeout = 0;
   private boolean enableCombineDescribe = false;
   private int httpClientConnectionTimeout = 60000; // milliseconds
   private int httpClientSocketTimeout = DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT; // milliseconds
@@ -166,7 +167,9 @@ public QueryStatus getQueryStatus(String queryID) throws SQLException {
       try {
         get.setHeader("Content-type", "application/json");
         get.setHeader("Authorization", "Snowflake Token=\"" + this.sessionToken + "\"");
-        response = HttpUtil.executeGeneralRequest(get, loginTimeout, getHttpClientKey());
+        response =
+            HttpUtil.executeGeneralRequest(
+                get, loginTimeout, authTimeout, httpClientSocketTimeout, 0, getHttpClientKey());
         jsonNode = OBJECT_MAPPER.readTree(response);
       } catch (Exception e) {
         throw new SnowflakeSQLLoggedException(
@@ -432,6 +435,7 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
         .setOKTAUserName((String) connectionPropertiesMap.get(SFSessionProperty.OKTA_USERNAME))
         .setAccountName((String) connectionPropertiesMap.get(SFSessionProperty.ACCOUNT))
         .setLoginTimeout(loginTimeout)
+        .setAuthTimeout(authTimeout)
         .setUserName((String) connectionPropertiesMap.get(SFSessionProperty.USER))
         .setPassword((String) connectionPropertiesMap.get(SFSessionProperty.PASSWORD))
         .setToken((String) connectionPropertiesMap.get(SFSessionProperty.TOKEN))
@@ -457,6 +461,7 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
         SessionUtil.openSession(loginInput, connectionPropertiesMap, tracingLevel.toString());
     isClosed = false;
 
+    authTimeout = loginInput.getAuthTimeout();
     sessionToken = loginOutput.getSessionToken();
     masterToken = loginOutput.getMasterToken();
     idToken = loginOutput.getIdToken();
@@ -716,7 +721,13 @@ protected void heartbeat() throws SFException, SQLException {
         // per https://support-snowflake.zendesk.com/agent/tickets/6629
         int SF_HEARTBEAT_TIMEOUT = 300;
         String theResponse =
-            HttpUtil.executeGeneralRequest(postRequest, SF_HEARTBEAT_TIMEOUT, getHttpClientKey());
+            HttpUtil.executeGeneralRequest(
+                postRequest,
+                SF_HEARTBEAT_TIMEOUT,
+                authTimeout,
+                httpClientSocketTimeout,
+                0,
+                getHttpClientKey());
 
         JsonNode rootNode;
 
@@ -789,6 +800,18 @@ public int getNetworkTimeoutInMilli() {
     return networkTimeoutInMilli;
   }
 
+  public int getAuthTimeout() {
+    return authTimeout;
+  }
+
+  public int getHttpClientSocketTimeout() {
+    return httpClientSocketTimeout;
+  }
+
+  public int getHttpClientConnectionTimeout() {
+    return httpClientConnectionTimeout;
+  }
+
   public boolean isClosed() {
     return isClosed;
   }