Skip to content

GH-3683: Support programmatic KMS client factories #3785

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
stevenwarejones wants to merge 3 commits into
base: master
Choose a base branch
from
Open

GH-3683: Support programmatic KMS client factories #3785

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
cbf290f
GH-3683: Support programmatic KMS client factories
stevenwarejones Sep 9, 2026
b3da958
GH-3683: Isolate factory caches and add cleanup
stevenwarejones Sep 9, 2026
446345d
GH-3683: Pass KMS factory creation context
stevenwarejones Sep 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 7 additions & 6 deletions parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyUnwrapper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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(
Expand Down Expand Up @@ -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.");
Expand Down
14 changes: 7 additions & 7 deletions parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyWrapper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -73,25 +70,28 @@ 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(
KeyToolkit.KMS_INSTANCE_ID_PROPERTY_NAME, KmsClient.KMS_INSTANCE_ID_DEFAULT);
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();
kmsClient = kmsClientAndDetails.getKmsClient();
}

if (doubleWrapping) {
KEK_WRITE_CACHE_PER_TOKEN.checkCacheForExpiredTokens(cacheEntryLifetime);
KEKPerMasterKeyID = KEK_WRITE_CACHE_PER_TOKEN.getOrCreateInternalCache(accessToken, cacheEntryLifetime);
TwoLevelCacheWithExpiration<KeyEncryptionKey> kekWriteCache = cacheContext.getKekWriteCache();
kekWriteCache.checkCacheForExpiredTokens(cacheEntryLifetime);
KEKPerMasterKeyID = kekWriteCache.getOrCreateInternalCache(accessToken, cacheEntryLifetime);
int kekLengthBits =
configuration.getInt(KeyToolkit.KEK_LENGTH_PROPERTY_NAME, KeyToolkit.KEK_LENGTH_DEFAULT);
if (Arrays.binarySearch(ACCEPTABLE_KEK_LENGTHS, kekLengthBits) < 0) {
Expand Down
196 changes: 177 additions & 19 deletions parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,10 @@

import java.io.IOException;
import java.util.Base64;
import java.util.Collections;
import java.util.IdentityHashMap;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ConcurrentMap;
import org.apache.hadoop.conf.Configuration;
Expand Down Expand Up @@ -116,6 +120,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, KEK_READ_CACHE_PER_TOKEN);

// Programmatically supplied factories and their caches, scoped to the exact Configuration object.
// Callers must remove registrations when the Configuration is no longer in use.
private static final Map<Configuration, KmsClientCacheContext> KMS_CLIENT_FACTORY_REGISTRATIONS =
Collections.synchronizedMap(new IdentityHashMap<>());

private enum KmsClientCache {
INSTANCE;
private final TwoLevelCacheWithExpiration<KmsClient> cache = new TwoLevelCacheWithExpiration<>();
Expand Down Expand Up @@ -143,6 +155,56 @@ private TwoLevelCacheWithExpiration<byte[]> getCache() {
}
}

static final class KmsClientCacheContext {
private final KmsClientFactory factory;
private final TwoLevelCacheWithExpiration<KmsClient> kmsClientCache;
private final TwoLevelCacheWithExpiration<KeyEncryptionKey> kekWriteCache;
private final TwoLevelCacheWithExpiration<byte[]> kekReadCache;

private KmsClientCacheContext(KmsClientFactory factory) {
this(
factory,
new TwoLevelCacheWithExpiration<>(),
new TwoLevelCacheWithExpiration<>(),
new TwoLevelCacheWithExpiration<>());
}

private KmsClientCacheContext(
KmsClientFactory factory,
TwoLevelCacheWithExpiration<KmsClient> kmsClientCache,
TwoLevelCacheWithExpiration<KeyEncryptionKey> kekWriteCache,
TwoLevelCacheWithExpiration<byte[]> kekReadCache) {
this.factory = factory;
this.kmsClientCache = kmsClientCache;
this.kekWriteCache = kekWriteCache;
this.kekReadCache = kekReadCache;
}

TwoLevelCacheWithExpiration<KmsClient> getKmsClientCache() {
return kmsClientCache;
}

TwoLevelCacheWithExpiration<KeyEncryptionKey> getKekWriteCache() {
return kekWriteCache;
}

TwoLevelCacheWithExpiration<byte[]> 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;
Expand Down Expand Up @@ -220,7 +282,7 @@ public static void rotateMasterKeys(String folderPath, Configuration hadoopConfi
long currentTime = System.currentTimeMillis();
synchronized (lastCacheCleanForKeyRotationTimeLock) {
if (currentTime - lastCacheCleanForKeyRotationTime > CACHE_CLEAN_PERIOD_FOR_KEY_ROTATION) {
KEK_WRITE_CACHE_PER_TOKEN.clear();
clearKekWriteCaches();
lastCacheCleanForKeyRotationTime = currentTime;
}
}
Expand Down Expand Up @@ -281,15 +343,68 @@ public static void rotateMasterKeys(String folderPath, Configuration hadoopConfi
* @param accessToken access token
*/
public static void removeCacheEntriesForToken(String accessToken) {
KMS_CLIENT_CACHE_PER_TOKEN.removeCacheEntriesForToken(accessToken);
KEK_WRITE_CACHE_PER_TOKEN.removeCacheEntriesForToken(accessToken);
KEK_READ_CACHE_PER_TOKEN.removeCacheEntriesForToken(accessToken);
DEFAULT_KMS_CLIENT_CACHE_CONTEXT.removeCacheEntriesForToken(accessToken);
synchronized (KMS_CLIENT_FACTORY_REGISTRATIONS) {
for (KmsClientCacheContext cacheContext : KMS_CLIENT_FACTORY_REGISTRATIONS.values()) {
cacheContext.removeCacheEntriesForToken(accessToken);
}
}
}

public static void removeCacheEntriesForAllTokens() {
KMS_CLIENT_CACHE_PER_TOKEN.clear();
KEK_WRITE_CACHE_PER_TOKEN.clear();
KEK_READ_CACHE_PER_TOKEN.clear();
DEFAULT_KMS_CLIENT_CACHE_CONTEXT.clear();
synchronized (KMS_CLIENT_FACTORY_REGISTRATIONS) {
for (KmsClientCacheContext cacheContext : KMS_CLIENT_FACTORY_REGISTRATIONS.values()) {
cacheContext.clear();
}
}
}

/**
* Sets the factory used to create KMS clients for the supplied configuration.
*
* <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 is associated with the exact {@code Configuration} object, so changing
* that object's properties does not affect the registration. The factory receives the current
* configuration and resolved KMS details when it creates a client.
*
* <p>The association is not serialized, and configuration copies must register their own
* factory. The caller must invoke {@link #removeKmsClientFactory(Configuration)} after all
* readers and writers using the configuration have closed. Replacing a factory clears the
* previous registration and its caches.
*
* @param configuration Hadoop configuration associated with the factory
* @param kmsClientFactory factory used to create KMS clients
*/
public static void setKmsClientFactory(Configuration configuration, KmsClientFactory kmsClientFactory) {
Objects.requireNonNull(configuration, "configuration");
Objects.requireNonNull(kmsClientFactory, "kmsClientFactory");
KmsClientCacheContext previous =
KMS_CLIENT_FACTORY_REGISTRATIONS.put(configuration, new KmsClientCacheContext(kmsClientFactory));

@ggershinsky ggershinsky Sep 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What if a shell/notebook user adds some parameter to the config during a session?
They'll need to call removeKmsClientFactory and then setKmsClientFactory each time?
Are there usecases where config changes are hard to trace?
Maybe there is an alternative approach? (eg using something similar to the kms instance, say "parquet.encryption.kms.factory.instance")

@stevenwarejones stevenwarejones Sep 10, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mutating the same Configuration does not require deregistration. Registrations use object identity, so changes to its properties do not affect lookup.

I added a test that mutates the configuration after registration and verifies that the factory remains registered and receives the updated configuration. If a client has already been cached, changing a property will not recreate it; that is also the behavior of the existing reflective path. Calling setKmsClientFactory again replaces the registration and clears its caches, so a separate remove call is not needed.

A parquet.encryption.kms.factory.instance string would require a static ID-to-object registry because Configuration cannot contain the live factory itself. That would recreate the side channel this API is intended to remove and would have ambiguous behavior when a copied configuration is serialized to another JVM. Configuration copies therefore still need their own explicit registration.

@ggershinsky ggershinsky Sep 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good, thanks. Indeed, the Configuration.equals implementation ignores the content.

@ggershinsky ggershinsky Sep 11, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One more question on this. What if Spark/Flink/etc copies a Configuration into another object (today or in the future), some time during a session?

@stevenwarejones stevenwarejones Sep 11, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point. That would be a copied Configuration, so the identity-based registration would not be found.

One option is to store an opaque registration ID in the configuration. That would support copies within the same JVM, but it could not carry the live factory into another JVM. Transparent cross-JVM support would require Spark/Flink-side integration to register or create the factory on each worker.

Would same-JVM copy support, with an explicit error when no local factory is registered, be the right scope here? Or would you suggest a different approach?

@ggershinsky ggershinsky Sep 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nope, this option sounds good to me. As for an error, can you check that an exception is indeed thrown? To make sure a dataframe is not written unencrypted silently, when no factory or kms class parameter are found.

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");
KmsClientCacheContext registration = KMS_CLIENT_FACTORY_REGISTRATIONS.remove(configuration);
if (registration != null) {
registration.clear();
}
}

/**
Expand Down Expand Up @@ -335,32 +450,75 @@ static KmsClient getKmsClient(
String accessToken,
long cacheEntryLifetime) {

return getKmsClient(
kmsInstanceID,
kmsInstanceURL,
configuration,
accessToken,
cacheEntryLifetime,
getKmsClientCacheContext(configuration));
}

static KmsClient getKmsClient(
String kmsInstanceID,
String kmsInstanceURL,
Configuration configuration,
String accessToken,
long cacheEntryLifetime,
KmsClientCacheContext cacheContext) {

ConcurrentMap<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) {
KmsClientCacheContext cacheContext = KMS_CLIENT_FACTORY_REGISTRATIONS.get(configuration);
return cacheContext == null ? DEFAULT_KMS_CLIENT_CACHE_CONTEXT : cacheContext;
}

private static void clearKekWriteCaches() {
KEK_WRITE_CACHE_PER_TOKEN.clear();
synchronized (KMS_CLIENT_FACTORY_REGISTRATIONS) {
for (KmsClientCacheContext cacheContext : KMS_CLIENT_FACTORY_REGISTRATIONS.values()) {
cacheContext.getKekWriteCache().clear();
}
}
}

private static KmsClient createAndInitKmsClient(
Configuration configuration, String kmsInstanceID, String kmsInstanceURL, String accessToken) {
Configuration configuration,
String kmsInstanceID,
String kmsInstanceURL,
String accessToken,
KmsClientFactory factory) {

Class<?> kmsClientClass = null;
KmsClient kmsClient = null;

try {
kmsClientClass = ConfigurationUtil.getClassFromConfig(
configuration, KMS_CLIENT_CLASS_PROPERTY_NAME, KmsClient.class);

if (null == kmsClientClass) {
throw new ParquetCryptoRuntimeException("Unspecified " + KMS_CLIENT_CLASS_PROPERTY_NAME);
if (factory != null) {
kmsClient = factory.createKmsClient(configuration, kmsInstanceID, kmsInstanceURL, accessToken);
if (kmsClient == null) {
throw new ParquetCryptoRuntimeException("KmsClientFactory returned null");
}
} else {
try {
kmsClientClass = ConfigurationUtil.getClassFromConfig(
configuration, KMS_CLIENT_CLASS_PROPERTY_NAME, KmsClient.class);

if (null == kmsClientClass) {
throw new ParquetCryptoRuntimeException("Unspecified " + KMS_CLIENT_CLASS_PROPERTY_NAME);
}
kmsClient = (KmsClient) kmsClientClass.newInstance();
} catch (InstantiationException | IllegalAccessException | BadConfigurationException e) {
throw new ParquetCryptoRuntimeException("Could not instantiate KmsClient class: " + kmsClientClass, e);
}
kmsClient = (KmsClient) kmsClientClass.newInstance();
} catch (InstantiationException | IllegalAccessException | BadConfigurationException e) {
throw new ParquetCryptoRuntimeException("Could not instantiate KmsClient class: " + kmsClientClass, e);
}

kmsClient.initialize(configuration, kmsInstanceID, kmsInstanceURL, accessToken);
Expand Down
42 changes: 42 additions & 0 deletions parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KmsClientFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.
*
* <p>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);
}
Loading
Loading