Skip to content

Commit

Permalink
Add test scenarios
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-dheyman
sfc-gh-dheyman committed Dec 3, 2024
1 parent 4acd8df commit 214c98c
Show file tree
Hide file tree
Showing 5 changed files with 130 additions and 18 deletions.
6 changes: 3 additions & 3 deletions parent-pom.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
<apache.httpcore.version>4.4.16</apache.httpcore.version>
<zstd-jni.version>1.5.6-5</zstd-jni.version>
<arrow.version>17.0.0</arrow.version>
<asm.version>9.3</asm.version>
<asm.version>9.6</asm.version>
<avro.version>1.8.1</avro.version>
<awaitility.version>4.2.0</awaitility.version>
<awssdk.version>1.12.655</awssdk.version>
Expand Down Expand Up @@ -60,7 +60,7 @@
<javax.servlet.version>3.1.0</javax.servlet.version>
<jna.version>5.13.0</jna.version>
<joda.time.version>2.8.1</joda.time.version>
<json.smart.version>2.4.9</json.smart.version>
<json.smart.version>2.5.1</json.smart.version>
<junit4.version>4.13.2</junit4.version>
<junit.version>5.11.1</junit.version>
<junit.platform.version>1.11.1</junit.platform.version>
Expand All @@ -69,7 +69,7 @@
<metrics.version>2.2.0</metrics.version>
<mockito.version>4.11.0</mockito.version>
<netty.version>4.1.115.Final</netty.version>
<nimbusds.version>9.37.3</nimbusds.version>
<nimbusds.version>9.40</nimbusds.version>
<nimbusds.oauth2.version>11.20.1</nimbusds.oauth2.version>
<opencensus.version>0.31.1</opencensus.version>
<plexus.container.version>1.0-alpha-9-stable-1</plexus.container.version>
Expand Down
5 changes: 4 additions & 1 deletion src/main/java/net/snowflake/client/core/SessionUtil.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,8 @@ public class SessionUtil {
public static int DEFAULT_CLIENT_PREFETCH_THREADS = 4;
public static int MIN_CLIENT_CHUNK_SIZE = 48;
public static int MAX_CLIENT_CHUNK_SIZE = 160;
private static final int DEFAULT_BROWSER_AUTHORIZATION_TIMEOUT_SECONDS = 180;

public static Map<String, String> JVM_PARAMS_TO_PARAMS =
Stream.of(
new String[][] {
Expand Down Expand Up @@ -281,7 +283,8 @@ static SFLoginOutput openSession(
"passing clientSecret is required for OAUTH_AUTHORIZATION_CODE_FLOW authentication");
OauthAccessTokenProvider accessTokenProvider =
new AuthorizationCodeFlowAccessTokenProvider(
new SessionUtilExternalBrowser.DefaultAuthExternalBrowserHandlers());
new SessionUtilExternalBrowser.DefaultAuthExternalBrowserHandlers(),
DEFAULT_BROWSER_AUTHORIZATION_TIMEOUT_SECONDS);
String oauthAccessToken = accessTokenProvider.getAccessToken(loginInput);
loginInput.setAuthenticator(AuthenticatorType.OAUTH.name());
loginInput.setToken(oauthAccessToken);
Expand Down
24 changes: 15 additions & 9 deletions ...n/java/net/snowflake/client/core/auth/oauth/AuthorizationCodeFlowAccessTokenProvider.java
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
package net.snowflake.client.core.auth.oauth;

import static net.snowflake.client.core.SessionUtilExternalBrowser.*;
import static net.snowflake.client.core.SessionUtilExternalBrowser.AuthExternalBrowserHandlers;

import com.amazonaws.util.StringUtils;
import com.fasterxml.jackson.databind.ObjectMapper;
Expand Down Expand Up @@ -52,13 +52,14 @@ public class AuthorizationCodeFlowAccessTokenProvider implements OauthAccessToke
private static final String REDIRECT_URI_ENDPOINT = "/snowflake/oauth-redirect";
public static final String SESSION_ROLE_SCOPE = "session:role";

public static int AUTHORIZE_REDIRECT_TIMEOUT_MINUTES = 2;

private final AuthExternalBrowserHandlers browserHandler;
private final ObjectMapper objectMapper = new ObjectMapper();
private final int browserAuthorizationTimeoutSeconds;

public AuthorizationCodeFlowAccessTokenProvider(AuthExternalBrowserHandlers browserHandler) {
public AuthorizationCodeFlowAccessTokenProvider(
AuthExternalBrowserHandlers browserHandler, int browserAuthorizationTimeoutSeconds) {
this.browserHandler = browserHandler;
this.browserAuthorizationTimeoutSeconds = browserAuthorizationTimeoutSeconds;
}

@Override
Expand All @@ -75,21 +76,25 @@ private AuthorizationCode requestAuthorizationCode(
URI authorizeRequestURI = request.toURI();
CompletableFuture<String> codeFuture =
setupRedirectURIServerForAuthorizationCode(loginInput.getRedirectUriPort());
logger.debug(
"Waiting for authorization code on "
+ buildRedirectURI(loginInput.getRedirectUriPort())
+ "...");
letUserAuthorizeViaBrowser(authorizeRequestURI);
String code = codeFuture.get(AUTHORIZE_REDIRECT_TIMEOUT_MINUTES, TimeUnit.MINUTES);
String code = codeFuture.get(this.browserAuthorizationTimeoutSeconds, TimeUnit.SECONDS);
return new AuthorizationCode(code);
} catch (Exception e) {
if (e instanceof TimeoutException) {
logger.error(
"Authorization request timed out. Did not receive authorization code back to the redirect URI");
throw new RuntimeException(
"Authorization request timed out. Snowflake driver did not receive authorization code back to the redirect URI. Verify your security integration and driver configuration.",
e);
}
throw new RuntimeException(e);
}
}

private String exchangeAuthorizationCodeForAccessToken(
SFLoginInput loginInput, AuthorizationCode authorizationCode, CodeVerifier pkceVerifier)
throws SFException {
SFLoginInput loginInput, AuthorizationCode authorizationCode, CodeVerifier pkceVerifier) {
try {
TokenRequest request = buildTokenRequest(loginInput, authorizationCode, pkceVerifier);
String tokenResponse =
Expand Down Expand Up @@ -126,6 +131,7 @@ private static CompletableFuture<String> setupRedirectURIServerForAuthorizationC
String authorizationCode =
extractAuthorizationCodeFromQueryParameters(exchange.getRequestURI().getQuery());
if (!StringUtils.isNullOrEmpty(authorizationCode)) {
logger.debug("Received authorization code on redirect URI");
accessTokenFuture.complete(authorizationCode);
httpServer.stop(0);
}
Expand Down
109 changes: 106 additions & 3 deletions src/test/java/net/snowflake/client/jdbc/OauthAuthorizationCodeFlowLatestIT.java
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
package net.snowflake.client.jdbc;

import static net.snowflake.client.core.SessionUtilExternalBrowser.*;
import static net.snowflake.client.core.SessionUtilExternalBrowser.AuthExternalBrowserHandlers;

import java.io.IOException;
import java.net.URI;
Expand All @@ -26,7 +26,7 @@
@Tag(TestTags.CORE)
public class OauthAuthorizationCodeFlowLatestIT extends BaseWiremockTest {

public static final String SUCCESSFUL_FLOW_SCENARIO_MAPPINGS =
private static final String SUCCESSFUL_FLOW_SCENARIO_MAPPINGS =
"{\n"
+ " \"mappings\": [\n"
+ " {\n"
Expand Down Expand Up @@ -81,6 +81,81 @@ public class OauthAuthorizationCodeFlowLatestIT extends BaseWiremockTest {
+ " }\n"
+ "}";

public static final String BROWSER_TIMEOUT_SCENARIO_MAPPING =
"{\n"
+ " \"mappings\": [\n"
+ " {\n"
+ " \"scenarioName\": \"Browser Authorization timeout\",\n"
+ " \"request\": {\n"
+ " \"urlPathPattern\": \"/oauth/authorize.*\",\n"
+ " \"method\": \"GET\"\n"
+ " },\n"
+ " \"response\": {\n"
+ " \"status\": 200,\n"
+ " \"fixedDelayMilliseconds\": 5000\n"
+ " }\n"
+ " }\n"
+ " ],\n"
+ " \"importOptions\": {\n"
+ " \"duplicatePolicy\": \"IGNORE\",\n"
+ " \"deleteAllNotInImport\": true\n"
+ " }\n"
+ "}";

public static final String TOKEN_REQUEST_ERROR_SCENARIO_MAPPING =
"{\n"
+ " \"mappings\": [\n"
+ " {\n"
+ " \"scenarioName\": \"OAuth token request error\",\n"
+ " \"requiredScenarioState\": \"Started\",\n"
+ " \"newScenarioState\": \"Authorized\",\n"
+ " \"request\": {\n"
+ " \"urlPathPattern\": \"/oauth/authorize.*\",\n"
+ " \"method\": \"GET\"\n"
+ " },\n"
+ " \"response\": {\n"
+ " \"status\": 200\n"
+ " },\n"
+ " \"serveEventListeners\": [\n"
+ " {\n"
+ " \"name\": \"webhook\",\n"
+ " \"parameters\": {\n"
+ " \"method\": \"GET\",\n"
+ " \"url\": \"http://localhost:8001/snowflake/oauth-redirect?code=123\"\n"
+ " }\n"
+ " }\n"
+ " ]\n"
+ " },\n"
+ " {\n"
+ " \"scenarioName\": \"OAuth token request error\",\n"
+ " \"requiredScenarioState\": \"Authorized\",\n"
+ " \"newScenarioState\": \"Token request error\",\n"
+ " \"request\": {\n"
+ " \"urlPathPattern\": \"/oauth/token-request.*\",\n"
+ " \"method\": \"POST\",\n"
+ " \"headers\": {\n"
+ " \"Authorization\": {\n"
+ " \"contains\": \"Basic\"\n"
+ " },\n"
+ " \"Content-Type\": {\n"
+ " \"contains\": \"application/x-www-form-urlencoded; charset=UTF-8\"\n"
+ " }\n"
+ " },\n"
+ " \"bodyPatterns\": [{\n"
+ " \"contains\": \"grant_type=authorization_code&code=123&redirect_uri=http%3A%2F%2Flocalhost%3A8001%2Fsnowflake%2Foauth-redirect&code_verifier=\"\n"
+ " }]\n"
+ " },\n"
+ " \"response\": {\n"
+ " \"status\": 400\n"
+ " }\n"
+ " }\n"
+ " ],\n"
+ " \"importOptions\": {\n"
+ " \"duplicatePolicy\": \"IGNORE\",\n"
+ " \"deleteAllNotInImport\": true\n"
+ " }\n"
+ "}";

private static final Logger log =
LoggerFactory.getLogger(OauthAuthorizationCodeFlowLatestIT.class);

Expand All @@ -93,13 +168,41 @@ public void successfulFlowScenario() throws SFException {
SFLoginInput loginInput = createLoginInputStub();

OauthAccessTokenProvider provider =
new AuthorizationCodeFlowAccessTokenProvider(wiremockProxyRequestBrowserHandler);
new AuthorizationCodeFlowAccessTokenProvider(wiremockProxyRequestBrowserHandler, 30);
String accessToken = provider.getAccessToken(loginInput);

Assertions.assertTrue(StringUtils.isNotBlank(accessToken));
Assertions.assertEquals("access-token-123", accessToken);
}

@Test
public void browserTimeoutFlowScenario() {
importMapping(BROWSER_TIMEOUT_SCENARIO_MAPPING);
SFLoginInput loginInput = createLoginInputStub();

OauthAccessTokenProvider provider =
new AuthorizationCodeFlowAccessTokenProvider(wiremockProxyRequestBrowserHandler, 1);
RuntimeException e =
Assertions.assertThrows(RuntimeException.class, () -> provider.getAccessToken(loginInput));
Assertions.assertEquals(
"Authorization request timed out. Snowflake driver did not receive authorization code back to the redirect URI. Verify your security integration and driver configuration.",
e.getMessage());
}

@Test
public void tokenRequestErrorFlowScenario() {
importMapping(TOKEN_REQUEST_ERROR_SCENARIO_MAPPING);
SFLoginInput loginInput = createLoginInputStub();

OauthAccessTokenProvider provider =
new AuthorizationCodeFlowAccessTokenProvider(wiremockProxyRequestBrowserHandler, 30);
RuntimeException e =
Assertions.assertThrows(RuntimeException.class, () -> provider.getAccessToken(loginInput));
Assertions.assertEquals(
"net.snowflake.client.jdbc.SnowflakeSQLException: JDBC driver encountered communication error. Message: HTTP status=400.",
e.getMessage());
}

private SFLoginInput createLoginInputStub() {
SFLoginInput loginInputStub = new SFLoginInput();
loginInputStub.setServerUrl(String.format("http://%s:%d/", WIREMOCK_HOST, wiremockHttpPort));
Expand Down
4 changes: 2 additions & 2 deletions thin_public_pom.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,11 +55,11 @@
<javax.servlet.version>3.1.0</javax.servlet.version>
<jna.version>5.13.0</jna.version>
<joda.time.version>2.8.1</joda.time.version>
<json.smart.version>2.4.9</json.smart.version>
<json.smart.version>2.5.1</json.smart.version>
<jsoup.version>1.15.3</jsoup.version>
<metrics.version>2.2.0</metrics.version>
<netty.version>4.1.115.Final</netty.version>
<nimbusds.version>9.37.3</nimbusds.version>
<nimbusds.version>9.40.0</nimbusds.version>
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
<project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
<slf4j.version>2.0.13</slf4j.version>
Expand Down

0 comments on commit 214c98c

Please sign in to comment.