Skip to content

Commit 94ac08c

Browse files
sfc-gh-pmansoursfc-gh-pmansour
authored
SNOW-1955810 Support client-side opt-in of Refresh Token Rotation in Snowflake OAuth (#2166)
1 parent 34326f2 commit 94ac08c

6 files changed

Lines changed: 150 additions & 17 deletions

File tree

‎src/main/java/net/snowflake/client/core/SFOauthLoginInput.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ public class SFOauthLoginInput {
99
private final String authorizationUrl;
1010
private final String tokenRequestUrl;
1111
private final String scope;
12+
private final boolean enableSingleUseRefreshTokens;
1213

1314
public SFOauthLoginInput(
1415
String clientId,
@@ -17,12 +18,24 @@ public SFOauthLoginInput(
1718
String authorizationUrl,
1819
String tokenRequestUrl,
1920
String scope) {
21+
this(clientId, clientSecret, redirectUri, authorizationUrl, tokenRequestUrl, scope, false);
22+
}
23+
24+
public SFOauthLoginInput(
25+
String clientId,
26+
String clientSecret,
27+
String redirectUri,
28+
String authorizationUrl,
29+
String tokenRequestUrl,
30+
String scope,
31+
boolean enableSingleUseRefreshTokens) {
2032
this.redirectUri = redirectUri;
2133
this.clientId = clientId;
2234
this.clientSecret = clientSecret;
2335
this.authorizationUrl = authorizationUrl;
2436
this.tokenRequestUrl = tokenRequestUrl;
2537
this.scope = scope;
38+
this.enableSingleUseRefreshTokens = enableSingleUseRefreshTokens;
2639
}
2740

2841
public String getRedirectUri() {
@@ -48,4 +61,8 @@ public String getTokenRequestUrl() {
4861
public String getScope() {
4962
return scope;
5063
}
64+
65+
public boolean getEnableSingleUseRefreshTokens() {
66+
return enableSingleUseRefreshTokens;
67+
}
5168
}

‎src/main/java/net/snowflake/client/core/SFSession.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -680,7 +680,10 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
680680
(String) connectionPropertiesMap.get(SFSessionProperty.OAUTH_REDIRECT_URI),
681681
(String) connectionPropertiesMap.get(SFSessionProperty.OAUTH_AUTHORIZATION_URL),
682682
(String) connectionPropertiesMap.get(SFSessionProperty.OAUTH_TOKEN_REQUEST_URL),
683-
(String) connectionPropertiesMap.get(SFSessionProperty.OAUTH_SCOPE));
683+
(String) connectionPropertiesMap.get(SFSessionProperty.OAUTH_SCOPE),
684+
getBooleanValue(
685+
connectionPropertiesMap.get(
686+
SFSessionProperty.OAUTH_ENABLE_SINGLE_USE_REFRESH_TOKENS)));
684687

685688
loginInput
686689
.setServerUrl((String) connectionPropertiesMap.get(SFSessionProperty.SERVER_URL))

‎src/main/java/net/snowflake/client/core/SFSessionProperty.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ public enum SFSessionProperty {
2626
OAUTH_SCOPE("oauthScope", false, String.class),
2727
OAUTH_AUTHORIZATION_URL("oauthAuthorizationUrl", false, String.class),
2828
OAUTH_TOKEN_REQUEST_URL("oauthTokenRequestUrl", false, String.class),
29+
OAUTH_ENABLE_SINGLE_USE_REFRESH_TOKENS("oauthEnableSingleUseRefreshTokens", false, Boolean.class),
2930
WORKLOAD_IDENTITY_PROVIDER("workloadIdentityProvider", false, String.class),
3031
WORKLOAD_IDENTITY_ENTRA_RESOURCE("workloadIdentityEntraResource", false, String.class),
3132
WAREHOUSE("warehouse", false, String.class),

‎src/main/java/net/snowflake/client/core/auth/oauth/OAuthAuthorizationCodeAccessTokenProvider.java‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -215,13 +215,17 @@ private HttpRequestBase buildTokenRequest(
215215
new Secret(loginInput.getOauthLoginInput().getClientSecret()));
216216
Scope scope =
217217
new Scope(OAuthUtil.getScope(loginInput.getOauthLoginInput(), loginInput.getRole()));
218-
TokenRequest tokenRequest =
219-
new TokenRequest(
220-
OAuthUtil.getTokenRequestUrl(
221-
loginInput.getOauthLoginInput(), loginInput.getServerUrl()),
222-
clientAuthentication,
223-
codeGrant,
224-
scope);
218+
TokenRequest.Builder tokenRequestBuilder =
219+
new TokenRequest.Builder(
220+
OAuthUtil.getTokenRequestUrl(
221+
loginInput.getOauthLoginInput(), loginInput.getServerUrl()),
222+
clientAuthentication,
223+
codeGrant)
224+
.scope(scope);
225+
if (loginInput.getOauthLoginInput().getEnableSingleUseRefreshTokens()) {
226+
tokenRequestBuilder.customParameter("enable_single_use_refresh_tokens", "true");
227+
}
228+
TokenRequest tokenRequest = tokenRequestBuilder.build();
225229
HTTPRequest tokenHttpRequest = tokenRequest.toHTTPRequest();
226230
HttpRequestBase convertedTokenRequest = OAuthUtil.convertToBaseRequest(tokenHttpRequest);
227231

‎src/test/java/net/snowflake/client/core/OAuthAuthorizationCodeFlowLatestIT.java‎

Lines changed: 37 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ public class OAuthAuthorizationCodeFlowLatestIT extends BaseWiremockTest {
3030
SCENARIOS_BASE_DIR + "/successful_flow.json";
3131
private static final String SUCCESSFUL_DPOP_FLOW_SCENARIO_MAPPINGS =
3232
SCENARIOS_BASE_DIR + "/successful_dpop_flow.json";
33+
private static final String SUCCESSFUL_FLOW_WITH_SINGLE_USE_REFRESH_TOKENS_SCENARIO_MAPPINGS =
34+
SCENARIOS_BASE_DIR + "/successful_flow_with_single_use_refresh_tokens.json";
3335
private static final String DPOP_NONCE_ERROR_SCENARIO_MAPPINGS =
3436
SCENARIOS_BASE_DIR + "/dpop_nonce_error_flow.json";
3537
private static final String BROWSER_TIMEOUT_SCENARIO_MAPPING =
@@ -59,7 +61,7 @@ public OAuthAuthorizationCodeFlowLatestIT() throws SFException {}
5961
public void successfulFlowScenario() throws SFException {
6062
importMappingFromResources(SUCCESSFUL_FLOW_SCENARIO_MAPPINGS);
6163
SFLoginInput loginInput =
62-
createLoginInputStub("http://localhost:8009/snowflake/oauth-redirect", null, null);
64+
createLoginInputStub("http://localhost:8009/snowflake/oauth-redirect", null, null, false);
6365

6466
TokenResponseDTO tokenResponse = provider.getAccessToken(loginInput);
6567
String accessToken = tokenResponse.getAccessToken();
@@ -96,14 +98,31 @@ public void successfulFlowDPoPScenarioWithNonce() throws SFException {
9698
Assertions.assertEquals("access-token-123", accessToken);
9799
}
98100

101+
@Test
102+
public void successfulFlowWithSingleUseRefreshTokensScenario() throws SFException {
103+
importMappingFromResources(SUCCESSFUL_FLOW_WITH_SINGLE_USE_REFRESH_TOKENS_SCENARIO_MAPPINGS);
104+
SFLoginInput loginInput =
105+
createLoginInputStub("http://localhost:8009/snowflake/oauth-redirect", null, null, true);
106+
107+
TokenResponseDTO tokenResponse = provider.getAccessToken(loginInput);
108+
String accessToken = tokenResponse.getAccessToken();
109+
String refreshToken = tokenResponse.getRefreshToken();
110+
111+
Assertions.assertFalse(StringUtils.isNullOrEmpty(accessToken));
112+
Assertions.assertFalse(StringUtils.isNullOrEmpty(refreshToken));
113+
Assertions.assertEquals("access-token-123", accessToken);
114+
Assertions.assertEquals("refresh-token-123", refreshToken);
115+
}
116+
99117
@Test
100118
public void customUrlsScenario() throws SFException {
101119
importMappingFromResources(CUSTOM_URLS_SCENARIO_MAPPINGS);
102120
SFLoginInput loginInput =
103121
createLoginInputStub(
104122
"http://localhost:8007/snowflake/oauth-redirect",
105123
String.format("http://%s:%d/authorization", WIREMOCK_HOST, wiremockHttpPort),
106-
String.format("http://%s:%d/tokenrequest", WIREMOCK_HOST, wiremockHttpPort));
124+
String.format("http://%s:%d/tokenrequest", WIREMOCK_HOST, wiremockHttpPort),
125+
false);
107126

108127
TokenResponseDTO tokenResponse = provider.getAccessToken(loginInput);
109128
String accessToken = tokenResponse.getAccessToken();
@@ -116,7 +135,7 @@ public void customUrlsScenario() throws SFException {
116135
public void browserTimeoutFlowScenario() throws SFException {
117136
importMappingFromResources(BROWSER_TIMEOUT_SCENARIO_MAPPING);
118137
SFLoginInput loginInput =
119-
createLoginInputStub("http://localhost:8004/snowflake/oauth-redirect", null, null);
138+
createLoginInputStub("http://localhost:8004/snowflake/oauth-redirect", null, null, false);
120139

121140
AccessTokenProvider provider =
122141
new OAuthAuthorizationCodeAccessTokenProvider(
@@ -133,7 +152,7 @@ public void browserTimeoutFlowScenario() throws SFException {
133152
public void invalidScopeFlowScenario() {
134153
importMappingFromResources(INVALID_SCOPE_SCENARIO_MAPPING);
135154
SFLoginInput loginInput =
136-
createLoginInputStub("http://localhost:8002/snowflake/oauth-redirect", null, null);
155+
createLoginInputStub("http://localhost:8002/snowflake/oauth-redirect", null, null, false);
137156
SFException e =
138157
Assertions.assertThrows(SFException.class, () -> provider.getAccessToken(loginInput));
139158
Assertions.assertTrue(
@@ -146,7 +165,7 @@ public void invalidScopeFlowScenario() {
146165
public void invalidStateFlowScenario() {
147166
importMappingFromResources(INVALID_STATE_SCENARIO_MAPPING);
148167
SFLoginInput loginInput =
149-
createLoginInputStub("http://localhost:8010/snowflake/oauth-redirect", null, null);
168+
createLoginInputStub("http://localhost:8010/snowflake/oauth-redirect", null, null, false);
150169
SFException e =
151170
Assertions.assertThrows(SFException.class, () -> provider.getAccessToken(loginInput));
152171
Assertions.assertTrue(
@@ -159,7 +178,7 @@ public void invalidStateFlowScenario() {
159178
public void tokenRequestErrorFlowScenario() {
160179
importMappingFromResources(TOKEN_REQUEST_ERROR_SCENARIO_MAPPING);
161180
SFLoginInput loginInput =
162-
createLoginInputStub("http://localhost:8003/snowflake/oauth-redirect", null, null);
181+
createLoginInputStub("http://localhost:8003/snowflake/oauth-redirect", null, null, false);
163182

164183
SFException e =
165184
Assertions.assertThrows(SFException.class, () -> provider.getAccessToken(loginInput));
@@ -169,12 +188,21 @@ public void tokenRequestErrorFlowScenario() {
169188
}
170189

171190
private SFLoginInput createLoginInputStub(
172-
String redirectUri, String authorizationUrl, String tokenRequestUrl) {
191+
String redirectUri,
192+
String authorizationUrl,
193+
String tokenRequestUrl,
194+
boolean enableSingleUseRefreshTokens) {
173195
SFLoginInput loginInputStub = new SFLoginInput();
174196
loginInputStub.setServerUrl(String.format("http://%s:%d/", WIREMOCK_HOST, wiremockHttpPort));
175197
loginInputStub.setOauthLoginInput(
176198
new SFOauthLoginInput(
177-
"123", "123", redirectUri, authorizationUrl, tokenRequestUrl, "session:role:ANALYST"));
199+
"123",
200+
"123",
201+
redirectUri,
202+
authorizationUrl,
203+
tokenRequestUrl,
204+
"session:role:ANALYST",
205+
enableSingleUseRefreshTokens));
178206
loginInputStub.setSocketTimeout(Duration.ofMinutes(5));
179207
loginInputStub.setHttpClientSettingsKey(new HttpClientSettingsKey(OCSPMode.FAIL_OPEN));
180208

@@ -184,7 +212,7 @@ private SFLoginInput createLoginInputStub(
184212
private SFLoginInput createLoginInputStubWithDPoPEnabled(
185213
String redirectUri, String authorizationUrl, String tokenRequestUrl) {
186214
SFLoginInput loginInputStub =
187-
createLoginInputStub(redirectUri, authorizationUrl, tokenRequestUrl);
215+
createLoginInputStub(redirectUri, authorizationUrl, tokenRequestUrl, false);
188216
loginInputStub.setDPoPEnabled(true);
189217
return loginInputStub;
190218
}

‎src/test/resources/wiremock/mappings/oauth/authorization_code/successful_flow_with_single_use_refresh_tokens.json‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
{
2+
"mappings": [
3+
{
4+
"scenarioName": "Successful OAuth authorization code flow",
5+
"requiredScenarioState": "Started",
6+
"newScenarioState": "Authorized",
7+
"request": {
8+
"urlPathPattern": "/oauth/authorize",
9+
"queryParameters": {
10+
"response_type": {
11+
"equalTo": "code"
12+
},
13+
"scope": {
14+
"equalTo": "session:role:ANALYST"
15+
},
16+
"code_challenge_method": {
17+
"equalTo": "S256"
18+
},
19+
"redirect_uri": {
20+
"equalTo": "http://localhost:8009/snowflake/oauth-redirect"
21+
},
22+
"code_challenge": {
23+
"matches": ".*"
24+
},
25+
"state": {
26+
"matches": ".*"
27+
},
28+
"client_id": {
29+
"equalTo": "123"
30+
}
31+
},
32+
"method": "GET"
33+
},
34+
"response": {
35+
"status": 302,
36+
"headers": {
37+
"Location": "http://localhost:8009/snowflake/oauth-redirect?code=123&state=abc123"
38+
}
39+
}
40+
},
41+
{
42+
"scenarioName": "Successful OAuth authorization code flow",
43+
"requiredScenarioState": "Authorized",
44+
"newScenarioState": "Acquired access token",
45+
"request": {
46+
"urlPathPattern": "/oauth/token-request.*",
47+
"method": "POST",
48+
"headers": {
49+
"Authorization": {
50+
"contains": "Basic"
51+
},
52+
"Content-Type": {
53+
"contains": "application/x-www-form-urlencoded; charset=UTF-8"
54+
}
55+
},
56+
"bodyPatterns": [
57+
{
58+
"contains": "grant_type=authorization_code&code=123&redirect_uri=http%3A%2F%2Flocalhost%3A8009%2Fsnowflake%2Foauth-redirect"
59+
},
60+
{
61+
"contains": "&enable_single_use_refresh_tokens=true"
62+
}
63+
]
64+
},
65+
"response": {
66+
"status": 200,
67+
"jsonBody": {
68+
"access_token": "access-token-123",
69+
"refresh_token": "refresh-token-123",
70+
"token_type": "Bearer",
71+
"username": "user",
72+
"scope": "refresh_token session:role:ANALYST",
73+
"expires_in": 600,
74+
"refresh_token_expires_in": 86399,
75+
"idpInitiated": false
76+
}
77+
}
78+
}
79+
]
80+
}

0 commit comments

Comments
 (0)