HADOOP-19660: Add support for custom ClientAssertionProvider in WorkloadIdentityTokenProvider (#7901)

Contributed by Kunal Sevkani.
Reviewed by Anmol Asrani

Author: kunalmnnit

hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/AbfsConfiguration.java

@@ -55,6 +55,7 @@
 import org.apache.hadoop.fs.azurebfs.extensions.EncryptionContextProvider;
 import org.apache.hadoop.fs.azurebfs.extensions.SASTokenProvider;
 import org.apache.hadoop.fs.azurebfs.oauth2.AccessTokenProvider;
+import org.apache.hadoop.fs.azurebfs.oauth2.ClientAssertionProvider;
 import org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider;
 import org.apache.hadoop.fs.azurebfs.oauth2.CustomTokenProviderAdapter;
 import org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider;
@@ -1330,12 +1331,38 @@ public AccessTokenProvider getTokenProvider() throws TokenAccessProviderExceptio
               getMandatoryPasswordString(FS_AZURE_ACCOUNT_OAUTH_MSI_TENANT);
           String clientId =
               getMandatoryPasswordString(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ID);
-          String tokenFile =
-              getTrimmedPasswordString(FS_AZURE_ACCOUNT_OAUTH_TOKEN_FILE,
-              AuthConfigurations.DEFAULT_FS_AZURE_ACCOUNT_OAUTH_TOKEN_FILE);
-          tokenProvider = new WorkloadIdentityTokenProvider(
-              authority, tenantGuid, clientId, tokenFile);
-          LOG.trace("WorkloadIdentityTokenProvider initialized");
+
+          // Check if a custom ClientAssertionProvider is configured
+          String clientAssertionProviderType =
+              getPasswordString(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE);
+
+          if (clientAssertionProviderType != null && !clientAssertionProviderType.trim().isEmpty()) {
+            // Use custom ClientAssertionProvider
+            try {
+              Class<?> providerClass = Class.forName(clientAssertionProviderType.trim());
+              ClientAssertionProvider clientAssertionProvider =
+                  (ClientAssertionProvider) providerClass.getDeclaredConstructor().newInstance();
+
+              // Initialize the provider with configuration
+              clientAssertionProvider.initialize(rawConfig, accountName);
+
+              tokenProvider = new WorkloadIdentityTokenProvider(
+                  authority, tenantGuid, clientId, clientAssertionProvider);
+              LOG.trace("WorkloadIdentityTokenProvider initialized with custom ClientAssertionProvider: {}",
+                  clientAssertionProviderType);
+            } catch (Exception e) {
+              throw new TokenAccessProviderException(
+                  "Failed to initialize custom ClientAssertionProvider: " + clientAssertionProviderType, e);
+            }
+          } else {
+            // Use file-based approach (backward compatibility)
+            String tokenFile =
+                getTrimmedPasswordString(FS_AZURE_ACCOUNT_OAUTH_TOKEN_FILE,
+                AuthConfigurations.DEFAULT_FS_AZURE_ACCOUNT_OAUTH_TOKEN_FILE);
+            tokenProvider = new WorkloadIdentityTokenProvider(
+                authority, tenantGuid, clientId, tokenFile);
+            LOG.trace("WorkloadIdentityTokenProvider initialized with file-based token");
+          }
         } else {
           throw new IllegalArgumentException("Failed to initialize " + tokenProviderClass);
         }

hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/constants/ConfigurationKeys.java

