Azure identity second try

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

@@ -8,6 +8,7 @@
 import com.microsoft.azure.kusto.data.http.CloseParentResourcesStream;
 import com.microsoft.azure.kusto.data.http.HttpRequestBuilder;
 import com.microsoft.azure.kusto.data.http.HttpStatus;
+import com.microsoft.azure.kusto.data.req.RequestUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.http.conn.EofSensorInputStream;
 import org.slf4j.Logger;
@@ -23,6 +24,7 @@
 
 public abstract class BaseClient implements Client, StreamingClient {
 
+    // TODO - this is never used?
     private static final int MAX_REDIRECT_COUNT = 1;
     private static final int EXTRA_TIMEOUT_FOR_CLIENT_SIDE = (int) TimeUnit.SECONDS.toMillis(30);
 
@@ -152,8 +154,7 @@ public static DataServiceException createExceptionFromResponse(String url, HttpR
     private static Context getContextTimeout(long timeoutMs) {
         int requestTimeout = timeoutMs > Integer.MAX_VALUE ? Integer.MAX_VALUE : Math.toIntExact(timeoutMs) + EXTRA_TIMEOUT_FOR_CLIENT_SIDE;
 
-        // See https://github.com/Azure/azure-sdk-for-java/blob/azure-core-http-netty_1.10.2/sdk/core/azure-core-http-netty/CHANGELOG.md#features-added
-        return Context.NONE.addData("azure-response-timeout", Duration.ofMillis(requestTimeout));
+        return RequestUtils.contextWithTimeout(Duration.ofMillis(requestTimeout));
     }
 
     private static void closeResourcesIfNeeded(boolean returnInputStream, HttpResponse httpResponse) {

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

@@ -397,7 +397,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 ((ClientSecretCredentialBuilder)builder)
+                .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,27 @@
+package com.microsoft.azure.kusto.data.auth;
+
+import com.microsoft.azure.kusto.data.exceptions.DataClientException;
+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 -> DataClientException.unwrapThrowable(clusterUrl, e));
+    }
+
+    @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,82 @@
+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(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);
+
+        }
+
+        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

@@ -3,38 +3,26 @@
 
 package com.microsoft.azure.kusto.data.auth;
 
-import com.azure.core.http.HttpClient;
 import com.microsoft.azure.kusto.data.exceptions.DataClientException;
 
-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;
 
 public class CallbackTokenProvider extends TokenProviderBase {
     public static final String CALLBACK_TOKEN_PROVIDER = "CallbackTokenProvider";
-    private final CallbackTokenProviderFunction tokenProvider;
+    private final @NotNull Callable<String> tokenProvider;
 
     CallbackTokenProvider(@NotNull String clusterUrl, @NotNull Callable<String> tokenProvider) throws URISyntaxException {
         super(clusterUrl, null);
-        this.tokenProvider = (httpClient) -> tokenProvider.call();
-    }
-
-    CallbackTokenProvider(@NotNull String clusterUrl, @NotNull CallbackTokenProviderFunction tokenProvider,
-            @Nullable HttpClient httpClient) throws URISyntaxException {
-        super(clusterUrl, httpClient);
         this.tokenProvider = tokenProvider;
     }
 
     @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).onErrorMap(e -> DataClientException.unwrapThrowable(clusterUrl, e));
     }
 
     @Override