add oauth and okta automated tests

snowflakedb · Dec 11, 2024 · 648ee37 · 648ee37
1 parent d979454
commit 648ee37
Show file tree

Hide file tree

Showing 4 changed files with 184 additions and 38 deletions.
diff --git a/src/test/java/net/snowflake/client/authentication/AuthConnectionParameters.java b/src/test/java/net/snowflake/client/authentication/AuthConnectionParameters.java
@@ -34,4 +34,20 @@ static Properties getStoreIDTokenConnectionParameters() {
     properties.put("CLIENT_STORE_TEMPORARY_CREDENTIAL", true);
     return properties;
   }
+
+  static Properties getOktaConnectionParameters() {
+    Properties properties = getBaseConnectionParameters();
+    properties.put("user", SSO_USER);
+    properties.put("password", SSO_PASSWORD);
+    properties.put("authenticator", systemGetEnv("SNOWFLAKE_AUTH_TEST_OAUTH_URL"));
+    return properties;
+  }
+
+  static Properties getOauthConnectionParameters(String token) {
+    Properties properties = getBaseConnectionParameters();
+    properties.put("user", SSO_USER);
+    properties.put("authenticator", "OAUTH");
+    properties.put("token", token);
+    return properties;
+  }
 }
diff --git a/src/test/java/net/snowflake/client/authentication/OauthIT.java b/src/test/java/net/snowflake/client/authentication/OauthIT.java
@@ -0,0 +1,94 @@
+package net.snowflake.client.authentication;
+
+import static net.snowflake.client.authentication.AuthConnectionParameters.getOauthConnectionParameters;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.Base64;
+import java.util.List;
+import java.util.Properties;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import net.snowflake.client.category.TestTags;
+import org.apache.http.client.entity.UrlEncodedFormEntity;
+import org.apache.http.client.methods.CloseableHttpResponse;
+import org.apache.http.client.methods.HttpPost;
+import org.apache.http.impl.client.CloseableHttpClient;
+import org.apache.http.impl.client.HttpClients;
+import org.apache.http.message.BasicNameValuePair;
+import org.apache.http.util.EntityUtils;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+
+@Tag(TestTags.AUTHENTICATION)
+public class OauthIT {
+
+  AuthTest authTest;
+
+  @BeforeEach
+  public void setUp() throws IOException {
+    authTest = new AuthTest();
+  }
+
+  @Test
+  void shouldAuthenticateUsingOauth() throws IOException {
+    authTest.connectAndExecuteSimpleQuery(getOauthConnectionParameters(getToken()), null);
+    authTest.verifyExceptionIsNotThrown();
+  }
+
+  @Test
+  void shouldThrowErrorForInvalidToken() {
+    authTest.connectAndExecuteSimpleQuery(getOauthConnectionParameters("invalidToken"), null);
+    authTest.verifyExceptionIsThrown("Invalid OAuth access token. ");
+  }
+
+  @Test
+  void shouldThrowErrorForMismatchedOauthUsername() throws IOException {
+    Properties properties = getOauthConnectionParameters(getToken());
+    properties.put("user", "differentUsername");
+    authTest.connectAndExecuteSimpleQuery(properties, null);
+    authTest.verifyExceptionIsThrown(
+        "The user you were trying to authenticate as differs from the user tied to the access token.");
+  }
+
+  private String getToken() throws IOException {
+    List<BasicNameValuePair> data =
+        Stream.of(
+                new BasicNameValuePair("username", System.getenv("SNOWFLAKE_AUTH_TEST_OKTA_USER")),
+                new BasicNameValuePair("password", System.getenv("SNOWFLAKE_AUTH_TEST_OKTA_PASS")),
+                new BasicNameValuePair("grant_type", "password"),
+                new BasicNameValuePair(
+                    "scope",
+                    "session:role:" + System.getenv("SNOWFLAKE_AUTH_TEST_ROLE").toLowerCase()))
+            .collect(Collectors.toList());
+
+    String auth =
+        System.getenv("SNOWFLAKE_AUTH_TEST_OAUTH_CLIENT_ID")
+            + ":"
+            + System.getenv("SNOWFLAKE_AUTH_TEST_OAUTH_CLIENT_SECRET");
+    String encodedAuth = Base64.getEncoder().encodeToString(auth.getBytes(StandardCharsets.UTF_8));
+
+    try (CloseableHttpClient httpClient = HttpClients.createDefault()) {
+      HttpPost httpPost = new HttpPost(System.getenv("SNOWFLAKE_AUTH_TEST_OAUTH_URL"));
+      httpPost.setHeader("Content-Type", "application/x-www-form-urlencoded;charset=UTF-8");
+      httpPost.setHeader("Authorization", "Basic " + encodedAuth);
+      httpPost.setEntity(new UrlEncodedFormEntity(data, StandardCharsets.UTF_8));
+
+      try (CloseableHttpResponse response = httpClient.execute(httpPost)) {
+        if (response.getStatusLine().getStatusCode() != 200) {
+          throw new IOException(
+              "Failed to get access token, response code: "
+                  + response.getStatusLine().getStatusCode());
+        }
+
+        String responseBody = EntityUtils.toString(response.getEntity(), StandardCharsets.UTF_8);
+        ObjectMapper mapper = new ObjectMapper();
+        JsonNode jsonNode = mapper.readTree(responseBody);
+        return jsonNode.get("access_token").asText();
+      }
+    }
+  }
+}
diff --git a/src/test/java/net/snowflake/client/authentication/OktaAuthIT.java b/src/test/java/net/snowflake/client/authentication/OktaAuthIT.java
@@ -0,0 +1,74 @@
+package net.snowflake.client.authentication;
+
+import static net.snowflake.client.authentication.AuthConnectionParameters.SSO_USER;
+import static net.snowflake.client.authentication.AuthConnectionParameters.getOktaConnectionParameters;
+
+import java.io.IOException;
+import java.util.Properties;
+import net.snowflake.client.category.TestTags;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+
+@Tag(TestTags.AUTHENTICATION)
+class OktaAuthIT {
+
+  AuthTest authTest;
+
+  @BeforeEach
+  public void setUp() throws IOException {
+    authTest = new AuthTest();
+  }
+
+  @Test
+  void shouldAuthenticateUsingOkta() {
+    authTest.connectAndExecuteSimpleQuery(getOktaConnectionParameters(), null);
+    authTest.verifyExceptionIsNotThrown();
+  }
+
+  @Test
+  void shouldAuthenticateUsingOktaWithOktaUsernameParam() {
+    Properties properties = getOktaConnectionParameters();
+    properties.replace("user", "differentUsername");
+    authTest.connectAndExecuteSimpleQuery(properties, "oktausername=" + SSO_USER);
+    authTest.verifyExceptionIsNotThrown();
+  }
+
+  @Test
+  void shouldThrowErrorForWrongOktaCredentials() {
+    Properties properties = getOktaConnectionParameters();
+    properties.put("user", "invalidUsername");
+    properties.put("password", "fakepassword");
+    authTest.connectAndExecuteSimpleQuery(properties, null);
+    authTest.verifyExceptionIsThrown(
+        "JDBC driver encountered communication error. Message: HTTP status=401.");
+  }
+
+  @Test
+  void shouldThrowErrorForWrongOktaCredentialsInOktaUsernameParam() {
+    Properties properties = getOktaConnectionParameters();
+    properties.replace("user", "differentUsername");
+    authTest.connectAndExecuteSimpleQuery(properties, "oktausername=invalidUser");
+    authTest.verifyExceptionIsThrown(
+        "JDBC driver encountered communication error. Message: HTTP status=401.");
+  }
+
+  @Test
+  void shouldThrowErrorForWrongOktaUrl() {
+    Properties properties = getOktaConnectionParameters();
+    properties.put("authenticator", "https://invalid.okta.com/");
+    authTest.connectAndExecuteSimpleQuery(properties, null);
+    authTest.verifyExceptionIsThrown(
+        "The specified authenticator is not accepted by your Snowflake account configuration.  Please contact your local system administrator to get the correct URL to use.");
+  }
+
+  @Test
+  @Disabled // todo SNOW-1852279 implement error handling for invalid URL
+  void shouldThrowErrorForWrongUrlWithoutOktaPath() {
+    Properties properties = getOktaConnectionParameters();
+    properties.put("authenticator", "https://invalid.abc.com/");
+    authTest.connectAndExecuteSimpleQuery(properties, null);
+    authTest.verifyExceptionIsThrown("todo");
+  }
+}
diff --git a/src/test/java/net/snowflake/client/jdbc/ConnectionIT.java b/src/test/java/net/snowflake/client/jdbc/ConnectionIT.java
@@ -44,7 +44,6 @@
 import java.util.concurrent.Executors;
 import net.snowflake.client.TestUtil;
 import net.snowflake.client.annotations.DontRunOnGithubActions;
