SNOW-1831103/SNOW-1853435: Refresh token & OAuth tokens caching support #2009

sfc-gh-dheyman · 2024-12-18T20:54:22Z

Overview

SNOW-1831103/SNOW-1853435: Refresh token & OAuth tokens caching support

Pre-review self checklist

PR branch is updated with all the changes from master branch
The code is correctly formatted (run mvn -P check-style validate)
New public API is not unnecessary exposed (run mvn verify and inspect target/japicmp/japicmp.html)
The pull request name is prefixed with SNOW-XXXX:
Code is in compliance with internal logging requirements

External contributors - please answer these questions before submitting a pull request. Thanks!

What GitHub issue is this PR addressing? Make sure that there is an accompanying issue to your PR.

Issue: #NNNN
Fill out the following pre-review checklist:
- I am adding a new automated test(s) to verify correctness of my new code
- I am adding new logging messages
- I am modifying authorization mechanisms
- I am adding new credentials
- I am modifying OCSP code
- I am adding a new dependency or upgrading an existing one
- I am adding new public/protected component not marked with @SnowflakeJdbcInternalApi (note that public/protected methods/fields in classes marked with this annotation are already internal)
Please describe how your code solves the related issue.

Please write a short description of how your code change solves the related issue.

sfc-gh-dprzybysz · 2024-12-19T06:26:56Z

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

