apache · wombatu-kun · May 23, 2026
diff --git a/...integration/java/org/apache/iceberg/azure/keymanagement/TestAzureKeyManagementClient.java b/...integration/java/org/apache/iceberg/azure/keymanagement/TestAzureKeyManagementClient.java
@@ -20,15 +20,20 @@
 
 import static org.apache.iceberg.azure.AzureProperties.AZURE_KEYVAULT_URL;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRequestContext;
 import com.azure.identity.DefaultAzureCredentialBuilder;
 import com.azure.security.keyvault.keys.KeyClient;
 import com.azure.security.keyvault.keys.KeyClientBuilder;
 import com.azure.security.keyvault.keys.models.KeyType;
 import java.nio.ByteBuffer;
 import java.time.Duration;
 import org.apache.iceberg.TestHelpers;
+import org.apache.iceberg.azure.AzureProperties;
 import org.apache.iceberg.encryption.KeyManagementClient;
+import org.apache.iceberg.exceptions.ValidationException;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
@@ -48,17 +53,21 @@ public class TestAzureKeyManagementClient {
 
   private static KeyManagementClient azureKeyManagementClient;
   private static KeyClient keyClient;
+  private static String vaultToken;
 
   @BeforeAll
   public static void beforeClass() {
-    keyClient =
-        new KeyClientBuilder()
-            .vaultUrl(KEY_VAULT_URI)
-            .credential(new DefaultAzureCredentialBuilder().build())
-            .buildClient();
+    TokenCredential credential = new DefaultAzureCredentialBuilder().build();
+    keyClient = new KeyClientBuilder().vaultUrl(KEY_VAULT_URI).credential(credential).buildClient();
     keyClient.createKey(ICEBERG_TEST_KEY_NAME, KeyType.RSA);
-    azureKeyManagementClient = new AzureKeyManagementClient();
-    azureKeyManagementClient.initialize(ImmutableMap.of(AZURE_KEYVAULT_URL, KEY_VAULT_URI));
+    // The KMS client no longer falls back to ambient credentials, so vend it a Key Vault-scoped
+    // token explicitly (as a catalog would) using the same ambient identity this test runs with.
+    vaultToken =
+        credential
+            .getToken(new TokenRequestContext().addScopes("https://vault.azure.net/.default"))
+            .block()
+            .getToken();
+    azureKeyManagementClient = newClient();
   }
 
   @AfterAll
@@ -69,6 +78,22 @@ public static void afterClass() {
     }
   }
 
