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.

Sign up for GitHub

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

Added retry logic for authentication with Okta #1661

Closed
wants to merge 11 commits into from
Closed

Added retry logic for authentication with Okta #1661

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
8891c79
Added retry logic for authentication with Okta
sfc-gh-ext-simba-jf Mar 7, 2024
32b42b3
Added retry logic for authentication with Okta
sfc-gh-ext-simba-jf Mar 7, 2024
653635c
Add log message and testing
sfc-gh-ext-simba-jf Mar 14, 2024
dcdf8d0
Merge branch 'Okta-429-Retry-Fix' of https://github.com/snowflakedb/s…
sfc-gh-ext-simba-jf Mar 14, 2024
8f9be45
Revert "SNOW-974575 Structured type array map (#1654)"
sfc-gh-ext-simba-jf Mar 14, 2024
6c66fa3
Revert "SNOW-1168175 Update CHANGELOG.rst (#1668)"
sfc-gh-ext-simba-jf Mar 14, 2024
58ceaf2
Revert "SNOW-1203558 Fix geo types values (#1663)"
sfc-gh-ext-simba-jf Mar 14, 2024
8b7817c
Revert "SNOW-1196041: Ignore default credentials in GCS client (#1656)"
sfc-gh-ext-simba-jf Mar 14, 2024
e85e7e5
Check-style fix
sfc-gh-ext-simba-jf Mar 14, 2024
76e4ff7
Added retry for the okta saml flow
sfc-gh-ext-simba-jf Apr 4, 2024
01fc2c9
Merge branch 'master' into Okta-429-Retry-Fix
sfc-gh-ext-simba-jf Apr 4, 2024
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
37 changes: 26 additions & 11 deletions src/main/java/net/snowflake/client/core/SessionUtil.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -651,16 +651,30 @@ private static SFLoginOutput newSession(
loginInput.getHttpClientSettingsKey());
} catch (SnowflakeSQLException ex) {
if (ex.getErrorCode() == ErrorCode.AUTHENTICATOR_REQUEST_TIMEOUT.getMessageCode()) {
if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT) {
SessionUtilKeyPair s =
new SessionUtilKeyPair(
loginInput.getPrivateKey(),
loginInput.getPrivateKeyFile(),
loginInput.getPrivateKeyFilePwd(),
loginInput.getAccountName(),
loginInput.getUserName());

data.put(ClientAuthnParameter.TOKEN.name(), s.issueJwtToken());
if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT
|| authenticatorType == ClientAuthnDTO.AuthenticatorType.OKTA) {

if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT) {
SessionUtilKeyPair s =
new SessionUtilKeyPair(
loginInput.getPrivateKey(),
loginInput.getPrivateKeyFile(),
loginInput.getPrivateKeyFilePwd(),
loginInput.getAccountName(),
loginInput.getUserName());

data.put(ClientAuthnParameter.TOKEN.name(), s.issueJwtToken());
} else if (authenticatorType == ClientAuthnDTO.AuthenticatorType.OKTA) {
// If we need to retry, we need to get a new Okta token
tokenOrSamlResponse = getSamlResponseUsingOkta(loginInput);
sfc-gh-dprzybysz marked this conversation as resolved.
Show resolved Hide resolved
data.put(ClientAuthnParameter.RAW_SAML_RESPONSE.name(), tokenOrSamlResponse);
authnData.setData(data);
String updatedJson = mapper.writeValueAsString(authnData);

StringEntity updatedInput = new StringEntity(updatedJson, StandardCharsets.UTF_8);
updatedInput.setContentType("application/json");
postRequest.setEntity(updatedInput);
}

long elapsedSeconds = ex.getElapsedSeconds();

Expand Down Expand Up @@ -690,7 +704,8 @@ private static SFLoginOutput newSession(
}
}

// JWT renew should not count as a retry, so we pass back the current retry count.
// JWT or Okta renew should not count as a retry, so we pass back the current retry
// count.
retryCount = ex.getRetryCount();

continue;
Expand Down
75 changes: 75 additions & 0 deletions src/test/java/net/snowflake/client/core/SessionUtilLatestIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -390,4 +390,79 @@ public void testOktaAuthGetFail() throws Throwable {
assertEquals(SqlState.IO_ERROR, e.getSQLState());
}
}

