SNOW-1874066: Adding Soteria tests for JDBC (#2026)

Co-authored-by: Dawid Heyman <[email protected]> Co-authored-by: Adam Kolodziejczyk <[email protected]> Co-authored-by: Przemyslaw Motacki <[email protected]>
snowflakedb · Feb 25, 2025 · 403c651 · 403c651
1 parent 485e690
commit 403c651
Show file tree

Hide file tree

Showing 18 changed files with 579 additions and 29 deletions.
diff --git a/.github/workflows/parameters_aws_auth_tests.json.gpg b/.github/workflows/parameters_aws_auth_tests.json.gpg
diff --git a/TestOnly/pom.xml b/TestOnly/pom.xml
@@ -414,6 +414,8 @@
             <exclude>**/ResultSetAsyncIT.java</exclude>
             <!-- ResultSet serialization is new -->
             <exclude>**/SnowflakeResultSetSerializable*IT.java</exclude>
+            <!-- AuthTestHelper is a new class for authentication tests -->
+            <exclude>**/AuthTestHelper.java</exclude>
           </testExcludes>
           <useIncrementalCompilation>false</useIncrementalCompilation>
         </configuration>

diff --git a/ci/container/test_authentication.sh b/ci/container/test_authentication.sh
@@ -12,6 +12,8 @@ eval $(jq -r '.authtestparams | to_entries | map("export \(.key)=\(.value|tostri
 export SF_ENABLE_EXPERIMENTAL_AUTHENTICATION=true
 
 $MVNW_EXE -DjenkinsIT \
+    -Dnet.snowflake.jdbc.temporaryCredentialCacheDir=/mnt/workspace/abc \
+    -Dnet.snowflake.jdbc.ocspResponseCacheDir=/mnt/workspace/abc \
     -Djava.io.tmpdir=$WORKSPACE \
     -Djacoco.skip.instrument=true \
     -Dskip.unitTests=true \

diff --git a/ci/test_authentication.sh b/ci/test_authentication.sh
@@ -1,16 +1,29 @@
 #!/bin/bash -e
 
 set -o pipefail
-THIS_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+export THIS_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
 export WORKSPACE=${WORKSPACE:-/tmp}
 export INTERNAL_REPO=nexus.int.snowflakecomputing.com:8086
 
-source $THIS_DIR/scripts/login_internal_docker.sh
+
+CI_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+if [[ -n "$JENKINS_HOME" ]]; then
+  ROOT_DIR="$(cd "${CI_DIR}/.." && pwd)"
+  export WORKSPACE=${WORKSPACE:-/tmp}
+
+  source $CI_DIR/_init.sh
+  source $CI_DIR/scripts/login_internal_docker.sh
+
+  echo "Use /sbin/ip"
+  IP_ADDR=$(/sbin/ip -4 addr show scope global dev eth0 | grep inet | awk '{print $2}' | cut -d / -f 1)
+
+fi
+
 gpg --quiet --batch --yes --decrypt --passphrase="$PARAMETERS_SECRET" --output $THIS_DIR/../.github/workflows/parameters_aws_auth_tests.json "$THIS_DIR/../.github/workflows/parameters_aws_auth_tests.json.gpg"
 
 docker run \
   -v $(cd $THIS_DIR/.. && pwd):/mnt/host \
   -v $WORKSPACE:/mnt/workspace \
   --rm \
-  nexus.int.snowflakecomputing.com:8086/docker/snowdrivers-test-external-browser-jdbc:1 \
+  nexus.int.snowflakecomputing.com:8086/docker/snowdrivers-test-external-browser-jdbc:2 \
   "/mnt/host/ci/container/test_authentication.sh"
diff --git a/src/main/java/net/snowflake/client/core/CredentialManager.java b/src/main/java/net/snowflake/client/core/CredentialManager.java
@@ -280,28 +280,32 @@ static void deleteMfaTokenCache(String host, String user) {
     getInstance().deleteTemporaryCredential(host, user, CachedCredentialType.MFA_TOKEN);
   }
 
+  /** Delete the Oauth access token cache */
+  static void deleteOAuthAccessTokenCache(String host, String user) {
+    logger.debug(
+        "Removing cached mfa token from a secure storage for user: {}, host: {}", user, host);
+    getInstance().deleteTemporaryCredential(host, user, CachedCredentialType.OAUTH_ACCESS_TOKEN);
+  }
+
   /** Delete the OAuth access token cache */
   static void deleteOAuthAccessTokenCache(SFLoginInput loginInput) throws SFException {
     String host = getHostForOAuthCacheKey(loginInput);
-    logger.debug(
-        "Removing cached OAuth access token from a secure storage for user: {}, host: {}",
-        loginInput.getUserName(),
-        host);
-    getInstance()
-        .deleteTemporaryCredential(
-            host, loginInput.getUserName(), CachedCredentialType.OAUTH_ACCESS_TOKEN);
+    deleteOAuthAccessTokenCache(host, loginInput.getUserName());
   }
 
   /** Delete the OAuth refresh token cache */
   static void deleteOAuthRefreshTokenCache(SFLoginInput loginInput) throws SFException {
     String host = getHostForOAuthCacheKey(loginInput);
+    deleteOAuthRefreshTokenCache(host, loginInput.getUserName());
+  }
+
+  /** Delete the Oauth refresh token cache */
+  static void deleteOAuthRefreshTokenCache(String host, String user) {
     logger.debug(
         "Removing cached OAuth refresh token from a secure storage for user: {}, host: {}",
-        loginInput.getUserName(),
+        user,
         host);
-    getInstance()
-        .deleteTemporaryCredential(
-            host, loginInput.getUserName(), CachedCredentialType.OAUTH_REFRESH_TOKEN);
+    getInstance().deleteTemporaryCredential(host, user, CachedCredentialType.OAUTH_REFRESH_TOKEN);
   }
 
   /**

diff --git a/src/main/java/net/snowflake/client/core/SFLoginInput.java b/src/main/java/net/snowflake/client/core/SFLoginInput.java
@@ -59,6 +59,11 @@ public class SFLoginInput {
   private boolean enableClientStoreTemporaryCredential;
   private boolean enableClientRequestMfaToken;
 
+  // OAuth
+  private int redirectUriPort = -1;
+  private String clientId;
+  private String clientSecret;
+
   private Duration browserResponseTimeout;
 
   // Additional headers to add for Snowsight.
@@ -426,13 +431,13 @@ SFLoginInput setHttpClientSettingsKey(HttpClientSettingsKey key) {
     this.httpClientKey = key;
     return this;
   }
-  // Opaque string sent for Snowsight account activation
 
+  // Opaque string sent for Snowsight account activation
   String getInFlightCtx() {
     return inFlightCtx;
   }
-  // Opaque string sent for Snowsight account activation
 
+  // Opaque string sent for Snowsight account activation
   SFLoginInput setInFlightCtx(String inFlightCtx) {
     this.inFlightCtx = inFlightCtx;
     return this;
@@ -447,6 +452,33 @@ SFLoginInput setDisableSamlURLCheck(boolean disableSamlURLCheck) {
     return this;
   }
 
+  public int getRedirectUriPort() {
+    return redirectUriPort;
+  }
+
+  public SFLoginInput setRedirectUriPort(int redirectUriPort) {
+    this.redirectUriPort = redirectUriPort;
+    return this;
+  }
+
+  public String getClientId() {
+    return clientId;
+  }
+
+  public SFLoginInput setClientId(String clientId) {
+    this.clientId = clientId;
+    return this;
+  }
+
+  public String getClientSecret() {
+    return clientSecret;
+  }
+
+  public SFLoginInput setClientSecret(String clientSecret) {
+    this.clientSecret = clientSecret;
+    return this;
+  }
+
   Map<String, String> getAdditionalHttpHeadersForSnowsight() {
     return additionalHttpHeadersForSnowsight;
   }

diff --git a/src/main/java/net/snowflake/client/core/SFSession.java b/src/main/java/net/snowflake/client/core/SFSession.java
@@ -1259,6 +1259,11 @@ public String getIdToken() {
     return idToken;
   }
 
+  @SnowflakeJdbcInternalApi
+  public String getAccessToken() {
+    return oauthAccessToken;
+  }
+
   public String getMfaToken() {
     return mfaToken;
   }

diff --git a/src/main/java/net/snowflake/client/core/SessionUtil.java b/src/main/java/net/snowflake/client/core/SessionUtil.java
@@ -1090,6 +1090,28 @@ public static void deleteIdTokenCache(String host, String user) {
     CredentialManager.deleteIdTokenCache(host, user);
   }
 
+  /**
+   * Delete the Oauth access token cache
+   *
+   * @param host The host string
+   * @param user The user
+   */
+  @SnowflakeJdbcInternalApi
+  public static void deleteOAuthAccessTokenCache(String host, String user) {
+    CredentialManager.deleteOAuthAccessTokenCache(host, user);
+  }
+
+  /**
+   * Delete the Oauth refresh token cache
+   *
+   * @param host The host string
+   * @param user The user
+   */
+  @SnowflakeJdbcInternalApi
+  public static void deleteOAuthRefreshTokenCache(String host, String user) {
+    CredentialManager.deleteOAuthRefreshTokenCache(host, user);
+  }
+
   /**
    * Delete the mfa token cache
    *

diff --git a/src/test/java/net/snowflake/client/authentication/AuthConnectionParameters.java b/src/test/java/net/snowflake/client/authentication/AuthConnectionParameters.java
@@ -9,6 +9,9 @@ public class AuthConnectionParameters {
   static final String SSO_USER = systemGetEnv("SNOWFLAKE_AUTH_TEST_BROWSER_USER");
   static final String HOST = systemGetEnv("SNOWFLAKE_AUTH_TEST_HOST");
   static final String SSO_PASSWORD = systemGetEnv("SNOWFLAKE_AUTH_TEST_OKTA_PASS");
+  static final String OKTA = systemGetEnv("SNOWFLAKE_AUTH_TEST_OKTA_NAME");
+  static final String OAUTH_PASSWORD =
+      systemGetEnv("SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_USER_PASSWORD");
 
   static Properties getBaseConnectionParameters() {
     Properties properties = new Properties();
@@ -19,6 +22,7 @@ static Properties getBaseConnectionParameters() {
     properties.put("db", systemGetEnv("SNOWFLAKE_AUTH_TEST_DATABASE"));
     properties.put("schema", systemGetEnv("SNOWFLAKE_AUTH_TEST_SCHEMA"));
     properties.put("warehouse", systemGetEnv("SNOWFLAKE_AUTH_TEST_WAREHOUSE"));
+    properties.put("CLIENT_STORE_TEMPORARY_CREDENTIAL", false);
     return properties;
   }
 
@@ -50,4 +54,65 @@ static Properties getOauthConnectionParameters(String token) {
     properties.put("token", token);
     return properties;
   }
+
+  static Properties getOAuthExternalAuthorizationCodeConnectionParameters() {
+    Properties properties = getBaseConnectionParameters();
+    properties.put("authenticator", "OAUTH_AUTHORIZATION_CODE");
+    properties.put(
+        "oauthClientId", systemGetEnv("SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_ID"));
+    properties.put(
+        "oauthClientSecret", systemGetEnv("SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_SECRET"));
+    properties.put(
+        "oauthRedirectURI", systemGetEnv("SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_REDIRECT_URI"));
+    properties.put(
+        "oauthAuthorizationUrl", systemGetEnv("SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_AUTH_URL"));
+    properties.put(
+        "oauthTokenRequestUrl", systemGetEnv("SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_TOKEN"));
+    properties.put("user", SSO_USER);
+
+    return properties;
+  }
+
+  static Properties getOAuthSnowflakeAuthorizationCodeConnectionParameters() {
+    Properties properties = getBaseConnectionParameters();
+    properties.put("authenticator", "OAUTH_AUTHORIZATION_CODE");
+    properties.put(
+        "oauthClientId", systemGetEnv("SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_CLIENT_ID"));
+    properties.put(
+        "oauthClientSecret",
+        systemGetEnv("SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_CLIENT_SECRET"));
+    properties.put(
+        "oauthRedirectURI",
+        systemGetEnv("SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_REDIRECT_URI"));
+    properties.put("role", systemGetEnv("SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_ROLE"));
+    properties.put("user", systemGetEnv("SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_ID"));
+    return properties;
+  }
+
+  static Properties getOAuthSnowflakeWildcardsAuthorizationCodeConnectionParameters() {
+    Properties properties = getBaseConnectionParameters();
+    properties.put("authenticator", "OAUTH_AUTHORIZATION_CODE");
+    properties.put(
+        "oauthClientId",
+        systemGetEnv("SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_WILDCARDS_CLIENT_ID"));
+    properties.put(
+        "oauthClientSecret",
+        systemGetEnv("SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_WILDCARDS_CLIENT_SECRET"));
+    properties.put("role", systemGetEnv("SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_ROLE"));
+    properties.put("user", systemGetEnv("SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_ID"));
+    return properties;
+  }
+
+  static Properties getOAuthSnowflakeClientCredentialParameters() {
+    Properties properties = getBaseConnectionParameters();
+    properties.put("authenticator", "OAUTH_CLIENT_CREDENTIALS");
+    properties.put(
+        "oauthClientId", systemGetEnv("SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_ID"));
+    properties.put(
+        "oauthClientSecret", systemGetEnv("SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_SECRET"));
+    properties.put(
+        "oauthTokenRequestUrl", systemGetEnv("SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_TOKEN"));
+    properties.put("user", systemGetEnv("SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_ID"));
+    return properties;
+  }
 }
diff --git a/src/test/java/net/snowflake/client/authentication/AuthTestHelper.java b/src/test/java/net/snowflake/client/authentication/AuthTestHelper.java
@@ -22,6 +22,7 @@ public class AuthTestHelper {
 
   private Exception exception;
   private String idToken;
+  private String accessToken;
   private final boolean runAuthTestsManually;
 
   public AuthTestHelper() {
@@ -87,17 +88,36 @@ public static void deleteIdToken() {
         AuthConnectionParameters.HOST, AuthConnectionParameters.SSO_USER);
   }
 
+  public static void deleteIdToken(String host, String user) {
+    SessionUtil.deleteIdTokenCache(host, user);
+  }
+
+  public static void deleteOauthToken() {
+    SessionUtil.deleteOAuthAccessTokenCache(
+        AuthConnectionParameters.OKTA, AuthConnectionParameters.SSO_USER);
+  }
+
+  public static void deleteOauthToken(String host, String user) {
+    SessionUtil.deleteOAuthAccessTokenCache(host, user);
+  }
+
+  public static void deleteOauthRefreshToken(String host, String user) {
+    SessionUtil.deleteOAuthRefreshTokenCache(host, user);
+  }
+
   public void connectAndExecuteSimpleQuery(Properties props, String sessionParameters) {
-    String url = String.format("jdbc:snowflake://%s:%s", props.get("host"), props.get("port"));
+    String url = String.format("jdbc:snowflake://%s", props.get("host"));
     if (sessionParameters != null) {
       url += "?" + sessionParameters;
     }
     try (Connection con = DriverManager.getConnection(url, props);
         Statement stmt = con.createStatement();
         ResultSet rs = stmt.executeQuery("select 1")) {
       assertTrue(rs.next());
-      assertEquals(1, rs.getInt(1));
+      int value = rs.getInt(1);
+      assertEquals(1, value);
       saveToken(con);
+      saveAccessToken(con);
     } catch (SQLException e) {
       this.exception = e;
     }
@@ -108,7 +128,16 @@ private void saveToken(Connection con) throws SnowflakeSQLException {
     this.idToken = sfcon.getSfSession().getIdToken();
   }
 
+  private void saveAccessToken(Connection con) throws SnowflakeSQLException {
+    SnowflakeConnectionV1 sfcon = (SnowflakeConnectionV1) con;
+    this.accessToken = sfcon.getSfSession().getAccessToken();
+  }
+
   public String getIdToken() {
     return idToken;
   }
+
+  public String getAccessToken() {
+    return accessToken;
+  }
 }
diff --git a/src/test/java/net/snowflake/client/authentication/ExternalBrowserLatestIT.java b/src/test/java/net/snowflake/client/authentication/ExternalBrowserLatestIT.java
@@ -16,10 +16,12 @@ class ExternalBrowserLatestIT {
   String login = AuthConnectionParameters.SSO_USER;
   String password = AuthConnectionParameters.SSO_PASSWORD;
   AuthTestHelper authTestHelper = new AuthTestHelper();
+  Properties properties;
 
   @BeforeEach
   public void setUp() throws IOException {
     AuthTestHelper.deleteIdToken();
+    properties = getExternalBrowserConnectionParameters();
   }
 
   @AfterEach
@@ -32,17 +34,14 @@ public void tearDown() {
   void shouldAuthenticateUsingExternalBrowser() throws InterruptedException {
     Thread provideCredentialsThread =
         new Thread(() -> authTestHelper.provideCredentials("success", login, password));
-    Thread connectThread =
-        authTestHelper.getConnectAndExecuteSimpleQueryThread(
-            getExternalBrowserConnectionParameters());
+    Thread connectThread = authTestHelper.getConnectAndExecuteSimpleQueryThread(properties);
 
     authTestHelper.connectAndProvideCredentials(provideCredentialsThread, connectThread);
     authTestHelper.verifyExceptionIsNotThrown();
   }
 
   @Test
   void shouldThrowErrorForMismatchedUsername() throws InterruptedException {
-    Properties properties = getExternalBrowserConnectionParameters();
     properties.put("user", "differentUsername");
     Thread provideCredentialsThread =
         new Thread(() -> authTestHelper.provideCredentials("success", login, password));
@@ -61,7 +60,7 @@ void shouldThrowErrorForWrongCredentials() throws InterruptedException {
         new Thread(() -> authTestHelper.provideCredentials("fail", login, password));
     Thread connectThread =
         authTestHelper.getConnectAndExecuteSimpleQueryThread(
-            getExternalBrowserConnectionParameters(), "BROWSER_RESPONSE_TIMEOUT=10");
+            properties, "BROWSER_RESPONSE_TIMEOUT=10");
 
     authTestHelper.connectAndProvideCredentials(provideCredentialsThread, connectThread);
     authTestHelper.verifyExceptionIsThrown(
@@ -74,7 +73,7 @@ void shouldThrowErrorForBrowserTimeout() throws InterruptedException {
         new Thread(() -> authTestHelper.provideCredentials("timeout", login, password));
     Thread connectThread =
         authTestHelper.getConnectAndExecuteSimpleQueryThread(
-            getExternalBrowserConnectionParameters(), "BROWSER_RESPONSE_TIMEOUT=1");
+            properties, "BROWSER_RESPONSE_TIMEOUT=1");
 
     authTestHelper.connectAndProvideCredentials(provideCredentialsThread, connectThread);
     authTestHelper.verifyExceptionIsThrown(