Skip to content

Commit

Permalink
SNOW-1825471: Port and remove snowflake-common auth related codebase
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-dheyman committed Nov 29, 2024
1 parent a20f2cf commit 9c728ec
Show file tree
Hide file tree
Showing 11 changed files with 159 additions and 66 deletions.
2 changes: 1 addition & 1 deletion 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.ClientAuthnDTO;
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
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,50 @@
/*
* 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,

/*
* Authenticator to support PAT
*/
PROGRAMMATIC_ACCESS_TOKEN,
}
27 changes: 27 additions & 0 deletions src/main/java/net/snowflake/client/core/auth/ClientAuthnDTO.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
/*
* Copyright (c) 2024 Snowflake Computing Inc. All right reserved.
*/
package net.snowflake.client.core.auth;

import java.util.Map;
import javax.annotation.Nullable;
import net.snowflake.client.core.SnowflakeJdbcInternalApi;

@SnowflakeJdbcInternalApi
public class ClientAuthnDTO {

// contains all the required data for current authn step
public final Map<String, Object> data;

/*
* current state
* tokenized string with all current parameters and the authn step
*/
public final String inFlightCtx;

/** Required by Jackson */
public ClientAuthnDTO(Map<String, Object> data, @Nullable String inFlightCtx) {
this.data = data;
this.inFlightCtx = inFlightCtx;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
/*
* Copyright (c) 2024 Snowflake Computing Inc. All right reserved.
*/
package net.snowflake.client.core.auth;

import net.snowflake.client.core.SnowflakeJdbcInternalApi;

@SnowflakeJdbcInternalApi
public enum ClientAuthnParameter {
LOGIN_NAME,
PASSWORD,
RAW_SAML_RESPONSE,
ACCOUNT_NAME,
CLIENT_APP_ID,
CLIENT_APP_VERSION,
EXT_AUTHN_DUO_METHOD,
PASSCODE,
CLIENT_ENVIRONMENT,
AUTHENTICATOR,
BROWSER_MODE_REDIRECT_PORT,
SESSION_PARAMETERS,
PROOF_KEY,
TOKEN
}
Loading

0 comments on commit 9c728ec

Please sign in to comment.