This is an automated email from the ASF dual-hosted git repository.
ggershinsky pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/parquet-java.git
The following commit(s) were added to refs/heads/master by this push:
new 1c9367fd6 GH-3683: Support programmatic KMS client factories (#3785)
1c9367fd6 is described below
commit 1c9367fd6d61146efe530af749cd57b6b002fc79
Author: Steven Ware Jones <[email protected]>
AuthorDate: Tue Sep 22 03:12:02 2026 -0400
GH-3683: Support programmatic KMS client factories (#3785)
* GH-3683: Support programmatic KMS client factories
* GH-3683: Isolate factory caches and add cleanup
* GH-3683: Pass KMS factory creation context
* GH-3683: Preserve KMS factories across configuration copies
* GH-3683: Scope KEK caches by KMS instance
* GH-3683: Use KMS URL default in assertion
* GH-3683: Limit KMS-scoped KEK cache to factories
* GH-3683: Document and test KMS client factories
---
parquet-hadoop/README.md | 22 +
.../parquet/crypto/keytools/FileKeyUnwrapper.java | 13 +-
.../parquet/crypto/keytools/FileKeyWrapper.java | 12 +-
.../apache/parquet/crypto/keytools/KeyToolkit.java | 243 +++++++++-
.../parquet/crypto/keytools/KmsClientFactory.java | 45 ++
.../org/apache/parquet/crypto/TestKmsUrlRead.java | 142 +++++-
.../parquet/crypto/keytools/KeyToolkitTest.java | 496 +++++++++++++++++++++
7 files changed, 918 insertions(+), 55 deletions(-)
diff --git a/parquet-hadoop/README.md b/parquet-hadoop/README.md
index cb4cf3622..13b513270 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 b681187de..9bb0ffcfa 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<byte[]> kekReadCache =
cacheContext.getKekReadCache();
+ kekReadCache.checkCacheForExpiredTokens(cacheEntryLifetime);
+ kekPerKekID = kekReadCache.getOrCreateInternalCache(accessToken,
cacheEntryLifetime);
if (LOG.isDebugEnabled()) {
LOG.debug(
@@ -168,7 +169,7 @@ public class FileKeyUnwrapper implements
DecryptionKeyRetriever {
}
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 195a02424..7c7683eda 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,7 @@ public class FileKeyWrapper {
}
if (doubleWrapping) {
- KEK_WRITE_CACHE_PER_TOKEN.checkCacheForExpiredTokens(cacheEntryLifetime);
- KEKPerMasterKeyID =
KEK_WRITE_CACHE_PER_TOKEN.getOrCreateInternalCache(accessToken,
cacheEntryLifetime);
+ KEKPerMasterKeyID = cacheContext.getOrCreateKekWriteCache(accessToken,
kmsInstanceID, 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 854976d37..9d3da61ec 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,13 @@ package org.apache.parquet.crypto.keytools;
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.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileStatus;
@@ -44,6 +50,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).
*/
@@ -116,6 +125,14 @@ public class KeyToolkit {
// KEK two level cache for unwrapping: token -> KEK_ID -> KEK bytes
static final TwoLevelCacheWithExpiration<byte[]> 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, null,
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<String, KmsClientCacheContext>
KMS_CLIENT_FACTORY_REGISTRATIONS =
+ Collections.synchronizedMap(new HashMap<>());
+
private enum KmsClientCache {
INSTANCE;
private final TwoLevelCacheWithExpiration<KmsClient> cache = new
TwoLevelCacheWithExpiration<>();
@@ -143,6 +160,80 @@ public class KeyToolkit {
}
}
+ static final class KmsClientCacheContext {
+ private final KmsClientFactory factory;
+ private final TwoLevelCacheWithExpiration<KmsClient> kmsClientCache;
+ private final TwoLevelCacheWithExpiration<KeyEncryptionKey>
defaultKekWriteCache;
+ private final TwoLevelCacheWithExpiration<ConcurrentMap<String,
KeyEncryptionKey>> factoryKekWriteCache;
+ private final TwoLevelCacheWithExpiration<byte[]> kekReadCache;
+
+ private KmsClientCacheContext(KmsClientFactory factory) {
+ this(
+ factory,
+ new TwoLevelCacheWithExpiration<>(),
+ null,
+ new TwoLevelCacheWithExpiration<>(),
+ new TwoLevelCacheWithExpiration<>());
+ }
+
+ private KmsClientCacheContext(
+ KmsClientFactory factory,
+ TwoLevelCacheWithExpiration<KmsClient> kmsClientCache,
+ TwoLevelCacheWithExpiration<KeyEncryptionKey> defaultKekWriteCache,
+ TwoLevelCacheWithExpiration<ConcurrentMap<String, KeyEncryptionKey>>
factoryKekWriteCache,
+ TwoLevelCacheWithExpiration<byte[]> kekReadCache) {
+ this.factory = factory;
+ this.kmsClientCache = kmsClientCache;
+ this.defaultKekWriteCache = defaultKekWriteCache;
+ this.factoryKekWriteCache = factoryKekWriteCache;
+ this.kekReadCache = kekReadCache;
+ }
+
+ TwoLevelCacheWithExpiration<KmsClient> getKmsClientCache() {
+ return kmsClientCache;
+ }
+
+ ConcurrentMap<String, KeyEncryptionKey> getOrCreateKekWriteCache(
+ String accessToken, String kmsInstanceID, long cacheEntryLifetime) {
+ if (defaultKekWriteCache != null) {
+ defaultKekWriteCache.checkCacheForExpiredTokens(cacheEntryLifetime);
+ return defaultKekWriteCache.getOrCreateInternalCache(accessToken,
cacheEntryLifetime);
+ }
+ factoryKekWriteCache.checkCacheForExpiredTokens(cacheEntryLifetime);
+ ConcurrentMap<String, ConcurrentMap<String, KeyEncryptionKey>>
kekPerKmsInstanceID =
+ factoryKekWriteCache.getOrCreateInternalCache(accessToken,
cacheEntryLifetime);
+ return kekPerKmsInstanceID.computeIfAbsent(kmsInstanceID, ignored -> new
ConcurrentHashMap<>());
+ }
+
+ TwoLevelCacheWithExpiration<byte[]> getKekReadCache() {
+ return kekReadCache;
+ }
+
+ void removeCacheEntriesForToken(String accessToken) {
+ kmsClientCache.removeCacheEntriesForToken(accessToken);
+ if (defaultKekWriteCache != null) {
+ defaultKekWriteCache.removeCacheEntriesForToken(accessToken);
+ } else {
+ factoryKekWriteCache.removeCacheEntriesForToken(accessToken);
+ }
+ kekReadCache.removeCacheEntriesForToken(accessToken);
+ }
+
+ void clear() {
+ kmsClientCache.clear();
+ clearKekWriteCache();
+ kekReadCache.clear();
+ }
+
+ void clearKekWriteCache() {
+ if (defaultKekWriteCache != null) {
+ defaultKekWriteCache.clear();
+ } else {
+ factoryKekWriteCache.clear();
+ }
+ }
+ }
+
static class KeyWithMasterID {
private final byte[] keyBytes;
private final String masterID;
@@ -220,7 +311,7 @@ public class KeyToolkit {
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 +372,79 @@ public class KeyToolkit {
* @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.
+ *
+ * <p>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.
+ *
+ * <p>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.
+ *
+ * <p>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.
+ *
+ * <p>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 +490,82 @@ public class KeyToolkit {
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<String, KmsClient> 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() {
+ DEFAULT_KMS_CLIENT_CACHE_CONTEXT.clearKekWriteCache();
+ synchronized (KMS_CLIENT_FACTORY_REGISTRATIONS) {
+ for (KmsClientCacheContext cacheContext :
KMS_CLIENT_FACTORY_REGISTRATIONS.values()) {
+ cacheContext.clearKekWriteCache();
+ }
+ }
+ }
+
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 000000000..2f5de3117
--- /dev/null
+++
b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KmsClientFactory.java
@@ -0,0 +1,45 @@
+/*
+ * 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.
+ *
+ * <p>Each invocation must return a distinct, uninitialized client.
+ *
+ * <p>This method may be invoked concurrently for different access-token and
KMS-instance
+ * combinations. Implementations must be thread-safe.
+ *
+ * @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 5d3d3e04f..2467293d3 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 static
org.assertj.core.api.Assertions.assertThatThrownBy;
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,29 @@ public class TestKmsUrlRead {
}
}
+ private static class ConstructorInjectedKmsClient extends UnitestUrlReadKMS {
+ private final String dependency;
+ private int initializeCalls;
+ private int wrapCalls;
+
+ 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);
+ }
+
+ @Override
+ public synchronized String wrapKey(byte[] keyBytes, String
masterKeyIdentifier) {
+ wrapCalls++;
+ return super.wrapKey(keyBytes, masterKeyIdentifier);
+ }
+ }
+
@BeforeAll
public static void writeEncryptedFile() throws IOException {
Configuration writeConf = new Configuration();
@@ -88,29 +112,7 @@ public class TestKmsUrlRead {
filePath = new Path(Files.createTempFile("test-kms-url_", ".parquet")
.toAbsolutePath()
.toString());
-
- MessageType schema = SingleRow.getSchema();
- SimpleGroupFactory f = new SimpleGroupFactory(schema);
-
- try (ParquetWriter<Group> writer = ExampleParquetWriter.builder(filePath)
- .withConf(writeConf)
- .withWriteMode(OVERWRITE)
- .withType(schema)
- .build()) {
-
- for (SingleRow singleRow : DATA) {
- writer.write(f.newGroup()
- .append(SingleRow.BOOLEAN_FIELD_NAME, singleRow.boolean_field)
- .append(SingleRow.INT32_FIELD_NAME, singleRow.int32_field)
- .append(SingleRow.FLOAT_FIELD_NAME, singleRow.float_field)
- .append(SingleRow.DOUBLE_FIELD_NAME, singleRow.double_field)
- .append(SingleRow.BINARY_FIELD_NAME,
Binary.fromConstantByteArray(singleRow.ba_field))
- .append(
- SingleRow.FIXED_LENGTH_BINARY_FIELD_NAME,
- Binary.fromConstantByteArray(singleRow.flba_field))
- .append(SingleRow.PLAINTEXT_INT32_FIELD_NAME,
singleRow.plaintext_int32_field));
- }
- }
+ writeEncryptedFile(writeConf, filePath, DATA);
}
@Test
@@ -179,6 +181,74 @@ public class TestKmsUrlRead {
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<ConstructorInjectedKmsClient> kmsClients = new ArrayList<>();
+ KeyToolkit.setKmsClientFactory(
+ readConf, (ignoredConfiguration, ignoredKmsInstanceID,
ignoredKmsInstanceURL, ignoredAccessToken) -> {
+ ConstructorInjectedKmsClient kmsClient = new
ConstructorInjectedKmsClient("dependency");
+ kmsClients.add(kmsClient);
+ return kmsClient;
+ });
+
+ try {
+ try (ParquetReader<Group> 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);
+ }
+ }
+
+ @Test
+ public void testProgrammaticKmsClientFactoryForWrite() throws IOException {
+ Configuration writeConf = new Configuration();
+ writeConf.set(
+ EncryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME,
+ PropertiesDrivenCryptoFactory.class.getName());
+ writeConf.set(PropertiesDrivenCryptoFactory.UNIFORM_KEY_PROPERTY_NAME,
UNIFORM_MASTER_KEY_ID);
+ writeConf.set(InMemoryKMS.KEY_LIST_PROPERTY_NAME, KEY_LIST);
+ writeConf.set(KeyToolkit.KEY_ACCESS_TOKEN_PROPERTY_NAME,
"factory-writer-token");
+ Path factoryFilePath = new Path(Files.createTempFile("test-kms-factory_",
".parquet")
+ .toAbsolutePath()
+ .toString());
+ List<ConstructorInjectedKmsClient> kmsClients = new ArrayList<>();
+ KeyToolkit.setKmsClientFactory(
+ writeConf, (ignoredConfiguration, ignoredKmsInstanceID,
ignoredKmsInstanceURL, ignoredAccessToken) -> {
+ ConstructorInjectedKmsClient kmsClient = new
ConstructorInjectedKmsClient("dependency");
+ kmsClients.add(kmsClient);
+ return kmsClient;
+ });
+
+ try {
+ writeEncryptedFile(writeConf, factoryFilePath,
Collections.singletonList(DATA.get(0)));
+
+ assertThat(kmsClients).hasSize(1);
+ assertThat(kmsClients.get(0).dependency).isEqualTo("dependency");
+ assertThat(kmsClients.get(0).initializeCalls).isEqualTo(1);
+ assertThat(kmsClients.get(0).wrapCalls).isEqualTo(1);
+ try (ParquetReader<Group> reader = ParquetReader.builder(new
GroupReadSupport(), factoryFilePath)
+ .withConf(new Configuration())
+ .build()) {
+ assertThatThrownBy(reader::read)
+ .isInstanceOf(ParquetCryptoRuntimeException.class)
+ .hasMessageContaining("Trying to read file with encrypted footer.
No keys available");
+ }
+ } finally {
+ KeyToolkit.removeKmsClientFactory(writeConf);
+ factoryFilePath.getFileSystem(new
Configuration()).delete(factoryFilePath, false);
+ }
+ }
+
@AfterAll
public static void deleteFile() throws IOException {
filePath.getFileSystem(new Configuration()).delete(filePath, false);
@@ -194,4 +264,30 @@ public class TestKmsUrlRead {
return readConf;
}
+
+ private static void writeEncryptedFile(Configuration writeConf, Path
outputPath, List<SingleRow> rows)
+ throws IOException {
+ MessageType schema = SingleRow.getSchema();
+ SimpleGroupFactory groupFactory = new SimpleGroupFactory(schema);
+
+ try (ParquetWriter<Group> writer = ExampleParquetWriter.builder(outputPath)
+ .withConf(writeConf)
+ .withWriteMode(OVERWRITE)
+ .withType(schema)
+ .build()) {
+ for (SingleRow singleRow : rows) {
+ writer.write(groupFactory
+ .newGroup()
+ .append(SingleRow.BOOLEAN_FIELD_NAME, singleRow.boolean_field)
+ .append(SingleRow.INT32_FIELD_NAME, singleRow.int32_field)
+ .append(SingleRow.FLOAT_FIELD_NAME, singleRow.float_field)
+ .append(SingleRow.DOUBLE_FIELD_NAME, singleRow.double_field)
+ .append(SingleRow.BINARY_FIELD_NAME,
Binary.fromConstantByteArray(singleRow.ba_field))
+ .append(
+ SingleRow.FIXED_LENGTH_BINARY_FIELD_NAME,
+ Binary.fromConstantByteArray(singleRow.flba_field))
+ .append(SingleRow.PLAINTEXT_INT32_FIELD_NAME,
singleRow.plaintext_int32_field));
+ }
+ }
+ }
}
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 000000000..dc34f9b3f
--- /dev/null
+++
b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/keytools/KeyToolkitTest.java
@@ -0,0 +1,496 @@
+/*
+ * 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<Configuration> 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<Configuration> factoryConfigurations = new ArrayList<>();
+ List<String> factoryValues = new ArrayList<>();
+ List<String> factoryKmsInstanceIDs = new ArrayList<>();
+ List<String> factoryKmsInstanceURLs = new ArrayList<>();
+ List<String> 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
missingFactoryAndKmsClientClassFailsEncryptionPropertiesCreation() {
+ Configuration configuration = new Configuration(false);
+ configuration.set(PropertiesDrivenCryptoFactory.UNIFORM_KEY_PROPERTY_NAME,
MASTER_KEY_ID);
+
+ assertThatThrownBy(() -> new PropertiesDrivenCryptoFactory()
+ .getFileEncryptionProperties(configuration, new
Path("encrypted.parquet"), null))
+ .isInstanceOf(ParquetCryptoRuntimeException.class)
+ .hasMessage("Unspecified " +
KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME);
+ }
+
+ @Test
+ public void createsDistinctKmsClientsForDifferentAccessTokens() {
+ Configuration configuration = new Configuration(false);
+ List<ConstructorInjectedKmsClient> 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
+ .getOrCreateKekWriteCache("token", "instance", 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.getOrCreateKekWriteCache("token", "instance",
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);
+ 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();
+ }
+ }
+}
