Refactor

snowflakedb · Nov 30, 2024 · 195ff36 · 195ff36
1 parent 1741680
commit 195ff36
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 11 deletions.
diff --git a/...n/java/net/snowflake/client/core/auth/oauth/AuthorizationCodeFlowAccessTokenProvider.java b/...n/java/net/snowflake/client/core/auth/oauth/AuthorizationCodeFlowAccessTokenProvider.java
@@ -47,17 +47,15 @@ public String getAccessToken(SFLoginInput loginInput) throws SFException {
     private AuthorizationCode requestAuthorizationCode(SFLoginInput loginInput) throws SFException {
         try {
             AuthorizationRequest request = buildAuthorizationRequest(loginInput);
-
             URI requestURI = request.toURI();
             HttpUtil.executeGeneralRequest(new HttpGet(requestURI),
                     loginInput.getLoginTimeout(),
                     loginInput.getAuthTimeout(),
                     loginInput.getSocketTimeoutInMillis(),
                     0,
                     loginInput.getHttpClientSettingsKey());
-            CompletableFuture<String> f = getAuthorizationCodeFromRedirectURI();
-            f.join();
-            return new AuthorizationCode(f.get());
+            String code = getAuthorizationCodeFromRedirectURI().join();
+            return new AuthorizationCode(code);
         } catch (Exception e) {
             throw new SFException(e, ErrorCode.INTERNAL_ERROR);
         }
@@ -92,9 +90,9 @@ private static CompletableFuture<String> getAuthorizationCodeFromRedirectURI() t
     private static AuthorizationRequest buildAuthorizationRequest(SFLoginInput loginInput) throws URISyntaxException {
         URI authorizeEndpoint = new URI(String.format("%s/oauth/authorize", loginInput.getServerUrl()));
         ClientID clientID = new ClientID("123");
-        Scope scope = new Scope("read", "write");
+        Scope scope = new Scope(String.format("session:role:%s", loginInput.getRole()));
         URI callback = buildRedirectURI();
-        State state = new State();
+        State state = new State(256);
         return new AuthorizationRequest.Builder(
                 new ResponseType(ResponseType.Value.CODE), clientID)
                 .scope(scope)
@@ -111,10 +109,9 @@ private static URI buildRedirectURI() throws URISyntaxException {
     private static TokenRequest buildTokenRequest(SFLoginInput loginInput, AuthorizationCode authorizationCode) throws URISyntaxException {
         URI callback = buildRedirectURI();
         AuthorizationGrant codeGrant = new AuthorizationCodeGrant(authorizationCode, callback);
-        ClientID clientID = new ClientID("123");
-        Secret clientSecret = new Secret("123");
-        ClientAuthentication clientAuthentication = new ClientSecretBasic(clientID, clientSecret);
-        URI tokenEndpoint = new URI(String.format("%s/oauth/token", loginInput.getServerUrl()));
-        return new TokenRequest(tokenEndpoint, clientAuthentication, codeGrant, new Scope());
+        ClientAuthentication clientAuthentication = new ClientSecretBasic(new ClientID("123"), new Secret("123"));
+        URI tokenEndpoint = new URI(String.format("%s/oauth/token-request", loginInput.getServerUrl()));
+        Scope scope = new Scope("session:role", loginInput.getRole());
+        return new TokenRequest(tokenEndpoint, clientAuthentication, codeGrant, scope);
     }
 }
diff --git a/src/main/java/net/snowflake/client/core/auth/oauth/OauthAccessTokenProvider.java b/src/main/java/net/snowflake/client/core/auth/oauth/OauthAccessTokenProvider.java
@@ -2,7 +2,9 @@
 
 import net.snowflake.client.core.SFException;
 import net.snowflake.client.core.SFLoginInput;
+import net.snowflake.client.core.SnowflakeJdbcInternalApi;
 
+@SnowflakeJdbcInternalApi
 public interface OauthAccessTokenProvider {
 
     String getAccessToken(SFLoginInput loginInput) throws SFException;