Skip to content

Commit 06693df

Browse files
sfc-gh-dheymansfc-gh-dheyman
committed
SNOW-1902211: Add initial DPoP support
1 parent 464540d commit 06693df

23 files changed

+1456
-134
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,8 @@ 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_PUBLIC_KEY("DPOP_PUBLIC_KEY");
89

910
private final String value;
1011

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

Lines changed: 61 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,22 @@ static void fillCachedOAuthRefreshToken(SFLoginInput loginInput) throws SFExcept
124124
loginInput, host, loginInput.getUserName(), CachedCredentialType.OAUTH_REFRESH_TOKEN);
125125
}
126126

127+
/**
128+
* Reuse the cached OAuth DPoP public kley stored locally
129+
*
130+
* @param loginInput login input to attach refresh token
131+
*/
132+
static void fillCachedDPoPPublicKey(SFLoginInput loginInput) throws SFException {
133+
String host = getHostForOAuthCacheKey(loginInput);
134+
logger.debug(
135+
"Looking for cached DPoP public key for user: {}, host: {}",
136+
loginInput.getUserName(),
137+
host);
138+
getInstance()
139+
.fillCachedCredential(
140+
loginInput, host, loginInput.getUserName(), CachedCredentialType.DPOP_PUBLIC_KEY);
141+
}
142+
127143
/** Reuse the cached token stored locally */
128144
synchronized void fillCachedCredential(
129145
SFLoginInput loginInput, String host, String username, CachedCredentialType credType)
@@ -168,6 +184,9 @@ synchronized void fillCachedCredential(
168184
case OAUTH_REFRESH_TOKEN:
169185
loginInput.setOauthRefreshToken(cred);
170186
break;
187+
case DPOP_PUBLIC_KEY:
188+
loginInput.setDPoPPublicKeyBase64(cred);
189+
break;
171190
default:
172191
throw new SFException(
173192
ErrorCode.INTERNAL_ERROR, "Unrecognized type {} for local cached credential", credType);
@@ -238,6 +257,25 @@ static void writeOAuthRefreshToken(SFLoginInput loginInput) throws SFException {
238257
CachedCredentialType.OAUTH_REFRESH_TOKEN);
239258
}
240259

260+
/**
261+
* Store OAuth DPoP Public Key
262+
*
263+
* @param loginInput loginInput to denote to the cache
264+
*/
265+
static void writeDPoPPublicKey(SFLoginInput loginInput) throws SFException {
266+
String host = getHostForOAuthCacheKey(loginInput);
267+
logger.debug(
268+
"Caching DPoP public key in a secure storage for user: {}, host: {}",
269+
loginInput.getUserName(),
270+
host);
271+
getInstance()
272+
.writeTemporaryCredential(
273+
host,
274+
loginInput.getUserName(),
275+
loginInput.getDPoPPublicKeyBase64(),
276+
CachedCredentialType.DPOP_PUBLIC_KEY);
277+
}
278+
241279
/** Store the temporary credential */
242280
synchronized void writeTemporaryCredential(
243281
String host, String user, String cred, CachedCredentialType credType) {
@@ -279,10 +317,28 @@ static void deleteMfaTokenCache(String host, String user) {
279317
/** Delete the Oauth access token cache */
280318
static void deleteOAuthAccessTokenCache(String host, String user) {
281319
logger.debug(
282-
"Removing cached mfa token from a secure storage for user: {}, host: {}", user, host);
320+
"Removing cached oauth access token from a secure storage for user: {}, host: {}",
321+
user,
322+
host);
283323
getInstance().deleteTemporaryCredential(host, user, CachedCredentialType.OAUTH_ACCESS_TOKEN);
284324
}
285325

326+
/** Delete the Oauth refresh token cache */
327+
static void deleteOAuthRefreshTokenCache(String host, String user) {
328+
logger.debug(
329+
"Removing cached OAuth refresh token from a secure storage for user: {}, host: {}",
330+
user,
331+
host);
332+
getInstance().deleteTemporaryCredential(host, user, CachedCredentialType.OAUTH_REFRESH_TOKEN);
333+
}
334+
335+
/** Delete the Oauth access token cache */
336+
static void deleteDPoPPublicKeyCache(String host, String user) {
337+
logger.debug(
338+
"Removing cached DPoP public key from a secure storage for user: {}, host: {}", user, host);
339+
getInstance().deleteTemporaryCredential(host, user, CachedCredentialType.DPOP_PUBLIC_KEY);
340+
}
341+
286342
/** Delete the OAuth access token cache */
287343
static void deleteOAuthAccessTokenCache(SFLoginInput loginInput) throws SFException {
288344
String host = getHostForOAuthCacheKey(loginInput);
@@ -295,13 +351,10 @@ static void deleteOAuthRefreshTokenCache(SFLoginInput loginInput) throws SFExcep
295351
deleteOAuthRefreshTokenCache(host, loginInput.getUserName());
296352
}
297353

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);
354+
/** Delete the DPoP refresh token cache */
355+
static void deleteDPoPPublicKeyCache(SFLoginInput loginInput) throws SFException {
356+
String host = getHostForOAuthCacheKey(loginInput);
357+
deleteDPoPPublicKeyCache(host, loginInput.getUserName());
305358
}
306359

