From cbf290fc134fb02b650877149276cd754ead23bd Mon Sep 17 00:00:00 2001 From: Steven Jones Date: Wed, 9 Sep 2026 15:41:11 +0000 Subject: [PATCH 1/6] GH-3683: Support programmatic KMS client factories --- .../parquet/crypto/keytools/KeyToolkit.java | 89 ++++++- .../crypto/keytools/KmsClientFactory.java | 32 +++ .../apache/parquet/crypto/TestKmsUrlRead.java | 25 ++ .../crypto/keytools/KeyToolkitTest.java | 237 ++++++++++++++++++ 4 files changed, 371 insertions(+), 12 deletions(-) create mode 100644 parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KmsClientFactory.java create mode 100644 parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java index 854976d37b..c2719dc938 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java @@ -21,7 +21,11 @@ import java.io.IOException; import java.util.Base64; +import java.util.Collections; +import java.util.Map; +import java.util.Objects; import java.util.Set; +import java.util.WeakHashMap; import java.util.concurrent.ConcurrentMap; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.FileStatus; @@ -109,6 +113,10 @@ public class KeyToolkit { // KMS client two level cache: token -> KMSInstanceId -> KmsClient static final TwoLevelCacheWithExpiration KMS_CLIENT_CACHE_PER_TOKEN = KmsClientCache.INSTANCE.getCache(); + // A weak association keeps programmatically supplied factories scoped to their Configuration. + private static final Map KMS_CLIENT_FACTORY_REGISTRATIONS = + Collections.synchronizedMap(new WeakHashMap<>()); + // KEK two level cache for wrapping: token -> MEK_ID -> KeyEncryptionKey static final TwoLevelCacheWithExpiration KEK_WRITE_CACHE_PER_TOKEN = KEKWriteCache.INSTANCE.getCache(); @@ -143,6 +151,15 @@ private TwoLevelCacheWithExpiration getCache() { } } + private static final class KmsClientFactoryRegistration { + private final KmsClientFactory factory; + private final TwoLevelCacheWithExpiration cache = new TwoLevelCacheWithExpiration<>(); + + private KmsClientFactoryRegistration(KmsClientFactory factory) { + this.factory = factory; + } + } + static class KeyWithMasterID { private final byte[] keyBytes; private final String masterID; @@ -282,16 +299,44 @@ public static void rotateMasterKeys(String folderPath, Configuration hadoopConfi */ public static void removeCacheEntriesForToken(String accessToken) { KMS_CLIENT_CACHE_PER_TOKEN.removeCacheEntriesForToken(accessToken); + synchronized (KMS_CLIENT_FACTORY_REGISTRATIONS) { + for (KmsClientFactoryRegistration registration : KMS_CLIENT_FACTORY_REGISTRATIONS.values()) { + registration.cache.removeCacheEntriesForToken(accessToken); + } + } KEK_WRITE_CACHE_PER_TOKEN.removeCacheEntriesForToken(accessToken); KEK_READ_CACHE_PER_TOKEN.removeCacheEntriesForToken(accessToken); } public static void removeCacheEntriesForAllTokens() { KMS_CLIENT_CACHE_PER_TOKEN.clear(); + synchronized (KMS_CLIENT_FACTORY_REGISTRATIONS) { + for (KmsClientFactoryRegistration registration : KMS_CLIENT_FACTORY_REGISTRATIONS.values()) { + registration.cache.clear(); + } + } KEK_WRITE_CACHE_PER_TOKEN.clear(); KEK_READ_CACHE_PER_TOKEN.clear(); } + /** + * Sets the factory used to create KMS clients for the supplied configuration. + * + *

The factory is local to this JVM and must be set before constructing a reader or writer. + * Reflection through {@link #KMS_CLIENT_CLASS_PROPERTY_NAME} remains the default for other + * configurations. Clients returned by the factory are initialized and cached by {@link + * KeyToolkit} in the same way as reflectively constructed clients. + * The association is not serialized; configuration copies must register their own factory. + * + * @param configuration Hadoop configuration associated with the factory + * @param kmsClientFactory factory used to create KMS clients + */ + public static void setKmsClientFactory(Configuration configuration, KmsClientFactory kmsClientFactory) { + Objects.requireNonNull(configuration, "configuration"); + Objects.requireNonNull(kmsClientFactory, "kmsClientFactory"); + KMS_CLIENT_FACTORY_REGISTRATIONS.put(configuration, new KmsClientFactoryRegistration(kmsClientFactory)); + } + /** * Encrypts "key" with "masterKey", using AES-GCM and the "AAD" * @@ -335,32 +380,52 @@ static KmsClient getKmsClient( String accessToken, long cacheEntryLifetime) { + KmsClientFactoryRegistration registration = KMS_CLIENT_FACTORY_REGISTRATIONS.get(configuration); + TwoLevelCacheWithExpiration cache = + registration == null ? KMS_CLIENT_CACHE_PER_TOKEN : registration.cache; + if (registration != null) { + // The callers already expire the default cache. Factory registrations have independent + // caches, so expire the selected registration here. + cache.checkCacheForExpiredTokens(cacheEntryLifetime); + } + ConcurrentMap kmsClientPerKmsInstanceCache = - KMS_CLIENT_CACHE_PER_TOKEN.getOrCreateInternalCache(accessToken, cacheEntryLifetime); + cache.getOrCreateInternalCache(accessToken, cacheEntryLifetime); KmsClient kmsClient = kmsClientPerKmsInstanceCache.computeIfAbsent( kmsInstanceID, - (k) -> createAndInitKmsClient(configuration, kmsInstanceID, kmsInstanceURL, accessToken)); + (k) -> createAndInitKmsClient(configuration, kmsInstanceID, kmsInstanceURL, accessToken, registration)); return kmsClient; } private static KmsClient createAndInitKmsClient( - Configuration configuration, String kmsInstanceID, String kmsInstanceURL, String accessToken) { + Configuration configuration, + String kmsInstanceID, + String kmsInstanceURL, + String accessToken, + KmsClientFactoryRegistration registration) { Class kmsClientClass = null; KmsClient kmsClient = null; - try { - kmsClientClass = ConfigurationUtil.getClassFromConfig( - configuration, KMS_CLIENT_CLASS_PROPERTY_NAME, KmsClient.class); - - if (null == kmsClientClass) { - throw new ParquetCryptoRuntimeException("Unspecified " + KMS_CLIENT_CLASS_PROPERTY_NAME); + if (registration != null) { + kmsClient = registration.factory.createKmsClient(); + if (kmsClient == null) { + throw new ParquetCryptoRuntimeException("KmsClientFactory returned null"); + } + } else { + try { + kmsClientClass = ConfigurationUtil.getClassFromConfig( + configuration, KMS_CLIENT_CLASS_PROPERTY_NAME, KmsClient.class); + + if (null == kmsClientClass) { + throw new ParquetCryptoRuntimeException("Unspecified " + KMS_CLIENT_CLASS_PROPERTY_NAME); + } + kmsClient = (KmsClient) kmsClientClass.newInstance(); + } catch (InstantiationException | IllegalAccessException | BadConfigurationException e) { + throw new ParquetCryptoRuntimeException("Could not instantiate KmsClient class: " + kmsClientClass, e); } - kmsClient = (KmsClient) kmsClientClass.newInstance(); - } catch (InstantiationException | IllegalAccessException | BadConfigurationException e) { - throw new ParquetCryptoRuntimeException("Could not instantiate KmsClient class: " + kmsClientClass, e); } kmsClient.initialize(configuration, kmsInstanceID, kmsInstanceURL, accessToken); diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KmsClientFactory.java b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KmsClientFactory.java new file mode 100644 index 0000000000..ea5ac33fb6 --- /dev/null +++ b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KmsClientFactory.java @@ -0,0 +1,32 @@ +/* + * 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.parquet.crypto.keytools; + +/** Factory for creating {@link KmsClient} instances with programmatically supplied dependencies. */ +@FunctionalInterface +public interface KmsClientFactory { + + /** + * Creates a KMS client. {@link KeyToolkit} initializes the returned client before using it. + * + * @return a new or pre-built KMS client + */ + KmsClient createKmsClient(); +} diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java index 5d3d3e04fc..f8ed69b3b4 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java @@ -73,6 +73,14 @@ static String getStaticKmsURL() { } } + private static class ConstructorInjectedKmsClient extends UnitestUrlReadKMS { + private final String dependency; + + private ConstructorInjectedKmsClient(String dependency) { + this.dependency = dependency; + } + } + @BeforeAll public static void writeEncryptedFile() throws IOException { Configuration writeConf = new Configuration(); @@ -179,6 +187,23 @@ public void testSetKmsUrl() throws IOException { assertThat(readerSetURL.equals(UnitestUrlReadKMS.getStaticKmsURL())); } + @Test + public void testProgrammaticKmsClientFactory() throws IOException { + Configuration readConf = basicDecryptionConfig(); + readConf.set(KeyToolkit.KEY_ACCESS_TOKEN_PROPERTY_NAME, "factory-token"); + ConstructorInjectedKmsClient kmsClient = new ConstructorInjectedKmsClient("dependency"); + KeyToolkit.setKmsClientFactory(readConf, () -> kmsClient); + + try (ParquetReader reader = ParquetReader.builder(new GroupReadSupport(), filePath) + .withConf(readConf) + .build()) { + assertThat(reader.read()).isNotNull(); + } + + assertThat(kmsClient.dependency).isEqualTo("dependency"); + assertThat(UnitestUrlReadKMS.getStaticKmsURL()).isEqualTo(KmsClient.KMS_INSTANCE_ID_DEFAULT); + } + @AfterAll public static void deleteFile() throws IOException { filePath.getFileSystem(new Configuration()).delete(filePath, false); diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java new file mode 100644 index 0000000000..036745e118 --- /dev/null +++ b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java @@ -0,0 +1,237 @@ +/* + * 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.parquet.crypto.keytools; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.hadoop.conf.Configuration; +import org.apache.parquet.crypto.ParquetCryptoRuntimeException; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +public class KeyToolkitTest { + + private static final long CACHE_LIFETIME_MILLIS = 60_000; + + @AfterEach + public void clearCaches() { + KeyToolkit.removeCacheEntriesForAllTokens(); + } + + @Test + public void prefersConfiguredKmsClientFactory() { + Configuration configuration = new Configuration(false); + configuration.set(KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME, ReflectiveKmsClient.class.getName()); + ConstructorInjectedKmsClient client = new ConstructorInjectedKmsClient("dependency"); + AtomicInteger factoryCalls = new AtomicInteger(); + KeyToolkit.setKmsClientFactory(configuration, () -> { + factoryCalls.incrementAndGet(); + return client; + }); + + KmsClient first = KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS); + KmsClient second = KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS); + + assertThat(first).isSameAs(client); + assertThat(second).isSameAs(client); + assertThat(factoryCalls).hasValue(1); + assertThat(client.configuration).isSameAs(configuration); + assertThat(client.kmsInstanceID).isEqualTo("instance"); + assertThat(client.kmsInstanceURL).isEqualTo("url"); + assertThat(client.accessToken).isEqualTo("token"); + assertThat(client.initializeCalls).isEqualTo(1); + } + + @Test + public void scopesKmsClientFactoryAndCacheToConfiguration() { + Configuration firstConfiguration = new Configuration(false); + Configuration secondConfiguration = new Configuration(false); + ConstructorInjectedKmsClient firstClient = new ConstructorInjectedKmsClient("first"); + ConstructorInjectedKmsClient secondClient = new ConstructorInjectedKmsClient("second"); + KeyToolkit.setKmsClientFactory(firstConfiguration, () -> firstClient); + KeyToolkit.setKmsClientFactory(secondConfiguration, () -> secondClient); + + KmsClient first = + KeyToolkit.getKmsClient("DEFAULT", "DEFAULT", firstConfiguration, "DEFAULT", CACHE_LIFETIME_MILLIS); + KmsClient second = + KeyToolkit.getKmsClient("DEFAULT", "DEFAULT", secondConfiguration, "DEFAULT", CACHE_LIFETIME_MILLIS); + + assertThat(first).isSameAs(firstClient); + assertThat(second).isSameAs(secondClient); + } + + @Test + public void factoryRegistrationDoesNotReuseCachedReflectiveClient() { + Configuration reflectiveConfiguration = new Configuration(false); + reflectiveConfiguration.set(KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME, ReflectiveKmsClient.class.getName()); + KmsClient reflectiveClient = + KeyToolkit.getKmsClient("instance", "url", reflectiveConfiguration, "token", CACHE_LIFETIME_MILLIS); + + Configuration factoryConfiguration = new Configuration(false); + ConstructorInjectedKmsClient factoryClient = new ConstructorInjectedKmsClient("dependency"); + KeyToolkit.setKmsClientFactory(factoryConfiguration, () -> factoryClient); + KmsClient actual = + KeyToolkit.getKmsClient("instance", "url", factoryConfiguration, "token", CACHE_LIFETIME_MILLIS); + + assertThat(reflectiveClient).isInstanceOf(ReflectiveKmsClient.class); + assertThat(actual).isSameAs(factoryClient); + } + + @Test + public void replacingKmsClientFactoryDiscardsCachedClient() { + Configuration configuration = new Configuration(false); + ConstructorInjectedKmsClient firstClient = new ConstructorInjectedKmsClient("first"); + ConstructorInjectedKmsClient replacementClient = new ConstructorInjectedKmsClient("replacement"); + KeyToolkit.setKmsClientFactory(configuration, () -> firstClient); + KmsClient first = KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS); + + KeyToolkit.setKmsClientFactory(configuration, () -> replacementClient); + KmsClient replacement = + KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS); + + assertThat(first).isSameAs(firstClient); + assertThat(replacement).isSameAs(replacementClient); + } + + @Test + public void removeCacheEntriesForTokenClearsOnlyMatchingFactoryClients() { + Configuration configuration = new Configuration(false); + AtomicInteger factoryCalls = new AtomicInteger(); + KeyToolkit.setKmsClientFactory( + configuration, + () -> new ConstructorInjectedKmsClient(Integer.toString(factoryCalls.incrementAndGet()))); + KmsClient firstTokenClient = + KeyToolkit.getKmsClient("instance", "url", configuration, "first-token", CACHE_LIFETIME_MILLIS); + KmsClient otherTokenClient = + KeyToolkit.getKmsClient("instance", "url", configuration, "other-token", CACHE_LIFETIME_MILLIS); + + KeyToolkit.removeCacheEntriesForToken("first-token"); + + KmsClient refreshedFirstTokenClient = + KeyToolkit.getKmsClient("instance", "url", configuration, "first-token", CACHE_LIFETIME_MILLIS); + KmsClient cachedOtherTokenClient = + KeyToolkit.getKmsClient("instance", "url", configuration, "other-token", CACHE_LIFETIME_MILLIS); + assertThat(refreshedFirstTokenClient).isNotSameAs(firstTokenClient); + assertThat(cachedOtherTokenClient).isSameAs(otherTokenClient); + assertThat(factoryCalls).hasValue(3); + } + + @Test + public void removeCacheEntriesForAllTokensClearsFactoryClients() { + Configuration configuration = new Configuration(false); + AtomicInteger factoryCalls = new AtomicInteger(); + KeyToolkit.setKmsClientFactory( + configuration, + () -> new ConstructorInjectedKmsClient(Integer.toString(factoryCalls.incrementAndGet()))); + KmsClient firstTokenClient = + KeyToolkit.getKmsClient("instance", "url", configuration, "first-token", CACHE_LIFETIME_MILLIS); + KmsClient secondTokenClient = + KeyToolkit.getKmsClient("instance", "url", configuration, "second-token", CACHE_LIFETIME_MILLIS); + + KeyToolkit.removeCacheEntriesForAllTokens(); + + KmsClient refreshedFirstTokenClient = + KeyToolkit.getKmsClient("instance", "url", configuration, "first-token", CACHE_LIFETIME_MILLIS); + KmsClient refreshedSecondTokenClient = + KeyToolkit.getKmsClient("instance", "url", configuration, "second-token", CACHE_LIFETIME_MILLIS); + assertThat(refreshedFirstTokenClient).isNotSameAs(firstTokenClient); + assertThat(refreshedSecondTokenClient).isNotSameAs(secondTokenClient); + assertThat(factoryCalls).hasValue(4); + } + + @Test + public void rejectsNullKmsClientFromFactory() { + Configuration configuration = new Configuration(false); + KeyToolkit.setKmsClientFactory(configuration, () -> null); + + assertThatThrownBy( + () -> KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS)) + .isInstanceOf(ParquetCryptoRuntimeException.class) + .hasMessage("KmsClientFactory returned null"); + } + + @Test + public void usesConfiguredClassWhenFactoryIsNotSet() { + Configuration configuration = new Configuration(false); + configuration.set(KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME, ReflectiveKmsClient.class.getName()); + + KmsClient client = KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS); + + assertThat(client).isInstanceOf(ReflectiveKmsClient.class); + assertThat(((ReflectiveKmsClient) client).initializeCalls).isEqualTo(1); + } + + private static class ConstructorInjectedKmsClient implements KmsClient { + private final String dependency; + private Configuration configuration; + private String kmsInstanceID; + private String kmsInstanceURL; + private String accessToken; + private int initializeCalls; + + private ConstructorInjectedKmsClient(String dependency) { + this.dependency = dependency; + } + + @Override + public void initialize( + Configuration configuration, String kmsInstanceID, String kmsInstanceURL, String accessToken) { + this.configuration = configuration; + this.kmsInstanceID = kmsInstanceID; + this.kmsInstanceURL = kmsInstanceURL; + this.accessToken = accessToken; + initializeCalls++; + } + + @Override + public String wrapKey(byte[] keyBytes, String masterKeyIdentifier) { + return dependency; + } + + @Override + public byte[] unwrapKey(String wrappedKey, String masterKeyIdentifier) { + return dependency.getBytes(); + } + } + + public static class ReflectiveKmsClient implements KmsClient { + private int initializeCalls; + + public ReflectiveKmsClient() {} + + @Override + public void initialize( + Configuration configuration, String kmsInstanceID, String kmsInstanceURL, String accessToken) { + initializeCalls++; + } + + @Override + public String wrapKey(byte[] keyBytes, String masterKeyIdentifier) { + throw new UnsupportedOperationException(); + } + + @Override + public byte[] unwrapKey(String wrappedKey, String masterKeyIdentifier) { + throw new UnsupportedOperationException(); + } + } +} From b3da958b1944eac983413557d1ca0e6bcdaffedf Mon Sep 17 00:00:00 2001 From: Steven Jones Date: Wed, 9 Sep 2026 17:12:27 +0000 Subject: [PATCH 2/6] GH-3683: Isolate factory caches and add cleanup --- .../crypto/keytools/FileKeyUnwrapper.java | 13 +- .../crypto/keytools/FileKeyWrapper.java | 14 +- .../parquet/crypto/keytools/KeyToolkit.java | 159 ++++++++++++++---- .../apache/parquet/crypto/TestKmsUrlRead.java | 18 +- .../crypto/keytools/KeyToolkitTest.java | 141 +++++++++++++++- 5 files changed, 281 insertions(+), 64 deletions(-) diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyUnwrapper.java b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyUnwrapper.java index b681187def..9bb0ffcfa6 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyUnwrapper.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyUnwrapper.java @@ -19,8 +19,6 @@ package org.apache.parquet.crypto.keytools; -import static org.apache.parquet.crypto.keytools.KeyToolkit.KEK_READ_CACHE_PER_TOKEN; -import static org.apache.parquet.crypto.keytools.KeyToolkit.KMS_CLIENT_CACHE_PER_TOKEN; import static org.apache.parquet.crypto.keytools.KeyToolkit.stringIsEmpty; import java.io.IOException; @@ -47,6 +45,7 @@ public class FileKeyUnwrapper implements DecryptionKeyRetriever { private final Path parquetFilePath; private final String accessToken; private final long cacheEntryLifetime; + private final KeyToolkit.KmsClientCacheContext cacheContext; FileKeyUnwrapper(Configuration hadoopConfiguration, Path filePath) { this.hadoopConfiguration = hadoopConfiguration; @@ -58,11 +57,13 @@ public class FileKeyUnwrapper implements DecryptionKeyRetriever { accessToken = hadoopConfiguration.getTrimmed( KeyToolkit.KEY_ACCESS_TOKEN_PROPERTY_NAME, KmsClient.KEY_ACCESS_TOKEN_DEFAULT); + cacheContext = KeyToolkit.getKmsClientCacheContext(hadoopConfiguration); // Check cache upon each file reading (clean once in cacheEntryLifetime) - KMS_CLIENT_CACHE_PER_TOKEN.checkCacheForExpiredTokens(cacheEntryLifetime); - KEK_READ_CACHE_PER_TOKEN.checkCacheForExpiredTokens(cacheEntryLifetime); - kekPerKekID = KEK_READ_CACHE_PER_TOKEN.getOrCreateInternalCache(accessToken, cacheEntryLifetime); + cacheContext.getKmsClientCache().checkCacheForExpiredTokens(cacheEntryLifetime); + TwoLevelCacheWithExpiration kekReadCache = cacheContext.getKekReadCache(); + kekReadCache.checkCacheForExpiredTokens(cacheEntryLifetime); + kekPerKekID = kekReadCache.getOrCreateInternalCache(accessToken, cacheEntryLifetime); if (LOG.isDebugEnabled()) { LOG.debug( @@ -168,7 +169,7 @@ KeyToolkit.KmsClientAndDetails getKmsClientFromConfigOrKeyMaterial(KeyMaterial k } KmsClient kmsClient = KeyToolkit.getKmsClient( - kmsInstanceID, kmsInstanceURL, hadoopConfiguration, accessToken, cacheEntryLifetime); + kmsInstanceID, kmsInstanceURL, hadoopConfiguration, accessToken, cacheEntryLifetime, cacheContext); if (null == kmsClient) { throw new ParquetCryptoRuntimeException( "KMSClient was not successfully created for reading encrypted data."); diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyWrapper.java b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyWrapper.java index 195a024247..6d2421abc4 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyWrapper.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyWrapper.java @@ -19,9 +19,6 @@ package org.apache.parquet.crypto.keytools; -import static org.apache.parquet.crypto.keytools.KeyToolkit.KEK_WRITE_CACHE_PER_TOKEN; -import static org.apache.parquet.crypto.keytools.KeyToolkit.KMS_CLIENT_CACHE_PER_TOKEN; - import java.nio.charset.StandardCharsets; import java.security.SecureRandom; import java.util.Arrays; @@ -73,8 +70,10 @@ public class FileKeyWrapper { accessToken = hadoopConfiguration.getTrimmed( KeyToolkit.KEY_ACCESS_TOKEN_PROPERTY_NAME, KmsClient.KEY_ACCESS_TOKEN_DEFAULT); + KeyToolkit.KmsClientCacheContext cacheContext = KeyToolkit.getKmsClientCacheContext(configuration); + // Check caches upon each file writing (clean once in cacheEntryLifetime) - KMS_CLIENT_CACHE_PER_TOKEN.checkCacheForExpiredTokens(cacheEntryLifetime); + cacheContext.getKmsClientCache().checkCacheForExpiredTokens(cacheEntryLifetime); if (null == kmsClientAndDetails) { kmsInstanceID = hadoopConfiguration.getTrimmed( @@ -82,7 +81,7 @@ public class FileKeyWrapper { kmsInstanceURL = hadoopConfiguration.getTrimmed( KeyToolkit.KMS_INSTANCE_URL_PROPERTY_NAME, KmsClient.KMS_INSTANCE_URL_DEFAULT); kmsClient = KeyToolkit.getKmsClient( - kmsInstanceID, kmsInstanceURL, configuration, accessToken, cacheEntryLifetime); + kmsInstanceID, kmsInstanceURL, configuration, accessToken, cacheEntryLifetime, cacheContext); } else { kmsInstanceID = kmsClientAndDetails.getKmsInstanceID(); kmsInstanceURL = kmsClientAndDetails.getKmsInstanceURL(); @@ -90,8 +89,9 @@ public class FileKeyWrapper { } if (doubleWrapping) { - KEK_WRITE_CACHE_PER_TOKEN.checkCacheForExpiredTokens(cacheEntryLifetime); - KEKPerMasterKeyID = KEK_WRITE_CACHE_PER_TOKEN.getOrCreateInternalCache(accessToken, cacheEntryLifetime); + TwoLevelCacheWithExpiration kekWriteCache = cacheContext.getKekWriteCache(); + kekWriteCache.checkCacheForExpiredTokens(cacheEntryLifetime); + KEKPerMasterKeyID = kekWriteCache.getOrCreateInternalCache(accessToken, cacheEntryLifetime); int kekLengthBits = configuration.getInt(KeyToolkit.KEK_LENGTH_PROPERTY_NAME, KeyToolkit.KEK_LENGTH_DEFAULT); if (Arrays.binarySearch(ACCEPTABLE_KEK_LENGTHS, kekLengthBits) < 0) { diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java index c2719dc938..a37adbd3d5 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java @@ -22,10 +22,10 @@ import java.io.IOException; import java.util.Base64; import java.util.Collections; +import java.util.IdentityHashMap; import java.util.Map; import java.util.Objects; import java.util.Set; -import java.util.WeakHashMap; import java.util.concurrent.ConcurrentMap; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.FileStatus; @@ -113,10 +113,6 @@ public class KeyToolkit { // KMS client two level cache: token -> KMSInstanceId -> KmsClient static final TwoLevelCacheWithExpiration KMS_CLIENT_CACHE_PER_TOKEN = KmsClientCache.INSTANCE.getCache(); - // A weak association keeps programmatically supplied factories scoped to their Configuration. - private static final Map KMS_CLIENT_FACTORY_REGISTRATIONS = - Collections.synchronizedMap(new WeakHashMap<>()); - // KEK two level cache for wrapping: token -> MEK_ID -> KeyEncryptionKey static final TwoLevelCacheWithExpiration KEK_WRITE_CACHE_PER_TOKEN = KEKWriteCache.INSTANCE.getCache(); @@ -124,6 +120,14 @@ public class KeyToolkit { // KEK two level cache for unwrapping: token -> KEK_ID -> KEK bytes static final TwoLevelCacheWithExpiration KEK_READ_CACHE_PER_TOKEN = KEKReadCache.INSTANCE.getCache(); + private static final KmsClientCacheContext DEFAULT_KMS_CLIENT_CACHE_CONTEXT = new KmsClientCacheContext( + null, KMS_CLIENT_CACHE_PER_TOKEN, KEK_WRITE_CACHE_PER_TOKEN, KEK_READ_CACHE_PER_TOKEN); + + // Programmatically supplied factories and their caches, scoped to the exact Configuration object. + // Callers must remove registrations when the Configuration is no longer in use. + private static final Map KMS_CLIENT_FACTORY_REGISTRATIONS = + Collections.synchronizedMap(new IdentityHashMap<>()); + private enum KmsClientCache { INSTANCE; private final TwoLevelCacheWithExpiration cache = new TwoLevelCacheWithExpiration<>(); @@ -151,12 +155,53 @@ private TwoLevelCacheWithExpiration getCache() { } } - private static final class KmsClientFactoryRegistration { + static final class KmsClientCacheContext { private final KmsClientFactory factory; - private final TwoLevelCacheWithExpiration cache = new TwoLevelCacheWithExpiration<>(); + private final TwoLevelCacheWithExpiration kmsClientCache; + private final TwoLevelCacheWithExpiration kekWriteCache; + private final TwoLevelCacheWithExpiration kekReadCache; + + private KmsClientCacheContext(KmsClientFactory factory) { + this( + factory, + new TwoLevelCacheWithExpiration<>(), + new TwoLevelCacheWithExpiration<>(), + new TwoLevelCacheWithExpiration<>()); + } - private KmsClientFactoryRegistration(KmsClientFactory factory) { + private KmsClientCacheContext( + KmsClientFactory factory, + TwoLevelCacheWithExpiration kmsClientCache, + TwoLevelCacheWithExpiration kekWriteCache, + TwoLevelCacheWithExpiration kekReadCache) { this.factory = factory; + this.kmsClientCache = kmsClientCache; + this.kekWriteCache = kekWriteCache; + this.kekReadCache = kekReadCache; + } + + TwoLevelCacheWithExpiration getKmsClientCache() { + return kmsClientCache; + } + + TwoLevelCacheWithExpiration getKekWriteCache() { + return kekWriteCache; + } + + TwoLevelCacheWithExpiration getKekReadCache() { + return kekReadCache; + } + + void removeCacheEntriesForToken(String accessToken) { + kmsClientCache.removeCacheEntriesForToken(accessToken); + kekWriteCache.removeCacheEntriesForToken(accessToken); + kekReadCache.removeCacheEntriesForToken(accessToken); + } + + void clear() { + kmsClientCache.clear(); + kekWriteCache.clear(); + kekReadCache.clear(); } } @@ -237,7 +282,7 @@ public static void rotateMasterKeys(String folderPath, Configuration hadoopConfi long currentTime = System.currentTimeMillis(); synchronized (lastCacheCleanForKeyRotationTimeLock) { if (currentTime - lastCacheCleanForKeyRotationTime > CACHE_CLEAN_PERIOD_FOR_KEY_ROTATION) { - KEK_WRITE_CACHE_PER_TOKEN.clear(); + clearKekWriteCaches(); lastCacheCleanForKeyRotationTime = currentTime; } } @@ -298,25 +343,21 @@ public static void rotateMasterKeys(String folderPath, Configuration hadoopConfi * @param accessToken access token */ public static void removeCacheEntriesForToken(String accessToken) { - KMS_CLIENT_CACHE_PER_TOKEN.removeCacheEntriesForToken(accessToken); + DEFAULT_KMS_CLIENT_CACHE_CONTEXT.removeCacheEntriesForToken(accessToken); synchronized (KMS_CLIENT_FACTORY_REGISTRATIONS) { - for (KmsClientFactoryRegistration registration : KMS_CLIENT_FACTORY_REGISTRATIONS.values()) { - registration.cache.removeCacheEntriesForToken(accessToken); + for (KmsClientCacheContext cacheContext : KMS_CLIENT_FACTORY_REGISTRATIONS.values()) { + cacheContext.removeCacheEntriesForToken(accessToken); } } - KEK_WRITE_CACHE_PER_TOKEN.removeCacheEntriesForToken(accessToken); - KEK_READ_CACHE_PER_TOKEN.removeCacheEntriesForToken(accessToken); } public static void removeCacheEntriesForAllTokens() { - KMS_CLIENT_CACHE_PER_TOKEN.clear(); + DEFAULT_KMS_CLIENT_CACHE_CONTEXT.clear(); synchronized (KMS_CLIENT_FACTORY_REGISTRATIONS) { - for (KmsClientFactoryRegistration registration : KMS_CLIENT_FACTORY_REGISTRATIONS.values()) { - registration.cache.clear(); + for (KmsClientCacheContext cacheContext : KMS_CLIENT_FACTORY_REGISTRATIONS.values()) { + cacheContext.clear(); } } - KEK_WRITE_CACHE_PER_TOKEN.clear(); - KEK_READ_CACHE_PER_TOKEN.clear(); } /** @@ -325,8 +366,13 @@ public static void removeCacheEntriesForAllTokens() { *

The factory is local to this JVM and must be set before constructing a reader or writer. * Reflection through {@link #KMS_CLIENT_CLASS_PROPERTY_NAME} remains the default for other * configurations. Clients returned by the factory are initialized and cached by {@link - * KeyToolkit} in the same way as reflectively constructed clients. - * The association is not serialized; configuration copies must register their own factory. + * KeyToolkit} in the same way as reflectively constructed clients. The KMS client and key + * encryption key caches are isolated from registrations for other configurations. + * + *

The association is not serialized, and configuration copies must register their own + * factory. The caller must invoke {@link #removeKmsClientFactory(Configuration)} after all + * readers and writers using the configuration have closed. Replacing a factory clears the + * previous registration and its caches. * * @param configuration Hadoop configuration associated with the factory * @param kmsClientFactory factory used to create KMS clients @@ -334,7 +380,27 @@ public static void removeCacheEntriesForAllTokens() { public static void setKmsClientFactory(Configuration configuration, KmsClientFactory kmsClientFactory) { Objects.requireNonNull(configuration, "configuration"); Objects.requireNonNull(kmsClientFactory, "kmsClientFactory"); - KMS_CLIENT_FACTORY_REGISTRATIONS.put(configuration, new KmsClientFactoryRegistration(kmsClientFactory)); + KmsClientCacheContext previous = + KMS_CLIENT_FACTORY_REGISTRATIONS.put(configuration, new KmsClientCacheContext(kmsClientFactory)); + if (previous != null) { + previous.clear(); + } + } + + /** + * Removes the KMS client factory for the supplied configuration and clears all of its caches. + * + *

This method must be called only after all readers and writers using the configuration have + * closed. + * + * @param configuration Hadoop configuration associated with the factory + */ + public static void removeKmsClientFactory(Configuration configuration) { + Objects.requireNonNull(configuration, "configuration"); + KmsClientCacheContext registration = KMS_CLIENT_FACTORY_REGISTRATIONS.remove(configuration); + if (registration != null) { + registration.clear(); + } } /** @@ -380,37 +446,60 @@ static KmsClient getKmsClient( String accessToken, long cacheEntryLifetime) { - KmsClientFactoryRegistration registration = KMS_CLIENT_FACTORY_REGISTRATIONS.get(configuration); - TwoLevelCacheWithExpiration cache = - registration == null ? KMS_CLIENT_CACHE_PER_TOKEN : registration.cache; - if (registration != null) { - // The callers already expire the default cache. Factory registrations have independent - // caches, so expire the selected registration here. - cache.checkCacheForExpiredTokens(cacheEntryLifetime); - } + return getKmsClient( + kmsInstanceID, + kmsInstanceURL, + configuration, + accessToken, + cacheEntryLifetime, + getKmsClientCacheContext(configuration)); + } + + static KmsClient getKmsClient( + String kmsInstanceID, + String kmsInstanceURL, + Configuration configuration, + String accessToken, + long cacheEntryLifetime, + KmsClientCacheContext cacheContext) { ConcurrentMap kmsClientPerKmsInstanceCache = - cache.getOrCreateInternalCache(accessToken, cacheEntryLifetime); + cacheContext.getKmsClientCache().getOrCreateInternalCache(accessToken, cacheEntryLifetime); KmsClient kmsClient = kmsClientPerKmsInstanceCache.computeIfAbsent( kmsInstanceID, - (k) -> createAndInitKmsClient(configuration, kmsInstanceID, kmsInstanceURL, accessToken, registration)); + (k) -> createAndInitKmsClient( + configuration, kmsInstanceID, kmsInstanceURL, accessToken, cacheContext.factory)); return kmsClient; } + static KmsClientCacheContext getKmsClientCacheContext(Configuration configuration) { + KmsClientCacheContext cacheContext = KMS_CLIENT_FACTORY_REGISTRATIONS.get(configuration); + return cacheContext == null ? DEFAULT_KMS_CLIENT_CACHE_CONTEXT : cacheContext; + } + + private static void clearKekWriteCaches() { + KEK_WRITE_CACHE_PER_TOKEN.clear(); + synchronized (KMS_CLIENT_FACTORY_REGISTRATIONS) { + for (KmsClientCacheContext cacheContext : KMS_CLIENT_FACTORY_REGISTRATIONS.values()) { + cacheContext.getKekWriteCache().clear(); + } + } + } + private static KmsClient createAndInitKmsClient( Configuration configuration, String kmsInstanceID, String kmsInstanceURL, String accessToken, - KmsClientFactoryRegistration registration) { + KmsClientFactory factory) { Class kmsClientClass = null; KmsClient kmsClient = null; - if (registration != null) { - kmsClient = registration.factory.createKmsClient(); + if (factory != null) { + kmsClient = factory.createKmsClient(); if (kmsClient == null) { throw new ParquetCryptoRuntimeException("KmsClientFactory returned null"); } diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java index f8ed69b3b4..761022bfd2 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java @@ -194,14 +194,18 @@ public void testProgrammaticKmsClientFactory() throws IOException { ConstructorInjectedKmsClient kmsClient = new ConstructorInjectedKmsClient("dependency"); KeyToolkit.setKmsClientFactory(readConf, () -> kmsClient); - try (ParquetReader reader = ParquetReader.builder(new GroupReadSupport(), filePath) - .withConf(readConf) - .build()) { - assertThat(reader.read()).isNotNull(); - } + try { + try (ParquetReader reader = ParquetReader.builder(new GroupReadSupport(), filePath) + .withConf(readConf) + .build()) { + assertThat(reader.read()).isNotNull(); + } - assertThat(kmsClient.dependency).isEqualTo("dependency"); - assertThat(UnitestUrlReadKMS.getStaticKmsURL()).isEqualTo(KmsClient.KMS_INSTANCE_ID_DEFAULT); + assertThat(kmsClient.dependency).isEqualTo("dependency"); + assertThat(UnitestUrlReadKMS.getStaticKmsURL()).isEqualTo(KmsClient.KMS_INSTANCE_ID_DEFAULT); + } finally { + KeyToolkit.removeKmsClientFactory(readConf); + } } @AfterAll diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java index 036745e118..37d6635721 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java @@ -22,8 +22,12 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; import java.util.concurrent.atomic.AtomicInteger; import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.Path; import org.apache.parquet.crypto.ParquetCryptoRuntimeException; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; @@ -31,19 +35,30 @@ public class KeyToolkitTest { private static final long CACHE_LIFETIME_MILLIS = 60_000; + private static final String MASTER_KEY_ID = "shared-master-key"; + + private final List configurationsWithFactories = new ArrayList<>(); @AfterEach public void clearCaches() { + for (Configuration configuration : configurationsWithFactories) { + KeyToolkit.removeKmsClientFactory(configuration); + } KeyToolkit.removeCacheEntriesForAllTokens(); } + private void setKmsClientFactory(Configuration configuration, KmsClientFactory factory) { + KeyToolkit.setKmsClientFactory(configuration, factory); + configurationsWithFactories.add(configuration); + } + @Test public void prefersConfiguredKmsClientFactory() { Configuration configuration = new Configuration(false); configuration.set(KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME, ReflectiveKmsClient.class.getName()); ConstructorInjectedKmsClient client = new ConstructorInjectedKmsClient("dependency"); AtomicInteger factoryCalls = new AtomicInteger(); - KeyToolkit.setKmsClientFactory(configuration, () -> { + setKmsClientFactory(configuration, () -> { factoryCalls.incrementAndGet(); return client; }); @@ -67,8 +82,8 @@ public void scopesKmsClientFactoryAndCacheToConfiguration() { Configuration secondConfiguration = new Configuration(false); ConstructorInjectedKmsClient firstClient = new ConstructorInjectedKmsClient("first"); ConstructorInjectedKmsClient secondClient = new ConstructorInjectedKmsClient("second"); - KeyToolkit.setKmsClientFactory(firstConfiguration, () -> firstClient); - KeyToolkit.setKmsClientFactory(secondConfiguration, () -> secondClient); + setKmsClientFactory(firstConfiguration, () -> firstClient); + setKmsClientFactory(secondConfiguration, () -> secondClient); KmsClient first = KeyToolkit.getKmsClient("DEFAULT", "DEFAULT", firstConfiguration, "DEFAULT", CACHE_LIFETIME_MILLIS); @@ -88,7 +103,7 @@ public void factoryRegistrationDoesNotReuseCachedReflectiveClient() { Configuration factoryConfiguration = new Configuration(false); ConstructorInjectedKmsClient factoryClient = new ConstructorInjectedKmsClient("dependency"); - KeyToolkit.setKmsClientFactory(factoryConfiguration, () -> factoryClient); + setKmsClientFactory(factoryConfiguration, () -> factoryClient); KmsClient actual = KeyToolkit.getKmsClient("instance", "url", factoryConfiguration, "token", CACHE_LIFETIME_MILLIS); @@ -101,10 +116,10 @@ public void replacingKmsClientFactoryDiscardsCachedClient() { Configuration configuration = new Configuration(false); ConstructorInjectedKmsClient firstClient = new ConstructorInjectedKmsClient("first"); ConstructorInjectedKmsClient replacementClient = new ConstructorInjectedKmsClient("replacement"); - KeyToolkit.setKmsClientFactory(configuration, () -> firstClient); + setKmsClientFactory(configuration, () -> firstClient); KmsClient first = KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS); - KeyToolkit.setKmsClientFactory(configuration, () -> replacementClient); + setKmsClientFactory(configuration, () -> replacementClient); KmsClient replacement = KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS); @@ -116,7 +131,7 @@ public void replacingKmsClientFactoryDiscardsCachedClient() { public void removeCacheEntriesForTokenClearsOnlyMatchingFactoryClients() { Configuration configuration = new Configuration(false); AtomicInteger factoryCalls = new AtomicInteger(); - KeyToolkit.setKmsClientFactory( + setKmsClientFactory( configuration, () -> new ConstructorInjectedKmsClient(Integer.toString(factoryCalls.incrementAndGet()))); KmsClient firstTokenClient = @@ -139,7 +154,7 @@ public void removeCacheEntriesForTokenClearsOnlyMatchingFactoryClients() { public void removeCacheEntriesForAllTokensClearsFactoryClients() { Configuration configuration = new Configuration(false); AtomicInteger factoryCalls = new AtomicInteger(); - KeyToolkit.setKmsClientFactory( + setKmsClientFactory( configuration, () -> new ConstructorInjectedKmsClient(Integer.toString(factoryCalls.incrementAndGet()))); KmsClient firstTokenClient = @@ -161,7 +176,7 @@ public void removeCacheEntriesForAllTokensClearsFactoryClients() { @Test public void rejectsNullKmsClientFromFactory() { Configuration configuration = new Configuration(false); - KeyToolkit.setKmsClientFactory(configuration, () -> null); + setKmsClientFactory(configuration, () -> null); assertThatThrownBy( () -> KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS)) @@ -169,6 +184,74 @@ public void rejectsNullKmsClientFromFactory() { .hasMessage("KmsClientFactory returned null"); } + @Test + public void removeKmsClientFactoryRemovesRegistrationForClientRetainingConfiguration() { + Configuration configuration = new Configuration(false); + configuration.set(KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME, ReflectiveKmsClient.class.getName()); + ConstructorInjectedKmsClient factoryClient = new ConstructorInjectedKmsClient("dependency"); + setKmsClientFactory(configuration, () -> factoryClient); + KmsClient registered = + KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS); + KeyToolkit.KmsClientCacheContext cacheContext = KeyToolkit.getKmsClientCacheContext(configuration); + cacheContext + .getKekWriteCache() + .getOrCreateInternalCache("token", CACHE_LIFETIME_MILLIS) + .put("master-key", new KeyToolkit.KeyEncryptionKey(new byte[16], new byte[16], "wrapped")); + cacheContext + .getKekReadCache() + .getOrCreateInternalCache("token", CACHE_LIFETIME_MILLIS) + .put("kek", new byte[16]); + + KeyToolkit.removeKmsClientFactory(configuration); + + KmsClient fallback = KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS); + assertThat(registered).isSameAs(factoryClient); + assertThat(factoryClient.configuration).isSameAs(configuration); + assertThat(fallback).isInstanceOf(ReflectiveKmsClient.class); + assertThat(cacheContext.getKmsClientCache().getOrCreateInternalCache("token", CACHE_LIFETIME_MILLIS)) + .isEmpty(); + assertThat(cacheContext.getKekWriteCache().getOrCreateInternalCache("token", CACHE_LIFETIME_MILLIS)) + .isEmpty(); + assertThat(cacheContext.getKekReadCache().getOrCreateInternalCache("token", CACHE_LIFETIME_MILLIS)) + .isEmpty(); + } + + @Test + public void isolatesDoubleWrappingWriteCacheByFactoryRegistration() { + TrackingKmsClient firstClient = new TrackingKmsClient("0123456789012346", false); + TrackingKmsClient secondClient = new TrackingKmsClient("6543210987654321", false); + Configuration firstConfiguration = newFactoryConfiguration(firstClient); + Configuration secondConfiguration = newFactoryConfiguration(secondClient); + byte[] dataKey = new byte[16]; + + new FileKeyWrapper(firstConfiguration, null).getEncryptionKeyMetadata(dataKey, MASTER_KEY_ID, true); + new FileKeyWrapper(secondConfiguration, null).getEncryptionKeyMetadata(dataKey, MASTER_KEY_ID, true); + + assertThat(firstClient.wrapCalls).hasValue(1); + assertThat(secondClient.wrapCalls).hasValue(1); + } + + @Test + public void isolatesDoubleWrappingReadCacheByFactoryRegistration() { + TrackingKmsClient permittedClient = new TrackingKmsClient("0123456789012346", false); + Configuration permittedConfiguration = newFactoryConfiguration(permittedClient); + byte[] dataKey = new byte[16]; + byte[] keyMetadata = + new FileKeyWrapper(permittedConfiguration, null).getEncryptionKeyMetadata(dataKey, MASTER_KEY_ID, true); + FileKeyUnwrapper permittedUnwrapper = + new FileKeyUnwrapper(permittedConfiguration, new Path("encrypted.parquet")); + assertThat(permittedUnwrapper.getKey(keyMetadata)).isEqualTo(dataKey); + + TrackingKmsClient deniedClient = new TrackingKmsClient("6543210987654321", true); + Configuration deniedConfiguration = newFactoryConfiguration(deniedClient); + FileKeyUnwrapper deniedUnwrapper = new FileKeyUnwrapper(deniedConfiguration, new Path("encrypted.parquet")); + + assertThatThrownBy(() -> deniedUnwrapper.getKey(keyMetadata)) + .isInstanceOf(ParquetCryptoRuntimeException.class) + .hasMessage("KMS access denied"); + assertThat(deniedClient.unwrapCalls).hasValue(1); + } + @Test public void usesConfiguredClassWhenFactoryIsNotSet() { Configuration configuration = new Configuration(false); @@ -180,6 +263,13 @@ public void usesConfiguredClassWhenFactoryIsNotSet() { assertThat(((ReflectiveKmsClient) client).initializeCalls).isEqualTo(1); } + private Configuration newFactoryConfiguration(KmsClient kmsClient) { + Configuration configuration = new Configuration(false); + configuration.setBoolean(KeyToolkit.DOUBLE_WRAPPING_PROPERTY_NAME, true); + setKmsClientFactory(configuration, () -> kmsClient); + return configuration; + } + private static class ConstructorInjectedKmsClient implements KmsClient { private final String dependency; private Configuration configuration; @@ -213,6 +303,39 @@ public byte[] unwrapKey(String wrappedKey, String masterKeyIdentifier) { } } + private static class TrackingKmsClient implements KmsClient { + private final byte[] masterKey; + private final boolean denyUnwrap; + private final AtomicInteger wrapCalls = new AtomicInteger(); + private final AtomicInteger unwrapCalls = new AtomicInteger(); + + private TrackingKmsClient(String masterKey, boolean denyUnwrap) { + this.masterKey = masterKey.getBytes(StandardCharsets.UTF_8); + this.denyUnwrap = denyUnwrap; + } + + @Override + public void initialize( + Configuration configuration, String kmsInstanceID, String kmsInstanceURL, String accessToken) {} + + @Override + public String wrapKey(byte[] keyBytes, String masterKeyIdentifier) { + wrapCalls.incrementAndGet(); + return KeyToolkit.encryptKeyLocally( + keyBytes, masterKey, masterKeyIdentifier.getBytes(StandardCharsets.UTF_8)); + } + + @Override + public byte[] unwrapKey(String wrappedKey, String masterKeyIdentifier) { + unwrapCalls.incrementAndGet(); + if (denyUnwrap) { + throw new ParquetCryptoRuntimeException("KMS access denied"); + } + return KeyToolkit.decryptKeyLocally( + wrappedKey, masterKey, masterKeyIdentifier.getBytes(StandardCharsets.UTF_8)); + } + } + public static class ReflectiveKmsClient implements KmsClient { private int initializeCalls; From 446345d35b83920766844133c137660c74917f26 Mon Sep 17 00:00:00 2001 From: Steven Jones Date: Thu, 10 Sep 2026 15:53:17 +0000 Subject: [PATCH 3/6] GH-3683: Pass KMS factory creation context --- .../parquet/crypto/keytools/KeyToolkit.java | 6 +- .../crypto/keytools/KmsClientFactory.java | 16 +++- .../apache/parquet/crypto/TestKmsUrlRead.java | 22 +++++- .../crypto/keytools/KeyToolkitTest.java | 79 ++++++++++++++++--- 4 files changed, 105 insertions(+), 18 deletions(-) diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java index a37adbd3d5..df4be72c68 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java @@ -369,6 +369,10 @@ public static void removeCacheEntriesForAllTokens() { * KeyToolkit} in the same way as reflectively constructed clients. The KMS client and key * encryption key caches are isolated from registrations for other configurations. * + *

The registration is associated with the exact {@code Configuration} object, so changing + * that object's properties does not affect the registration. The factory receives the current + * configuration and resolved KMS details when it creates a client. + * *

The association is not serialized, and configuration copies must register their own * factory. The caller must invoke {@link #removeKmsClientFactory(Configuration)} after all * readers and writers using the configuration have closed. Replacing a factory clears the @@ -499,7 +503,7 @@ private static KmsClient createAndInitKmsClient( KmsClient kmsClient = null; if (factory != null) { - kmsClient = factory.createKmsClient(); + kmsClient = factory.createKmsClient(configuration, kmsInstanceID, kmsInstanceURL, accessToken); if (kmsClient == null) { throw new ParquetCryptoRuntimeException("KmsClientFactory returned null"); } diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KmsClientFactory.java b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KmsClientFactory.java index ea5ac33fb6..ecfbff169a 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KmsClientFactory.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KmsClientFactory.java @@ -19,14 +19,24 @@ package org.apache.parquet.crypto.keytools; +import org.apache.hadoop.conf.Configuration; + /** Factory for creating {@link KmsClient} instances with programmatically supplied dependencies. */ @FunctionalInterface public interface KmsClientFactory { /** - * Creates a KMS client. {@link KeyToolkit} initializes the returned client before using it. + * Creates a new KMS client. {@link KeyToolkit} invokes this method for each uncached combination + * of access token and KMS instance ID, then initializes the returned client before using it. + * + *

Each invocation must return a distinct, uninitialized client. * - * @return a new or pre-built KMS client + * @param configuration current Hadoop configuration + * @param kmsInstanceID ID of the KMS instance + * @param kmsInstanceURL URL of the KMS instance + * @param accessToken KMS access token + * @return a new KMS client */ - KmsClient createKmsClient(); + KmsClient createKmsClient( + Configuration configuration, String kmsInstanceID, String kmsInstanceURL, String accessToken); } diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java index 761022bfd2..07fa5174ea 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java @@ -26,6 +26,7 @@ import java.io.IOException; import java.nio.charset.StandardCharsets; import java.nio.file.Files; +import java.util.ArrayList; import java.util.Base64; import java.util.Collections; import java.util.List; @@ -75,10 +76,18 @@ static String getStaticKmsURL() { private static class ConstructorInjectedKmsClient extends UnitestUrlReadKMS { private final String dependency; + private int initializeCalls; private ConstructorInjectedKmsClient(String dependency) { this.dependency = dependency; } + + @Override + public synchronized void initialize( + Configuration configuration, String kmsInstanceID, String kmsInstanceURL, String accessToken) { + initializeCalls++; + super.initialize(configuration, kmsInstanceID, kmsInstanceURL, accessToken); + } } @BeforeAll @@ -191,8 +200,13 @@ public void testSetKmsUrl() throws IOException { public void testProgrammaticKmsClientFactory() throws IOException { Configuration readConf = basicDecryptionConfig(); readConf.set(KeyToolkit.KEY_ACCESS_TOKEN_PROPERTY_NAME, "factory-token"); - ConstructorInjectedKmsClient kmsClient = new ConstructorInjectedKmsClient("dependency"); - KeyToolkit.setKmsClientFactory(readConf, () -> kmsClient); + List kmsClients = new ArrayList<>(); + KeyToolkit.setKmsClientFactory( + readConf, (ignoredConfiguration, ignoredKmsInstanceID, ignoredKmsInstanceURL, ignoredAccessToken) -> { + ConstructorInjectedKmsClient kmsClient = new ConstructorInjectedKmsClient("dependency"); + kmsClients.add(kmsClient); + return kmsClient; + }); try { try (ParquetReader reader = ParquetReader.builder(new GroupReadSupport(), filePath) @@ -201,7 +215,9 @@ public void testProgrammaticKmsClientFactory() throws IOException { assertThat(reader.read()).isNotNull(); } - assertThat(kmsClient.dependency).isEqualTo("dependency"); + assertThat(kmsClients).hasSize(1); + assertThat(kmsClients.get(0).dependency).isEqualTo("dependency"); + assertThat(kmsClients.get(0).initializeCalls).isEqualTo(1); assertThat(UnitestUrlReadKMS.getStaticKmsURL()).isEqualTo(KmsClient.KMS_INSTANCE_ID_DEFAULT); } finally { KeyToolkit.removeKmsClientFactory(readConf); diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java index 37d6635721..d7203018bf 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java @@ -58,7 +58,7 @@ public void prefersConfiguredKmsClientFactory() { configuration.set(KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME, ReflectiveKmsClient.class.getName()); ConstructorInjectedKmsClient client = new ConstructorInjectedKmsClient("dependency"); AtomicInteger factoryCalls = new AtomicInteger(); - setKmsClientFactory(configuration, () -> { + setKmsClientFactory(configuration, (conf, kmsId, kmsUrl, token) -> { factoryCalls.incrementAndGet(); return client; }); @@ -76,14 +76,69 @@ public void prefersConfiguredKmsClientFactory() { assertThat(client.initializeCalls).isEqualTo(1); } + @Test + public void factoryRegistrationSurvivesConfigurationMutationAndReceivesCurrentContext() { + Configuration configuration = new Configuration(false); + ConstructorInjectedKmsClient client = new ConstructorInjectedKmsClient("dependency"); + List factoryConfigurations = new ArrayList<>(); + List factoryValues = new ArrayList<>(); + List factoryKmsInstanceIDs = new ArrayList<>(); + List factoryKmsInstanceURLs = new ArrayList<>(); + List factoryAccessTokens = new ArrayList<>(); + setKmsClientFactory(configuration, (currentConfiguration, kmsInstanceID, kmsInstanceURL, accessToken) -> { + factoryConfigurations.add(currentConfiguration); + factoryValues.add(currentConfiguration.get("custom.factory.parameter")); + factoryKmsInstanceIDs.add(kmsInstanceID); + factoryKmsInstanceURLs.add(kmsInstanceURL); + factoryAccessTokens.add(accessToken); + return client; + }); + + configuration.set("custom.factory.parameter", "updated"); + + KmsClient actual = KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS); + + assertThat(actual).isSameAs(client); + assertThat(factoryConfigurations).containsExactly(configuration); + assertThat(factoryValues).containsExactly("updated"); + assertThat(factoryKmsInstanceIDs).containsExactly("instance"); + assertThat(factoryKmsInstanceURLs).containsExactly("url"); + assertThat(factoryAccessTokens).containsExactly("token"); + } + + @Test + public void createsDistinctKmsClientsForDifferentAccessTokens() { + Configuration configuration = new Configuration(false); + List clients = new ArrayList<>(); + setKmsClientFactory(configuration, (conf, kmsId, kmsUrl, token) -> { + ConstructorInjectedKmsClient client = new ConstructorInjectedKmsClient("client-" + clients.size()); + clients.add(client); + return client; + }); + + KmsClient first = + KeyToolkit.getKmsClient("instance", "url", configuration, "first-token", CACHE_LIFETIME_MILLIS); + KmsClient second = + KeyToolkit.getKmsClient("instance", "url", configuration, "second-token", CACHE_LIFETIME_MILLIS); + + assertThat(clients).hasSize(2); + assertThat(first).isSameAs(clients.get(0)); + assertThat(second).isSameAs(clients.get(1)); + assertThat(first).isNotSameAs(second); + assertThat(clients.get(0).accessToken).isEqualTo("first-token"); + assertThat(clients.get(1).accessToken).isEqualTo("second-token"); + assertThat(clients.get(0).initializeCalls).isEqualTo(1); + assertThat(clients.get(1).initializeCalls).isEqualTo(1); + } + @Test public void scopesKmsClientFactoryAndCacheToConfiguration() { Configuration firstConfiguration = new Configuration(false); Configuration secondConfiguration = new Configuration(false); ConstructorInjectedKmsClient firstClient = new ConstructorInjectedKmsClient("first"); ConstructorInjectedKmsClient secondClient = new ConstructorInjectedKmsClient("second"); - setKmsClientFactory(firstConfiguration, () -> firstClient); - setKmsClientFactory(secondConfiguration, () -> secondClient); + setKmsClientFactory(firstConfiguration, (conf, kmsId, kmsUrl, token) -> firstClient); + setKmsClientFactory(secondConfiguration, (conf, kmsId, kmsUrl, token) -> secondClient); KmsClient first = KeyToolkit.getKmsClient("DEFAULT", "DEFAULT", firstConfiguration, "DEFAULT", CACHE_LIFETIME_MILLIS); @@ -103,7 +158,7 @@ public void factoryRegistrationDoesNotReuseCachedReflectiveClient() { Configuration factoryConfiguration = new Configuration(false); ConstructorInjectedKmsClient factoryClient = new ConstructorInjectedKmsClient("dependency"); - setKmsClientFactory(factoryConfiguration, () -> factoryClient); + setKmsClientFactory(factoryConfiguration, (conf, kmsId, kmsUrl, token) -> factoryClient); KmsClient actual = KeyToolkit.getKmsClient("instance", "url", factoryConfiguration, "token", CACHE_LIFETIME_MILLIS); @@ -116,10 +171,10 @@ public void replacingKmsClientFactoryDiscardsCachedClient() { Configuration configuration = new Configuration(false); ConstructorInjectedKmsClient firstClient = new ConstructorInjectedKmsClient("first"); ConstructorInjectedKmsClient replacementClient = new ConstructorInjectedKmsClient("replacement"); - setKmsClientFactory(configuration, () -> firstClient); + setKmsClientFactory(configuration, (conf, kmsId, kmsUrl, token) -> firstClient); KmsClient first = KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS); - setKmsClientFactory(configuration, () -> replacementClient); + setKmsClientFactory(configuration, (conf, kmsId, kmsUrl, token) -> replacementClient); KmsClient replacement = KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS); @@ -133,7 +188,8 @@ public void removeCacheEntriesForTokenClearsOnlyMatchingFactoryClients() { AtomicInteger factoryCalls = new AtomicInteger(); setKmsClientFactory( configuration, - () -> new ConstructorInjectedKmsClient(Integer.toString(factoryCalls.incrementAndGet()))); + (conf, kmsId, kmsUrl, token) -> + new ConstructorInjectedKmsClient(Integer.toString(factoryCalls.incrementAndGet()))); KmsClient firstTokenClient = KeyToolkit.getKmsClient("instance", "url", configuration, "first-token", CACHE_LIFETIME_MILLIS); KmsClient otherTokenClient = @@ -156,7 +212,8 @@ public void removeCacheEntriesForAllTokensClearsFactoryClients() { AtomicInteger factoryCalls = new AtomicInteger(); setKmsClientFactory( configuration, - () -> new ConstructorInjectedKmsClient(Integer.toString(factoryCalls.incrementAndGet()))); + (conf, kmsId, kmsUrl, token) -> + new ConstructorInjectedKmsClient(Integer.toString(factoryCalls.incrementAndGet()))); KmsClient firstTokenClient = KeyToolkit.getKmsClient("instance", "url", configuration, "first-token", CACHE_LIFETIME_MILLIS); KmsClient secondTokenClient = @@ -176,7 +233,7 @@ public void removeCacheEntriesForAllTokensClearsFactoryClients() { @Test public void rejectsNullKmsClientFromFactory() { Configuration configuration = new Configuration(false); - setKmsClientFactory(configuration, () -> null); + setKmsClientFactory(configuration, (conf, kmsId, kmsUrl, token) -> null); assertThatThrownBy( () -> KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS)) @@ -189,7 +246,7 @@ public void removeKmsClientFactoryRemovesRegistrationForClientRetainingConfigura Configuration configuration = new Configuration(false); configuration.set(KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME, ReflectiveKmsClient.class.getName()); ConstructorInjectedKmsClient factoryClient = new ConstructorInjectedKmsClient("dependency"); - setKmsClientFactory(configuration, () -> factoryClient); + setKmsClientFactory(configuration, (conf, kmsId, kmsUrl, token) -> factoryClient); KmsClient registered = KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS); KeyToolkit.KmsClientCacheContext cacheContext = KeyToolkit.getKmsClientCacheContext(configuration); @@ -266,7 +323,7 @@ public void usesConfiguredClassWhenFactoryIsNotSet() { private Configuration newFactoryConfiguration(KmsClient kmsClient) { Configuration configuration = new Configuration(false); configuration.setBoolean(KeyToolkit.DOUBLE_WRAPPING_PROPERTY_NAME, true); - setKmsClientFactory(configuration, () -> kmsClient); + setKmsClientFactory(configuration, (conf, kmsId, kmsUrl, token) -> kmsClient); return configuration; } From d524e57c9bf403f72c3d390316c4123110271365 Mon Sep 17 00:00:00 2001 From: Steven Jones Date: Fri, 11 Sep 2026 18:54:31 +0000 Subject: [PATCH 4/6] GH-3683: Preserve KMS factories across configuration copies --- .../parquet/crypto/keytools/KeyToolkit.java | 54 +++++++++++++------ .../crypto/keytools/KeyToolkitTest.java | 28 ++++++++++ 2 files changed, 66 insertions(+), 16 deletions(-) diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java index df4be72c68..0b530b492f 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java @@ -22,10 +22,11 @@ import java.io.IOException; import java.util.Base64; import java.util.Collections; -import java.util.IdentityHashMap; +import java.util.HashMap; import java.util.Map; import java.util.Objects; import java.util.Set; +import java.util.UUID; import java.util.concurrent.ConcurrentMap; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.FileStatus; @@ -48,6 +49,9 @@ public class KeyToolkit { * KMS stands for “key management service”. */ public static final String KMS_CLIENT_CLASS_PROPERTY_NAME = "parquet.encryption.kms.client.class"; + + private static final String KMS_CLIENT_FACTORY_REGISTRATION_ID_PROPERTY_NAME = + "parquet.encryption.kms.client.factory.registration.id"; /** * ID of the KMS instance that will be used for encryption (if multiple KMS instances are available). */ @@ -123,10 +127,10 @@ public class KeyToolkit { private static final KmsClientCacheContext DEFAULT_KMS_CLIENT_CACHE_CONTEXT = new KmsClientCacheContext( null, KMS_CLIENT_CACHE_PER_TOKEN, KEK_WRITE_CACHE_PER_TOKEN, KEK_READ_CACHE_PER_TOKEN); - // Programmatically supplied factories and their caches, scoped to the exact Configuration object. - // Callers must remove registrations when the Configuration is no longer in use. - private static final Map KMS_CLIENT_FACTORY_REGISTRATIONS = - Collections.synchronizedMap(new IdentityHashMap<>()); + // Programmatically supplied factories and their caches, scoped by an ID copied with Configuration. + // Callers must remove registrations when the Configuration and its copies are no longer in use. + private static final Map KMS_CLIENT_FACTORY_REGISTRATIONS = + Collections.synchronizedMap(new HashMap<>()); private enum KmsClientCache { INSTANCE; @@ -369,14 +373,15 @@ public static void removeCacheEntriesForAllTokens() { * KeyToolkit} in the same way as reflectively constructed clients. The KMS client and key * encryption key caches are isolated from registrations for other configurations. * - *

The registration is associated with the exact {@code Configuration} object, so changing - * that object's properties does not affect the registration. The factory receives the current - * configuration and resolved KMS details when it creates a client. + *

The registration ID is stored in the {@code Configuration}, so copies made in the same JVM + * share the factory and caches. The factory receives the current configuration and resolved KMS + * details when it creates a client. * - *

The association is not serialized, and configuration copies must register their own - * factory. The caller must invoke {@link #removeKmsClientFactory(Configuration)} after all - * readers and writers using the configuration have closed. Replacing a factory clears the - * previous registration and its caches. + *

The factory itself is local to this JVM. A configuration deserialized in another JVM must + * register its factory before use. The caller must invoke {@link + * #removeKmsClientFactory(Configuration)} after all readers and writers using the configuration + * or its copies have closed. Replacing a factory clears the previous registration and its + * caches. * * @param configuration Hadoop configuration associated with the factory * @param kmsClientFactory factory used to create KMS clients @@ -384,8 +389,13 @@ public static void removeCacheEntriesForAllTokens() { public static void setKmsClientFactory(Configuration configuration, KmsClientFactory kmsClientFactory) { Objects.requireNonNull(configuration, "configuration"); Objects.requireNonNull(kmsClientFactory, "kmsClientFactory"); + String registrationId = configuration.getTrimmed(KMS_CLIENT_FACTORY_REGISTRATION_ID_PROPERTY_NAME); + if (stringIsEmpty(registrationId)) { + registrationId = UUID.randomUUID().toString(); + configuration.set(KMS_CLIENT_FACTORY_REGISTRATION_ID_PROPERTY_NAME, registrationId); + } KmsClientCacheContext previous = - KMS_CLIENT_FACTORY_REGISTRATIONS.put(configuration, new KmsClientCacheContext(kmsClientFactory)); + KMS_CLIENT_FACTORY_REGISTRATIONS.put(registrationId, new KmsClientCacheContext(kmsClientFactory)); if (previous != null) { previous.clear(); } @@ -401,7 +411,12 @@ public static void setKmsClientFactory(Configuration configuration, KmsClientFac */ public static void removeKmsClientFactory(Configuration configuration) { Objects.requireNonNull(configuration, "configuration"); - KmsClientCacheContext registration = KMS_CLIENT_FACTORY_REGISTRATIONS.remove(configuration); + String registrationId = configuration.getTrimmed(KMS_CLIENT_FACTORY_REGISTRATION_ID_PROPERTY_NAME); + if (stringIsEmpty(registrationId)) { + return; + } + configuration.unset(KMS_CLIENT_FACTORY_REGISTRATION_ID_PROPERTY_NAME); + KmsClientCacheContext registration = KMS_CLIENT_FACTORY_REGISTRATIONS.remove(registrationId); if (registration != null) { registration.clear(); } @@ -479,8 +494,15 @@ static KmsClient getKmsClient( } static KmsClientCacheContext getKmsClientCacheContext(Configuration configuration) { - KmsClientCacheContext cacheContext = KMS_CLIENT_FACTORY_REGISTRATIONS.get(configuration); - return cacheContext == null ? DEFAULT_KMS_CLIENT_CACHE_CONTEXT : cacheContext; + String registrationId = configuration.getTrimmed(KMS_CLIENT_FACTORY_REGISTRATION_ID_PROPERTY_NAME); + if (stringIsEmpty(registrationId)) { + return DEFAULT_KMS_CLIENT_CACHE_CONTEXT; + } + KmsClientCacheContext cacheContext = KMS_CLIENT_FACTORY_REGISTRATIONS.get(registrationId); + if (cacheContext == null) { + throw new ParquetCryptoRuntimeException("No KmsClientFactory is registered for this configuration"); + } + return cacheContext; } private static void clearKekWriteCaches() { diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java index d7203018bf..c5fab80d88 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java @@ -106,6 +106,34 @@ public void factoryRegistrationSurvivesConfigurationMutationAndReceivesCurrentCo assertThat(factoryAccessTokens).containsExactly("token"); } + @Test + public void configurationCopyUsesRegisteredKmsClientFactory() { + Configuration configuration = new Configuration(false); + ConstructorInjectedKmsClient client = new ConstructorInjectedKmsClient("dependency"); + setKmsClientFactory(configuration, (conf, kmsId, kmsUrl, token) -> client); + Configuration copy = new Configuration(configuration); + + KmsClient actual = KeyToolkit.getKmsClient("instance", "url", copy, "token", CACHE_LIFETIME_MILLIS); + + assertThat(actual).isSameAs(client); + assertThat(client.configuration).isSameAs(copy); + } + + @Test + public void missingFactoryForConfigurationCopyFailsEncryptionPropertiesCreation() { + Configuration configuration = new Configuration(false); + configuration.set(PropertiesDrivenCryptoFactory.UNIFORM_KEY_PROPERTY_NAME, MASTER_KEY_ID); + setKmsClientFactory( + configuration, (conf, kmsId, kmsUrl, token) -> new ConstructorInjectedKmsClient("dependency")); + Configuration copy = new Configuration(configuration); + KeyToolkit.removeKmsClientFactory(configuration); + + assertThatThrownBy(() -> new PropertiesDrivenCryptoFactory() + .getFileEncryptionProperties(copy, new Path("encrypted.parquet"), null)) + .isInstanceOf(ParquetCryptoRuntimeException.class) + .hasMessage("No KmsClientFactory is registered for this configuration"); + } + @Test public void createsDistinctKmsClientsForDifferentAccessTokens() { Configuration configuration = new Configuration(false); From 5f3051ecd8419092de104415d4e9d3e0c6a09d28 Mon Sep 17 00:00:00 2001 From: Steven Jones Date: Fri, 11 Sep 2026 19:16:10 +0000 Subject: [PATCH 5/6] GH-3683: Scope KEK caches by KMS instance --- parquet-hadoop/README.md | 22 ++++++++++ .../crypto/keytools/FileKeyWrapper.java | 9 +++- .../parquet/crypto/keytools/KeyToolkit.java | 15 +++---- .../crypto/keytools/KeyToolkitTest.java | 42 +++++++++++++++++++ 4 files changed, 79 insertions(+), 9 deletions(-) diff --git a/parquet-hadoop/README.md b/parquet-hadoop/README.md index cb4cf36225..13b5132701 100644 --- a/parquet-hadoop/README.md +++ b/parquet-hadoop/README.md @@ -455,6 +455,28 @@ If `false`, write files in encrypted footer mode, that fully encrypts the footer **Description:** Class implementing the KmsClient interface. "KMS" stands for “key management service”. The Client will interact with a KMS Server to wrap/unrwap encryption keys. **Default value:** None +KMS clients can also be supplied programmatically when they require constructor-injected dependencies: + +```java +KeyToolkit.setKmsClientFactory( + configuration, + (conf, kmsInstanceID, kmsInstanceURL, accessToken) -> new CustomKmsClient(dependency)); +try { + // Construct and close readers and writers using configuration or its copies. +} finally { + KeyToolkit.removeKmsClientFactory(configuration); +} +``` + +A registered factory takes precedence over `parquet.encryption.kms.client.class`. Each invocation must return a +distinct, uninitialized `KmsClient`; `KeyToolkit` initializes and caches it. The factory can be invoked concurrently +for different access-token and KMS-instance combinations, so it must be thread-safe. + +Copies of the `Configuration` in the same JVM share the registration and its caches. A configuration deserialized in +another JVM must register the factory there before use. Call `removeKmsClientFactory` only after all readers and writers +using the configuration and its copies have closed. Registering another factory for the configuration or one of its +copies replaces the previous factory and clears the registration's caches. + --- **Property:** `parquet.encryption.kms.instance.id` diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyWrapper.java b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyWrapper.java index 6d2421abc4..f658e7fd91 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyWrapper.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyWrapper.java @@ -22,6 +22,7 @@ import java.nio.charset.StandardCharsets; import java.security.SecureRandom; import java.util.Arrays; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import org.apache.hadoop.conf.Configuration; import org.apache.parquet.crypto.ParquetCryptoRuntimeException; @@ -89,9 +90,13 @@ public class FileKeyWrapper { } if (doubleWrapping) { - TwoLevelCacheWithExpiration kekWriteCache = cacheContext.getKekWriteCache(); + TwoLevelCacheWithExpiration> kekWriteCache = + cacheContext.getKekWriteCache(); kekWriteCache.checkCacheForExpiredTokens(cacheEntryLifetime); - KEKPerMasterKeyID = kekWriteCache.getOrCreateInternalCache(accessToken, cacheEntryLifetime); + ConcurrentMap> kekPerKmsInstanceID = + kekWriteCache.getOrCreateInternalCache(accessToken, cacheEntryLifetime); + KEKPerMasterKeyID = + kekPerKmsInstanceID.computeIfAbsent(kmsInstanceID, ignored -> new ConcurrentHashMap<>()); int kekLengthBits = configuration.getInt(KeyToolkit.KEK_LENGTH_PROPERTY_NAME, KeyToolkit.KEK_LENGTH_DEFAULT); if (Arrays.binarySearch(ACCEPTABLE_KEK_LENGTHS, kekLengthBits) < 0) { diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java index 0b530b492f..51f8383c7d 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java @@ -117,8 +117,8 @@ public class KeyToolkit { // KMS client two level cache: token -> KMSInstanceId -> KmsClient static final TwoLevelCacheWithExpiration KMS_CLIENT_CACHE_PER_TOKEN = KmsClientCache.INSTANCE.getCache(); - // KEK two level cache for wrapping: token -> MEK_ID -> KeyEncryptionKey - static final TwoLevelCacheWithExpiration KEK_WRITE_CACHE_PER_TOKEN = + // KEK cache for wrapping: token -> KMS instance ID -> master key ID -> KeyEncryptionKey + static final TwoLevelCacheWithExpiration> KEK_WRITE_CACHE_PER_TOKEN = KEKWriteCache.INSTANCE.getCache(); // KEK two level cache for unwrapping: token -> KEK_ID -> KEK bytes @@ -143,9 +143,10 @@ private TwoLevelCacheWithExpiration getCache() { private enum KEKWriteCache { INSTANCE; - private final TwoLevelCacheWithExpiration cache = new TwoLevelCacheWithExpiration<>(); + private final TwoLevelCacheWithExpiration> cache = + new TwoLevelCacheWithExpiration<>(); - private TwoLevelCacheWithExpiration getCache() { + private TwoLevelCacheWithExpiration> getCache() { return cache; } } @@ -162,7 +163,7 @@ private TwoLevelCacheWithExpiration getCache() { static final class KmsClientCacheContext { private final KmsClientFactory factory; private final TwoLevelCacheWithExpiration kmsClientCache; - private final TwoLevelCacheWithExpiration kekWriteCache; + private final TwoLevelCacheWithExpiration> kekWriteCache; private final TwoLevelCacheWithExpiration kekReadCache; private KmsClientCacheContext(KmsClientFactory factory) { @@ -176,7 +177,7 @@ private KmsClientCacheContext(KmsClientFactory factory) { private KmsClientCacheContext( KmsClientFactory factory, TwoLevelCacheWithExpiration kmsClientCache, - TwoLevelCacheWithExpiration kekWriteCache, + TwoLevelCacheWithExpiration> kekWriteCache, TwoLevelCacheWithExpiration kekReadCache) { this.factory = factory; this.kmsClientCache = kmsClientCache; @@ -188,7 +189,7 @@ TwoLevelCacheWithExpiration getKmsClientCache() { return kmsClientCache; } - TwoLevelCacheWithExpiration getKekWriteCache() { + TwoLevelCacheWithExpiration> getKekWriteCache() { return kekWriteCache; } diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java index c5fab80d88..5e7caeef67 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java @@ -25,6 +25,7 @@ import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicInteger; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.Path; @@ -281,6 +282,7 @@ public void removeKmsClientFactoryRemovesRegistrationForClientRetainingConfigura cacheContext .getKekWriteCache() .getOrCreateInternalCache("token", CACHE_LIFETIME_MILLIS) + .computeIfAbsent("instance", ignored -> new ConcurrentHashMap<>()) .put("master-key", new KeyToolkit.KeyEncryptionKey(new byte[16], new byte[16], "wrapped")); cacheContext .getKekReadCache() @@ -316,6 +318,46 @@ public void isolatesDoubleWrappingWriteCacheByFactoryRegistration() { assertThat(secondClient.wrapCalls).hasValue(1); } + @Test + public void isolatesDoubleWrappingWriteCacheByKmsInstanceForConfigurationCopies() { + String firstKmsInstanceID = "first-instance"; + String secondKmsInstanceID = "second-instance"; + TrackingKmsClient firstClient = new TrackingKmsClient("0123456789012346", false); + TrackingKmsClient secondClient = new TrackingKmsClient("6543210987654321", false); + Configuration firstConfiguration = new Configuration(false); + firstConfiguration.setBoolean(KeyToolkit.DOUBLE_WRAPPING_PROPERTY_NAME, true); + firstConfiguration.set(KeyToolkit.KEY_ACCESS_TOKEN_PROPERTY_NAME, "shared-token"); + firstConfiguration.set(KeyToolkit.KMS_INSTANCE_ID_PROPERTY_NAME, firstKmsInstanceID); + setKmsClientFactory( + firstConfiguration, + (conf, kmsId, kmsUrl, token) -> kmsId.equals(firstKmsInstanceID) ? firstClient : secondClient); + Configuration secondConfiguration = new Configuration(firstConfiguration); + secondConfiguration.set(KeyToolkit.KMS_INSTANCE_ID_PROPERTY_NAME, secondKmsInstanceID); + byte[] dataKey = new byte[16]; + + byte[] firstMetadata = + new FileKeyWrapper(firstConfiguration, null).getEncryptionKeyMetadata(dataKey, MASTER_KEY_ID, true); + byte[] secondMetadata = + new FileKeyWrapper(secondConfiguration, null).getEncryptionKeyMetadata(dataKey, MASTER_KEY_ID, true); + + assertThat(firstClient.wrapCalls).hasValue(1); + assertThat(secondClient.wrapCalls).hasValue(1); + assertThat(KeyMaterial.parse(new String(firstMetadata, StandardCharsets.UTF_8)) + .getKmsInstanceID()) + .isEqualTo(firstKmsInstanceID); + assertThat(KeyMaterial.parse(new String(secondMetadata, StandardCharsets.UTF_8)) + .getKmsInstanceID()) + .isEqualTo(secondKmsInstanceID); + + assertThat(new FileKeyUnwrapper(firstConfiguration, new Path("first.parquet")).getKey(firstMetadata)) + .isEqualTo(dataKey); + KeyToolkit.getKmsClientCacheContext(firstConfiguration) + .getKekReadCache() + .clear(); + assertThat(new FileKeyUnwrapper(secondConfiguration, new Path("second.parquet")).getKey(secondMetadata)) + .isEqualTo(dataKey); + } + @Test public void isolatesDoubleWrappingReadCacheByFactoryRegistration() { TrackingKmsClient permittedClient = new TrackingKmsClient("0123456789012346", false); From a311069882f532ded045699568d50a6206e95834 Mon Sep 17 00:00:00 2001 From: Steven Jones Date: Fri, 11 Sep 2026 19:45:56 +0000 Subject: [PATCH 6/6] GH-3683: Use KMS URL default in assertion --- .../src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java index 07fa5174ea..d30505a0f3 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestKmsUrlRead.java @@ -218,7 +218,7 @@ public void testProgrammaticKmsClientFactory() throws IOException { assertThat(kmsClients).hasSize(1); assertThat(kmsClients.get(0).dependency).isEqualTo("dependency"); assertThat(kmsClients.get(0).initializeCalls).isEqualTo(1); - assertThat(UnitestUrlReadKMS.getStaticKmsURL()).isEqualTo(KmsClient.KMS_INSTANCE_ID_DEFAULT); + assertThat(UnitestUrlReadKMS.getStaticKmsURL()).isEqualTo(KmsClient.KMS_INSTANCE_URL_DEFAULT); } finally { KeyToolkit.removeKmsClientFactory(readConf); }