Azure identity second try

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

@@ -394,7 +394,7 @@ private long determineTimeout(ClientRequestProperties properties, CommandType co
 
     private String getAuthorizationHeaderValue() throws DataServiceException, DataClientException {
         if (aadAuthenticationHelper != null) {
-            return String.format("Bearer %s", aadAuthenticationHelper.acquireAccessToken());
+            return String.format("Bearer %s", aadAuthenticationHelper.acquireAccessToken().block());
         }
         return null;
     }

data/src/main/java/com/microsoft/azure/kusto/data/auth/AccessTokenTokenProvider.java

@@ -6,9 +6,9 @@
 import java.net.URISyntaxException;
 
 import org.jetbrains.annotations.NotNull;
+import reactor.core.publisher.Mono;
 
 public class AccessTokenTokenProvider extends TokenProviderBase {
-    public static final String ACCESS_TOKEN_TOKEN_PROVIDER = "AccessTokenTokenProvider";
     private final String accessToken;
 
     AccessTokenTokenProvider(@NotNull String clusterUrl, @NotNull String accessToken) throws URISyntaxException {
@@ -17,12 +17,7 @@ public class AccessTokenTokenProvider extends TokenProviderBase {
     }
 
     @Override
-    protected String acquireAccessTokenImpl() {
-        return accessToken;
-    }
-
-    @Override
-    protected String getAuthMethod() {
-        return ACCESS_TOKEN_TOKEN_PROVIDER;
+    protected Mono<String> acquireAccessTokenImpl() {
+        return Mono.just(accessToken);
     }
 }

data/src/main/java/com/microsoft/azure/kusto/data/auth/ApplicationCertificateTokenProvider.java

@@ -13,8 +13,9 @@
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+// Azure identity doesn't provide a solution for all certificate types, so for now we still use MSAL for this.
+
 public class ApplicationCertificateTokenProvider extends ConfidentialAppTokenProviderBase {
-    public static final String APPLICATION_CERTIFICATE_TOKEN_PROVIDER = "ApplicationCertificateTokenProvider";
     private final IClientCertificate clientCertificate;
 
     ApplicationCertificateTokenProvider(@NotNull String clusterUrl, @NotNull String applicationClientId, @NotNull IClientCertificate clientCertificate,
@@ -34,9 +35,4 @@ protected IConfidentialClientApplication getClientApplication() throws Malformed
         return clientApplication = builder
                 .build();
     }
-
-    @Override
-    protected String getAuthMethod() {
-        return APPLICATION_CERTIFICATE_TOKEN_PROVIDER;
-    }
 }

data/src/main/java/com/microsoft/azure/kusto/data/auth/ApplicationKeyTokenProvider.java

@@ -3,38 +3,34 @@
 
 package com.microsoft.azure.kusto.data.auth;
 
-import com.microsoft.aad.msal4j.ConfidentialClientApplication;
-import com.microsoft.aad.msal4j.IClientSecret;
-import com.microsoft.aad.msal4j.IConfidentialClientApplication;
-import java.net.MalformedURLException;
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.http.HttpClient;
+import com.azure.identity.ClientSecretCredentialBuilder;
+import com.azure.identity.CredentialBuilderBase;
+
 import java.net.URISyntaxException;
 
-import com.azure.core.http.HttpClient;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-public class ApplicationKeyTokenProvider extends ConfidentialAppTokenProviderBase {
-    private final IClientSecret clientSecret;
-    public static final String APPLICATION_KEY_TOKEN_PROVIDER = "ApplicationKeyTokenProvider";
+public class ApplicationKeyTokenProvider extends AzureIdentityTokenProvider {
+    private final String clientSecret;
 
-    ApplicationKeyTokenProvider(@NotNull String clusterUrl, @NotNull String applicationClientId, @NotNull IClientSecret clientSecret,
-            String authorityId, @Nullable HttpClient httpClient) throws URISyntaxException {
-        super(clusterUrl, applicationClientId, authorityId, httpClient);
+    public ApplicationKeyTokenProvider(@NotNull String clusterUrl, String clientId, String clientSecret, String tenantId,
+            @Nullable HttpClient httpClient) throws URISyntaxException {
+        super(clusterUrl, tenantId, clientId, httpClient);
         this.clientSecret = clientSecret;
     }
 
     @Override
-    protected IConfidentialClientApplication getClientApplication() throws MalformedURLException {
-        ConfidentialClientApplication.Builder authority = ConfidentialClientApplication.builder(applicationClientId, clientSecret)
-                .authority(aadAuthorityUrl);
-        if (httpClient != null) {
-            authority.httpClient(new HttpClientWrapper(httpClient));
-        }
-        return authority.build();
+    protected TokenCredential createTokenCredential(CredentialBuilderBase<?> builder) {
+        return new ClientSecretCredentialBuilder()
+                .clientSecret(clientSecret)
+                .build();
     }
 
     @Override
-    protected String getAuthMethod() {
-        return APPLICATION_KEY_TOKEN_PROVIDER;
+    protected CredentialBuilderBase<?> initBuilder() {
+        return new ClientSecretCredentialBuilder();
     }
 }

data/src/main/java/com/microsoft/azure/kusto/data/auth/AsyncCallbackTokenProvider.java

@@ -0,0 +1,36 @@
+package com.microsoft.azure.kusto.data.auth;
+
+import com.microsoft.azure.kusto.data.exceptions.DataClientException;
+import com.microsoft.azure.kusto.data.exceptions.ExceptionsUtils;
+import org.jetbrains.annotations.NotNull;
+import reactor.core.publisher.Mono;
+
+import java.net.URISyntaxException;
+
+public class AsyncCallbackTokenProvider extends TokenProviderBase {
+    public static final String CALLBACK_TOKEN_PROVIDER = "CallbackTokenProvider";
+    private final Mono<String> tokenProvider;
+
+    AsyncCallbackTokenProvider(@NotNull String clusterUrl, @NotNull Mono<String> tokenProvider) throws URISyntaxException {
+        super(clusterUrl, null);
+        this.tokenProvider = tokenProvider;
+    }
+
+    @Override
+    protected Mono<String> acquireAccessTokenImpl() {
+        return tokenProvider
+                .onErrorMap(e -> {
+                    if (e instanceof Exception) {
+                        Exception ex = (Exception) e;
+                        return new DataClientException(clusterUrl, ExceptionsUtils.getMessageEx(ex), ex);
+                    } else {
+                        return new DataClientException(clusterUrl, e.toString(), null);
+                    }
+                });
+    }
+
+    @Override
+    protected String getAuthMethod() {
+        return CALLBACK_TOKEN_PROVIDER;
+    }
+}

data/src/main/java/com/microsoft/azure/kusto/data/auth/AzureCliTokenProvider.java

@@ -1,51 +1,26 @@
 package com.microsoft.azure.kusto.data.auth;
 
-import com.azure.core.credential.AccessToken;
-import com.azure.core.credential.TokenRequestContext;
+import com.azure.core.credential.TokenCredential;
 import com.azure.core.http.HttpClient;
-import com.azure.identity.AzureCliCredential;
 import com.azure.identity.AzureCliCredentialBuilder;
-import com.microsoft.azure.kusto.data.exceptions.DataClientException;
-import com.microsoft.azure.kusto.data.exceptions.DataServiceException;
+import com.azure.identity.CredentialBuilderBase;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.net.URISyntaxException;
 
-public class AzureCliTokenProvider extends CloudDependentTokenProviderBase {
-    public static final String AZURE_CLI_TOKEN_PROVIDER = "AzureCliTokenProvider";
-    private TokenRequestContext tokenRequestContext;
-    private AzureCliCredential azureCliCredential;
-
+public class AzureCliTokenProvider extends AzureIdentityTokenProvider {
     public AzureCliTokenProvider(@NotNull String clusterUrl, @Nullable HttpClient httpClient) throws URISyntaxException {
         super(clusterUrl, httpClient);
-
-    }
-
-    @Override
-    protected void initializeWithCloudInfo(CloudInfo cloudInfo) throws DataServiceException, DataClientException {
-        super.initializeWithCloudInfo(cloudInfo);
-        AzureCliCredentialBuilder builder = new AzureCliCredentialBuilder();
-
-        if (httpClient != null) {
-            builder = builder.httpClient(new HttpClientWrapper(httpClient));
-        }
-
-        this.azureCliCredential = builder.build();
-        tokenRequestContext = new TokenRequestContext().addScopes(scopes.toArray(new String[0]));
     }
 
     @Override
-    protected String acquireAccessTokenImpl() throws DataServiceException {
-        AccessToken accessToken = azureCliCredential.getToken(tokenRequestContext).block();
-        if (accessToken == null) {
-            throw new DataServiceException(clusterUrl, "Couldn't get token from Azure Identity", true);
-        }
-        return accessToken.getToken();
+    protected CredentialBuilderBase<?> initBuilder() {
+        return new AzureCliCredentialBuilder();
     }
 
     @Override
-    protected String getAuthMethod() {
-        return AZURE_CLI_TOKEN_PROVIDER;
+    protected TokenCredential createTokenCredential(CredentialBuilderBase<?> builder) {
+        return ((AzureCliCredentialBuilder) builder).build();
     }
 }

data/src/main/java/com/microsoft/azure/kusto/data/auth/AzureIdentityTokenProvider.java

@@ -0,0 +1,87 @@
+package com.microsoft.azure.kusto.data.auth;
+
+import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRequestContext;
+import com.azure.core.http.HttpClient;
+import com.azure.identity.AadCredentialBuilderBase;
+import com.azure.identity.CredentialBuilderBase;
+import com.microsoft.azure.kusto.data.UriUtils;
+import com.microsoft.azure.kusto.data.exceptions.DataClientException;
+import com.microsoft.azure.kusto.data.exceptions.DataServiceException;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import reactor.core.publisher.Mono;
+
+import java.net.URISyntaxException;
+
+public abstract class AzureIdentityTokenProvider extends CloudDependentTokenProviderBase {
+    private String clientId;
+    private final String tenantId;
+    TokenCredential cred;
+    TokenRequestContext tokenRequestContext;
+
+    protected static final String ORGANIZATION_URI_SUFFIX = "organizations";
+    protected static final String ERROR_INVALID_AUTHORITY_URL = "Error acquiring ApplicationAccessToken due to invalid Authority URL";
+
+    private String determineAadAuthorityUrl(CloudInfo cloudInfo) throws DataClientException {
+        String aadAuthorityUrlFromEnv = System.getenv("AadAuthorityUri");
+        String authorityIdToUse = tenantId != null ? tenantId : ORGANIZATION_URI_SUFFIX;
+        try {
+            return UriUtils.setPathForUri(aadAuthorityUrlFromEnv == null ? cloudInfo.getLoginEndpoint() : aadAuthorityUrlFromEnv, authorityIdToUse, true);
+        } catch (URISyntaxException e) {
+            throw new DataClientException(clusterUrl, ERROR_INVALID_AUTHORITY_URL, e);
+        }
+    }
+
+    AzureIdentityTokenProvider(@NotNull String clusterUrl, @Nullable String tenantId, @Nullable String clientId,
+            @Nullable HttpClient httpClient) throws URISyntaxException {
+        super(clusterUrl, httpClient);
+        this.tenantId = tenantId;
+        this.clientId = clientId;
+    }
+
+    AzureIdentityTokenProvider(@NotNull String clusterUrl, @Nullable HttpClient httpClient) throws URISyntaxException {
+        this(clusterUrl, null, null, httpClient);
+    }
+
+    @Override
+    protected final Mono<String> acquireAccessTokenImpl() {
+        return cred.getToken(tokenRequestContext).map(AccessToken::getToken);
+    }
+
+    @Override
+    protected final void initializeWithCloudInfo(CloudInfo cloudInfo) throws DataClientException, DataServiceException {
+        super.initializeWithCloudInfo(cloudInfo);
+        CredentialBuilderBase<?> builder = initBuilder();
+        if (httpClient != null) {
+            builder.httpClient(new HttpClientWrapper(httpClient));
+        }
+        if (builder instanceof AadCredentialBuilderBase<?>) {
+            AadCredentialBuilderBase<?> aadBuilder = (AadCredentialBuilderBase<?>) builder;
+            if (tenantId != null)
+                aadBuilder.tenantId(tenantId);
+
+            aadBuilder.authorityHost(determineAadAuthorityUrl(cloudInfo));
+
+            if (clientId == null) {
+                clientId = cloudInfo.getKustoClientAppId();
+            }
+
+            aadBuilder.clientId(clientId);
+
+            // TODO: do we need to port this code?
+            // if (account.homeAccountId() != null && account.homeAccountId().endsWith(PERSONAL_TENANT_IDV2_AAD)) {
+            // authorityUrl = firstPartyAuthorityUrl;
+            // }
+
+        }
+
+        cred = createTokenCredential(builder);
+        tokenRequestContext = new TokenRequestContext().addScopes(scopes.toArray(new String[0]));
+    }
+
+    protected abstract CredentialBuilderBase<?> initBuilder();
+
+    protected abstract TokenCredential createTokenCredential(CredentialBuilderBase<?> builder);
+}

data/src/main/java/com/microsoft/azure/kusto/data/auth/CallbackTokenProvider.java

@@ -9,6 +9,7 @@
 import com.microsoft.azure.kusto.data.exceptions.ExceptionsUtils;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import reactor.core.publisher.Mono;
 
 import java.net.URISyntaxException;
 import java.util.concurrent.Callable;
@@ -29,12 +30,17 @@ public class CallbackTokenProvider extends TokenProviderBase {
     }
 
     @Override
-    protected String acquireAccessTokenImpl() throws DataClientException {
-        try {
-            return tokenProvider.apply(httpClient);
-        } catch (Exception e) {
-            throw new DataClientException(clusterUrl, ExceptionsUtils.getMessageEx(e), e);
-        }
+    protected Mono<String> acquireAccessTokenImpl() {
+        return Mono.fromCallable(() -> tokenProvider.apply(httpClient))
+                // TODO - is this a better way?
+                .onErrorMap(e -> {
+                    if (e instanceof Exception) {
+                        Exception ex = (Exception) e;
+                        return new DataClientException(clusterUrl, ExceptionsUtils.getMessageEx(ex), ex);
+                    } else {
+                        return new DataClientException(clusterUrl, e.toString(), null);
+                    }
+                });
     }
 
     @Override

data/src/main/java/com/microsoft/azure/kusto/data/auth/CloudDependentTokenProviderBase.java

@@ -16,6 +16,7 @@
 
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import reactor.core.publisher.Mono;
 
 public abstract class CloudDependentTokenProviderBase extends TokenProviderBase {
     private static final String ERROR_INVALID_SERVICE_RESOURCE_URL = "Error determining scope due to invalid Kusto Service Resource URL";
@@ -28,17 +29,29 @@ public abstract class CloudDependentTokenProviderBase extends TokenProviderBase
     }
 
     @Override
-    synchronized void initialize() throws DataClientException, DataServiceException {
-        if (initialized) {
-            return;
-        }
-        // trace retrieveCloudInfo
-        cloudInfo = MonitoredActivity.invoke(
-                (SupplierTwoExceptions<CloudInfo, DataClientException, DataServiceException>) () -> CloudInfo.retrieveCloudInfoForCluster(clusterUrl,
-                        httpClient),
-                "CloudDependentTokenProviderBase.retrieveCloudInfo", getTracingAttributes());
-        initializeWithCloudInfo(cloudInfo);
-        initialized = true;
+    final Mono<Void> initialize() {
+        return Mono.fromCallable(() -> initialized)
+                .flatMap(initialized -> {
+                    if (initialized) {
+                        return Mono.empty();
+                    }
+                    // trace retrieveCloudInfo
+                    return MonitoredActivity.wrap(
+                            CloudInfo.retrieveCloudInfoForClusterAsync(clusterUrl,
+                                    httpClient),
+                            "CloudDependentTokenProviderBase.retrieveCloudInfo", getTracingAttributes())
+                            .flatMap(cloudInfo -> {
+                                this.cloudInfo = cloudInfo;
+
+                                try {
+                                    initializeWithCloudInfo(cloudInfo);
+                                } catch (DataClientException | DataServiceException e) {
+                                    return Mono.error(e);
+                                }
+                                this.initialized = true;
+                                return Mono.empty();
+                            });
+                });
     }
 
     protected void initializeWithCloudInfo(CloudInfo cloudInfo) throws DataClientException, DataServiceException {