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/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..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 @@ -19,12 +19,10 @@ 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; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import org.apache.hadoop.conf.Configuration; import org.apache.parquet.crypto.ParquetCryptoRuntimeException; @@ -73,8 +71,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 +82,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 +90,13 @@ 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); + 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 854976d37b..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 @@ -21,7 +21,12 @@ import java.io.IOException; import java.util.Base64; +import java.util.Collections; +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; @@ -44,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). */ @@ -109,13 +117,21 @@ 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 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 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; private final TwoLevelCacheWithExpiration cache = new TwoLevelCacheWithExpiration<>(); @@ -127,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; } } @@ -143,6 +160,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 +287,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 +348,79 @@ 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 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 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 + */ + 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(registrationId, 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"); + 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(); + } } /** @@ -335,32 +466,82 @@ 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) { + 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() { + 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..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 @@ -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_URL_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..5e7caeef67 --- /dev/null +++ b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java @@ -0,0 +1,487 @@ +/* + * 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.ConcurrentHashMap; +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 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); + 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) + .computeIfAbsent("instance", ignored -> new ConcurrentHashMap<>()) + .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 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); + 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(); + } + } +}