Merge pull request #168 from forcedotcom/develop

Authentication flow updated to use instance url (#167)
forcedotcom · Jun 7, 2024 · 3bdaee1 · 3bdaee1
2 parents c5ca7a4 + 6aa6383
commit 3bdaee1
Show file tree

Hide file tree

Showing 9 changed files with 92 additions and 9 deletions.
diff --git a/pom.xml b/pom.xml
@@ -6,7 +6,7 @@
 
     <groupId>com.queryService</groupId>
     <artifactId>Salesforce-CDP-jdbc</artifactId>
-    <version>1.19.6</version>
+    <version>1.19.7</version>
 
     <properties>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>

diff --git a/src/main/java/com/salesforce/cdp/queryservice/auth/TokenExchangeHelper.java b/src/main/java/com/salesforce/cdp/queryservice/auth/TokenExchangeHelper.java
@@ -30,17 +30,24 @@ public TokenExchangeHelper(Properties properties, OkHttpClient client) {
     }
 
     public OffcoreToken exchangeToken(CoreToken coreToken) throws TokenException {
-        return exchangeToken(properties.getProperty(Constants.LOGIN_URL),
+        log.info("Exchanging core token for cdp token");
+        OffcoreToken offcoreToken = exchangeToken(coreToken.getInstanceUrl(),
                 coreToken.getAccessToken(), properties.getProperty(Constants.DATASPACE));
+        log.info("Finished exchanging core token for CDP token");
+        return offcoreToken;
     }
 
     private OffcoreToken exchangeToken(String url, String coreToken, String dataspace) throws TokenException {
+        log.info("Token exchange url : {}", url);
         String token_url = url + Constants.TOKEN_EXCHANGE_URL;
         Map<String, String> requestBody = new HashMap<>();
         requestBody.put(Constants.GRANT_TYPE_NAME, Constants.GRANT_TYPE);
         requestBody.put(Constants.SUBJECT_TOKEN_TYPE_NAME, Constants.SUBJECT_TOKEN_TYPE);
         requestBody.put(Constants.SUBJECT_TOKEN, coreToken);
-        if(StringUtils.isNotBlank(dataspace)) {requestBody.put(Constants.DATASPACE,dataspace);}
+        if(StringUtils.isNotBlank(dataspace)) {
+            log.info("Token exchange dataspace : {}", dataspace);
+            requestBody.put(Constants.DATASPACE,dataspace);
+        }
         Calendar expireTime = Calendar.getInstance();
         Response response = null;
         try {
@@ -52,6 +59,7 @@ private OffcoreToken exchangeToken(String url, String coreToken, String dataspac
                 String message = token.getErrorDescription();
                 throw new TokenException(message);
             }
+            log.info("Setting token expiry time");
             expireTime.add(Calendar.SECOND, token.getExpiresIn());
             token.setExpireTime(expireTime);
             return token;
@@ -69,7 +77,10 @@ private Response login(Map<String, String> requestBody, String url) throws Token
         log.info(requestBody.toString());
         try {
             Request request = HttpHelper.buildRequest(Constants.POST, url, formBody.build(), headers);
+            long startTime = System.currentTimeMillis();
+            log.info("Sending call for token exchange");
             Response response = client.newCall(request).execute();
+            log.info("Time taken for token exchange : {}", System.currentTimeMillis() - startTime);
             if (!response.isSuccessful()) {
                 log.error("login failed with status code {}", response.code());
                 HttpHelper.handleErrorResponse(response, Constants.ERROR_DESCRIPTION);

diff --git a/src/main/java/com/salesforce/cdp/queryservice/auth/TokenManager.java b/src/main/java/com/salesforce/cdp/queryservice/auth/TokenManager.java
@@ -3,10 +3,12 @@
 import com.salesforce.cdp.queryservice.auth.token.CoreToken;
 import com.salesforce.cdp.queryservice.auth.token.OffcoreToken;
 import com.salesforce.cdp.queryservice.util.TokenException;
+import lombok.extern.slf4j.Slf4j;
 import okhttp3.OkHttpClient;
 
 import java.util.Properties;
 
+@Slf4j
 public class TokenManager {
 
     private final TokenProvider tokenProvider;
@@ -19,11 +21,19 @@ public TokenManager(TokenProvider tokenProvider) {
         this.tokenProvider = tokenProvider;
     }
 
-    public CoreToken getCoreToken() throws TokenException {
-        return tokenProvider.getCoreToken();
+    public synchronized CoreToken getCoreToken() throws TokenException {
+        log.info("Begin fetching core token");
+        long startTime = System.currentTimeMillis();
+        CoreToken coreToken = tokenProvider.getCoreToken();
+        log.info("Fetched core token in {}ms", System.currentTimeMillis() - startTime);
+        return coreToken;
     }
 
-    public OffcoreToken getOffcoreToken() throws TokenException {
-        return tokenProvider.getOffcoreToken();
+    public synchronized OffcoreToken getOffcoreToken() throws TokenException {
+        log.info("Begin fetching offcore token");
+        long startTime = System.currentTimeMillis();
+        OffcoreToken offcoreToken = tokenProvider.getOffcoreToken();
+        log.info("Fetched offcore core token in {}ms", System.currentTimeMillis() - startTime);
+        return offcoreToken;
     }
 }
diff --git a/src/main/java/com/salesforce/cdp/queryservice/auth/TokenUtils.java b/src/main/java/com/salesforce/cdp/queryservice/auth/TokenUtils.java
@@ -19,7 +19,14 @@ public class TokenUtils {
     public static boolean isValid(OffcoreToken offcoreToken) {
         if (offcoreToken != null) {
             Calendar now = Calendar.getInstance();
-            return now.compareTo(offcoreToken.getExpireTime()) < 1;
+            if (now.compareTo(offcoreToken.getExpireTime()) < 1) {
+                log.info("Offcore token is still valid");
+                return true;
+            } else {
+                log.info("Current token is expired");
+            }
+        } else {
+            log.info("No existing tokens available");
         }
         return false;
     }
@@ -36,6 +43,7 @@ public static void invalidateCoreToken(String url, String coreToken, OkHttpClien
             Request request = HttpHelper.buildRequest(Constants.POST, tokenRevokeUrl, formBody, headers);
             // Response is not needed for this call.
             client.newCall(request).execute();
+            log.info("Successfully invalidated the core token");
         } catch (Exception e) {
             log.error("Revoking the core token failed", e);
         }

diff --git a/src/main/java/com/salesforce/cdp/queryservice/auth/jwt/JwtLoginClient.java b/src/main/java/com/salesforce/cdp/queryservice/auth/jwt/JwtLoginClient.java
@@ -39,6 +39,7 @@ public CoreToken keyPairAuthLogin(
             String grantType, String clientId,
             String userName, String privateKey, String audience, String tokenUrl
     ) throws TokenException {
+        log.info("Fetching core token using pvt key from url : {}", tokenUrl);
         byte[] grantTypeSegment = (Constants.GRANT_TYPE_NAME + Constants.TOKEN_ASSIGNMENT + grantType)
                 .getBytes(StandardCharsets.UTF_8);
         byte[] jwsSegment = (Constants.ASSERTION + Constants.TOKEN_ASSIGNMENT +

diff --git a/src/main/java/com/salesforce/cdp/queryservice/auth/refresh/RefreshTokenClient.java b/src/main/java/com/salesforce/cdp/queryservice/auth/refresh/RefreshTokenClient.java
@@ -29,6 +29,7 @@ public RefreshTokenClient(OkHttpClient client) {
 
     public CoreToken getCoreToken(String url, String refreshToken, String clientId, String secret) throws TokenException {
         String token_url = url + Constants.CORE_TOKEN_URL;
+        log.info("Fetching Refresh token using URL : {}", token_url);
         Map<String, String> requestBody = new HashMap<>();
         requestBody.put(Constants.GRANT_TYPE_NAME, Constants.REFRESH_TOKEN_GRANT_TYPE);
         requestBody.put(Constants.CLIENT_ID_NAME, clientId);

diff --git a/src/main/java/com/salesforce/cdp/queryservice/auth/unpwd/UnPwdAuthClient.java b/src/main/java/com/salesforce/cdp/queryservice/auth/unpwd/UnPwdAuthClient.java
@@ -47,10 +47,14 @@ public CoreToken un_pw_login(
                 separator, userNameSegment, separator, passwordSegment, passwordBytes
         );
         try {
+            log.info("Fetching the core token using UN Password Flow");
             RequestBody requestBody = RequestBody.create(body, MediaType.parse(Constants.URL_ENCODED_CONTENT));
             Map<String, String> headers = Collections.singletonMap(Constants.CONTENT_TYPE, Constants.URL_ENCODED_CONTENT);
             Request request = HttpHelper.buildRequest(Constants.POST, tokenUrl, requestBody, headers);
+            long l = System.currentTimeMillis();
+            log.info("Starting request to get core token using username password flow");
             Response response = client.newCall(request).execute();
+            log.info("Finished request to get core token using username password flow in {}ms", System.currentTimeMillis() - l);
             if (!response.isSuccessful()) {
                 log.error("login with user credentials failed with status code {}", response.code());
                 HttpHelper.handleErrorResponse(response, Constants.ERROR_DESCRIPTION);

diff --git a/src/main/java/com/salesforce/cdp/queryservice/auth/unpwd/UnPwdTokenFlow.java b/src/main/java/com/salesforce/cdp/queryservice/auth/unpwd/UnPwdTokenFlow.java
@@ -30,8 +30,10 @@ public UnPwdTokenFlow(Properties properties, UnPwdAuthClient client, TokenExchan
 
     @Override
     public CoreToken getCoreToken() throws TokenException {
+        log.info("Getting core token using username password flow");
         validateProperties();
         String token_url = properties.getProperty(Constants.LOGIN_URL) + Constants.CORE_TOKEN_URL;
+        log.info("Fetching core token from url: {}", token_url);
         byte[] passwordBytes = null;
         byte[] clientSecret = null;
         // Convert the password and client secret to byte arrays so, we can empty them at will.
@@ -60,7 +62,8 @@ public CoreToken getCoreToken() throws TokenException {
 
     @Override
     public OffcoreToken getOffcoreToken() throws TokenException {
-        if (TokenUtils.isValid(offcoreToken)) return offcoreToken;
+        if (TokenUtils.isValid(offcoreToken))
+            return offcoreToken;
         CoreToken coreToken = getCoreToken();
         offcoreToken = tokenExchangeHelper.exchangeToken(coreToken);
         return offcoreToken;

diff --git a/src/test/java/com/salesforce/cdp/queryservice/auth/TokenExchangeHelperTest.java b/src/test/java/com/salesforce/cdp/queryservice/auth/TokenExchangeHelperTest.java
@@ -0,0 +1,45 @@
+package com.salesforce.cdp.queryservice.auth;
+
+import com.salesforce.cdp.queryservice.auth.token.CoreToken;
+import com.salesforce.cdp.queryservice.util.TokenException;
+import okhttp3.OkHttpClient;
+import okhttp3.Request;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
+
+import java.util.Properties;
+
+import static com.salesforce.cdp.queryservice.util.Constants.LOGIN_URL;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.*;
+
+public class TokenExchangeHelperTest {
+
+    @Test
+    public void testInstanceUrlIsUsedForExchange() throws TokenException {
+        OkHttpClient mockHttpClient = mock(OkHttpClient.class);
+        Properties mockProperties = mock(Properties.class);
+        when(mockProperties.getProperty(eq(LOGIN_URL))).thenReturn("https://login.salesforce.com");
+        TokenExchangeHelper tokenExchangeHelper = new TokenExchangeHelper(mockProperties, mockHttpClient);
+
+        CoreToken mockCoreToken = mock(CoreToken.class);
+        when(mockCoreToken.getInstanceUrl()).thenReturn("https://tenant-instance.com");
+
+        try {
+            tokenExchangeHelper.exchangeToken(mockCoreToken);
+        } catch (Exception e) {
+            // Not worried about the exception as we are checking only
+            // whether the instance url is used for exchange
+        }
+
+        verify(mockCoreToken, Mockito.times(1)).getInstanceUrl();
+        ArgumentCaptor<Request> captor = ArgumentCaptor.forClass(Request.class);
+        verify(mockHttpClient, Mockito.times(1)).newCall(captor.capture());
+
+        assertThat(captor.getValue().url().toString())
+                .isEqualTo("https://tenant-instance.com/services/a360/token");
+    }
+
+}