cache-lastAuthenticationResult

data/src/main/java/com/microsoft/azure/kusto/data/AadAuthenticationHelper.java

@@ -14,12 +14,16 @@
 import java.net.URISyntaxException;
 import java.security.PrivateKey;
 import java.security.cert.X509Certificate;
+import java.util.Date;
 import java.util.concurrent.*;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 
 public class AadAuthenticationHelper {
 
     private final static String DEFAULT_AAD_TENANT = "common";
     private final static String CLIENT_ID = "db662dc1-0cfe-4e1c-a843-19a68e65be58";
+    public final static long ONE_MINUTE_IN_MILLIS = 60000;
 
     private ClientCredential clientCredential;
     private String userUsername;
@@ -29,6 +33,11 @@ public class AadAuthenticationHelper {
     private X509Certificate x509Certificate;
     private PrivateKey privateKey;
     private AuthenticationType authenticationType;
+    private AuthenticationResult lastAuthenticationResult;
+    private ExecutorService service = Executors.newSingleThreadExecutor();
+    private AuthenticationContext context;
+    private Lock lastAuthenticationResultLock = new ReentrantLock();
+    private String applicationClientId;
 
     private enum AuthenticationType {
         AAD_USERNAME_PASSWORD,
@@ -50,7 +59,7 @@ public AadAuthenticationHelper(@NotNull ConnectionStringBuilder csb) throws URIS
         } else if (csb.getX509Certificate() != null && csb.getPrivateKey() != null) {
             x509Certificate = csb.getX509Certificate();
             privateKey = csb.getPrivateKey();
-            clientCredential = new ClientCredential(csb.getApplicationClientId(), null);
+            applicationClientId = csb.getApplicationClientId();
             authenticationType = AuthenticationType.AAD_APPLICATION_CERTIFICATE;
         } else {
             authenticationType = AuthenticationType.AAD_DEVICE_LOGIN;
@@ -61,44 +70,39 @@ public AadAuthenticationHelper(@NotNull ConnectionStringBuilder csb) throws URIS
         aadAuthorityUri = String.format("https://login.microsoftonline.com/%s", aadAuthorityId);
     }
 
-    String acquireAccessToken() throws DataServiceException  {
-        try {
-            switch (authenticationType) {
-                case AAD_APPLICATION_KEY:
-                    return acquireAadApplicationAccessToken().getAccessToken();
-                case AAD_USERNAME_PASSWORD:
-                    return acquireAadUserAccessToken().getAccessToken();
-                case AAD_DEVICE_LOGIN:
-                    return acquireAccessTokenUsingDeviceCodeFlow().getAccessToken();
-                case AAD_APPLICATION_CERTIFICATE:
-                    return acquireWithClientCertificate().getAccessToken();
-                default:
-                    throw new DataServiceException("Authentication type: " + authenticationType.name() + " is invalid");
+    String acquireAccessToken() throws DataServiceException {
+        if (lastAuthenticationResult == null) {
+            aquireTokenByType();
+        } else {
+            if (isTokenExpired()) {
+                if (lastAuthenticationResult.getRefreshToken() == null) {
+                    aquireTokenByType();
+                } else {
+                    lastAuthenticationResultLock.lock();
+                    if (isTokenExpired()) {
+                        lastAuthenticationResult = refreshToken();
+                    }
+                    lastAuthenticationResultLock.unlock();
+                }
             }
-        } catch (Exception e) {
-            throw new DataServiceException(e.getMessage());
         }
 
+        return lastAuthenticationResult.getAccessToken();
     }
 
-    private AuthenticationResult acquireAadUserAccessToken() throws DataServiceException, DataClientException {
-        AuthenticationContext context;
-        AuthenticationResult result;
-        ExecutorService service = null;
-        try {
-            service = Executors.newSingleThreadExecutor();
+    private AuthenticationResult acquireAadUserAccessToken() throws DataServiceException, DataClientException, MalformedURLException {
+        if (context == null) {
             context = new AuthenticationContext(aadAuthorityUri, true, service);
+        }
 
+        AuthenticationResult result;
+        try {
             Future<AuthenticationResult> future = context.acquireToken(
                     clusterUrl, CLIENT_ID, userUsername, userPassword,
                     null);
             result = future.get();
-        } catch (InterruptedException | ExecutionException | MalformedURLException e) {
+        } catch (InterruptedException | ExecutionException e) {
             throw new DataClientException(clusterUrl, "Error in acquiring UserAccessToken", e);
-        } finally {
-            if (service != null) {
-                service.shutdown();
-            }
         }
 
         if (result == null) {
@@ -107,21 +111,17 @@ private AuthenticationResult acquireAadUserAccessToken() throws DataServiceExcep
         return result;
     }
 
-    private AuthenticationResult acquireAadApplicationAccessToken() throws DataServiceException, DataClientException {
-        AuthenticationContext context;
+    private AuthenticationResult acquireAadApplicationAccessToken() throws DataServiceException, DataClientException, MalformedURLException {
+        if (context == null) {
+            context = new AuthenticationContext(aadAuthorityUri, true, service);
+        }
+
         AuthenticationResult result;
-        ExecutorService service = null;
         try {
-            service = Executors.newSingleThreadExecutor();
-            context = new AuthenticationContext(aadAuthorityUri, true, service);
             Future<AuthenticationResult> future = context.acquireToken(clusterUrl, clientCredential, null);
             result = future.get();
-        } catch (InterruptedException | ExecutionException | MalformedURLException e) {
+        } catch (InterruptedException | ExecutionException e) {
             throw new DataClientException(clusterUrl, "Error in acquiring ApplicationAccessToken", e);
-        } finally {
-            if (service != null) {
-                service.shutdown();
-            }
         }
 
         if (result == null) {
@@ -131,71 +131,108 @@ private AuthenticationResult acquireAadApplicationAccessToken() throws DataServi
     }
 
     private AuthenticationResult acquireAccessTokenUsingDeviceCodeFlow() throws Exception {
-        AuthenticationContext context = null;
-        AuthenticationResult result = null;
-        ExecutorService service = null;
-        try {
-            service = Executors.newSingleThreadExecutor();
-            context = new AuthenticationContext( aadAuthorityUri, true, service);
-            Future<DeviceCode> future = context.acquireDeviceCode(CLIENT_ID, clusterUrl, null);
-            DeviceCode deviceCode = future.get();
-            System.out.println(deviceCode.getMessage());
-            if (Desktop.isDesktopSupported()) {
-                Desktop.getDesktop().browse(new URI(deviceCode.getVerificationUrl()));
-            }
-            result = waitAndAcquireTokenByDeviceCode(deviceCode, context);
-
+        if (context == null) {
+            context = new AuthenticationContext(aadAuthorityUri, true, service);
+        }
 
-        } finally {
-            if (service != null) {
-                service.shutdown();
-            }
+        AuthenticationResult result;
+        Future<DeviceCode> future = context.acquireDeviceCode(CLIENT_ID, clusterUrl, null);
+        DeviceCode deviceCode = future.get();
+        System.out.println(deviceCode.getMessage() + " " + Thread.currentThread());
+        if (Desktop.isDesktopSupported()) {
+            Desktop.getDesktop().browse(new URI(deviceCode.getVerificationUrl()));
         }
+        result = waitAndAcquireTokenByDeviceCode(deviceCode);
+
         if (result == null) {
             throw new ServiceUnavailableException("authentication result was null");
         }
         return result;
     }
 
-    private AuthenticationResult waitAndAcquireTokenByDeviceCode(DeviceCode deviceCode, AuthenticationContext context)
-            throws  InterruptedException{
+    private AuthenticationResult waitAndAcquireTokenByDeviceCode(DeviceCode deviceCode)
+            throws InterruptedException {
         int timeout = 15 * 1000;
         AuthenticationResult result = null;
-        while (timeout > 0){
-            try{
+        while (timeout > 0) {
+            try {
                 Future<AuthenticationResult> futureResult = context.acquireTokenByDeviceCode(deviceCode, null);
                 return futureResult.get();
             } catch (ExecutionException e) {
-                    Thread.sleep(1000);
-                    timeout -= 1000;
+                Thread.sleep(1000);
+                timeout -= 1000;
             }
         }
         return result;
     }
 
     AuthenticationResult acquireWithClientCertificate()
-            throws IOException, InterruptedException, ExecutionException, ServiceUnavailableException{
-
-        AuthenticationContext context;
+            throws IOException, InterruptedException, ExecutionException, ServiceUnavailableException {
+        if (context == null) {
+            context = new AuthenticationContext(aadAuthorityUri, true, service);
+        }
         AuthenticationResult result;
-        ExecutorService service = null;
 
-        try {
-            service = Executors.newSingleThreadExecutor();
-            context = new AuthenticationContext(aadAuthorityUri, false, service);
-            AsymmetricKeyCredential asymmetricKeyCredential = AsymmetricKeyCredential.create(clientCredential.getClientId(),
-                    privateKey, x509Certificate);
-            // pass null value for optional callback function and acquire access token
-            result = context.acquireToken(clusterUrl, asymmetricKeyCredential, null).get();
-        } finally {
-            if (service != null) {
-                service.shutdown();
-            }
-        }
+        AsymmetricKeyCredential asymmetricKeyCredential = AsymmetricKeyCredential.create(applicationClientId,
+                privateKey, x509Certificate);
+        // pass null value for optional callback function and acquire access token
+        result = context.acquireToken(clusterUrl, asymmetricKeyCredential, null).get();
+
         if (result == null) {
             throw new ServiceUnavailableException("authentication result was null");
         }
         return result;
     }
 
+    private void aquireTokenByType() throws DataServiceException {
+        lastAuthenticationResultLock.lock();
+        if(lastAuthenticationResult == null || isTokenExpired()) {
+            try {
+                switch (authenticationType) {
+                    case AAD_APPLICATION_KEY:
+                        lastAuthenticationResult = acquireAadApplicationAccessToken();
+                        break;
+                    case AAD_USERNAME_PASSWORD:
+                        lastAuthenticationResult = acquireAadUserAccessToken();
+                        break;
+                    case AAD_DEVICE_LOGIN:
+                        lastAuthenticationResult = acquireAccessTokenUsingDeviceCodeFlow();
+                        break;
+                    case AAD_APPLICATION_CERTIFICATE:
+                        lastAuthenticationResult = acquireWithClientCertificate();
+                        break;
+                    default:
+                        throw new DataServiceException("Authentication type: " + authenticationType.name() + " is invalid");
+                }
+            } catch (Exception e) {
+                throw new DataServiceException(e.getMessage());
+            }
+        }
+        lastAuthenticationResultLock.unlock();
+    }
+
+    private boolean isTokenExpired(){
+        return lastAuthenticationResult.getExpiresOnDate().before(dateInAMinute());
+    }
+
+    AuthenticationResult refreshToken() throws DataServiceException {
+        try {
+            switch (authenticationType) {
+                case AAD_APPLICATION_KEY:
+                case AAD_APPLICATION_CERTIFICATE:
+                    return context.acquireTokenByRefreshToken(lastAuthenticationResult.getRefreshToken(), clientCredential, null).get();
+                case AAD_USERNAME_PASSWORD:
+                case AAD_DEVICE_LOGIN:
+                    return context.acquireTokenByRefreshToken(lastAuthenticationResult.getRefreshToken(), CLIENT_ID, clusterUrl, null).get();
+                default:
+                    throw new DataServiceException("Authentication type: " + authenticationType.name() + " is invalid");
+            }
+        } catch (Exception e) {
+            throw new DataServiceException(e.getMessage());
+        }
+    }
+
+    Date dateInAMinute() {
+        return new Date(System.currentTimeMillis() + ONE_MINUTE_IN_MILLIS);
+    }
 }

data/src/test/java/com/microsoft/azure/kusto/data/AadAuthenticationHelperTest.java

@@ -1,6 +1,9 @@
 package com.microsoft.azure.kusto.data;
 
 
+import com.microsoft.aad.adal4j.AuthenticationResult;
+import com.microsoft.aad.adal4j.UserInfo;
+import com.microsoft.azure.kusto.data.exceptions.DataServiceException;
 import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
 import org.bouncycastle.cert.X509CertificateHolder;
 import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
@@ -25,10 +28,21 @@
 import java.security.Security;
 import java.security.cert.CertificateException;
 import java.security.cert.X509Certificate;
+import java.util.Date;
 import java.util.concurrent.ExecutionException;
 
+import javax.naming.ServiceUnavailableException;
+
+import static com.microsoft.azure.kusto.data.AadAuthenticationHelper.ONE_MINUTE_IN_MILLIS;
+import static org.mockito.Mockito.*;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
+
 public class AadAuthenticationHelperTest {
 
+
+
     @Test
     @DisplayName("validate auth with certificate throws exception when missing or invalid parameters")
     void acquireWithClientCertificateNullKey() throws CertificateException, OperatorCreationException,
@@ -47,7 +61,6 @@ void acquireWithClientCertificateNullKey() throws CertificateException, Operator
 
         Assertions.assertThrows(ExecutionException.class,
                 () -> aadAuthenticationHelper.acquireWithClientCertificate());
-
     }
 
     static KeyCert readPem(String path, String password)
@@ -83,4 +96,37 @@ static KeyCert readPem(String path, String password)
         return keycert;
     }
 
+    @Test
+    @DisplayName("validate cached token. Refresh if needed. Call regularly if no refresh token")
+    void useCachedTokenAndRefreshWhenNeeded() throws InterruptedException, ExecutionException, ServiceUnavailableException, IOException, DataServiceException, URISyntaxException, CertificateException, OperatorCreationException, PKCSException {
+        String certFilePath = Paths.get("src","test","resources", "cert.cer").toString();
+        String privateKeyPath = Paths.get("src","test","resources","key.pem").toString();
+
+        X509Certificate x509Certificate = readPem(certFilePath, "basic").getCertificate();
+        PrivateKey privateKey = readPem(privateKeyPath, "basic").getKey();
+
+        ConnectionStringBuilder csb = ConnectionStringBuilder
+                .createWithAadApplicationCertificate("resource.uri", "client-id", x509Certificate, privateKey);
+
+        AadAuthenticationHelper aadAuthenticationHelperSpy = spy(new AadAuthenticationHelper(csb));
+
+        AuthenticationResult authenticationResult = new AuthenticationResult("testType", "firstToken","refreshToken",0,"id", mock(UserInfo.class),false);
+        AuthenticationResult authenticationResultFromRefresh = new AuthenticationResult("testType", "fromRefresh",null,90,"id", mock(UserInfo.class),false);
+        AuthenticationResult authenticationResultNullRefreshTokenResult = new AuthenticationResult("testType", "nullRefreshResult",null,0,"id", mock(UserInfo.class),false);
+
+        doReturn(authenticationResultFromRefresh).when(aadAuthenticationHelperSpy).refreshToken();
+        doReturn(authenticationResult).when(aadAuthenticationHelperSpy).acquireWithClientCertificate();
+
+        assertEquals("firstToken", aadAuthenticationHelperSpy.acquireAccessToken());
+        assertEquals("fromRefresh", aadAuthenticationHelperSpy.acquireAccessToken());
+        assertEquals("fromRefresh", aadAuthenticationHelperSpy.acquireAccessToken());
+
+        // If no refresh token - it will authenticate again,
+        doReturn(new Date(System.currentTimeMillis() + ONE_MINUTE_IN_MILLIS * 2)).when(aadAuthenticationHelperSpy).dateInAMinute();
+        doReturn(authenticationResultNullRefreshTokenResult).when(aadAuthenticationHelperSpy).acquireWithClientCertificate();
+
+        assertEquals("nullRefreshResult", aadAuthenticationHelperSpy.acquireAccessToken());
+        // Null refreshToken
+
+    }
 }