-      logger.debug("Unrecognized type {} for local cached credential", credType);
+    switch (credType) {
+      case ID_TOKEN:
+        logger.debug(


the messages are different only in the part after {} - could we make the login before and once?

sfc-gh-dprzybysz · 2024-12-19T06:28:28Z

src/main/java/net/snowflake/client/core/SFLoginOutput.java

@@ -98,6 +104,14 @@ String getMfaToken() {
    return mfaToken;
  }

+  public String getOauthAccessToken() {


new public methods may be internal?

Removed public modifier, it was actually unnecessary

sfc-gh-dprzybysz · 2024-12-19T06:29:23Z

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

-    preNewSession(loginInput);
+    readCachedTokens(loginInput);
+
+    if (AccessTokenProviderFactory.isEligible(getAuthenticator(loginInput))) {


can we extract methods from newly added parts in openSession?

sfc-gh-dprzybysz · 2024-12-19T06:30:45Z

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

+
+    if (AccessTokenProviderFactory.isEligible(getAuthenticator(loginInput))) {
+      if (loginInput.getOauthAccessToken() != null) { // Access Token cached
+        loginInput.setAuthenticator(AuthenticatorType.OAUTH.name());


will the new authenticators work with session renewal also?

Yes, but I'm still thinking how to write automated test for it.

sfc-gh-dprzybysz · 2024-12-19T06:31:17Z

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

+    TokenResponseDTO tokenResponse = accessTokenProvider.getAccessToken(loginInput);
+    loginInput.setAuthenticator(AuthenticatorType.OAUTH.name());
+    loginInput.setToken(tokenResponse.getAccessToken());
+    loginInput.setOauthAccessToken(tokenResponse.getAccessToken());


why do we pass access token twice?

Token field is generic token field which is related to /login-request body and oauthToken on the other hand is a field that is being used to store only oauth access token and to interact with token cache

sfc-gh-dprzybysz · 2024-12-19T06:33:15Z

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

@@ -998,6 +1064,14 @@ public static void deleteMfaTokenCache(String host, String user) {
    CredentialManager.getInstance().deleteMfaTokenCache(host, user);
  }

+  private static void deleteOAuthAccessTokenCache(String host, String user) {
+    CredentialManager.getInstance().deleteOAuthAccessTokenCache(host, user);


cannot we expose static methods directly from CredentialManager? under the hood it's instance could be used and we don't need to provide such helper methods?

sfc-gh-dprzybysz · 2024-12-19T06:34:12Z

src/main/java/net/snowflake/client/core/auth/oauth/OAuthAccessTokenForRefreshTokenProvider.java

+  private static final SFLogger logger =
+      SFLoggerFactory.getLogger(OAuthClientCredentialsAccessTokenProvider.class);
+
+  private final ObjectMapper objectMapper = new ObjectMapper();


don't we want to use ObjectMapperFactory?

No really since it's configuration causes problems with deserialization via constructor

ObjectMapper is thread safe, maybe it could be a static field?

Changed to static field.

sfc-gh-dprzybysz · 2024-12-19T06:35:24Z

src/main/java/net/snowflake/client/core/auth/oauth/OAuthAccessTokenForRefreshTokenProvider.java

+        objectMapper.readValue(tokenResponse, TokenResponseDTO.class);
+    logger.debug(
+        "Received new OAuth access token from: {}",
+        requestUri.getAuthority() + requestUri.getPath());


let's use separate parameters without concatenation

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

src/main/java/net/snowflake/client/core/SFStatement.java

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

sfc-gh-pfus · 2024-12-20T07:03:23Z

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

    try {
      cred =
          secureStorageManager.getCredential(
-              loginInput.getHostFromServerUrl(), loginInput.getUserName(), credType);
+              loginInput.getHostFromServerUrl(), loginInput.getUserName(), credType.getValue());


I'm not sure if it is correct for oauth. Guessing from the code above getHostFromServerUrl returns snowflake host, right? For OAuth it should be IDP address. Otherwise we could accidentally leak tokens between IDPs.

Refactored according to your suggestion.

sfc-gh-pfus · 2024-12-20T07:29:07Z

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

+  }
+
+  /** Delete the OAuth access token cache */
+  static void deleteOAuthAccessTokenCache(String host, String user) {


...FromCache?

I'd rather leave it as is since another functions are named in similar convention

sfc-gh-pfus · 2024-12-20T07:30:55Z

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

-   * @return true if authenticator type is EXTERNALBROWSER
-   */
-  boolean isExternalbrowserAuthenticator() {
+  boolean isExternalbrowserOrOAuthFullFlowAuthenticator() {


Full flow is a new name? WDYT about isBrowserBasedAuthenticator?

CLIENT_CREDENTIALS flow is not browser based.

Full flow as in opposed to OAUTH authenticator :)

sfc-gh-pfus · 2024-12-20T07:34:27Z

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

+      if (tokenResponse.getRefreshToken() != null) {
+        loginInput.setOauthRefreshToken(tokenResponse.getRefreshToken());
+      }
+    } catch (Throwable e) {


Do not catch throwable. What if OOM or stack overflow is thrown here?

True, fixed.

sfc-gh-pfus · 2024-12-20T07:35:43Z

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

-        && asBoolean(loginInput.getSessionParameters().get(CLIENT_STORE_TEMPORARY_CREDENTIAL))) {
-      CredentialManager.getInstance().writeIdToken(loginInput, ret);
+    if (asBoolean(loginInput.getSessionParameters().get(CLIENT_STORE_TEMPORARY_CREDENTIAL))) {
+      if (consentCacheIdToken) {


What the consent is here?

not sure, this code was already there. I prefer not to touch it :).

sfc-gh-pfus · 2024-12-20T07:36:33Z

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

+      if (consentCacheIdToken) {
+        CredentialManager.writeIdToken(loginInput, ret);
+      }
+      if (loginInput.getOauthAccessToken() != null) {


This is a good assertion, but maybe it would be better placed in CredentialManager? It would make the code more robust.

I'd prefer to leave it here. For me it's counterintuitive - it would suggests that we always save this token

sfc-gh-pfus · 2024-12-20T07:37:49Z

src/main/java/net/snowflake/client/core/auth/oauth/OAuthAccessTokenForRefreshTokenProvider.java

+  private static final SFLogger logger =
+      SFLoggerFactory.getLogger(OAuthClientCredentialsAccessTokenProvider.class);
+
+  private final ObjectMapper objectMapper = new ObjectMapper();


ObjectMapper is thread safe, maybe it could be a static field?

sfc-gh-pfus · 2024-12-20T07:39:48Z

src/main/java/net/snowflake/client/core/auth/oauth/OAuthAccessTokenForRefreshTokenProvider.java

+        requestUri.getAuthority(),
+        requestUri.getPath());
+    String tokenResponse =
+        HttpUtil.executeGeneralRequest(


Should we use this utility for non-snowflake requests? It contains some of snowflake-specific behaviour, like telemetry, adding request id etc.

This is also used in case of external browser saml request. It supports configured retry stategy, timeouts and proxy.

sfc-gh-pfus · 2024-12-20T07:40:44Z

src/main/java/net/snowflake/client/core/auth/oauth/OAuthAccessTokenProviderFactory.java

@@ -18,10 +18,10 @@
 import net.snowflake.client.log.SFLoggerFactory;

 @SnowflakeJdbcInternalApi
-public class AccessTokenProviderFactory {
+public class OAuthAccessTokenProviderFactory {


Even more precise - OAuth21, right?

I think this is too detailed. Documentation can point this out but that's about it IMHO.

sfc-gh-dheyman added 2 commits December 12, 2024 16:47

Add oauth token cache support

aad5632

Add cache & refresh token support + wiremock tests

9aeeed0

sfc-gh-dheyman requested a review from a team as a code owner December 18, 2024 20:54

sfc-gh-dprzybysz reviewed Dec 19, 2024

View reviewed changes

sfc-gh-dheyman added 5 commits December 19, 2024 16:06

CR suggestions applied

e10594a

Extracted openSession code to separate methods

561abfb

Update session renewal

d07da78

Remove dead code

d1f3114

update comment

68e468c

sfc-gh-astachowski reviewed Dec 20, 2024

View reviewed changes

sfc-gh-pfus reviewed Dec 20, 2024

View reviewed changes

sfc-gh-dheyman added 3 commits December 20, 2024 09:59

CR suggestions applied

284ed4f

Fixed oauth cache + added state check

c99757b

Fix tests

257b224

sfc-gh-dheyman merged commit 0312556 into oauth-code-flow Dec 20, 2024
4 of 7 checks passed

sfc-gh-dheyman deleted the oauth-token-cache branch December 20, 2024 20:44

github-actions bot locked and limited conversation to collaborators Dec 20, 2024

SNOW-1831103/SNOW-1853435: Refresh token & OAuth tokens caching support #2009

SNOW-1831103/SNOW-1853435: Refresh token & OAuth tokens caching support #2009

Conversation

sfc-gh-dheyman commented Dec 18, 2024 • edited Loading

Overview

Pre-review self checklist

External contributors - please answer these questions before submitting a pull request. Thanks!

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sfc-gh-dheyman Dec 20, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sfc-gh-dheyman commented Dec 18, 2024 •

edited

Loading

sfc-gh-dheyman Dec 20, 2024 •

edited

Loading