@@ -339,6 +339,8 @@ public final class ConfigurationKeys {
   public static final String FS_AZURE_ACCOUNT_OAUTH_REFRESH_TOKEN_ENDPOINT = "fs.azure.account.oauth2.refresh.token.endpoint";
   /** Key for oauth AAD workload identity token file path: {@value}. */
   public static final String FS_AZURE_ACCOUNT_OAUTH_TOKEN_FILE = "fs.azure.account.oauth2.token.file";
+  /** Key for custom client assertion provider class for WorkloadIdentityTokenProvider */
+  public static final String FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE = "fs.azure.account.oauth2.client.assertion.provider.type";
   /** Key for enabling the tracking of ABFS API latency and sending the latency numbers to the ABFS API service */
   public static final String FS_AZURE_ABFS_LATENCY_TRACK = "fs.azure.abfs.latency.track";
 

hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/oauth2/ClientAssertionProvider.java

@@ -0,0 +1,72 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hadoop.fs.azurebfs.oauth2;
+
+import java.io.IOException;
+
+import org.apache.hadoop.classification.InterfaceAudience;
+import org.apache.hadoop.classification.InterfaceStability;
+import org.apache.hadoop.conf.Configuration;
+
+/**
+ * Interface for providing client assertions for Azure Workload Identity authentication.
+ *
+ * This interface allows custom implementations to provide JWT tokens through various mechanisms:
+ * - Kubernetes Token Request API
+ * - HashiCorp Vault
+ * - Custom token services
+ * - File-based tokens with custom logic
+ *
+ * Implementations should be thread-safe as they may be called concurrently.
+ */
+@InterfaceAudience.Public
+@InterfaceStability.Evolving
+public interface ClientAssertionProvider {
+
+  /**
+   * Initializes the provider with the given configuration.
+   * This method is called once after the provider is instantiated via reflection.
+   *
+   * @param configuration Hadoop configuration containing provider-specific settings
+   * @param accountName Azure storage account name for account-specific configuration
+   * @throws IOException if initialization fails
+   */
+  void initialize(Configuration configuration, String accountName) throws IOException;
+
+  /**
+   * Retrieves a client assertion (JWT token) for Azure Workload Identity authentication.
+   *
+   * The returned string should be a valid JWT token that can be used as a client assertion
+   * in OAuth 2.0 client credentials flow with JWT bearer assertion.
+   *
+   * @return JWT token as a string
+   * @throws IOException if token retrieval fails
+   */
+  String getClientAssertion() throws IOException;
+
+  /**
+   * Optional: Cleanup resources when the provider is no longer needed.
+   * Default implementation does nothing.
+   *
+   * @throws IOException if cleanup fails
+   */
+  default void close() throws IOException {
+    // Default: no-op
+  }
+}

hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/oauth2/WorkloadIdentityTokenProvider.java

@@ -20,14 +20,19 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.nio.charset.StandardCharsets;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+
 import org.apache.commons.io.FileUtils;
 import org.apache.hadoop.classification.VisibleForTesting;
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.thirdparty.com.google.common.base.Strings;
 import org.apache.hadoop.util.Preconditions;
 
+import static org.apache.hadoop.fs.azurebfs.constants.AbfsHttpConstants.EMPTY_STRING;
+
 /**
  * Provides tokens based on Azure AD Workload Identity.
  */
@@ -38,11 +43,73 @@ public class WorkloadIdentityTokenProvider extends AccessTokenProvider {
   private static final String EMPTY_TOKEN_FILE_ERROR = "Empty token file found at specified path: ";
   private static final String TOKEN_FILE_READ_ERROR = "Error reading token file at specified path: ";
 
+  /**
+   * Internal implementation of ClientAssertionProvider for file-based token reading.
+   * This provides backward compatibility for the file-based constructor.
+   */
+  private static class FileBasedClientAssertionProvider implements ClientAssertionProvider {
+    private final String tokenFile;
+
+    FileBasedClientAssertionProvider(String tokenFile) {
+      this.tokenFile = tokenFile;
+    }
+
+    @Override
+    public void initialize(Configuration configuration, String accountName) throws IOException {
+      // No initialization needed for file-based provider
+    }
+
+    @Override
+    public String getClientAssertion() throws IOException {
+      String clientAssertion = EMPTY_STRING;
+      try {
+        File file = new File(tokenFile);
+        clientAssertion = FileUtils.readFileToString(file, StandardCharsets.UTF_8);
+      } catch (Exception e) {
+        throw new IOException(TOKEN_FILE_READ_ERROR + tokenFile, e);
+      }
+      clientAssertion = clientAssertion.trim();
+      if (Strings.isNullOrEmpty(clientAssertion)) {
+        throw new IOException(EMPTY_TOKEN_FILE_ERROR + tokenFile);
+      }
+      return clientAssertion;
+    }
+  }
+
   private final String authEndpoint;
   private final String clientId;
-  private final String tokenFile;
+  private final ClientAssertionProvider clientAssertionProvider;
   private long tokenFetchTime = -1;
 
+  /**
+   * Constructor with custom ClientAssertionProvider.
+   * Use this for custom token retrieval mechanisms like Kubernetes Token Request API.
+   *
+   * @param authority OAuth authority URL
+   * @param tenantId Azure AD tenant ID
+   * @param clientId Azure AD client ID
+   * @param clientAssertionProvider Custom provider for client assertions
+   */
+  public WorkloadIdentityTokenProvider(final String authority, final String tenantId,
+                                       final String clientId, ClientAssertionProvider clientAssertionProvider) {
+    Preconditions.checkNotNull(authority, "authority");
+    Preconditions.checkNotNull(tenantId, "tenantId");
+    Preconditions.checkNotNull(clientId, "clientId");
+    Preconditions.checkNotNull(clientAssertionProvider, "clientAssertionProvider");
+
+    this.authEndpoint = authority + tenantId + OAUTH2_TOKEN_PATH;
+    this.clientId = clientId;
+    this.clientAssertionProvider = clientAssertionProvider;
+  }
+
+  /**
+   * Constructor with file-based token reading (backward compatibility).
+   *
+   * @param authority OAuth authority URL
+   * @param tenantId Azure AD tenant ID
+   * @param clientId Azure AD client ID
+   * @param tokenFile Path to file containing the JWT token
+   */
   public WorkloadIdentityTokenProvider(final String authority, final String tenantId,
       final String clientId, final String tokenFile) {
     Preconditions.checkNotNull(authority, "authority");
@@ -52,13 +119,13 @@ public WorkloadIdentityTokenProvider(final String authority, final String tenant
 
     this.authEndpoint = authority + tenantId + OAUTH2_TOKEN_PATH;
     this.clientId = clientId;
-    this.tokenFile = tokenFile;
+    this.clientAssertionProvider = new FileBasedClientAssertionProvider(tokenFile);
   }
 
   @Override
   protected AzureADToken refreshToken() throws IOException {
     LOG.debug("AADToken: refreshing token from JWT Assertion");
-    String clientAssertion = getClientAssertion();
+    String clientAssertion = clientAssertionProvider.getClientAssertion();
     AzureADToken token = getTokenUsingJWTAssertion(clientAssertion);
     tokenFetchTime = System.currentTimeMillis();
     return token;
@@ -90,31 +157,6 @@ protected boolean isTokenAboutToExpire() {
     return expiring;
   }
 
-  /**
-   * Gets the client assertion from the token file.
-   * The token file should contain the client assertion in JWT format.
-   * It should be a String containing Base64Url encoded JSON Web Token (JWT).
-   * See <a href="https://azure.github.io/azure-workload-identity/docs/faq.html#does-workload-identity-work-in-disconnected-environments">
-   * Azure Workload Identity FAQ</a>.
-   *
-   * @return the client assertion.
-   * @throws IOException if the token file is empty.
-   */
-  private String getClientAssertion()
-      throws IOException {
-    String clientAssertion = "";
-    try {
-      File file = new File(tokenFile);
-      clientAssertion = FileUtils.readFileToString(file, "UTF-8");
-    } catch (Exception e) {
-      throw new IOException(TOKEN_FILE_READ_ERROR + tokenFile, e);
-    }
-    if (Strings.isNullOrEmpty(clientAssertion)) {
-      throw new IOException(EMPTY_TOKEN_FILE_ERROR + tokenFile);
-    }
-    return clientAssertion;
-  }
-
   /**
    * Gets the Azure AD token from a client assertion in JWT format.
    * This method exists to make unit testing possible.