ggershinsky commented on code in PR #3785: URL: https://github.com/apache/parquet-java/pull/3785#discussion_r3999209647
########## 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<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() Review Comment: can you add a basic unitest that asserts an exception when neither factory nor client property are set? ########## parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/KeyToolkit.java: ########## @@ -109,13 +117,21 @@ public class KeyToolkit { // KMS client two level cache: token -> KMSInstanceId -> KmsClient static final TwoLevelCacheWithExpiration<KmsClient> KMS_CLIENT_CACHE_PER_TOKEN = KmsClientCache.INSTANCE.getCache(); - // KEK two level cache for wrapping: token -> MEK_ID -> KeyEncryptionKey - static final TwoLevelCacheWithExpiration<KeyEncryptionKey> KEK_WRITE_CACHE_PER_TOKEN = + // KEK cache for wrapping: token -> KMS instance ID -> master key ID -> KeyEncryptionKey + static final TwoLevelCacheWithExpiration<ConcurrentMap<String, KeyEncryptionKey>> KEK_WRITE_CACHE_PER_TOKEN = Review Comment: I'm somewhat uneasy about changing the existing functionality. Is this required for the new feature? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
