Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SNOW-1825471: PAT authentication support #1995

Merged
merged 1 commit into from
Dec 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it fine to reduce the visibility here?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we can reduce on internal methods

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
Loading