-import net.snowflake.client.annotations.RunOnTestaccountNotOnGithubActions;
 import net.snowflake.client.category.TestTags;
 import net.snowflake.client.core.SFSession;
 import net.snowflake.common.core.SqlState;
@@ -837,43 +836,6 @@ public void testResultSetsClosedByStatement() throws SQLException {
     assertTrue(rs4.isClosed());
   }
 
-  @Test
-  @RunOnTestaccountNotOnGithubActions
-  public void testOKTAConnection() throws Throwable {
-    Map<String, String> params = getConnectionParameters();
-    Properties properties = new Properties();
-    properties.put("user", params.get("ssoUser"));
-    properties.put("password", params.get("ssoPassword"));
-    properties.put("ssl", params.get("ssl"));
-    properties.put("authenticator", "https://snowflakecomputing.okta.com/");
-
-    DriverManager.getConnection(
-        String.format(
-            "jdbc:snowflake://%s.reg.snowflakecomputing.com:%s/",
-            params.get("account"), params.get("port")),
-        properties);
-  }
-
-  @Test
-  @RunOnTestaccountNotOnGithubActions
-  public void testOKTAConnectionWithOktauserParam() throws Throwable {
-    Map<String, String> params = getConnectionParameters();
-    Properties properties = new Properties();
-    properties.put("user", "test");
-    properties.put("password", params.get("ssoPassword"));
-    properties.put("ssl", params.get("ssl"));
-    properties.put(
-        "authenticator",
-        String.format(
-            "https://snowflakecomputing.okta.com;oktausername=%s;", params.get("ssoUser")));
-
-    DriverManager.getConnection(
-        String.format(
-            "jdbc:snowflake://%s.reg.snowflakecomputing.com:%s/",
-            params.get("account"), params.get("port")),
-        properties);
-  }
-
   @Test
   public void testValidateDefaultParameters() throws Throwable {
     Map<String, String> params = getConnectionParameters();