Skip to content

Commit 514b391

Browse files
sfc-gh-mkubiksfc-gh-mkubik
authored
Merge branch 'master' into SNOW-1437655-set-auth-timeout-to-0-in-not-auth-endpoint-calls
2 parents f88bb08 + 8d999f0 commit 514b391

24 files changed

+1620
-184
lines changed

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,10 @@ enum CachedCredentialType {
44
ID_TOKEN("ID_TOKEN"),
55
MFA_TOKEN("MFATOKEN"),
66
OAUTH_ACCESS_TOKEN("OAUTH_ACCESS_TOKEN"),
7-
OAUTH_REFRESH_TOKEN("OAUTH_REFRESH_TOKEN");
7+
OAUTH_REFRESH_TOKEN("OAUTH_REFRESH_TOKEN"),
8+
DPOP_BUNDLED_ACCESS_TOKEN(
9+
"DPOP_BUNDLED_ACCESS_TOKEN"); // contains '.' separated, base64 encoded access token and DPoP
10+
// public key
811

912
private final String value;
1013

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

Lines changed: 115 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33
import com.amazonaws.util.StringUtils;
44
import com.google.common.base.Strings;
55
import java.net.URI;
6+
import java.nio.charset.StandardCharsets;
7+
import java.util.Base64;
68
import net.snowflake.client.jdbc.ErrorCode;
79
import net.snowflake.client.log.SFLogger;
810
import net.snowflake.client.log.SFLoggerFactory;
@@ -124,6 +126,25 @@ static void fillCachedOAuthRefreshToken(SFLoginInput loginInput) throws SFExcept
124126
loginInput, host, loginInput.getUserName(), CachedCredentialType.OAUTH_REFRESH_TOKEN);
125127
}
126128

129+
/**
130+
* Reuse the cached OAuth access token & DPoP public key tied to it
131+
*
132+
* @param loginInput login input to attach refresh token
133+
*/
134+
static void fillCachedDPoPBundledAccessToken(SFLoginInput loginInput) throws SFException {
135+
String host = getHostForOAuthCacheKey(loginInput);
136+
logger.debug(
137+
"Looking for cached DPoP public key for user: {}, host: {}",
138+
loginInput.getUserName(),
139+
host);
140+
getInstance()
141+
.fillCachedCredential(
142+
loginInput,
143+
host,
144+
loginInput.getUserName(),
145+
CachedCredentialType.DPOP_BUNDLED_ACCESS_TOKEN);
146+
}
147+
127148
/** Reuse the cached token stored locally */
128149
synchronized void fillCachedCredential(
129150
SFLoginInput loginInput, String host, String username, CachedCredentialType credType)
@@ -137,24 +158,34 @@ synchronized void fillCachedCredential(
137158
return;
138159
}
139160

140-
String cred;
161+
String base64EncodedCred, cred = null;
141162
try {
142-
cred = secureStorageManager.getCredential(host, username, credType.getValue());
163+
base64EncodedCred = secureStorageManager.getCredential(host, username, credType.getValue());
143164
} catch (NoClassDefFoundError error) {
144165
logMissingJnaJarForSecureLocalStorage();
145166
return;
146167
}
147168