307360
/**

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

Lines changed: 34 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,26 @@ private static String executeRequestInternal(
950964
return theString;
951965
}
952966

967+
private static void checkForDPoPNonceError(HttpResponse response) throws IOException {
968+
StringWriter writer = new StringWriter();
969+
try (InputStream ins = response.getEntity().getContent()) {
970+
IOUtils.copy(ins, writer, "UTF-8");
971+
}
972+
String errorResponse = writer.toString();
973+
if (!Strings.isNullOrEmpty(errorResponse)) {
974+
ObjectMapper objectMapper = ObjectMapperFactory.getObjectMapper();
975+
JsonNode rootNode = objectMapper.readTree(errorResponse);
976+
JsonNode errorNode = rootNode.get(ERROR_FIELD_NAME);
977+
if (errorNode != null
978+
&& errorNode.isValueNode()
979+
&& errorNode.isTextual()
980+
&& errorNode.textValue().equals(ERROR_USE_DPOP_NONCE)) {
981+
throw new SnowflakeUseDPoPNonceException(
982+
response.getFirstHeader(DPOP_NONCE_HEADER_NAME).getValue());
983+
}
984+
}
985+
}
986+
953987
// This is a workaround for JDK-7036144.
954988
//
955989
// 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 dpopPublicKeyBase64;
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 getDPoPPublicKeyBase64() {
347+
return dpopPublicKeyBase64;
348+
}
349+
350+
SFLoginInput setDPoPPublicKeyBase64(String dpopPublicKeyBase64) {
351+
this.dpopPublicKeyBase64 = dpopPublicKeyBase64;
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
}

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

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import net.snowflake.client.core.auth.ClientAuthnDTO;
2929
import net.snowflake.client.core.auth.ClientAuthnParameter;
3030
import net.snowflake.client.core.auth.oauth.AccessTokenProvider;
31+
import net.snowflake.client.core.auth.oauth.DPoPUtil;
3132
import net.snowflake.client.core.auth.oauth.OAuthAccessTokenForRefreshTokenProvider;
3233
import net.snowflake.client.core.auth.oauth.OAuthAccessTokenProviderFactory;
3334
import net.snowflake.client.core.auth.oauth.TokenResponseDTO;
@@ -329,7 +330,7 @@ static SFLoginOutput openSession(
329330

330331
convertSessionParameterStringValueToBooleanIfGiven(loginInput, CLIENT_REQUEST_MFA_TOKEN);
331332

332-
readCachedTokensIfPossible(loginInput);
333+
readCachedCredentialsIfPossible(loginInput);
333334
if (OAuthAccessTokenProviderFactory.isEligible(getAuthenticator(loginInput))) {
334335
obtainAuthAccessTokenAndUpdateInput(loginInput);
335336
}
@@ -403,6 +404,7 @@ private static void fetchOAuthAccessTokenAndUpdateInput(SFLoginInput loginInput)
403404
loginInput.setToken(tokenResponse.getAccessToken());
404405
loginInput.setOauthAccessToken(tokenResponse.getAccessToken());
405406
loginInput.setOauthRefreshToken(tokenResponse.getRefreshToken());
407+
loginInput.setDPoPPublicKeyBase64(accessTokenProvider.getDPoPPublicKeyBase64());
406408
}
407409

408410
private static void refreshOAuthAccessTokenAndUpdateInput(SFLoginInput loginInput)
@@ -437,12 +439,13 @@ private static void convertSessionParameterStringValueToBooleanIfGiven(
437439
}
438440
}
439441

440-
private static void readCachedTokensIfPossible(SFLoginInput loginInput) throws SFException {
442+
private static void readCachedCredentialsIfPossible(SFLoginInput loginInput) throws SFException {
441443
if (!StringUtils.isNullOrEmpty(loginInput.getUserName())) {
442444
if (asBoolean(loginInput.getSessionParameters().get(CLIENT_STORE_TEMPORARY_CREDENTIAL))) {
443445
CredentialManager.fillCachedIdToken(loginInput);
444446
CredentialManager.fillCachedOAuthAccessToken(loginInput);
445447
CredentialManager.fillCachedOAuthRefreshToken(loginInput);
448+
CredentialManager.fillCachedDPoPPublicKey(loginInput);
446449
}
447450

448451
if (asBoolean(loginInput.getSessionParameters().get(CLIENT_REQUEST_MFA_TOKEN))) {
@@ -768,6 +771,10 @@ static SFLoginOutput newSession(
768771

769772
postRequest.addHeader("accept", "application/json");
770773
postRequest.addHeader("Accept-Encoding", "");
774+
if (loginInput.isDPoPEnabled()) {
775+
new DPoPUtil(loginInput.getDPoPPublicKeyBase64())
776+
.addDPoPProofHeaderToRequest(postRequest, null);
777+
}
771778

772779
/*
773780
* HttpClient should take authorization header from char[] instead of
@@ -896,11 +903,13 @@ static SFLoginOutput newSession(
896903
logger.debug("OAuth Access Token Invalid: {}", errorCode);
897904
loginInput.setOauthAccessToken(null);
898905
CredentialManager.deleteOAuthAccessTokenCache(loginInput);
906+
CredentialManager.deleteDPoPPublicKeyCache(loginInput);
899907
}
900908

901909
if (errorCode == Constants.OAUTH_ACCESS_TOKEN_EXPIRED_GS_CODE) {
902910
loginInput.setOauthAccessToken(null);
903911
CredentialManager.deleteOAuthAccessTokenCache(loginInput);
912+
CredentialManager.deleteDPoPPublicKeyCache(loginInput);
904913

905914
logger.debug("OAuth Access Token Expired: {}", errorCode);
906915
SnowflakeUtil.checkErrorAndThrowExceptionIncludingReauth(jsonNode);
@@ -1053,6 +1062,9 @@ static SFLoginOutput newSession(
10531062
if (loginInput.getOauthRefreshToken() != null) {
10541063
CredentialManager.writeOAuthRefreshToken(loginInput);
10551064
}
1065+
if (loginInput.getDPoPPublicKeyBase64() != null && loginInput.isDPoPEnabled()) {
1066+
CredentialManager.writeDPoPPublicKey(loginInput);
1067+
}
10561068
}
10571069

10581070
if (asBoolean(loginInput.getSessionParameters().get(CLIENT_REQUEST_MFA_TOKEN))) {

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,4 +8,6 @@
88
public interface AccessTokenProvider {
99

1010
TokenResponseDTO getAccessToken(SFLoginInput loginInput) throws SFException;
11+
12+
String getDPoPPublicKeyBase64();
1113
}

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

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
package net.snowflake.client.core.auth.oauth;
2+
3+
import com.nimbusds.jose.JOSEException;
4+
import com.nimbusds.jose.JWSAlgorithm;
5+
import com.nimbusds.jose.jwk.Curve;
6+
import com.nimbusds.jose.jwk.ECKey;
7+
import com.nimbusds.jose.jwk.gen.ECKeyGenerator;
8+
import com.nimbusds.jwt.SignedJWT;
9+
import com.nimbusds.oauth2.sdk.dpop.DPoPProofFactory;
10+
import com.nimbusds.oauth2.sdk.dpop.DefaultDPoPProofFactory;
11+
import com.nimbusds.oauth2.sdk.dpop.JWKThumbprintConfirmation;
12+
import com.nimbusds.openid.connect.sdk.Nonce;
13+
import java.net.URI;
14+
import java.net.URISyntaxException;
15+
import java.util.Base64;
16+
import net.snowflake.client.core.SFException;
17+
import net.snowflake.client.core.SnowflakeJdbcInternalApi;
18+
import net.snowflake.client.jdbc.ErrorCode;
19+
import org.apache.http.client.methods.HttpRequestBase;
20+
21+
@SnowflakeJdbcInternalApi
22+
public class DPoPUtil {
23+
24+
private final ECKey jwk;
25+
26+
DPoPUtil() throws SFException {
27+
try {
28+
jwk = new ECKeyGenerator(Curve.P_256).generate();
29+
} catch (JOSEException e) {
30+
throw new SFException(
31+
ErrorCode.INTERNAL_ERROR, "Error during DPoP JWK initialization: " + e.getMessage());
32+
}
33+
}
34+
35+
public DPoPUtil(String jsonKeyBase64) throws SFException {
36+
try {
37+
String jsonKey = new String(Base64.getDecoder().decode(jsonKeyBase64));
38+
jwk = ECKey.parse(jsonKey);
39+
} catch (Exception e) {
40+
throw new SFException(
41+
ErrorCode.INTERNAL_ERROR, "Error during DPoP JWK initialization: " + e.getMessage());
42+
}
43+
}
44+
45+
String getPublicKeyInBase64() {
46+
String jsonString = jwk.toJSONString();
47+
return Base64.getEncoder().encodeToString(jsonString.getBytes());
48+
}
49+
50+
JWKThumbprintConfirmation getThumbprint() throws SFException {
51+
try {
52+
return JWKThumbprintConfirmation.of(this.jwk);
53+
} catch (JOSEException e) {
54+
throw new SFException(
55+
ErrorCode.INTERNAL_ERROR, "Error during JWK thumbprint generation: " + e.getMessage());
56+
}
57+
}
58+
59+
public void addDPoPProofHeaderToRequest(HttpRequestBase httpRequest, String nonce)
60+
throws SFException {
61+
SignedJWT signedJWT = generateDPoPProof(httpRequest, nonce);
62+
httpRequest.setHeader("DPoP", signedJWT.serialize());
63+
}
64+
65+
private SignedJWT generateDPoPProof(HttpRequestBase httpRequest, String nonce)
66+
throws SFException {
67+
try {
68+
DPoPProofFactory proofFactory = new DefaultDPoPProofFactory(jwk, JWSAlgorithm.ES256);
69+
if (nonce != null) {
70+
return proofFactory.createDPoPJWT(
71+
httpRequest.getMethod(), httpRequest.getURI(), new Nonce(nonce));
72+
} else {
73+
return proofFactory.createDPoPJWT(
74+
httpRequest.getMethod(), getUriWithoutQuery(httpRequest.getURI()));
75+
}
76+
} catch (Exception e) {
77+
throw new SFException(
78+
ErrorCode.INTERNAL_ERROR, " Error during DPoP proof generation: " + e.getMessage());
79+
}
80+
}
81+
82+
private URI getUriWithoutQuery(URI uri) throws URISyntaxException {
83+
return new URI(uri.getScheme(), uri.getAuthority(), uri.getPath(), null, null);
84+
}
85+
}

0 commit comments

Comments
 (0)