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: Port and remove snowflake-common auth related codebase #1976

Merged
merged 1 commit into from
Nov 29, 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
10 changes: 4 additions & 6 deletions src/main/java/net/snowflake/client/core/SFSession.java
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
import java.util.concurrent.atomic.AtomicInteger;
import java.util.logging.Level;
import net.snowflake.client.config.SFClientConfig;
import net.snowflake.client.core.auth.AuthenticatorType;
import net.snowflake.client.jdbc.DefaultSFConnectionHandler;
import net.snowflake.client.jdbc.ErrorCode;
import net.snowflake.client.jdbc.QueryStatusV2;
Expand All @@ -49,7 +50,6 @@
import net.snowflake.client.log.SFLoggerFactory;
import net.snowflake.client.log.SFLoggerUtil;
import net.snowflake.client.util.Stopwatch;
import net.snowflake.common.core.ClientAuthnDTO;
import net.snowflake.common.core.SqlState;
import org.apache.http.HttpHeaders;
import org.apache.http.client.methods.HttpGet;
Expand Down Expand Up @@ -804,7 +804,7 @@ private boolean isSnowflakeAuthenticator() {
&& privateKey == null
&& privateKeyFileLocation == null
&& privateKeyBase64 == null)
|| ClientAuthnDTO.AuthenticatorType.SNOWFLAKE.name().equalsIgnoreCase(authenticator);
|| AuthenticatorType.SNOWFLAKE.name().equalsIgnoreCase(authenticator);
}

