Skip to content

Commit

Permalink
Add PAT support
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-dheyman committed Dec 11, 2024
1 parent e8dc943 commit 801728c
Show file tree
Hide file tree
Showing 17 changed files with 261 additions and 70 deletions.
6 changes: 2 additions & 4 deletions src/main/java/net/snowflake/client/core/SFLoginInput.java
Original file line number Diff line number Diff line change
Expand Up @@ -249,8 +249,7 @@ public int getSocketTimeoutInMillis() {
return (int) socketTimeout.toMillis();
}

@SnowflakeJdbcInternalApi
public SFLoginInput setSocketTimeout(Duration socketTimeout) {
SFLoginInput setSocketTimeout(Duration socketTimeout) {
this.socketTimeout = socketTimeout;
return this;
}
Expand Down Expand Up @@ -401,8 +400,7 @@ public HttpClientSettingsKey getHttpClientSettingsKey() {
return httpClientKey;
}

@SnowflakeJdbcInternalApi
public SFLoginInput setHttpClientSettingsKey(HttpClientSettingsKey key) {
SFLoginInput setHttpClientSettingsKey(HttpClientSettingsKey key) {
this.httpClientKey = key;
return this;
}
Expand Down
22 changes: 14 additions & 8 deletions src/main/java/net/snowflake/client/core/SessionUtil.java
Original file line number Diff line number Diff line change
Expand Up @@ -231,6 +231,10 @@ private static AuthenticatorType getAuthenticator(SFLoginInput loginInput) {
} else if (loginInput.getAuthenticator().equalsIgnoreCase(AuthenticatorType.OAUTH.name())) {
// OAuth access code Authentication
return AuthenticatorType.OAUTH;
} else if (loginInput
.getAuthenticator()
.equalsIgnoreCase(AuthenticatorType.PROGRAMMATIC_ACCESS_TOKEN.name())) {
return AuthenticatorType.PROGRAMMATIC_ACCESS_TOKEN;
} else if (loginInput
.getAuthenticator()
.equalsIgnoreCase(AuthenticatorType.SNOWFLAKE_JWT.name())) {
Expand Down Expand Up @@ -290,15 +294,16 @@ static SFLoginOutput openSession(
}

final AuthenticatorType authenticator = getAuthenticator(loginInput);
if (!authenticator.equals(AuthenticatorType.OAUTH)) {
// OAuth does not require a username
AssertUtil.assertTrue(
loginInput.getUserName() != null, "missing user name for opening session");
} else {
// OAUTH needs either token or password
if (authenticator.equals(AuthenticatorType.OAUTH)
|| authenticator.equals(AuthenticatorType.PROGRAMMATIC_ACCESS_TOKEN)) {
// OAUTH and PAT needs either token or password
AssertUtil.assertTrue(
loginInput.getToken() != null || loginInput.getPassword() != null,
"missing token or password for opening session");
} else {
// OAuth and PAT do not require a username
AssertUtil.assertTrue(
loginInput.getUserName() != null, "missing user name for opening session");
}
if (authenticator.equals(AuthenticatorType.EXTERNALBROWSER)) {
if ((Constants.getOS() == Constants.OS.MAC || Constants.getOS() == Constants.OS.WINDOWS)
Expand Down Expand Up @@ -363,7 +368,7 @@ private static boolean asBoolean(Object value) {
return false;
}

private static SFLoginOutput newSession(
static SFLoginOutput newSession(
SFLoginInput loginInput,
Map<SFSessionProperty, Object> connectionPropertiesMap,
String tracingLevel)
Expand Down Expand Up @@ -506,7 +511,8 @@ private static SFLoginOutput newSession(
}
} else if (authenticatorType == AuthenticatorType.OKTA) {
data.put(ClientAuthnParameter.RAW_SAML_RESPONSE.name(), tokenOrSamlResponse);
} else if (authenticatorType == AuthenticatorType.OAUTH) {
} else if (authenticatorType == AuthenticatorType.OAUTH
|| authenticatorType == AuthenticatorType.PROGRAMMATIC_ACCESS_TOKEN) {
data.put(ClientAuthnParameter.AUTHENTICATOR.name(), authenticatorType.name());

// Fix for HikariCP refresh token issue:SNOW-533673.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -51,5 +51,10 @@ public enum AuthenticatorType {
/*
* Client credentials flow with clientId and clientSecret as input
*/
OAUTH_CLIENT_CREDENTIALS
OAUTH_CLIENT_CREDENTIALS,

/*
* Authenticator to support PAT created in Snowflake
*/
PROGRAMMATIC_ACCESS_TOKEN
}
3 changes: 0 additions & 3 deletions src/test/java/net/snowflake/client/AbstractDriverIT.java
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.annotation.Nullable;
import net.snowflake.client.core.auth.AuthenticatorType;

/** Base test class with common constants, data structures and methods */
public class AbstractDriverIT {
Expand Down Expand Up @@ -325,8 +324,6 @@ public static Connection getConnection(
properties.put("internal", Boolean.TRUE.toString()); // TODO: do we need this?
properties.put("insecureMode", false); // use OCSP for all tests.

properties.put("authenticator", AuthenticatorType.OAUTH_CLIENT_CREDENTIALS.name());

if (injectSocketTimeout > 0) {
properties.put("injectSocketTimeout", String.valueOf(injectSocketTimeout));
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,21 +2,17 @@
* Copyright (c) 2024 Snowflake Computing Inc. All rights reserved.
*/

package net.snowflake.client.jdbc;
package net.snowflake.client.core;

import static net.snowflake.client.core.SessionUtilExternalBrowser.AuthExternalBrowserHandlers;

import com.amazonaws.util.StringUtils;
import java.net.URI;
import java.time.Duration;
import net.snowflake.client.category.TestTags;
import net.snowflake.client.core.HttpClientSettingsKey;
import net.snowflake.client.core.OCSPMode;
import net.snowflake.client.core.SFException;
import net.snowflake.client.core.SFLoginInput;
import net.snowflake.client.core.SFOauthLoginInput;
import net.snowflake.client.core.auth.oauth.AccessTokenProvider;
import net.snowflake.client.core.auth.oauth.OAuthAuthorizationCodeAccessTokenProvider;
import net.snowflake.client.jdbc.BaseWiremockTest;
import org.apache.http.HttpResponse;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpPost;
Expand All @@ -31,7 +27,7 @@
@Tag(TestTags.CORE)
public class OAuthAuthorizationCodeFlowLatestIT extends BaseWiremockTest {

private static final String SCENARIOS_BASE_DIR = "/oauth/authorization_code";
private static final String SCENARIOS_BASE_DIR = MAPPINGS_BASE_DIR + "/oauth/authorization_code";
private static final String SUCCESSFUL_FLOW_SCENARIO_MAPPINGS =
SCENARIOS_BASE_DIR + "/successful_scenario_mapping.json";
private static final String BROWSER_TIMEOUT_SCENARIO_MAPPING =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,44 +2,27 @@
* Copyright (c) 2024 Snowflake Computing Inc. All rights reserved.
*/

package net.snowflake.client.jdbc;

import static net.snowflake.client.core.SessionUtilExternalBrowser.AuthExternalBrowserHandlers;
package net.snowflake.client.core;

import com.amazonaws.util.StringUtils;
import java.net.URI;
import java.time.Duration;
import net.snowflake.client.category.TestTags;
import net.snowflake.client.core.HttpClientSettingsKey;
import net.snowflake.client.core.OCSPMode;
import net.snowflake.client.core.SFException;
import net.snowflake.client.core.SFLoginInput;
import net.snowflake.client.core.SFOauthLoginInput;
import net.snowflake.client.core.auth.oauth.AccessTokenProvider;
import net.snowflake.client.core.auth.oauth.OAuthClientCredentialsAccessTokenProvider;
import org.apache.http.HttpResponse;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import net.snowflake.client.jdbc.BaseWiremockTest;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@Tag(TestTags.CORE)
public class OAuthClientCredentialsFlowLatestIT extends BaseWiremockTest {

private static final String SCENARIOS_BASE_DIR = "/oauth/client_credentials";
private static final String SCENARIOS_BASE_DIR = MAPPINGS_BASE_DIR + "/oauth/client_credentials";
private static final String SUCCESSFUL_FLOW_SCENARIO_MAPPINGS =
SCENARIOS_BASE_DIR + "/successful_scenario_mapping.json";
private static final String TOKEN_REQUEST_ERROR_SCENARIO_MAPPING =
SCENARIOS_BASE_DIR + "/token_request_error_scenario_mapping.json";

private static final Logger logger =
LoggerFactory.getLogger(OAuthClientCredentialsFlowLatestIT.class);

@Test
public void successfulFlowScenario() throws SFException {
importMappingFromResources(SUCCESSFUL_FLOW_SCENARIO_MAPPINGS);
Expand Down Expand Up @@ -83,30 +66,4 @@ private SFLoginInput createLoginInputStub(String redirectUri) {

return loginInputStub;
}

static class WiremockProxyRequestBrowserHandler implements AuthExternalBrowserHandlers {
@Override
public HttpPost build(URI uri) {
// do nothing
return null;
}

@Override
public void openBrowser(String ssoUrl) {
try (CloseableHttpClient client = HttpClients.createDefault()) {
logger.debug("executing browser request to redirect uri: {}", ssoUrl);
HttpResponse response = client.execute(new HttpGet(ssoUrl));
if (response.getStatusLine().getStatusCode() != 200) {
throw new RuntimeException("Invalid response from " + ssoUrl);
}
} catch (Exception e) {
throw new RuntimeException(e);
}
}

@Override
public void output(String msg) {
// do nothing
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
/*
* Copyright (c) 2024 Snowflake Computing Inc. All rights reserved.
*/

package net.snowflake.client.core;

import java.util.HashMap;
import net.snowflake.client.category.TestTags;
import net.snowflake.client.core.auth.AuthenticatorType;
import net.snowflake.client.jdbc.BaseWiremockTest;
import net.snowflake.client.jdbc.SnowflakeSQLException;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;

@Tag(TestTags.CORE)
public class ProgrammaticAccessTokenAuthFlowLatestIT extends BaseWiremockTest {

private static final String SCENARIOS_BASE_DIR = MAPPINGS_BASE_DIR + "/pat";
private static final String SUCCESSFUL_FLOW_SCENARIO_MAPPINGS =
SCENARIOS_BASE_DIR + "/successful_scenario_mapping.json";
private static final String INVALID_TOKEN_SCENARIO_MAPPINGS =
SCENARIOS_BASE_DIR + "/invalid_token_scenario_mapping.json";

@Test
public void successfulFlowScenarioPatAsToken() throws SFException, SnowflakeSQLException {
importMappingFromResources(SUCCESSFUL_FLOW_SCENARIO_MAPPINGS);
SFLoginInput loginInputWithPatAsToken = createLoginInputStub("MOCK_TOKEN", null);
SFLoginOutput loginOutput =
SessionUtil.newSession(loginInputWithPatAsToken, new HashMap<>(), "INFO");
assertSuccessfulLoginOutput(loginOutput);
}

@Test
public void successfulFlowScenarioPatAsPassword() throws SFException, SnowflakeSQLException {
importMappingFromResources(SUCCESSFUL_FLOW_SCENARIO_MAPPINGS);
SFLoginInput loginInputWithPatAsPassword = createLoginInputStub(null, "MOCK_TOKEN");
SFLoginOutput loginOutput =
SessionUtil.newSession(loginInputWithPatAsPassword, new HashMap<>(), "INFO");
assertSuccessfulLoginOutput(loginOutput);
}

@Test
public void invalidTokenScenario() {
importMappingFromResources(INVALID_TOKEN_SCENARIO_MAPPINGS);
SnowflakeSQLException e =
Assertions.assertThrows(
SnowflakeSQLException.class,
() ->
SessionUtil.newSession(
createLoginInputStub("MOCK_TOKEN", null), new HashMap<>(), "INFO"));
Assertions.assertEquals("Programmatic access token is invalid.", e.getMessage());
}

private void assertSuccessfulLoginOutput(SFLoginOutput loginOutput) {
Assertions.assertNotNull(loginOutput);
Assertions.assertEquals("session token", loginOutput.getSessionToken());
Assertions.assertEquals("master token", loginOutput.getMasterToken());
Assertions.assertEquals(14400, loginOutput.getMasterTokenValidityInSeconds());
Assertions.assertEquals("8.48.0", loginOutput.getDatabaseVersion());
Assertions.assertEquals("TEST_DHEYMAN", loginOutput.getSessionDatabase());
Assertions.assertEquals("TEST_JDBC", loginOutput.getSessionSchema());
Assertions.assertEquals("ANALYST", loginOutput.getSessionRole());
Assertions.assertEquals("TEST_XSMALL", loginOutput.getSessionWarehouse());
Assertions.assertEquals("1172562260498", loginOutput.getSessionId());
Assertions.assertEquals(1, loginOutput.getCommonParams().size());
Assertions.assertEquals(4, loginOutput.getCommonParams().get("CLIENT_PREFETCH_THREADS"));
}

private SFLoginInput createLoginInputStub(String token, String password) {
SFLoginInput input = new SFLoginInput();
input.setAuthenticator(AuthenticatorType.PROGRAMMATIC_ACCESS_TOKEN.name());
input.setServerUrl(String.format("http://%s:%d/", WIREMOCK_HOST, wiremockHttpPort));
input.setUserName("MOCK_USERNAME");
input.setAccountName("MOCK_ACCOUNT_NAME");
input.setAppId("MOCK_APP_ID");
input.setAppVersion("MOCK_APP_VERSION");
input.setToken(token);
input.setPassword(password);
input.setOCSPMode(OCSPMode.FAIL_OPEN);
input.setHttpClientSettingsKey(new HttpClientSettingsKey(OCSPMode.FAIL_OPEN));
input.setLoginTimeout(1000);
input.setSessionParameters(new HashMap<>());
return input;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -34,14 +34,15 @@
import org.junit.jupiter.api.Assumptions;
import org.junit.jupiter.api.BeforeAll;

abstract class BaseWiremockTest {
public abstract class BaseWiremockTest {

protected static final SFLogger logger = SFLoggerFactory.getLogger(BaseWiremockTest.class);
protected static final String WIREMOCK_HOME_DIR = ".wiremock";
protected static final String WIREMOCK_M2_PATH =
"/.m2/repository/org/wiremock/wiremock-standalone/3.8.0/wiremock-standalone-3.8.0.jar";
protected static final String WIREMOCK_HOST = "localhost";
protected static final String TRUST_STORE_PROPERTY = "javax.net.ssl.trustStore";
protected static final String MAPPINGS_BASE_DIR = "/wiremock/mappings";
protected static int wiremockHttpPort;
protected static int wiremockHttpsPort;
private static String originalTrustStorePath;
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
{
"mappings": [
{
"scenarioName": "Successful PAT authentication flow",
"requiredScenarioState": "Started",
"newScenarioState": "Authenticated",
"request": {
"urlPathPattern": "/session/v1/login-request.*",
"method": "POST",
"headers": {
"CLIENT_APP_ID": {
"equalTo": "MOCK_APP_ID"
},
"CLIENT_APP_VERSION": {
"equalTo": "MOCK_APP_VERSION"
},
"Authorization": {
"equalTo": "Basic"
},
"accept": {
"equalTo": "application/json"
}
},
"bodyPatterns": [
{
"equalToJson" : {
"data": {
"ACCOUNT_NAME": "MOCK_ACCOUNT_NAME",
"CLIENT_APP_ID": "MOCK_APP_ID",
"CLIENT_ENVIRONMENT": {
"tracing": "INFO",
"OCSP_MODE": "FAIL_OPEN"
},
"CLIENT_APP_VERSION": "MOCK_APP_VERSION",
"TOKEN": "MOCK_TOKEN",
"LOGIN_NAME": "MOCK_USERNAME",
"AUTHENTICATOR": "PROGRAMMATIC_ACCESS_TOKEN"
}
},
"ignoreExtraElements" : true
}
]
},
"response": {
"status": 200,
"jsonBody": {
"data": {
"nextAction": "RETRY_LOGIN",
"authnMethod": "PAT",
"signInOptions": {}
},
"code": "394400",
"message": "Programmatic access token is invalid.",
"success": false,
"headers": null
}
}
}
]
}
Loading

0 comments on commit 801728c

Please sign in to comment.