+  @Test
+  public void wrapKeyFailsWithoutCatalogProvidedCredential() {
+    // Configuring only the vault URL used to authenticate via the ambient DefaultAzureCredential.
+    // After the security fix that path is rejected: the client requires a catalog-provided
+    // credential, so this previously-working configuration now fails fast.
+    AzureKeyManagementClient client = new AzureKeyManagementClient();
+    client.initialize(ImmutableMap.of(AZURE_KEYVAULT_URL, KEY_VAULT_URI));
+
+    ByteBuffer key = ByteBuffer.wrap("table-master-key".getBytes());
+
+    assertThatThrownBy(() -> client.wrapKey(key, ICEBERG_TEST_KEY_NAME))
+        .isInstanceOf(ValidationException.class)
+        .hasMessageContaining("ambient")
+        .hasMessageContaining(AzureProperties.ADLS_TOKEN);
+  }
+
   @Test
   public void keyWrapping() {
     ByteBuffer key = ByteBuffer.wrap("table-master-key".getBytes());
@@ -90,9 +115,7 @@ public void keyGenerationNotSupported() {
   public void testSerialization(
       TestHelpers.RoundTripSerializer<AzureKeyManagementClient> roundTripSerializer)
       throws Exception {
-    try (AzureKeyManagementClient keyManagementClient = new AzureKeyManagementClient()) {
-      keyManagementClient.initialize(ImmutableMap.of(AZURE_KEYVAULT_URL, KEY_VAULT_URI));
-
+    try (AzureKeyManagementClient keyManagementClient = newClient()) {
       AzureKeyManagementClient result = roundTripSerializer.apply(keyManagementClient);
 
       ByteBuffer key = ByteBuffer.wrap("super-secret-table-master-key".getBytes());
@@ -102,4 +125,11 @@ public void testSerialization(
       assertThat(result.unwrapKey(encryptedKey, ICEBERG_TEST_KEY_NAME)).isEqualTo(key);
     }
   }
+
+  private static AzureKeyManagementClient newClient() {
+    AzureKeyManagementClient client = new AzureKeyManagementClient();
+    client.initialize(
+        ImmutableMap.of(AZURE_KEYVAULT_URL, KEY_VAULT_URI, AzureProperties.ADLS_TOKEN, vaultToken));
+    return client;
+  }
 }
diff --git a/azure/src/main/java/org/apache/iceberg/azure/AzureProperties.java b/azure/src/main/java/org/apache/iceberg/azure/AzureProperties.java
@@ -178,17 +178,7 @@ public void applyClientConfiguration(String account, DataLakeFileSystemClientBui
         builder.credential(
             new StorageSharedKeyCredential(namedKeyCreds.getKey(), namedKeyCreds.getValue()));
       } else if (token != null && !token.isEmpty()) {
-        // Use TokenCredential with the provided token
-        TokenCredential tokenCredential =
-            new TokenCredential() {
-              @Override
-              public Mono<AccessToken> getToken(TokenRequestContext request) {
-                // Assume the token is valid for 1 hour from the current time
-                return Mono.just(
-                    new AccessToken(token, OffsetDateTime.now(ZoneOffset.UTC).plusHours(1)));
-              }
-            };
-        builder.credential(tokenCredential);
+        builder.credential(tokenCredential(token));
       } else {
         AdlsTokenCredentialProvider credentialProvider =
             AdlsTokenCredentialProviders.from(allProperties);
@@ -212,4 +202,38 @@ public String keyWrapAlgorithm() {
   public Optional<String> keyVaultUrl() {
     return Optional.ofNullable(this.keyVaultUrl);
   }
+
+  /**
+   * Returns the credential to use when authenticating to Azure Key Vault, derived only from
+   * credentials explicitly supplied through configuration: a catalog-vended {@link #ADLS_TOKEN}
+   * bearer token, or an explicitly configured {@link #ADLS_TOKEN_CREDENTIAL_PROVIDER}.
+   *
+   * <p>Ambient credentials such as {@link com.azure.identity.DefaultAzureCredential} are
+   * intentionally not used for Key Vault: combined with a misconfigured or malicious vault URL they
+   * could be used to exfiltrate the client's ambient identity. The supplied token must be scoped
+   * for Key Vault (for the public cloud, {@code https://vault.azure.net}).
+   *
+   * @return the configured Key Vault token credential, or empty if none was supplied
+   */
+  public Optional<TokenCredential> keyVaultTokenCredential() {
+    if (token != null && !token.isEmpty()) {
+      return Optional.of(tokenCredential(token));
+    }
+
+    if (allProperties.containsKey(ADLS_TOKEN_CREDENTIAL_PROVIDER)) {
+      return Optional.of(AdlsTokenCredentialProviders.from(allProperties).credential());
+    }
+
+    return Optional.empty();
+  }
+
+  private static TokenCredential tokenCredential(String token) {
+    return new TokenCredential() {
+      @Override
+      public Mono<AccessToken> getToken(TokenRequestContext request) {
+        // Assume the token is valid for 1 hour from the current time
+        return Mono.just(new AccessToken(token, OffsetDateTime.now(ZoneOffset.UTC).plusHours(1)));
+      }
+    };
+  }
 }
diff --git a/azure/src/main/java/org/apache/iceberg/azure/keymanagement/AzureKeyManagementClient.java b/azure/src/main/java/org/apache/iceberg/azure/keymanagement/AzureKeyManagementClient.java
@@ -18,16 +18,17 @@
  */
 package org.apache.iceberg.azure.keymanagement;
 
+import com.azure.core.credential.TokenCredential;
 import com.azure.security.keyvault.keys.KeyClient;
 import com.azure.security.keyvault.keys.KeyClientBuilder;
 import com.azure.security.keyvault.keys.cryptography.models.KeyWrapAlgorithm;
 import com.azure.security.keyvault.keys.cryptography.models.UnwrapResult;
 import com.azure.security.keyvault.keys.cryptography.models.WrapResult;
 import java.nio.ByteBuffer;
 import java.util.Map;
-import org.apache.iceberg.azure.AdlsTokenCredentialProviders;
 import org.apache.iceberg.azure.AzureProperties;
 import org.apache.iceberg.encryption.KeyManagementClient;
+import org.apache.iceberg.exceptions.ValidationException;
 import org.apache.iceberg.util.ByteBuffers;
 import org.apache.iceberg.util.SerializableMap;
 
@@ -74,12 +75,20 @@ private ClientState state() {
       synchronized (this) {
         if (state == null) {
           AzureProperties azureProperties = new AzureProperties(allProperties);
+          TokenCredential credential =
+              azureProperties
+                  .keyVaultTokenCredential()
+                  .orElseThrow(
+                      () ->
+                          new ValidationException(
+                              "Cannot authenticate to Azure Key Vault: set %s to a catalog-provided "
+                                  + "token or configure %s; ambient Azure credentials are not used "
+                                  + "for Key Vault",
+                              AzureProperties.ADLS_TOKEN,
+                              AzureProperties.ADLS_TOKEN_CREDENTIAL_PROVIDER));
           KeyClientBuilder keyClientBuilder = new KeyClientBuilder();
           azureProperties.keyVaultUrl().ifPresent(keyClientBuilder::vaultUrl);
-          KeyClient keyClient =
-              keyClientBuilder
-                  .credential(AdlsTokenCredentialProviders.from(allProperties).credential())
-                  .buildClient();
+          KeyClient keyClient = keyClientBuilder.credential(credential).buildClient();
           KeyWrapAlgorithm keyWrapAlgorithm =
               KeyWrapAlgorithm.fromString(azureProperties.keyWrapAlgorithm());
           state = new ClientState(keyClient, keyWrapAlgorithm);

diff --git a/azure/src/test/java/org/apache/iceberg/azure/TestAzureProperties.java b/azure/src/test/java/org/apache/iceberg/azure/TestAzureProperties.java
@@ -298,6 +298,41 @@ public void testCustomTokenCredentialProvider() {
     verify(clientBuilder, never()).credential(any(StorageSharedKeyCredential.class));
   }
 
+  @Test
+  public void keyVaultTokenCredentialFromAdlsToken() {
+    String token = "kv-token-value";
+    AzureProperties props = new AzureProperties(ImmutableMap.of(AzureProperties.ADLS_TOKEN, token));
+
+    Optional<TokenCredential> credential = props.keyVaultTokenCredential();
+
+    assertThat(credential).isPresent();
+    AccessToken accessToken = credential.get().getToken(new TokenRequestContext()).block();
+    assertThat(accessToken.getToken()).isEqualTo(token);
+  }
+
+  @Test
+  public void keyVaultTokenCredentialFromCustomProvider() {
+    AzureProperties props =
+        new AzureProperties(
+            ImmutableMap.of(
+                AzureProperties.ADLS_TOKEN_CREDENTIAL_PROVIDER,
+                DummyTokenCredentialProvider.class.getName()));
+
+    Optional<TokenCredential> credential = props.keyVaultTokenCredential();
+
+    assertThat(credential).isPresent();
+    assertThat(credential.get()).isInstanceOf(DummyTokenCredential.class);
+  }
+
+  @Test
+  public void keyVaultTokenCredentialAbsentWithoutExplicitConfig() {
+    AzureProperties props = new AzureProperties(ImmutableMap.of());
+
+    // Ambient DefaultAzureCredential must never be used for Key Vault: a misconfigured or malicious
+    // vault URL could otherwise be used to exfiltrate the client's ambient identity.
+    assertThat(props.keyVaultTokenCredential()).isEmpty();
+  }
+
   static class DummyTokenCredential implements TokenCredential {
     @Override
     public Mono<AccessToken> getToken(TokenRequestContext request) {