/**
Expand All @@ -815,7 +815,7 @@ private boolean isSnowflakeAuthenticator() {
boolean isExternalbrowserAuthenticator() {
Map<SFSessionProperty, Object> connectionPropertiesMap = getConnectionPropertiesMap();
String authenticator = (String) connectionPropertiesMap.get(SFSessionProperty.AUTHENTICATOR);
return ClientAuthnDTO.AuthenticatorType.EXTERNALBROWSER.name().equalsIgnoreCase(authenticator);
return AuthenticatorType.EXTERNALBROWSER.name().equalsIgnoreCase(authenticator);
}

/**
Expand All @@ -837,9 +837,7 @@ boolean isOKTAAuthenticator() {
boolean isUsernamePasswordMFAAuthenticator() {
Map<SFSessionProperty, Object> connectionPropertiesMap = getConnectionPropertiesMap();
String authenticator = (String) connectionPropertiesMap.get(SFSessionProperty.AUTHENTICATOR);
return ClientAuthnDTO.AuthenticatorType.USERNAME_PASSWORD_MFA
.name()
.equalsIgnoreCase(authenticator);
return AuthenticatorType.USERNAME_PASSWORD_MFA.name().equalsIgnoreCase(authenticator);
}

/**
Expand Down
88 changes: 41 additions & 47 deletions src/main/java/net/snowflake/client/core/SessionUtil.java
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,9 @@
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import net.snowflake.client.core.auth.AuthenticatorType;
import net.snowflake.client.core.auth.ClientAuthnDTO;
import net.snowflake.client.core.auth.ClientAuthnParameter;
import net.snowflake.client.jdbc.ErrorCode;
import net.snowflake.client.jdbc.SnowflakeDriver;
import net.snowflake.client.jdbc.SnowflakeReauthenticationRequest;
Expand All @@ -38,8 +41,6 @@
import net.snowflake.client.log.SFLoggerFactory;
import net.snowflake.client.util.SecretDetector;
import net.snowflake.client.util.Stopwatch;
import net.snowflake.common.core.ClientAuthnDTO;
import net.snowflake.common.core.ClientAuthnParameter;
import net.snowflake.common.core.SqlState;
import org.apache.http.HttpHeaders;
import org.apache.http.client.config.RequestConfig;
Expand Down Expand Up @@ -209,40 +210,38 @@ public class SessionUtil {
* @param loginInput login information
* @return Authenticator type
*/
private static ClientAuthnDTO.AuthenticatorType getAuthenticator(SFLoginInput loginInput) {
private static AuthenticatorType getAuthenticator(SFLoginInput loginInput) {
if (loginInput.getAuthenticator() != null) {
if (loginInput
.getAuthenticator()
.equalsIgnoreCase(ClientAuthnDTO.AuthenticatorType.EXTERNALBROWSER.name())) {
.equalsIgnoreCase(AuthenticatorType.EXTERNALBROWSER.name())) {
// SAML 2.0 compliant service/application
return ClientAuthnDTO.AuthenticatorType.EXTERNALBROWSER;
} else if (loginInput
.getAuthenticator()
.equalsIgnoreCase(ClientAuthnDTO.AuthenticatorType.OAUTH.name())) {
return AuthenticatorType.EXTERNALBROWSER;
} else if (loginInput.getAuthenticator().equalsIgnoreCase(AuthenticatorType.OAUTH.name())) {
// OAuth Authentication
return ClientAuthnDTO.AuthenticatorType.OAUTH;
return AuthenticatorType.OAUTH;
} else if (loginInput
.getAuthenticator()
.equalsIgnoreCase(ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT.name())) {
return ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT;
.equalsIgnoreCase(AuthenticatorType.SNOWFLAKE_JWT.name())) {
return AuthenticatorType.SNOWFLAKE_JWT;
} else if (loginInput
.getAuthenticator()
.equalsIgnoreCase(ClientAuthnDTO.AuthenticatorType.USERNAME_PASSWORD_MFA.name())) {
return ClientAuthnDTO.AuthenticatorType.USERNAME_PASSWORD_MFA;
.equalsIgnoreCase(AuthenticatorType.USERNAME_PASSWORD_MFA.name())) {
return AuthenticatorType.USERNAME_PASSWORD_MFA;
} else if (!loginInput
.getAuthenticator()
.equalsIgnoreCase(ClientAuthnDTO.AuthenticatorType.SNOWFLAKE.name())) {
.equalsIgnoreCase(AuthenticatorType.SNOWFLAKE.name())) {
// OKTA authenticator v1.
return ClientAuthnDTO.AuthenticatorType.OKTA;
return AuthenticatorType.OKTA;
}
}

// authenticator is null, then jdbc will decide authenticator depends on
// if privateKey is specified or not. If yes, authenticator type will be
// SNOWFLAKE_JWT, otherwise it will use SNOWFLAKE.
return loginInput.isPrivateKeyProvided()
? ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT
: ClientAuthnDTO.AuthenticatorType.SNOWFLAKE;
? AuthenticatorType.SNOWFLAKE_JWT
: AuthenticatorType.SNOWFLAKE;
}

/**
Expand All @@ -266,8 +265,8 @@ static SFLoginOutput openSession(
AssertUtil.assertTrue(
loginInput.getLoginTimeout() >= 0, "negative login timeout for opening session");

final ClientAuthnDTO.AuthenticatorType authenticator = getAuthenticator(loginInput);
if (!authenticator.equals(ClientAuthnDTO.AuthenticatorType.OAUTH)) {
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");
Expand All @@ -277,7 +276,7 @@ static SFLoginOutput openSession(
loginInput.getToken() != null || loginInput.getPassword() != null,
"missing token or password for opening session");
}
if (authenticator.equals(ClientAuthnDTO.AuthenticatorType.EXTERNALBROWSER)) {
if (authenticator.equals(AuthenticatorType.EXTERNALBROWSER)) {
if ((Constants.getOS() == Constants.OS.MAC || Constants.getOS() == Constants.OS.WINDOWS)
&& loginInput.isEnableClientStoreTemporaryCredential()) {
// force to set the flag for Mac/Windows users
Expand All @@ -299,7 +298,7 @@ static SFLoginOutput openSession(
}
}

if (authenticator.equals(ClientAuthnDTO.AuthenticatorType.USERNAME_PASSWORD_MFA)) {
if (authenticator.equals(AuthenticatorType.USERNAME_PASSWORD_MFA)) {
if ((Constants.getOS() == Constants.OS.MAC || Constants.getOS() == Constants.OS.WINDOWS)
&& loginInput.isEnableClientRequestMfaToken()) {
loginInput.getSessionParameters().put(CLIENT_REQUEST_MFA_TOKEN, true);
Expand Down Expand Up @@ -371,7 +370,7 @@ private static SFLoginOutput newSession(
int healthCheckInterval = DEFAULT_HEALTH_CHECK_INTERVAL;
int httpClientSocketTimeout = loginInput.getSocketTimeoutInMillis();
int httpClientConnectionTimeout = loginInput.getConnectionTimeoutInMillis();
final ClientAuthnDTO.AuthenticatorType authenticatorType = getAuthenticator(loginInput);
final AuthenticatorType authenticatorType = getAuthenticator(loginInput);
Map<String, Object> commonParams;

String oktaUsername = loginInput.getOKTAUserName();
Expand Down Expand Up @@ -406,7 +405,7 @@ private static SFLoginOutput newSession(
uriBuilder.addParameter(SF_QUERY_ROLE, loginInput.getRole());
}

if (authenticatorType == ClientAuthnDTO.AuthenticatorType.EXTERNALBROWSER) {
if (authenticatorType == AuthenticatorType.EXTERNALBROWSER) {
// try to reuse id_token if exists
if (loginInput.getIdToken() == null) {
// SAML 2.0 compliant service/application
Expand All @@ -416,10 +415,10 @@ private static SFLoginOutput newSession(
samlProofKey = s.getProofKey();
consentCacheIdToken = s.isConsentCacheIdToken();
}
} else if (authenticatorType == ClientAuthnDTO.AuthenticatorType.OKTA) {
} else if (authenticatorType == AuthenticatorType.OKTA) {
// okta authenticator v1
tokenOrSamlResponse = getSamlResponseUsingOkta(loginInput);
} else if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT) {
} else if (authenticatorType == AuthenticatorType.SNOWFLAKE_JWT) {
SessionUtilKeyPair s =
new SessionUtilKeyPair(
loginInput.getPrivateKey(),
Expand Down Expand Up @@ -453,9 +452,6 @@ private static SFLoginOutput newSession(
HttpPost postRequest = null;

try {
ClientAuthnDTO authnData = new ClientAuthnDTO();
authnData.setInFlightCtx(loginInput.getInFlightCtx());

Map<String, Object> data = new HashMap<>();
data.put(ClientAuthnParameter.CLIENT_APP_ID.name(), loginInput.getAppId());

Expand All @@ -472,22 +468,21 @@ private static SFLoginOutput newSession(
* authenticate with the IDP provider only, and GS should not have any
* trace for this information.
*/
if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE) {
if (authenticatorType == AuthenticatorType.SNOWFLAKE) {
data.put(ClientAuthnParameter.PASSWORD.name(), loginInput.getPassword());
} else if (authenticatorType == ClientAuthnDTO.AuthenticatorType.EXTERNALBROWSER) {
} else if (authenticatorType == AuthenticatorType.EXTERNALBROWSER) {
if (loginInput.getIdToken() != null) {
data.put(ClientAuthnParameter.AUTHENTICATOR.name(), ID_TOKEN_AUTHENTICATOR);
data.put(ClientAuthnParameter.TOKEN.name(), loginInput.getIdToken());
} else {
data.put(
ClientAuthnParameter.AUTHENTICATOR.name(),
ClientAuthnDTO.AuthenticatorType.EXTERNALBROWSER.name());
ClientAuthnParameter.AUTHENTICATOR.name(), AuthenticatorType.EXTERNALBROWSER.name());
data.put(ClientAuthnParameter.PROOF_KEY.name(), samlProofKey);
data.put(ClientAuthnParameter.TOKEN.name(), tokenOrSamlResponse);
}
} else if (authenticatorType == ClientAuthnDTO.AuthenticatorType.OKTA) {
} else if (authenticatorType == AuthenticatorType.OKTA) {
data.put(ClientAuthnParameter.RAW_SAML_RESPONSE.name(), tokenOrSamlResponse);
} else if (authenticatorType == ClientAuthnDTO.AuthenticatorType.OAUTH) {
} else if (authenticatorType == AuthenticatorType.OAUTH) {
data.put(ClientAuthnParameter.AUTHENTICATOR.name(), authenticatorType.name());

// Fix for HikariCP refresh token issue:SNOW-533673.
Expand All @@ -499,10 +494,10 @@ private static SFLoginOutput newSession(
data.put(ClientAuthnParameter.TOKEN.name(), loginInput.getPassword());
}

} else if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT) {
} else if (authenticatorType == AuthenticatorType.SNOWFLAKE_JWT) {
data.put(ClientAuthnParameter.AUTHENTICATOR.name(), authenticatorType.name());
data.put(ClientAuthnParameter.TOKEN.name(), loginInput.getToken());
} else if (authenticatorType == ClientAuthnDTO.AuthenticatorType.USERNAME_PASSWORD_MFA) {
} else if (authenticatorType == AuthenticatorType.USERNAME_PASSWORD_MFA) {
// No authenticator name should be added here, since this will be treated as snowflake
// default authenticator by backend
data.put(ClientAuthnParameter.PASSWORD.name(), loginInput.getPassword());
Expand Down Expand Up @@ -621,8 +616,7 @@ private static SFLoginOutput newSession(
}

data.put(ClientAuthnParameter.CLIENT_APP_VERSION.name(), loginInput.getAppVersion());

authnData.setData(data);
ClientAuthnDTO authnData = new ClientAuthnDTO(data, loginInput.getInFlightCtx());
String json = mapper.writeValueAsString(authnData);

postRequest = new HttpPost(loginURI);
Expand Down Expand Up @@ -672,10 +666,10 @@ private static SFLoginOutput newSession(
} catch (SnowflakeSQLException ex) {
lastRestException = ex;
if (ex.getErrorCode() == ErrorCode.AUTHENTICATOR_REQUEST_TIMEOUT.getMessageCode()) {
if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT
|| authenticatorType == ClientAuthnDTO.AuthenticatorType.OKTA) {
if (authenticatorType == AuthenticatorType.SNOWFLAKE_JWT
|| authenticatorType == AuthenticatorType.OKTA) {

if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT) {
if (authenticatorType == AuthenticatorType.SNOWFLAKE_JWT) {
SessionUtilKeyPair s =
new SessionUtilKeyPair(
loginInput.getPrivateKey(),
Expand All @@ -686,13 +680,14 @@ private static SFLoginOutput newSession(
loginInput.getUserName());

data.put(ClientAuthnParameter.TOKEN.name(), s.issueJwtToken());
} else if (authenticatorType == ClientAuthnDTO.AuthenticatorType.OKTA) {
} else if (authenticatorType == AuthenticatorType.OKTA) {
logger.debug("Retrieve new token for Okta authentication.");
// If we need to retry, we need to get a new Okta token
tokenOrSamlResponse = getSamlResponseUsingOkta(loginInput);
data.put(ClientAuthnParameter.RAW_SAML_RESPONSE.name(), tokenOrSamlResponse);
authnData.setData(data);
String updatedJson = mapper.writeValueAsString(authnData);
ClientAuthnDTO updatedAuthnData =
new ClientAuthnDTO(data, loginInput.getInFlightCtx());
String updatedJson = mapper.writeValueAsString(updatedAuthnData);

StringEntity updatedInput = new StringEntity(updatedJson, StandardCharsets.UTF_8);
updatedInput.setContentType("application/json");
Expand Down Expand Up @@ -785,7 +780,7 @@ private static SFLoginOutput newSession(
SnowflakeUtil.checkErrorAndThrowExceptionIncludingReauth(jsonNode);
}

if (authenticatorType == ClientAuthnDTO.AuthenticatorType.USERNAME_PASSWORD_MFA) {
if (authenticatorType == AuthenticatorType.USERNAME_PASSWORD_MFA) {
deleteMfaTokenCache(loginInput.getHostFromServerUrl(), loginInput.getUserName());
}

Expand Down Expand Up @@ -1378,8 +1373,7 @@ private static JsonNode federatedFlowStep1(SFLoginInput loginInput) throws Snowf
data.put(ClientAuthnParameter.CLIENT_APP_ID.name(), loginInput.getAppId());
data.put(ClientAuthnParameter.CLIENT_APP_VERSION.name(), loginInput.getAppVersion());

ClientAuthnDTO authnData = new ClientAuthnDTO();
authnData.setData(data);
ClientAuthnDTO authnData = new ClientAuthnDTO(data, null);
String json = mapper.writeValueAsString(authnData);

// attach the login info json body to the post request
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,12 +28,12 @@
import java.util.Locale;
import java.util.Map;
import java.util.TimeZone;
import net.snowflake.client.core.auth.ClientAuthnDTO;
import net.snowflake.client.core.auth.ClientAuthnParameter;
import net.snowflake.client.jdbc.ErrorCode;
import net.snowflake.client.jdbc.SnowflakeSQLException;
import net.snowflake.client.log.SFLogger;
import net.snowflake.client.log.SFLoggerFactory;
import net.snowflake.common.core.ClientAuthnDTO;
import net.snowflake.common.core.ClientAuthnParameter;
import net.snowflake.common.core.SqlState;
import org.apache.http.NameValuePair;
import org.apache.http.client.methods.HttpPost;
Expand Down Expand Up @@ -175,7 +175,6 @@ private String getSSOUrl(int port) throws SFException, SnowflakeSQLException {

HttpPost postRequest = this.handlers.build(fedUrlUri);

ClientAuthnDTO authnData = new ClientAuthnDTO();
Map<String, Object> data = new HashMap<>();

data.put(ClientAuthnParameter.AUTHENTICATOR.name(), authenticator);
Expand All @@ -185,7 +184,7 @@ private String getSSOUrl(int port) throws SFException, SnowflakeSQLException {
data.put(ClientAuthnParameter.CLIENT_APP_ID.name(), loginInput.getAppId());
data.put(ClientAuthnParameter.CLIENT_APP_VERSION.name(), loginInput.getAppVersion());

authnData.setData(data);
ClientAuthnDTO authnData = new ClientAuthnDTO(data, null);
String json = mapper.writeValueAsString(authnData);

// attach the login info json body to the post request
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
/*
* Copyright (c) 2024 Snowflake Computing Inc. All right reserved.
*/
package net.snowflake.client.core.auth;

import net.snowflake.client.core.SnowflakeJdbcInternalApi;

@SnowflakeJdbcInternalApi
public enum AuthenticatorType {
/*
* regular login username+password via Snowflake, may or may not have MFA
*/
SNOWFLAKE,

/*
* federated authentication, OKTA as IDP
*/
OKTA,

/*
* Web browser based authenticator for SAML 2.0 compliant
* service/application
*/
EXTERNALBROWSER,

/*
* OAUTH 2.0 flow
*/
OAUTH,

/*
* Snowflake local authentication using jwt token as a user credential
*/
SNOWFLAKE_JWT,

/*
* Internal authenticator to enable id_token for web browser based authenticator
*/
ID_TOKEN,

/*
* Authenticator to enable token for regular login with mfa
*/
USERNAME_PASSWORD_MFA
}
Loading