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 854976d37b..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 @@ -21,6 +21,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.concurrent.ConcurrentMap; import org.apache.hadoop.conf.Configuration; @@ -116,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<>(); @@ -143,6 +155,56 @@ private TwoLevelCacheWithExpiration getCache() { } } + static final class KmsClientCacheContext { + private final KmsClientFactory factory; + 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 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(); + } + } + static class KeyWithMasterID { private final byte[] keyBytes; private final String masterID; @@ -220,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; } } @@ -281,15 +343,68 @@ 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); - KEK_WRITE_CACHE_PER_TOKEN.removeCacheEntriesForToken(accessToken); - KEK_READ_CACHE_PER_TOKEN.removeCacheEntriesForToken(accessToken); + DEFAULT_KMS_CLIENT_CACHE_CONTEXT.removeCacheEntriesForToken(accessToken); + synchronized (KMS_CLIENT_FACTORY_REGISTRATIONS) { + for (KmsClientCacheContext cacheContext : KMS_CLIENT_FACTORY_REGISTRATIONS.values()) { + cacheContext.removeCacheEntriesForToken(accessToken); + } + } } public static void removeCacheEntriesForAllTokens() { - KMS_CLIENT_CACHE_PER_TOKEN.clear(); - KEK_WRITE_CACHE_PER_TOKEN.clear(); - KEK_READ_CACHE_PER_TOKEN.clear(); + DEFAULT_KMS_CLIENT_CACHE_CONTEXT.clear(); + synchronized (KMS_CLIENT_FACTORY_REGISTRATIONS) { + for (KmsClientCacheContext cacheContext : KMS_CLIENT_FACTORY_REGISTRATIONS.values()) { + cacheContext.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 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 + * previous registration and its caches. + * + * @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"); + 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(); + } } /** @@ -335,32 +450,75 @@ static KmsClient getKmsClient( String accessToken, long 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 = - KMS_CLIENT_CACHE_PER_TOKEN.getOrCreateInternalCache(accessToken, cacheEntryLifetime); + cacheContext.getKmsClientCache().getOrCreateInternalCache(accessToken, cacheEntryLifetime); KmsClient kmsClient = kmsClientPerKmsInstanceCache.computeIfAbsent( kmsInstanceID, - (k) -> createAndInitKmsClient(configuration, kmsInstanceID, kmsInstanceURL, accessToken)); + (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) { + Configuration configuration, + String kmsInstanceID, + String kmsInstanceURL, + String accessToken, + KmsClientFactory factory) { 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 (factory != null) { + kmsClient = factory.createKmsClient(configuration, kmsInstanceID, kmsInstanceURL, accessToken); + 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..ecfbff169a --- /dev/null +++ b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KmsClientFactory.java @@ -0,0 +1,42 @@ +/* + * 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 org.apache.hadoop.conf.Configuration; + +/** Factory for creating {@link KmsClient} instances with programmatically supplied dependencies. */ +@FunctionalInterface +public interface KmsClientFactory { + + /** + * 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. + * + * @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( + 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 5d3d3e04fc..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; @@ -73,6 +74,22 @@ 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 public static void writeEncryptedFile() throws IOException { Configuration writeConf = new Configuration(); @@ -179,6 +196,34 @@ 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"); + 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) + .withConf(readConf) + .build()) { + assertThat(reader.read()).isNotNull(); + } + + 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); + } + } + @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..d7203018bf --- /dev/null +++ b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java @@ -0,0 +1,417 @@ +/* + * 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.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; + +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(); + setKmsClientFactory(configuration, (conf, kmsId, kmsUrl, token) -> { + 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 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, (conf, kmsId, kmsUrl, token) -> firstClient); + setKmsClientFactory(secondConfiguration, (conf, kmsId, kmsUrl, token) -> 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"); + setKmsClientFactory(factoryConfiguration, (conf, kmsId, kmsUrl, token) -> 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"); + setKmsClientFactory(configuration, (conf, kmsId, kmsUrl, token) -> firstClient); + KmsClient first = KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS); + + setKmsClientFactory(configuration, (conf, kmsId, kmsUrl, token) -> 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(); + setKmsClientFactory( + configuration, + (conf, kmsId, kmsUrl, token) -> + 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(); + setKmsClientFactory( + configuration, + (conf, kmsId, kmsUrl, token) -> + 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); + setKmsClientFactory(configuration, (conf, kmsId, kmsUrl, token) -> null); + + assertThatThrownBy( + () -> KeyToolkit.getKmsClient("instance", "url", configuration, "token", CACHE_LIFETIME_MILLIS)) + .isInstanceOf(ParquetCryptoRuntimeException.class) + .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, (conf, kmsId, kmsUrl, token) -> 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); + 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 Configuration newFactoryConfiguration(KmsClient kmsClient) { + Configuration configuration = new Configuration(false); + configuration.setBoolean(KeyToolkit.DOUBLE_WRAPPING_PROPERTY_NAME, true); + setKmsClientFactory(configuration, (conf, kmsId, kmsUrl, token) -> kmsClient); + return configuration; + } + + 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(); + } + } + + 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; + + 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(); + } + } +}