Skip to content

Commit

Permalink
Snow 1925254 authentication logic refactor (#2073)
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-fpawlowski
sfc-gh-fpawlowski authored Feb 17, 2025
1 parent a2af702 commit 9005913
Show file tree
Hide file tree
Showing 3 changed files with 250 additions and 150 deletions.
8 changes: 4 additions & 4 deletions src/main/java/net/snowflake/client/core/HttpUtil.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -580,7 +580,7 @@ public static boolean isSocksProxyDisabled() {
}

/**
* Executes a HTTP request with the cookie spec set to IGNORE_COOKIES
* Executes an HTTP request with the cookie spec set to IGNORE_COOKIES
*
* @param httpRequest HttpRequestBase
* @param retryTimeout retry timeout
Expand Down Expand Up @@ -622,7 +622,7 @@ static String executeRequestWithoutCookies(
}

/**
* Executes a HTTP request for Snowflake.
* Executes an HTTP request for Snowflake.
*
* @param httpRequest HttpRequestBase
* @param retryTimeout retry timeout
Expand Down Expand Up @@ -658,7 +658,7 @@ public static String executeGeneralRequest(
}

/**
* Executes a HTTP request for Snowflake
* Executes an HTTP request for Snowflake
*
* @param httpRequest HttpRequestBase
* @param retryTimeout retry timeout
Expand Down Expand Up @@ -696,7 +696,7 @@ public static String executeGeneralRequest(
}

/**
* Executes a HTTP request for Snowflake.
* Executes an HTTP request for Snowflake.
*
* @param httpRequest HttpRequestBase
* @param retryTimeout retry timeout
Expand Down
237 changes: 151 additions & 86 deletions src/main/java/net/snowflake/client/core/SessionUtil.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import static net.snowflake.client.core.SFTrustManager.resetOCSPResponseCacherServerURL;
import static net.snowflake.client.jdbc.SnowflakeUtil.systemGetProperty;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Strings;
Expand Down Expand Up @@ -271,7 +272,7 @@ static SFLoginOutput openSession(
AssertUtil.assertTrue(
loginInput.getUserName() != null, "missing user name for opening session");
} else {
// OAUTH needs either token or passord
// OAUTH needs either token or password
AssertUtil.assertTrue(
loginInput.getToken() != null || loginInput.getPassword() != null,
"missing token or password for opening session");
Expand Down Expand Up @@ -707,7 +708,7 @@ private static SFLoginOutput newSession(
// In RestRequest.execute(), socket timeout is replaced with auth timeout
// so we can renew the request within auth timeout.
// auth timeout within socket timeout is thrown without backoff,
// and we need to update time remained in socket timeout here to control the
// and we need to update time remained in socket timeout here to control
// the actual socket timeout from customer setting.
if (loginInput.getSocketTimeoutInMillis() > 0) {
if (ex.issocketTimeoutNoBackoff()) {
Expand Down Expand Up @@ -737,29 +738,7 @@ private static SFLoginOutput newSession(
break;
}

if (theString == null) {
if (lastRestException != null) {
logger.error(
"Failed to open new session for user: {}, host: {}. Error: {}",
loginInput.getUserName(),
loginInput.getHostFromServerUrl(),
lastRestException);
throw lastRestException;
} else {
SnowflakeSQLException exception =
new SnowflakeSQLException(
NO_QUERY_ID,
"empty authentication response",
SqlState.CONNECTION_EXCEPTION,
ErrorCode.CONNECTION_ERROR.getMessageCode());
logger.error(
"Failed to open new session for user: {}, host: {}. Error: {}",
loginInput.getUserName(),
loginInput.getHostFromServerUrl(),
exception);
throw exception;
}
}
handleEmptyAuthResponse(theString, loginInput, lastRestException);

// general method, same as with data binding
JsonNode jsonNode = mapper.readTree(theString);
Expand Down Expand Up @@ -1201,22 +1180,11 @@ static void closeSession(SFLoginInput loginInput) throws SFException, SnowflakeS
private static String federatedFlowStep4(
SFLoginInput loginInput, String ssoUrl, String oneTimeToken) throws SnowflakeSQLException {
String responseHtml = "";
try {

final URL url = new URL(ssoUrl);
URI oktaGetUri =
new URIBuilder()
.setScheme(url.getProtocol())
.setHost(url.getHost())
.setPath(url.getPath())
.setParameter("RelayState", "%2Fsome%2Fdeep%2Flink")
.setParameter("onetimetoken", oneTimeToken)
.build();
HttpGet httpGet = new HttpGet(oktaGetUri);
try {

HeaderGroup headers = new HeaderGroup();
headers.addHeader(new BasicHeader(HttpHeaders.ACCEPT, "*/*"));
httpGet.setHeaders(headers.getAllHeaders());
HttpGet httpGet = new HttpGet();
prepareFederatedFlowStep4Request(httpGet, ssoUrl, oneTimeToken);

responseHtml =
HttpUtil.executeGeneralRequest(
Expand Down Expand Up @@ -1276,26 +1244,7 @@ private static String federatedFlowStep3(SFLoginInput loginInput, String tokenUr
URL url = new URL(tokenUrl);
URI tokenUri = url.toURI();
final HttpPost postRequest = new HttpPost(tokenUri);

String userName;
if (Strings.isNullOrEmpty(loginInput.getOKTAUserName())) {
userName = loginInput.getUserName();
} else {
userName = loginInput.getOKTAUserName();
}
StringEntity params =
new StringEntity(
"{\"username\":\""
+ userName
+ "\",\"password\":\""
+ loginInput.getPassword()
+ "\"}");
postRequest.setEntity(params);

HeaderGroup headers = new HeaderGroup();
headers.addHeader(new BasicHeader(HttpHeaders.ACCEPT, "application/json"));
headers.addHeader(new BasicHeader(HttpHeaders.CONTENT_TYPE, "application/json"));
postRequest.setHeaders(headers.getAllHeaders());
setFederatedFlowStep3PostRequestAuthData(postRequest, loginInput);

final String idpResponse =
HttpUtil.executeRequestWithoutCookies(
Expand Down Expand Up @@ -1363,29 +1312,8 @@ private static void federatedFlowStep2(SFLoginInput loginInput, String tokenUrl,
private static JsonNode federatedFlowStep1(SFLoginInput loginInput) throws SnowflakeSQLException {
JsonNode dataNode = null;
try {
URIBuilder fedUriBuilder = new URIBuilder(loginInput.getServerUrl());
fedUriBuilder.setPath(SF_PATH_AUTHENTICATOR_REQUEST);
URI fedUrlUri = fedUriBuilder.build();

Map<String, Object> data = new HashMap<>();
data.put(ClientAuthnParameter.ACCOUNT_NAME.name(), loginInput.getAccountName());
data.put(ClientAuthnParameter.AUTHENTICATOR.name(), loginInput.getAuthenticator());
data.put(ClientAuthnParameter.CLIENT_APP_ID.name(), loginInput.getAppId());
data.put(ClientAuthnParameter.CLIENT_APP_VERSION.name(), loginInput.getAppVersion());

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

// attach the login info json body to the post request
StringEntity input = new StringEntity(json, StandardCharsets.UTF_8);
input.setContentType("application/json");
HttpPost postRequest = new HttpPost(fedUrlUri);
postRequest.setEntity(input);
postRequest.addHeader("accept", "application/json");

// Add headers for driver name and version
postRequest.addHeader(SF_HEADER_CLIENT_APP_ID, loginInput.getAppId());
postRequest.addHeader(SF_HEADER_CLIENT_APP_VERSION, loginInput.getAppVersion());
StringEntity requestInput = prepareFederatedFlowStep1RequestInput(loginInput);
HttpPost postRequest = prepareFederatedFlowStep1PostRequest(loginInput, requestInput);

final String gsResponse =
HttpUtil.executeGeneralRequest(
Expand All @@ -1395,6 +1323,7 @@ private static JsonNode federatedFlowStep1(SFLoginInput loginInput) throws Snowf
loginInput.getSocketTimeoutInMillis(),
0,
loginInput.getHttpClientSettingsKey());

logger.debug("Authenticator-request response: {}", gsResponse);
JsonNode jsonNode = mapper.readTree(gsResponse);

Expand Down Expand Up @@ -1790,12 +1719,148 @@ public static boolean isNewRetryStrategyRequest(HttpRequestBase request) {
URI requestURI = request.getURI();
String requestPath = requestURI.getPath();
if (requestPath != null) {
if (requestPath.equals(SF_PATH_LOGIN_REQUEST)
return requestPath.equals(SF_PATH_LOGIN_REQUEST)
|| requestPath.equals(SF_PATH_AUTHENTICATOR_REQUEST)
|| requestPath.equals(SF_PATH_TOKEN_REQUEST)) {
return true;
}
|| requestPath.equals(SF_PATH_TOKEN_REQUEST);
}
return false;
}

/**
* Prepares an HTTP POST request for the first step of the federated authentication flow.
*
* @param loginInput The login information for the request.
* @param inputData The JSON input data to include in the request.
* @return An {@link HttpPost} object ready to execute the federated flow request.
* @throws URISyntaxException If the constructed URI is invalid.
*/
private static HttpPost prepareFederatedFlowStep1PostRequest(
SFLoginInput loginInput, StringEntity inputData) throws URISyntaxException {
URIBuilder fedUriBuilder = new URIBuilder(loginInput.getServerUrl());
// TODO: if loginInput.serverUrl contains port or additional segments - it will be ignored and
// overwritten here - to be fixed in SNOW-1922872
fedUriBuilder.setPath(SF_PATH_AUTHENTICATOR_REQUEST);
URI fedUrlUri = fedUriBuilder.build();

HttpPost postRequest = new HttpPost(fedUrlUri);
postRequest.setEntity(inputData);
postRequest.addHeader("accept", "application/json");

postRequest.addHeader(SF_HEADER_CLIENT_APP_ID, loginInput.getAppId());
postRequest.addHeader(SF_HEADER_CLIENT_APP_VERSION, loginInput.getAppVersion());

return postRequest;
}

/**
* Prepares the JSON input for the first step of the federated authentication flow.
*
* @param loginInput The login information for the request.
* @return A {@link StringEntity} containing the JSON input for the request.
* @throws JsonProcessingException If there is an error generating the JSON input.
*/
private static StringEntity prepareFederatedFlowStep1RequestInput(SFLoginInput loginInput)
throws JsonProcessingException {
Map<String, Object> data = new HashMap<>();
data.put(ClientAuthnParameter.ACCOUNT_NAME.name(), loginInput.getAccountName());
data.put(ClientAuthnParameter.AUTHENTICATOR.name(), loginInput.getAuthenticator());
data.put(ClientAuthnParameter.CLIENT_APP_ID.name(), loginInput.getAppId());
data.put(ClientAuthnParameter.CLIENT_APP_VERSION.name(), loginInput.getAppVersion());

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

StringEntity input = new StringEntity(json, StandardCharsets.UTF_8);
input.setContentType("application/json");
return input;
}

/**
* Sets the authentication data for the third step of the federated authentication flow.
*
* @param postRequest The {@link HttpPost} request to update with authentication data.
* @param loginInput The login information for the request.
* @throws SnowflakeSQLException If an error occurs while preparing the request.
*/
private static void setFederatedFlowStep3PostRequestAuthData(
HttpPost postRequest, SFLoginInput loginInput) throws SnowflakeSQLException {
String userName =
Strings.isNullOrEmpty(loginInput.getOKTAUserName())
? loginInput.getUserName()
: loginInput.getOKTAUserName();
try {
StringEntity params =
new StringEntity(
"{\"username\":\""
+ userName
+ "\",\"password\":\""
+ loginInput.getPassword()
+ "\"}");
postRequest.setEntity(params);

HeaderGroup headers = new HeaderGroup();
headers.addHeader(new BasicHeader(HttpHeaders.ACCEPT, "application/json"));
headers.addHeader(new BasicHeader(HttpHeaders.CONTENT_TYPE, "application/json"));
postRequest.setHeaders(headers.getAllHeaders());
} catch (IOException ex) {
handleFederatedFlowError(loginInput, ex);
}
}

/**
* Prepares an HTTP GET request for the fourth step of the federated authentication flow.
*
* @param retrieveSamlRequest The {@link HttpRequestBase} to update with the SAML request details.
* @param ssoUrl The SSO URL to use for the request.
* @param oneTimeToken The one-time token to include in the request.
* @throws MalformedURLException If the SSO URL is malformed.
* @throws URISyntaxException If the URI for the request cannot be built.
*/
private static void prepareFederatedFlowStep4Request(
HttpRequestBase retrieveSamlRequest, String ssoUrl, String oneTimeToken)
throws MalformedURLException, URISyntaxException {
final URL url = new URL(ssoUrl);
URI oktaGetUri =
new URIBuilder()
.setScheme(url.getProtocol())
.setHost(url.getHost())
.setPort(url.getPort())
.setPath(url.getPath())
.setParameter("RelayState", "%2Fsome%2Fdeep%2Flink")
.setParameter("onetimetoken", oneTimeToken)
.build();
retrieveSamlRequest.setURI(oktaGetUri);

HeaderGroup headers = new HeaderGroup();
headers.addHeader(new BasicHeader(HttpHeaders.ACCEPT, "*/*"));
retrieveSamlRequest.setHeaders(headers.getAllHeaders());
}

private static void handleEmptyAuthResponse(
String theString, SFLoginInput loginInput, Exception lastRestException)
throws Exception, SFException {
if (theString == null) {
if (lastRestException != null) {
logger.error(
"Failed to open new session for user: {}, host: {}. Error: {}",
loginInput.getUserName(),
loginInput.getHostFromServerUrl(),
lastRestException);
throw lastRestException;
} else {
SnowflakeSQLException exception =
new SnowflakeSQLException(
NO_QUERY_ID,
"empty authentication response",
SqlState.CONNECTION_EXCEPTION,
ErrorCode.CONNECTION_ERROR.getMessageCode());
logger.error(
"Failed to open new session for user: {}, host: {}. Error: {}",
loginInput.getUserName(),
loginInput.getHostFromServerUrl(),
exception);
throw exception;
}
}
}
}
Loading

0 comments on commit 9005913

Please sign in to comment.