From 112c8fee754df9be8cabb6095a6450e29b8b1203 Mon Sep 17 00:00:00 2001
From: cole <cole.snyder@gmork.tech>
Date: Mon, 28 Oct 2024 19:31:37 -0400
Subject: [PATCH] allows users to provide their own TokenCredential auth from
 azure-identity

---
 .../azure/kusto/data/ClientImpl.java          | 20 ++++++---
 .../data/auth/AzureIdentityTokenProvider.java | 45 +++++++++++++++++++
 2 files changed, 60 insertions(+), 5 deletions(-)
 create mode 100644 data/src/main/java/com/microsoft/azure/kusto/data/auth/AzureIdentityTokenProvider.java

diff --git a/data/src/main/java/com/microsoft/azure/kusto/data/ClientImpl.java b/data/src/main/java/com/microsoft/azure/kusto/data/ClientImpl.java
index 1ef0e39e..a8fc5567 100644
--- a/data/src/main/java/com/microsoft/azure/kusto/data/ClientImpl.java
+++ b/data/src/main/java/com/microsoft/azure/kusto/data/ClientImpl.java
@@ -3,13 +3,11 @@
 
 package com.microsoft.azure.kusto.data;
 
+import com.azure.core.credential.TokenCredential;
 import com.azure.core.http.HttpClient;
 import com.azure.core.http.HttpRequest;
 import com.azure.core.util.BinaryData;
-import com.microsoft.azure.kusto.data.auth.CloudInfo;
-import com.microsoft.azure.kusto.data.auth.ConnectionStringBuilder;
-import com.microsoft.azure.kusto.data.auth.TokenProviderBase;
-import com.microsoft.azure.kusto.data.auth.TokenProviderFactory;
+import com.microsoft.azure.kusto.data.auth.*;
 import com.microsoft.azure.kusto.data.auth.endpoints.KustoTrustedEndpoints;
 import com.microsoft.azure.kusto.data.exceptions.*;
 import com.microsoft.azure.kusto.data.http.*;
@@ -47,12 +45,24 @@ public ClientImpl(ConnectionStringBuilder csb, HttpClientProperties properties)
     }
 
     public ClientImpl(ConnectionStringBuilder csb, HttpClient httpClient) throws URISyntaxException {
+        this(csb, httpClient, null);
+    }
+
+    public ClientImpl(ConnectionStringBuilder csb, HttpClient httpClient, TokenCredential credentials) throws URISyntaxException {
         super(httpClient);
         String clusterURL = UriUtils.createClusterURLFrom(csb.getClusterUrl());
         csb.setClusterUrl(clusterURL);
 
         clusterUrl = csb.getClusterUrl();
-        aadAuthenticationHelper = clusterUrl.toLowerCase().startsWith(CloudInfo.LOCALHOST) ? null : TokenProviderFactory.createTokenProvider(csb, httpClient);
+
+        // Allow the user to provide their own TokenCredentials from Azure Identity
+        if (credentials != null) {
+            aadAuthenticationHelper = new AzureIdentityTokenProvider(clusterURL, httpClient, credentials);
+        } else {
+            aadAuthenticationHelper = clusterUrl.toLowerCase().startsWith(CloudInfo.LOCALHOST) ? null :
+                    TokenProviderFactory.createTokenProvider(csb, httpClient);
+        }
+
         clientDetails = new ClientDetails(csb.getApplicationNameForTracing(), csb.getUserNameForTracing(), csb.getClientVersionForTracing());
     }
 
diff --git a/data/src/main/java/com/microsoft/azure/kusto/data/auth/AzureIdentityTokenProvider.java b/data/src/main/java/com/microsoft/azure/kusto/data/auth/AzureIdentityTokenProvider.java
new file mode 100644
index 00000000..ace9f4b5
--- /dev/null
+++ b/data/src/main/java/com/microsoft/azure/kusto/data/auth/AzureIdentityTokenProvider.java
@@ -0,0 +1,45 @@
+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.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 java.net.URISyntaxException;
+
+public class AzureIdentityTokenProvider extends CloudDependentTokenProviderBase {
+    public static final String AZURE_ID_TOKEN_PROVIDER = "AzureIdentityTokenProvider";
+    private TokenRequestContext tokenRequestContext;
+    private final TokenCredential tokenCredential;
+
+    public AzureIdentityTokenProvider(@NotNull String clusterUrl, @Nullable HttpClient httpClient, @NotNull TokenCredential credentials) throws URISyntaxException {
+        super(clusterUrl, httpClient);
+        tokenCredential = credentials;
+    }
+
+    @Override
+    protected void initializeWithCloudInfo(CloudInfo cloudInfo) throws DataServiceException, DataClientException {
+        super.initializeWithCloudInfo(cloudInfo);
+        tokenRequestContext = new TokenRequestContext().addScopes(scopes.toArray(new String[0]));
+    }
+
+    @Override
+    protected String acquireAccessTokenImpl() throws DataServiceException {
+        // Fixme: Make this async
+        AccessToken accessToken = tokenCredential.getToken(tokenRequestContext).block();
+        if (accessToken == null) {
+            throw new DataServiceException(clusterUrl, "Couldn't get token from Azure Identity", true);
+        }
+        return accessToken.getToken();
+    }
+
+    @Override
+    protected String getAuthMethod() {
+        return AZURE_ID_TOKEN_PROVIDER;
+    }
+}