private SFLoginInput createOktaLoginInput() {
SFLoginInput input = new SFLoginInput();
input.setServerUrl("https://testauth.okta.com");
input.setUserName("MOCK_USERNAME");
input.setPassword("MOCK_PASSWORD");
input.setAccountName("MOCK_ACCOUNT_NAME");
input.setAppId("MOCK_APP_ID");
input.setOCSPMode(OCSPMode.FAIL_OPEN);
input.setHttpClientSettingsKey(new HttpClientSettingsKey(OCSPMode.FAIL_OPEN));
input.setLoginTimeout(1000);
input.setSessionParameters(new HashMap<>());
input.setAuthenticator("https://testauth.okta.com");
return input;
}

// Testing retry with Okta calls the service to get a new unique token. This is valid after
// version 3.15.0.
@Test
public void testOktaAuthRetry() throws Throwable {
SFLoginInput loginInput = createOktaLoginInput();
Map<SFSessionProperty, Object> connectionPropertiesMap = initConnectionPropertiesMap();
SnowflakeSQLException ex =
new SnowflakeSQLException(ErrorCode.AUTHENTICATOR_REQUEST_TIMEOUT, 0, true, 0);
Copy link
Collaborator

Choose a reason for hiding this comment

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

this exception is thrown when timeouts are reached - what about testing 429 error code? I think that there should be a test with NETWORK_ERROR and code 429 thrown during execution

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Everything around checking retry related to the 429 error code happens within RestRequest.execute(). When it detects a 429 error code it eventually throws the AUTHENTICATOR_REQUEST_TIMEOUT here which is what we're testing for in this case. I've tried to see if I could set up the test to show the whole flow from a 429 error being thrown then eventually leading to getting a new Okta token, but the mocking set up gets very quickly complicated because of the level of mocks you need to set up. So what I've done instead is I've added a new test in RestRequestTest to show that it throws the AUTHENTICATOR_REQUEST_TIMEOUT when it encounters a 429 error.

Does that seem like enough test support?

try (MockedStatic<HttpUtil> mockedHttpUtil = mockStatic(HttpUtil.class)) {
mockedHttpUtil
.when(
() ->
HttpUtil.executeGeneralRequest(
Mockito.any(HttpPost.class),
Mockito.anyInt(),
Mockito.anyInt(),
Mockito.anyInt(),
Mockito.anyInt(),
Mockito.nullable(HttpClientSettingsKey.class)))
.thenReturn(
"{\"data\":{\"tokenUrl\":\"https://testauth.okta.com/api/v1/authn\","
+ "\"ssoUrl\":\"https://testauth.okta.com/app/snowflake/abcdefghijklmnopqrstuvwxyz/sso/saml\","
+ "\"proofKey\":null},\"code\":null,\"message\":null,\"success\":true}")
.thenThrow(ex)
.thenReturn(
"{\"data\":{\"tokenUrl\":\"https://testauth.okta.com/api/v1/authn\","
+ "\"ssoUrl\":\"https://testauth.okta.com/app/snowflake/abcdefghijklmnopqrstuvwxyz/sso/saml\","
+ "\"proofKey\":null},\"code\":null,\"message\":null,\"success\":true}");

mockedHttpUtil
.when(
() ->
HttpUtil.executeRequestWithoutCookies(
Mockito.any(HttpRequestBase.class),
Mockito.anyInt(),
Mockito.anyInt(),
Mockito.anyInt(),
Mockito.anyInt(),
Mockito.anyInt(),
Mockito.nullable(AtomicBoolean.class),
Mockito.nullable(HttpClientSettingsKey.class)))
.thenReturn(
"{\"expiresAt\":\"2023-10-13T19:18:09.000Z\",\"status\":\"SUCCESS\",\"sessionToken\":\"testsessiontoken\"}");

mockedHttpUtil
.when(
() ->
HttpUtil.executeGeneralRequest(
Mockito.any(HttpGet.class),
Mockito.anyInt(),
Mockito.anyInt(),
Mockito.anyInt(),
Mockito.anyInt(),
Mockito.nullable(HttpClientSettingsKey.class)))
.thenReturn("<body><form action=\"https://testauth.okta.com\"></form></body>");

SessionUtil.openSession(loginInput, connectionPropertiesMap, "ALL");
}
}
}
Loading