148-
if (cred == null) {
169+
if (base64EncodedCred == null) {
149170
logger.debug("Retrieved {} is null", credType);
150171
}
151172

152173
logger.debug(
153174
"Setting {}{} token for user: {}, host: {}",
154-
cred == null ? "null " : "",
175+
base64EncodedCred == null ? "null " : "",
155176
credType.getValue(),
156177
username,
157178
host);
179+
180+
if (base64EncodedCred != null && credType != CachedCredentialType.DPOP_BUNDLED_ACCESS_TOKEN) {
181+
try {
182+
cred = new String(Base64.getDecoder().decode(base64EncodedCred));
183+
} catch (Exception e) {
184+
// handle legacy non-base64 encoded cache values (CredentialManager fails to decode)
185+
deleteTemporaryCredential(host, username, credType);
186+
return;
187+
}
188+
}
158189
switch (credType) {
159190
case ID_TOKEN:
160191
loginInput.setIdToken(cred);
@@ -168,12 +199,29 @@ synchronized void fillCachedCredential(
168199
case OAUTH_REFRESH_TOKEN:
169200
loginInput.setOauthRefreshToken(cred);
170201
break;
202+
case DPOP_BUNDLED_ACCESS_TOKEN:
203+
updateInputWithTokenAndPublicKey(base64EncodedCred, loginInput);
204+
break;
171205
default:
172206
throw new SFException(
173207
ErrorCode.INTERNAL_ERROR, "Unrecognized type {} for local cached credential", credType);
174208
}
175209
}
176210

211+
private void updateInputWithTokenAndPublicKey(String cred, SFLoginInput loginInput)
212+
throws SFException {
213+
if (Strings.isNullOrEmpty(cred)) {
214+
String[] values = cred.split("\\.");
215+
if (values.length != 2) {
216+
throw new SFException(
217+
ErrorCode.INTERNAL_ERROR, "Invalid DPoP bundled access token credential format");
218+
}
219+
Base64.Decoder decoder = Base64.getDecoder();
220+
loginInput.setOauthAccessToken(new String(decoder.decode(values[0])));
221+
loginInput.setDPoPPublicKey(new String(decoder.decode(values[1])));
222+
}
223+
}
224+
177225
static void writeIdToken(SFLoginInput loginInput, String idToken) throws SFException {
178226
logger.debug(
179227
"Caching id token in a secure storage for user: {}, host: {}",
@@ -238,6 +286,30 @@ static void writeOAuthRefreshToken(SFLoginInput loginInput) throws SFException {
238286
CachedCredentialType.OAUTH_REFRESH_TOKEN);
239287
}
240288

289+
/**
290+
* Store OAuth DPoP Public Key With Token
291+
*
292+
* @param loginInput loginInput to denote to the cache
293+
*/
294+
static void writeDPoPBundledAccessToken(SFLoginInput loginInput) throws SFException {
295+
String host = getHostForOAuthCacheKey(loginInput);
296+
logger.debug(
297+
"Caching DPoP public key in a secure storage for user: {}, host: {}",
298+
loginInput.getUserName(),
299+
host);
300+
Base64.Encoder encoder = Base64.getEncoder();
301+
String tokenBase64 =
302+
encoder.encodeToString(loginInput.getOauthAccessToken().getBytes(StandardCharsets.UTF_8));
303+
String publicKeyBase64 =
304+
encoder.encodeToString(loginInput.getDPoPPublicKey().getBytes(StandardCharsets.UTF_8));
305+
getInstance()
306+
.writeTemporaryCredential(
307+
host,
308+
loginInput.getUserName(),
309+
tokenBase64 + "." + publicKeyBase64,
310+
CachedCredentialType.DPOP_BUNDLED_ACCESS_TOKEN);
311+
}
312+
241313
/** Store the temporary credential */
242314
synchronized void writeTemporaryCredential(
243315
String host, String user, String cred, CachedCredentialType credType) {
@@ -256,52 +328,75 @@ synchronized void writeTemporaryCredential(
256328
}
257329

258330
try {
259-
secureStorageManager.setCredential(host, user, credType.getValue(), cred);
331+
if (credType == CachedCredentialType.DPOP_BUNDLED_ACCESS_TOKEN) {
332+
// DPOP_ACCESS_TOKEN is already preformatted and Base64 encoded
333+
secureStorageManager.setCredential(host, user, credType.getValue(), cred);
334+
} else {
335+
String base64EncodedCred =
336+
Base64.getEncoder().encodeToString(cred.getBytes(StandardCharsets.UTF_8));
337+
secureStorageManager.setCredential(host, user, credType.getValue(), base64EncodedCred);
338+
}
260339
} catch (NoClassDefFoundError error) {
261340
logMissingJnaJarForSecureLocalStorage();
262341
}
263342
}
264343

265344
/** Delete the id token cache */
266-
static void deleteIdTokenCache(String host, String user) {
345+
static void deleteIdTokenCacheEntry(String host, String user) {
267346
logger.debug(
268347
"Removing cached id token from a secure storage for user: {}, host: {}", user, host);
269348
getInstance().deleteTemporaryCredential(host, user, CachedCredentialType.ID_TOKEN);
270349
}
271350

272351
/** Delete the mfa token cache */
273-
static void deleteMfaTokenCache(String host, String user) {
352+
static void deleteMfaTokenCacheEntry(String host, String user) {
274353
logger.debug(
275354
"Removing cached mfa token from a secure storage for user: {}, host: {}", user, host);
276355
getInstance().deleteTemporaryCredential(host, user, CachedCredentialType.MFA_TOKEN);
277356
}
278357

279358
/** Delete the Oauth access token cache */
280-
static void deleteOAuthAccessTokenCache(String host, String user) {
359+
static void deleteOAuthAccessTokenCacheEntry(String host, String user) {
281360
logger.debug(
282-
"Removing cached mfa token from a secure storage for user: {}, host: {}", user, host);
361+
"Removing cached oauth access token from a secure storage for user: {}, host: {}",
362+
user,
363+
host);
283364
getInstance().deleteTemporaryCredential(host, user, CachedCredentialType.OAUTH_ACCESS_TOKEN);
284365
}
285366

367+
/** Delete the Oauth refresh token cache */
368+
static void deleteOAuthRefreshTokenCacheEntry(String host, String user) {
369+
logger.debug(
370+
"Removing cached OAuth refresh token from a secure storage for user: {}, host: {}",
371+
user,
372+
host);
373+
getInstance().deleteTemporaryCredential(host, user, CachedCredentialType.OAUTH_REFRESH_TOKEN);
374+
}
375+
376+
/** Delete the DPoP bundled access token cache */
377+
static void deleteDPoPBundledAccessTokenCacheEntry(String host, String user) {
378+
logger.debug(
379+
"Removing cached DPoP public key from a secure storage for user: {}, host: {}", user, host);
380+
getInstance()
381+
.deleteTemporaryCredential(host, user, CachedCredentialType.DPOP_BUNDLED_ACCESS_TOKEN);
382+
}
383+
286384
/** Delete the OAuth access token cache */
287-
static void deleteOAuthAccessTokenCache(SFLoginInput loginInput) throws SFException {
385+
static void deleteOAuthAccessTokenCacheEntry(SFLoginInput loginInput) throws SFException {
288386
String host = getHostForOAuthCacheKey(loginInput);
289-
deleteOAuthAccessTokenCache(host, loginInput.getUserName());
387+
deleteOAuthAccessTokenCacheEntry(host, loginInput.getUserName());
290388
}
291389

292390
/** Delete the OAuth refresh token cache */
293-
static void deleteOAuthRefreshTokenCache(SFLoginInput loginInput) throws SFException {
391+
static void deleteOAuthRefreshTokenCacheEntry(SFLoginInput loginInput) throws SFException {
294392
String host = getHostForOAuthCacheKey(loginInput);
295-
deleteOAuthRefreshTokenCache(host, loginInput.getUserName());
393+
deleteOAuthRefreshTokenCacheEntry(host, loginInput.getUserName());
296394
}
297395

298-
/** Delete the Oauth refresh token cache */
299-
static void deleteOAuthRefreshTokenCache(String host, String user) {
300-
logger.debug(
301-
"Removing cached OAuth refresh token from a secure storage for user: {}, host: {}",
302-
user,
303-
host);
304-
getInstance().deleteTemporaryCredential(host, user, CachedCredentialType.OAUTH_REFRESH_TOKEN);
396+
/** Delete the DPoP bundled access token cache */
397+
static void deleteDPoPBundledAccessTokenCacheEntry(SFLoginInput loginInput) throws SFException {
398+
String host = getHostForOAuthCacheKey(loginInput);
399+
deleteDPoPBundledAccessTokenCacheEntry(host, loginInput.getUserName());
305400
}
306401

307402
/**

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
import static org.apache.http.client.config.CookieSpecs.IGNORE_COOKIES;
66

77
import com.amazonaws.ClientConfiguration;
8+
import com.fasterxml.jackson.databind.JsonNode;
9+
import com.fasterxml.jackson.databind.ObjectMapper;
810
import com.google.common.annotations.VisibleForTesting;
911
import com.google.common.base.Strings;
1012
import com.microsoft.azure.storage.OperationContext;
@@ -31,6 +33,7 @@
3133
import net.snowflake.client.jdbc.RetryContextManager;
3234
import net.snowflake.client.jdbc.SnowflakeDriver;
3335
import net.snowflake.client.jdbc.SnowflakeSQLException;
36+
import net.snowflake.client.jdbc.SnowflakeUseDPoPNonceException;
3437
import net.snowflake.client.jdbc.SnowflakeUtil;
3538
import net.snowflake.client.jdbc.cloud.storage.S3HttpUtil;
3639
import net.snowflake.client.log.ArgSupplier;
@@ -42,6 +45,7 @@
4245
import net.snowflake.common.core.SqlState;
4346
import org.apache.commons.io.IOUtils;
4447
import org.apache.http.HttpHost;
48+
import org.apache.http.HttpResponse;
4549
import org.apache.http.auth.AuthScope;
4650
import org.apache.http.auth.Credentials;
4751
import org.apache.http.auth.UsernamePasswordCredentials;
@@ -66,6 +70,10 @@
6670
public class HttpUtil {
6771
private static final SFLogger logger = SFLoggerFactory.getLogger(HttpUtil.class);
6872

73+
static final String ERROR_FIELD_NAME = "error";
74+
static final String ERROR_USE_DPOP_NONCE = "use_dpop_nonce";
75+
static final String DPOP_NONCE_HEADER_NAME = "dpop-nonce";
76+
6977
static final int DEFAULT_MAX_CONNECTIONS = 300;
7078
static final int DEFAULT_MAX_CONNECTIONS_PER_ROUTE = 300;
7179
private static final int DEFAULT_HTTP_CLIENT_CONNECTION_TIMEOUT_IN_MS = 60000;
@@ -913,6 +921,12 @@ private static String executeRequestInternal(
913921
if (response == null || response.getStatusLine().getStatusCode() != 200) {
914922
logger.error("Error executing request: {}", requestInfoScrubbed);
915923

924+
if (response != null
925+
&& response.getStatusLine().getStatusCode() == 400
926+
&& response.getEntity() != null) {
927+
checkForDPoPNonceError(response);
928+
}
929+
916930
SnowflakeUtil.logResponseDetails(response, logger);
917931

918932
if (response != null) {
@@ -950,6 +964,22 @@ private static String executeRequestInternal(
950964
return theString;
951965
}
952966

967+
private static void checkForDPoPNonceError(HttpResponse response) throws IOException {
968+
String errorResponse = EntityUtils.toString(response.getEntity());
969+
if (!Strings.isNullOrEmpty(errorResponse)) {
970+
ObjectMapper objectMapper = ObjectMapperFactory.getObjectMapper();
971+
JsonNode rootNode = objectMapper.readTree(errorResponse);
972+
JsonNode errorNode = rootNode.get(ERROR_FIELD_NAME);
973+
if (errorNode != null
974+
&& errorNode.isValueNode()
975+
&& errorNode.isTextual()
976+
&& errorNode.textValue().equals(ERROR_USE_DPOP_NONCE)) {
977+
throw new SnowflakeUseDPoPNonceException(
978+
response.getFirstHeader(DPOP_NONCE_HEADER_NAME).getValue());
979+
}
980+
}
981+
}
982+
953983
// This is a workaround for JDK-7036144.
954984
//
955985
// The GZIPInputStream prematurely closes its input if a) it finds

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@ public class SFLoginInput {
4040
private String mfaToken;
4141
private String oauthAccessToken;
4242
private String oauthRefreshToken;
43+
private String dpopPublicKey;
44+
private boolean dpopEnabled = false;
4345
private String serviceName;
4446
private OCSPMode ocspMode;
4547
private HttpClientSettingsKey httpClientKey;
@@ -340,6 +342,27 @@ SFLoginInput setOauthRefreshToken(String oauthRefreshToken) {
340342
return this;
341343
}
342344

345+
@SnowflakeJdbcInternalApi
346+
public String getDPoPPublicKey() {
347+
return dpopPublicKey;
348+
}
349+
350+
SFLoginInput setDPoPPublicKey(String dpopPublicKey) {
351+
this.dpopPublicKey = dpopPublicKey;
352+
return this;
353+
}
354+
355+
@SnowflakeJdbcInternalApi
356+
public boolean isDPoPEnabled() {
357+
return dpopEnabled;
358+
}
359+
360+
// Currently only used for testing purpose
361+
@SnowflakeJdbcInternalApi
362+
public void setDPoPEnabled(boolean dpopEnabled) {
363+
this.dpopEnabled = dpopEnabled;
364+
}
365+
343366
Map<String, Object> getSessionParameters() {
344367
return sessionParameters;
345368
}

0 commit comments

Comments
 (0)