From b675a5c859fecc2052b524c872fc1c1d411efa1f Mon Sep 17 00:00:00 2001 From: Nevin Zheng Date: Tue, 21 Jul 2026 15:48:15 -0700 Subject: [PATCH 1/6] feat(kms): define key inspection API --- api/build.gradle.kts | 3 + .../gravitino/encryption/kms/KmsApi.java | 85 +++++++++++ .../kms/KmsAuthenticationException.java | 50 ++++++ .../gravitino/encryption/kms/KmsClient.java | 47 ++++++ .../encryption/kms/KmsClientFactory.java | 49 ++++++ .../kms/KmsConfigurationException.java | 49 ++++++ .../encryption/kms/KmsKeyProperties.java | 65 ++++++++ .../encryption/kms/KmsReference.java | 105 +++++++++++++ .../encryption/kms/TestFakeKmsClient.java | 44 ++++++ .../gravitino/encryption/kms/TestKmsApi.java | 47 ++++++ .../encryption/kms/TestKmsClient.java | 73 +++++++++ .../encryption/kms/TestKmsReference.java | 71 +++++++++ .../encryption/kms/FakeKmsClient.java | 142 ++++++++++++++++++ .../encryption/kms/TestKmsClientContract.java | 98 ++++++++++++ .../kms/TestKmsClientFactoryContract.java | 45 ++++++ 15 files changed, 973 insertions(+) create mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java create mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java create mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java create mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java create mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java create mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java create mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java create mode 100644 api/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java create mode 100644 api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsApi.java create mode 100644 api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java create mode 100644 api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsReference.java create mode 100644 api/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java create mode 100644 api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java create mode 100644 api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java diff --git a/api/build.gradle.kts b/api/build.gradle.kts index f0fe3ba5ee9..616703cb1eb 100644 --- a/api/build.gradle.kts +++ b/api/build.gradle.kts @@ -18,6 +18,7 @@ */ plugins { `maven-publish` + `java-test-fixtures` id("java") id("idea") } @@ -32,6 +33,8 @@ dependencies { testImplementation(libs.junit.jupiter.api) testImplementation(libs.junit.jupiter.params) testRuntimeOnly(libs.junit.jupiter.engine) + + testFixturesApi(libs.junit.jupiter.api) } tasks.build { diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java new file mode 100644 index 00000000000..e32e17d201f --- /dev/null +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java @@ -0,0 +1,85 @@ +/* + * 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.gravitino.encryption.kms; + +import java.util.Arrays; +import java.util.Locale; +import java.util.stream.Collectors; +import org.apache.gravitino.annotation.DeveloperApi; + +/** The KMS product API that defines key identifier and operation semantics. */ +@DeveloperApi +public enum KmsApi { + /** Amazon Web Services Key Management Service. */ + AWS_KMS("aws-kms"), + + /** Google Cloud Key Management Service. */ + GOOGLE_CLOUD_KMS("google-cloud-kms"), + + /** Microsoft Azure Key Vault. */ + AZURE_KEY_VAULT("azure-key-vault"), + + /** The OpenBao Transit secrets engine API. */ + OPENBAO_TRANSIT("openbao-transit"), + + /** The HashiCorp Vault Transit secrets engine API. */ + VAULT_TRANSIT("vault-transit"); + + private final String wireValue; + + KmsApi(String wireValue) { + this.wireValue = wireValue; + } + + /** + * Returns the stable configuration value for this API. + * + * @return the configuration value + */ + public String wireValue() { + return wireValue; + } + + /** + * Parses a configured API value. + * + * @param value configured API value + * @return the matching API + * @throws IllegalArgumentException if the value is null, blank, or unsupported + */ + public static KmsApi fromWireValue(String value) { + if (value == null || value.trim().isEmpty()) { + throw new IllegalArgumentException("KMS API cannot be blank"); + } + + String normalized = value.trim().toLowerCase(Locale.ROOT); + return Arrays.stream(values()) + .filter(api -> api.wireValue.equals(normalized)) + .findFirst() + .orElseThrow( + () -> + new IllegalArgumentException( + String.format( + "Unsupported KMS API '%s'; supported values are %s", + value, + Arrays.stream(values()) + .map(KmsApi::wireValue) + .collect(Collectors.joining(", "))))); + } +} diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java new file mode 100644 index 00000000000..d199c7cbd63 --- /dev/null +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java @@ -0,0 +1,50 @@ +/* + * 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.gravitino.encryption.kms; + +import com.google.errorprone.annotations.FormatMethod; +import com.google.errorprone.annotations.FormatString; +import org.apache.gravitino.exceptions.ConnectionFailedException; + +/** Indicates that a KMS backend rejected or could not resolve configured credentials. */ +public class KmsAuthenticationException extends ConnectionFailedException { + + /** + * Creates an authentication exception with the specified cause and detail message. + * + * @param cause the cause + * @param message the detail message + * @param args the arguments to the message + */ + @FormatMethod + public KmsAuthenticationException(Throwable cause, @FormatString String message, Object... args) { + super(cause, message, args); + } + + /** + * Creates an authentication exception with the specified detail message. + * + * @param message the detail message + * @param args the arguments to the message + */ + @FormatMethod + public KmsAuthenticationException(@FormatString String message, Object... args) { + super(message, args); + } +} diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java new file mode 100644 index 00000000000..bcc541f98dd --- /dev/null +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java @@ -0,0 +1,47 @@ +/* + * 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.gravitino.encryption.kms; + +import org.apache.gravitino.annotation.DeveloperApi; +import org.apache.gravitino.exceptions.ConnectionFailedException; + +/** + * Inspects keys managed by a configured KMS. + * + *

This is a server-side operation client, not a credential-vending API. Provider credentials + * authenticate calls made by the client and must never be returned to callers. This client does not + * perform cryptographic operations. + */ +@DeveloperApi +public interface KmsClient extends AutoCloseable { + + /** + * Reads the provider-reported properties of a key. + * + * @param reference key to inspect + * @return normalized key properties + * @throws IllegalArgumentException if the reference does not belong to this client + * @throws ConnectionFailedException if the provider cannot be queried + */ + KmsKeyProperties getKeyProperties(KmsReference reference); + + /** Releases resources owned by this client. */ + @Override + default void close() {} +} diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java new file mode 100644 index 00000000000..3cfc4898492 --- /dev/null +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java @@ -0,0 +1,49 @@ +/* + * 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.gravitino.encryption.kms; + +import java.util.Map; +import org.apache.gravitino.annotation.DeveloperApi; +import org.apache.gravitino.exceptions.ConnectionFailedException; + +/** Creates server-side KMS clients for one supported KMS API. */ +@DeveloperApi +public interface KmsClientFactory { + + /** + * Returns the KMS API implemented by this factory. + * + * @return the supported KMS API + */ + KmsApi api(); + + /** + * Creates a client bound to a configured KMS source. + * + *

Provider credentials are private implementation details of the returned client. They must + * not be exposed as Gravitino credentials or key properties. + * + * @param source logical name of the configured KMS instance + * @param properties provider-specific configuration + * @return the configured client + * @throws IllegalArgumentException if the source or configuration is invalid + * @throws ConnectionFailedException if required external initialization fails + */ + KmsClient create(String source, Map properties); +} diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java new file mode 100644 index 00000000000..ec1613666a6 --- /dev/null +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java @@ -0,0 +1,49 @@ +/* + * 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.gravitino.encryption.kms; + +import com.google.errorprone.annotations.FormatMethod; +import com.google.errorprone.annotations.FormatString; + +/** Indicates invalid KMS configuration detected during initialization. */ +public class KmsConfigurationException extends IllegalArgumentException { + + /** + * Creates a configuration exception with the specified detail message. + * + * @param message the detail message + * @param args the arguments to the message + */ + @FormatMethod + public KmsConfigurationException(@FormatString String message, Object... args) { + super(args.length == 0 ? message : String.format(message, args)); + } + + /** + * Creates a configuration exception with the specified cause and detail message. + * + * @param cause the cause + * @param message the detail message + * @param args the arguments to the message + */ + @FormatMethod + public KmsConfigurationException(Throwable cause, @FormatString String message, Object... args) { + super(args.length == 0 ? message : String.format(message, args), cause); + } +} diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java new file mode 100644 index 00000000000..ae1896a5fa8 --- /dev/null +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java @@ -0,0 +1,65 @@ +/* + * 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.gravitino.encryption.kms; + +import org.apache.gravitino.annotation.DeveloperApi; + +/** + * Common properties reported for a KMS key. + * + *

Capabilities describe downstream use; this API does not perform cryptographic operations. + */ +@DeveloperApi +public interface KmsKeyProperties { + + /** + * Returns the requested key reference. + * + * @return the key reference + */ + KmsReference reference(); + + /** + * Returns whether the provider reports that the key exists. + * + * @return whether the key exists + */ + boolean present(); + + /** + * Returns whether the provider reports that the key is enabled. + * + * @return whether the key is enabled + */ + boolean enabled(); + + /** + * Returns whether the key supports wrapping data-encryption keys. + * + * @return whether wrapping is supported + */ + boolean supportsWrapping(); + + /** + * Returns whether the key supports unwrapping data-encryption keys. + * + * @return whether unwrapping is supported + */ + boolean supportsUnwrapping(); +} diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java new file mode 100644 index 00000000000..9350e8844f5 --- /dev/null +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java @@ -0,0 +1,105 @@ +/* + * 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.gravitino.encryption.kms; + +import com.google.common.base.Preconditions; +import java.util.Objects; +import javax.annotation.Nullable; +import org.apache.commons.lang3.StringUtils; +import org.apache.gravitino.annotation.DeveloperApi; + +/** + * Identifies a key owned by a configured KMS source. + * + *

Contains no credentials or key material. + */ +@DeveloperApi +public final class KmsReference { + + private final KmsApi api; + private final String source; + private final String keyId; + + /** + * Creates a structurally valid key reference without contacting the provider. + * + * @param api explicitly selected KMS API + * @param source configured KMS client-instance name + * @param keyId provider-native key identifier + * @throws IllegalArgumentException if the API is null or either string is null or blank + */ + public KmsReference(KmsApi api, String source, String keyId) { + Preconditions.checkArgument(api != null, "KMS API cannot be null"); + Preconditions.checkArgument(StringUtils.isNotBlank(source), "KMS source cannot be blank"); + Preconditions.checkArgument(StringUtils.isNotBlank(keyId), "KMS key ID cannot be blank"); + + this.api = api; + this.source = source; + this.keyId = keyId; + } + + /** + * Returns the explicitly selected KMS API. + * + * @return the KMS API + */ + public KmsApi api() { + return api; + } + + /** + * Returns the configured KMS client-instance name. + * + * @return the source name + */ + public String source() { + return source; + } + + /** + * Returns the provider-native key identifier. + * + * @return the key identifier + */ + public String keyId() { + return keyId; + } + + @Override + public boolean equals(@Nullable Object other) { + if (this == other) { + return true; + } + if (!(other instanceof KmsReference)) { + return false; + } + KmsReference that = (KmsReference) other; + return api == that.api && source.equals(that.source) && keyId.equals(that.keyId); + } + + @Override + public int hashCode() { + return Objects.hash(api, source, keyId); + } + + @Override + public String toString() { + return String.format("KmsReference{api=%s, source='%s', keyId='%s'}", api, source, keyId); + } +} diff --git a/api/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java b/api/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java new file mode 100644 index 00000000000..7ad5e2acca8 --- /dev/null +++ b/api/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java @@ -0,0 +1,44 @@ +/* + * 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.gravitino.encryption.kms; + +public class TestFakeKmsClient extends TestKmsClientContract { + + private static final String SOURCE = "test"; + private static final String USABLE_KEY = "usable"; + private static final String MISSING_KEY = "missing"; + + private final FakeKmsClient client = + new FakeKmsClient(KmsApi.AWS_KMS, SOURCE).putKey(USABLE_KEY, true, true, true); + + @Override + protected KmsClient client() { + return client; + } + + @Override + protected KmsReference usableKey() { + return new KmsReference(KmsApi.AWS_KMS, SOURCE, USABLE_KEY); + } + + @Override + protected KmsReference missingKey() { + return new KmsReference(KmsApi.AWS_KMS, SOURCE, MISSING_KEY); + } +} diff --git a/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsApi.java b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsApi.java new file mode 100644 index 00000000000..cae91fb0303 --- /dev/null +++ b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsApi.java @@ -0,0 +1,47 @@ +/* + * 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.gravitino.encryption.kms; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class TestKmsApi { + + @Test + void testParsesStableWireValues() { + Assertions.assertEquals(KmsApi.AWS_KMS, KmsApi.fromWireValue("aws-kms")); + Assertions.assertEquals(KmsApi.GOOGLE_CLOUD_KMS, KmsApi.fromWireValue(" GOOGLE-CLOUD-KMS ")); + Assertions.assertEquals("azure-key-vault", KmsApi.AZURE_KEY_VAULT.wireValue()); + Assertions.assertEquals( + KmsApi.OPENBAO_TRANSIT, KmsApi.fromWireValue(KmsApi.OPENBAO_TRANSIT.wireValue())); + Assertions.assertEquals(KmsApi.VAULT_TRANSIT, KmsApi.fromWireValue("vault-transit")); + } + + @Test + void testRejectsUnknownValues() { + Assertions.assertThrows(IllegalArgumentException.class, () -> KmsApi.fromWireValue(null)); + Assertions.assertThrows(IllegalArgumentException.class, () -> KmsApi.fromWireValue(" ")); + + IllegalArgumentException exception = + Assertions.assertThrows( + IllegalArgumentException.class, () -> KmsApi.fromWireValue("custom")); + Assertions.assertTrue(exception.getMessage().contains("Unsupported KMS API 'custom'")); + Assertions.assertTrue(exception.getMessage().contains("openbao-transit")); + } +} diff --git a/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java new file mode 100644 index 00000000000..14fe5d8b4c2 --- /dev/null +++ b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java @@ -0,0 +1,73 @@ +/* + * 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.gravitino.encryption.kms; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class TestKmsClient { + + private static final KmsReference REFERENCE = + new KmsReference(KmsApi.AWS_KMS, "production", "key"); + + @Test + void testReturnsProviderProperties() { + KmsKeyProperties properties = new TestProperties(); + KmsClient client = reference -> properties; + + Assertions.assertSame(properties, client.getKeyProperties(REFERENCE)); + Assertions.assertTrue(properties.supportsWrapping()); + Assertions.assertFalse(properties.supportsUnwrapping()); + } + + @Test + void testDefaultClose() { + KmsClient client = reference -> new TestProperties(); + + Assertions.assertDoesNotThrow(client::close); + } + + private static final class TestProperties implements KmsKeyProperties { + + @Override + public KmsReference reference() { + return REFERENCE; + } + + @Override + public boolean present() { + return true; + } + + @Override + public boolean enabled() { + return true; + } + + @Override + public boolean supportsWrapping() { + return true; + } + + @Override + public boolean supportsUnwrapping() { + return false; + } + } +} diff --git a/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsReference.java b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsReference.java new file mode 100644 index 00000000000..fe6b566faa0 --- /dev/null +++ b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsReference.java @@ -0,0 +1,71 @@ +/* + * 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.gravitino.encryption.kms; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class TestKmsReference { + + @Test + void testStoresReferenceExactly() { + KmsReference reference = new KmsReference(KmsApi.AWS_KMS, "production", " alias/Customer-Key "); + + Assertions.assertEquals(KmsApi.AWS_KMS, reference.api()); + Assertions.assertEquals("production", reference.source()); + Assertions.assertEquals(" alias/Customer-Key ", reference.keyId()); + } + + @Test + void testRejectsMissingFields() { + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference(null, "production", "key")); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, null, "key")); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, "production", null)); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, "", "key")); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, " ", "key")); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, "production", "")); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, "production", " ")); + } + + @Test + void testValueSemantics() { + KmsReference first = new KmsReference(KmsApi.AWS_KMS, "production", "key"); + KmsReference same = new KmsReference(KmsApi.AWS_KMS, "production", "key"); + KmsReference differentApi = new KmsReference(KmsApi.GOOGLE_CLOUD_KMS, "production", "key"); + KmsReference differentSource = new KmsReference(KmsApi.AWS_KMS, "recovery", "key"); + KmsReference differentKey = new KmsReference(KmsApi.AWS_KMS, "production", "another-key"); + + Assertions.assertEquals(first, same); + Assertions.assertEquals(first.hashCode(), same.hashCode()); + Assertions.assertNotEquals(first, differentApi); + Assertions.assertNotEquals(first, differentSource); + Assertions.assertNotEquals(first, differentKey); + Assertions.assertNotEquals(first, null); + Assertions.assertNotEquals(first, "key"); + Assertions.assertEquals( + "KmsReference{api=AWS_KMS, source='production', keyId='key'}", first.toString()); + } +} diff --git a/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java b/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java new file mode 100644 index 00000000000..183ffb37b98 --- /dev/null +++ b/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java @@ -0,0 +1,142 @@ +/* + * 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.gravitino.encryption.kms; + +import java.util.HashMap; +import java.util.Map; + +/** In-memory KMS client for contract and consumer tests. */ +public final class FakeKmsClient implements KmsClient { + + private final KmsApi api; + private final String source; + private final Map keys = new HashMap<>(); + + /** + * Creates an empty fake client. + * + * @param api KMS API accepted by the client + * @param source configured source accepted by the client + */ + public FakeKmsClient(KmsApi api, String source) { + this.api = api; + this.source = source; + } + + /** + * Adds or replaces a key reported by this client. + * + * @param keyId provider-native key identifier + * @param enabled whether the key is enabled + * @param supportsWrapping whether the key supports wrapping + * @param supportsUnwrapping whether the key supports unwrapping + * @return this client + */ + public FakeKmsClient putKey( + String keyId, boolean enabled, boolean supportsWrapping, boolean supportsUnwrapping) { + keys.put(keyId, new KeyState(enabled, supportsWrapping, supportsUnwrapping)); + return this; + } + + /** {@inheritDoc} */ + @Override + public KmsKeyProperties getKeyProperties(KmsReference reference) { + requireReference(reference); + KeyState state = keys.get(reference.keyId()); + return state == null + ? new Properties(reference, false, false, false, false) + : new Properties( + reference, true, state.enabled, state.supportsWrapping, state.supportsUnwrapping); + } + + private void requireReference(KmsReference reference) { + if (reference == null) { + throw new IllegalArgumentException("KMS reference cannot be null"); + } + if (reference.api() != api) { + throw new IllegalArgumentException( + String.format("Expected KMS API %s but received %s", api, reference.api())); + } + if (!source.equals(reference.source())) { + throw new IllegalArgumentException( + String.format( + "KMS source %s does not match configured source %s", reference.source(), source)); + } + } + + private static final class KeyState { + + private final boolean enabled; + private final boolean supportsWrapping; + private final boolean supportsUnwrapping; + + private KeyState(boolean enabled, boolean supportsWrapping, boolean supportsUnwrapping) { + this.enabled = enabled; + this.supportsWrapping = supportsWrapping; + this.supportsUnwrapping = supportsUnwrapping; + } + } + + private static final class Properties implements KmsKeyProperties { + + private final KmsReference reference; + private final boolean present; + private final boolean enabled; + private final boolean supportsWrapping; + private final boolean supportsUnwrapping; + + private Properties( + KmsReference reference, + boolean present, + boolean enabled, + boolean supportsWrapping, + boolean supportsUnwrapping) { + this.reference = reference; + this.present = present; + this.enabled = enabled; + this.supportsWrapping = supportsWrapping; + this.supportsUnwrapping = supportsUnwrapping; + } + + @Override + public KmsReference reference() { + return reference; + } + + @Override + public boolean present() { + return present; + } + + @Override + public boolean enabled() { + return enabled; + } + + @Override + public boolean supportsWrapping() { + return supportsWrapping; + } + + @Override + public boolean supportsUnwrapping() { + return supportsUnwrapping; + } + } +} diff --git a/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java b/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java new file mode 100644 index 00000000000..855b4bdf6d9 --- /dev/null +++ b/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java @@ -0,0 +1,98 @@ +/* + * 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.gravitino.encryption.kms; + +import java.util.Arrays; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** Common inspection contract for implemented KMS clients. */ +public abstract class TestKmsClientContract { + + /** + * Returns the client under test. + * + * @return the client + */ + protected abstract KmsClient client(); + + /** + * Returns a reference to an enabled key that supports wrapping and unwrapping. + * + * @return the usable key reference + */ + protected abstract KmsReference usableKey(); + + /** + * Returns a reference to a key that does not exist. + * + * @return the missing key reference + */ + protected abstract KmsReference missingKey(); + + @Test + void testReportsUsableKeyProperties() { + KmsReference reference = usableKey(); + KmsKeyProperties properties = client().getKeyProperties(reference); + + Assertions.assertEquals(reference, properties.reference()); + Assertions.assertTrue(properties.present()); + Assertions.assertTrue(properties.enabled()); + Assertions.assertTrue(properties.supportsWrapping()); + Assertions.assertTrue(properties.supportsUnwrapping()); + } + + @Test + void testReportsMissingKey() { + KmsReference reference = missingKey(); + KmsKeyProperties properties = client().getKeyProperties(reference); + + Assertions.assertEquals(reference, properties.reference()); + Assertions.assertFalse(properties.present()); + } + + @Test + void testRejectsNullReference() { + Assertions.assertThrows(IllegalArgumentException.class, () -> client().getKeyProperties(null)); + } + + @Test + void testRejectsMismatchedSource() { + KmsReference reference = usableKey(); + KmsReference mismatched = + new KmsReference(reference.api(), reference.source() + "-other", reference.keyId()); + + Assertions.assertThrows( + IllegalArgumentException.class, () -> client().getKeyProperties(mismatched)); + } + + @Test + void testRejectsMismatchedApi() { + KmsReference reference = usableKey(); + KmsApi otherApi = + Arrays.stream(KmsApi.values()) + .filter(api -> api != reference.api()) + .findFirst() + .orElseThrow(IllegalStateException::new); + KmsReference mismatched = new KmsReference(otherApi, reference.source(), reference.keyId()); + + Assertions.assertThrows( + IllegalArgumentException.class, () -> client().getKeyProperties(mismatched)); + } +} diff --git a/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java b/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java new file mode 100644 index 00000000000..4136e5de41e --- /dev/null +++ b/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.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.gravitino.encryption.kms; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** Common contract for KMS client factories. */ +public abstract class TestKmsClientFactoryContract { + + /** + * Returns the factory under test. + * + * @return the factory + */ + protected abstract KmsClientFactory factory(); + + /** + * Returns the API expected from the factory. + * + * @return the expected API + */ + protected abstract KmsApi expectedApi(); + + @Test + void testReportsExpectedApi() { + Assertions.assertEquals(expectedApi(), factory().api()); + } +} From 6f6424df99f98f5afdc8627dc8bc3a418e7b7d60 Mon Sep 17 00:00:00 2001 From: Nevin Zheng Date: Wed, 22 Jul 2026 11:01:57 -0700 Subject: [PATCH 2/6] fix(kms): persist key references with canonical API values --- .../gravitino/encryption/kms/KmsApi.java | 4 ++++ .../gravitino/encryption/kms/KmsReference.java | 8 +++++++- .../apache/gravitino/json/TestJsonUtils.java | 18 ++++++++++++++++++ 3 files changed, 29 insertions(+), 1 deletion(-) diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java index e32e17d201f..28ef0d2dd46 100644 --- a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java @@ -18,6 +18,8 @@ */ package org.apache.gravitino.encryption.kms; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonValue; import java.util.Arrays; import java.util.Locale; import java.util.stream.Collectors; @@ -52,6 +54,7 @@ public enum KmsApi { * * @return the configuration value */ + @JsonValue public String wireValue() { return wireValue; } @@ -63,6 +66,7 @@ public String wireValue() { * @return the matching API * @throws IllegalArgumentException if the value is null, blank, or unsupported */ + @JsonCreator public static KmsApi fromWireValue(String value) { if (value == null || value.trim().isEmpty()) { throw new IllegalArgumentException("KMS API cannot be blank"); diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java index 9350e8844f5..49d7c798ff5 100644 --- a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java @@ -18,6 +18,8 @@ */ package org.apache.gravitino.encryption.kms; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; import com.google.common.base.Preconditions; import java.util.Objects; import javax.annotation.Nullable; @@ -44,7 +46,11 @@ public final class KmsReference { * @param keyId provider-native key identifier * @throws IllegalArgumentException if the API is null or either string is null or blank */ - public KmsReference(KmsApi api, String source, String keyId) { + @JsonCreator + public KmsReference( + @JsonProperty("api") KmsApi api, + @JsonProperty("source") String source, + @JsonProperty("keyId") String keyId) { Preconditions.checkArgument(api != null, "KMS API cannot be null"); Preconditions.checkArgument(StringUtils.isNotBlank(source), "KMS source cannot be blank"); Preconditions.checkArgument(StringUtils.isNotBlank(keyId), "KMS key ID cannot be blank"); diff --git a/common/src/test/java/org/apache/gravitino/json/TestJsonUtils.java b/common/src/test/java/org/apache/gravitino/json/TestJsonUtils.java index 84497a05055..a0214b99ad3 100644 --- a/common/src/test/java/org/apache/gravitino/json/TestJsonUtils.java +++ b/common/src/test/java/org/apache/gravitino/json/TestJsonUtils.java @@ -36,6 +36,8 @@ import org.apache.gravitino.dto.rel.partitions.ListPartitionDTO; import org.apache.gravitino.dto.rel.partitions.PartitionDTO; import org.apache.gravitino.dto.rel.partitions.RangePartitionDTO; +import org.apache.gravitino.encryption.kms.KmsApi; +import org.apache.gravitino.encryption.kms.KmsReference; import org.apache.gravitino.rel.indexes.Index; import org.apache.gravitino.rel.types.Type; import org.apache.gravitino.rel.types.Types; @@ -254,6 +256,22 @@ void testGetLong() throws Exception { assertEquals(1L, result); } + @Test + void testKmsReferenceAnyFieldMapperSerde() throws JsonProcessingException { + KmsReference reference = + new KmsReference(KmsApi.GOOGLE_CLOUD_KMS, "analytics-prod", "projects/p/keys/k"); + + String json = JsonUtils.anyFieldMapper().writeValueAsString(reference); + + assertEquals( + JsonUtils.anyFieldMapper() + .readTree( + "{\"api\":\"google-cloud-kms\",\"source\":\"analytics-prod\"," + + "\"keyId\":\"projects/p/keys/k\"}"), + JsonUtils.anyFieldMapper().readTree(json)); + assertEquals(reference, JsonUtils.anyFieldMapper().readValue(json, KmsReference.class)); + } + @Test void testPartitionDTOSerde() throws JsonProcessingException { String[] field1 = {"dt"}; From 391ec3d14410b4cef36d54936fba7c857f18abcb Mon Sep 17 00:00:00 2001 From: Nevin Zheng Date: Wed, 22 Jul 2026 13:11:24 -0700 Subject: [PATCH 3/6] refactor(kms): move provider SPI to common --- api/build.gradle.kts | 3 --- common/build.gradle.kts | 4 ++++ .../gravitino/encryption/kms/KmsAuthenticationException.java | 0 .../java/org/apache/gravitino/encryption/kms/KmsClient.java | 0 .../org/apache/gravitino/encryption/kms/KmsClientFactory.java | 0 .../gravitino/encryption/kms/KmsConfigurationException.java | 0 .../org/apache/gravitino/encryption/kms/KmsKeyProperties.java | 0 .../apache/gravitino/encryption/kms/TestFakeKmsClient.java | 0 .../org/apache/gravitino/encryption/kms/TestKmsClient.java | 0 .../org/apache/gravitino/encryption/kms/FakeKmsClient.java | 0 .../gravitino/encryption/kms/TestKmsClientContract.java | 0 .../encryption/kms/TestKmsClientFactoryContract.java | 0 12 files changed, 4 insertions(+), 3 deletions(-) rename {api => common}/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java (100%) rename {api => common}/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java (100%) rename {api => common}/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java (100%) rename {api => common}/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java (100%) rename {api => common}/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java (100%) rename {api => common}/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java (100%) rename {api => common}/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java (100%) rename {api => common}/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java (100%) rename {api => common}/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java (100%) rename {api => common}/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java (100%) diff --git a/api/build.gradle.kts b/api/build.gradle.kts index 616703cb1eb..f0fe3ba5ee9 100644 --- a/api/build.gradle.kts +++ b/api/build.gradle.kts @@ -18,7 +18,6 @@ */ plugins { `maven-publish` - `java-test-fixtures` id("java") id("idea") } @@ -33,8 +32,6 @@ dependencies { testImplementation(libs.junit.jupiter.api) testImplementation(libs.junit.jupiter.params) testRuntimeOnly(libs.junit.jupiter.engine) - - testFixturesApi(libs.junit.jupiter.api) } tasks.build { diff --git a/common/build.gradle.kts b/common/build.gradle.kts index e267e9a118e..3f20efd17a9 100644 --- a/common/build.gradle.kts +++ b/common/build.gradle.kts @@ -22,6 +22,7 @@ import java.util.Date plugins { `maven-publish` + `java-test-fixtures` id("java") id("idea") } @@ -50,6 +51,9 @@ dependencies { testImplementation(libs.junit.jupiter.params) testRuntimeOnly(libs.junit.jupiter.engine) + + testFixturesApi(project(":api")) + testFixturesApi(libs.junit.jupiter.api) } fun getGitCommitId(): String { diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java similarity index 100% rename from api/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java rename to common/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java similarity index 100% rename from api/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java rename to common/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java similarity index 100% rename from api/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java rename to common/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java similarity index 100% rename from api/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java rename to common/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java similarity index 100% rename from api/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java rename to common/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java diff --git a/api/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java b/common/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java similarity index 100% rename from api/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java rename to common/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java diff --git a/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java b/common/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java similarity index 100% rename from api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java rename to common/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java diff --git a/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java similarity index 100% rename from api/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java rename to common/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java diff --git a/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java similarity index 100% rename from api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java rename to common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java diff --git a/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java similarity index 100% rename from api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java rename to common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java From 10fb53adaa380377025f9229568ded90a0c7f7d4 Mon Sep 17 00:00:00 2001 From: Nevin Zheng Date: Wed, 22 Jul 2026 13:17:20 -0700 Subject: [PATCH 4/6] refactor(kms): use extensible API identifiers --- .../gravitino/encryption/kms/KmsApi.java | 89 ------------------- .../encryption/kms/KmsReference.java | 23 ++--- .../gravitino/encryption/kms/TestKmsApi.java | 47 ---------- .../encryption/kms/TestKmsReference.java | 34 +++---- .../encryption/kms/KmsClientFactory.java | 6 +- .../encryption/kms/TestFakeKmsClient.java | 7 +- .../encryption/kms/TestKmsClient.java | 3 +- .../apache/gravitino/json/TestJsonUtils.java | 10 ++- .../encryption/kms/FakeKmsClient.java | 14 +-- .../encryption/kms/TestKmsClientContract.java | 9 +- .../kms/TestKmsClientFactoryContract.java | 4 +- 11 files changed, 60 insertions(+), 186 deletions(-) delete mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java delete mode 100644 api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsApi.java diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java deleted file mode 100644 index 28ef0d2dd46..00000000000 --- a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java +++ /dev/null @@ -1,89 +0,0 @@ -/* - * 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.gravitino.encryption.kms; - -import com.fasterxml.jackson.annotation.JsonCreator; -import com.fasterxml.jackson.annotation.JsonValue; -import java.util.Arrays; -import java.util.Locale; -import java.util.stream.Collectors; -import org.apache.gravitino.annotation.DeveloperApi; - -/** The KMS product API that defines key identifier and operation semantics. */ -@DeveloperApi -public enum KmsApi { - /** Amazon Web Services Key Management Service. */ - AWS_KMS("aws-kms"), - - /** Google Cloud Key Management Service. */ - GOOGLE_CLOUD_KMS("google-cloud-kms"), - - /** Microsoft Azure Key Vault. */ - AZURE_KEY_VAULT("azure-key-vault"), - - /** The OpenBao Transit secrets engine API. */ - OPENBAO_TRANSIT("openbao-transit"), - - /** The HashiCorp Vault Transit secrets engine API. */ - VAULT_TRANSIT("vault-transit"); - - private final String wireValue; - - KmsApi(String wireValue) { - this.wireValue = wireValue; - } - - /** - * Returns the stable configuration value for this API. - * - * @return the configuration value - */ - @JsonValue - public String wireValue() { - return wireValue; - } - - /** - * Parses a configured API value. - * - * @param value configured API value - * @return the matching API - * @throws IllegalArgumentException if the value is null, blank, or unsupported - */ - @JsonCreator - public static KmsApi fromWireValue(String value) { - if (value == null || value.trim().isEmpty()) { - throw new IllegalArgumentException("KMS API cannot be blank"); - } - - String normalized = value.trim().toLowerCase(Locale.ROOT); - return Arrays.stream(values()) - .filter(api -> api.wireValue.equals(normalized)) - .findFirst() - .orElseThrow( - () -> - new IllegalArgumentException( - String.format( - "Unsupported KMS API '%s'; supported values are %s", - value, - Arrays.stream(values()) - .map(KmsApi::wireValue) - .collect(Collectors.joining(", "))))); - } -} diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java index 49d7c798ff5..93f8ec90d01 100644 --- a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java @@ -21,6 +21,7 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; import com.google.common.base.Preconditions; +import java.util.Locale; import java.util.Objects; import javax.annotation.Nullable; import org.apache.commons.lang3.StringUtils; @@ -34,38 +35,38 @@ @DeveloperApi public final class KmsReference { - private final KmsApi api; + private final String api; private final String source; private final String keyId; /** * Creates a structurally valid key reference without contacting the provider. * - * @param api explicitly selected KMS API + * @param api explicitly selected KMS API identifier * @param source configured KMS client-instance name * @param keyId provider-native key identifier - * @throws IllegalArgumentException if the API is null or either string is null or blank + * @throws IllegalArgumentException if any argument is null or blank */ @JsonCreator public KmsReference( - @JsonProperty("api") KmsApi api, + @JsonProperty("api") String api, @JsonProperty("source") String source, @JsonProperty("keyId") String keyId) { - Preconditions.checkArgument(api != null, "KMS API cannot be null"); + Preconditions.checkArgument(StringUtils.isNotBlank(api), "KMS API cannot be blank"); Preconditions.checkArgument(StringUtils.isNotBlank(source), "KMS source cannot be blank"); Preconditions.checkArgument(StringUtils.isNotBlank(keyId), "KMS key ID cannot be blank"); - this.api = api; + this.api = api.trim().toLowerCase(Locale.ROOT); this.source = source; this.keyId = keyId; } /** - * Returns the explicitly selected KMS API. + * Returns the explicitly selected KMS API identifier. * - * @return the KMS API + * @return the canonical KMS API identifier */ - public KmsApi api() { + public String api() { return api; } @@ -96,7 +97,7 @@ public boolean equals(@Nullable Object other) { return false; } KmsReference that = (KmsReference) other; - return api == that.api && source.equals(that.source) && keyId.equals(that.keyId); + return api.equals(that.api) && source.equals(that.source) && keyId.equals(that.keyId); } @Override @@ -106,6 +107,6 @@ public int hashCode() { @Override public String toString() { - return String.format("KmsReference{api=%s, source='%s', keyId='%s'}", api, source, keyId); + return String.format("KmsReference{api='%s', source='%s', keyId='%s'}", api, source, keyId); } } diff --git a/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsApi.java b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsApi.java deleted file mode 100644 index cae91fb0303..00000000000 --- a/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsApi.java +++ /dev/null @@ -1,47 +0,0 @@ -/* - * 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.gravitino.encryption.kms; - -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.Test; - -public class TestKmsApi { - - @Test - void testParsesStableWireValues() { - Assertions.assertEquals(KmsApi.AWS_KMS, KmsApi.fromWireValue("aws-kms")); - Assertions.assertEquals(KmsApi.GOOGLE_CLOUD_KMS, KmsApi.fromWireValue(" GOOGLE-CLOUD-KMS ")); - Assertions.assertEquals("azure-key-vault", KmsApi.AZURE_KEY_VAULT.wireValue()); - Assertions.assertEquals( - KmsApi.OPENBAO_TRANSIT, KmsApi.fromWireValue(KmsApi.OPENBAO_TRANSIT.wireValue())); - Assertions.assertEquals(KmsApi.VAULT_TRANSIT, KmsApi.fromWireValue("vault-transit")); - } - - @Test - void testRejectsUnknownValues() { - Assertions.assertThrows(IllegalArgumentException.class, () -> KmsApi.fromWireValue(null)); - Assertions.assertThrows(IllegalArgumentException.class, () -> KmsApi.fromWireValue(" ")); - - IllegalArgumentException exception = - Assertions.assertThrows( - IllegalArgumentException.class, () -> KmsApi.fromWireValue("custom")); - Assertions.assertTrue(exception.getMessage().contains("Unsupported KMS API 'custom'")); - Assertions.assertTrue(exception.getMessage().contains("openbao-transit")); - } -} diff --git a/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsReference.java b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsReference.java index fe6b566faa0..32b6b45b062 100644 --- a/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsReference.java +++ b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsReference.java @@ -24,10 +24,10 @@ public class TestKmsReference { @Test - void testStoresReferenceExactly() { - KmsReference reference = new KmsReference(KmsApi.AWS_KMS, "production", " alias/Customer-Key "); + void testStoresCanonicalApiAndPreservesProviderKey() { + KmsReference reference = new KmsReference(" AWS-KMS ", "production", " alias/Customer-Key "); - Assertions.assertEquals(KmsApi.AWS_KMS, reference.api()); + Assertions.assertEquals("aws-kms", reference.api()); Assertions.assertEquals("production", reference.source()); Assertions.assertEquals(" alias/Customer-Key ", reference.keyId()); } @@ -37,26 +37,30 @@ void testRejectsMissingFields() { Assertions.assertThrows( IllegalArgumentException.class, () -> new KmsReference(null, "production", "key")); Assertions.assertThrows( - IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, null, "key")); + IllegalArgumentException.class, () -> new KmsReference("", "production", "key")); Assertions.assertThrows( - IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, "production", null)); + IllegalArgumentException.class, () -> new KmsReference(" ", "production", "key")); Assertions.assertThrows( - IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, "", "key")); + IllegalArgumentException.class, () -> new KmsReference("aws-kms", null, "key")); Assertions.assertThrows( - IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, " ", "key")); + IllegalArgumentException.class, () -> new KmsReference("aws-kms", "production", null)); Assertions.assertThrows( - IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, "production", "")); + IllegalArgumentException.class, () -> new KmsReference("aws-kms", "", "key")); Assertions.assertThrows( - IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, "production", " ")); + IllegalArgumentException.class, () -> new KmsReference("aws-kms", " ", "key")); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference("aws-kms", "production", "")); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference("aws-kms", "production", " ")); } @Test void testValueSemantics() { - KmsReference first = new KmsReference(KmsApi.AWS_KMS, "production", "key"); - KmsReference same = new KmsReference(KmsApi.AWS_KMS, "production", "key"); - KmsReference differentApi = new KmsReference(KmsApi.GOOGLE_CLOUD_KMS, "production", "key"); - KmsReference differentSource = new KmsReference(KmsApi.AWS_KMS, "recovery", "key"); - KmsReference differentKey = new KmsReference(KmsApi.AWS_KMS, "production", "another-key"); + KmsReference first = new KmsReference("aws-kms", "production", "key"); + KmsReference same = new KmsReference(" AWS-KMS ", "production", "key"); + KmsReference differentApi = new KmsReference("google-cloud-kms", "production", "key"); + KmsReference differentSource = new KmsReference("aws-kms", "recovery", "key"); + KmsReference differentKey = new KmsReference("aws-kms", "production", "another-key"); Assertions.assertEquals(first, same); Assertions.assertEquals(first.hashCode(), same.hashCode()); @@ -66,6 +70,6 @@ void testValueSemantics() { Assertions.assertNotEquals(first, null); Assertions.assertNotEquals(first, "key"); Assertions.assertEquals( - "KmsReference{api=AWS_KMS, source='production', keyId='key'}", first.toString()); + "KmsReference{api='aws-kms', source='production', keyId='key'}", first.toString()); } } diff --git a/common/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java index 3cfc4898492..f533d4c2253 100644 --- a/common/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java +++ b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java @@ -27,11 +27,11 @@ public interface KmsClientFactory { /** - * Returns the KMS API implemented by this factory. + * Returns the stable, canonical KMS API identifier implemented by this factory. * - * @return the supported KMS API + * @return the supported KMS API identifier */ - KmsApi api(); + String api(); /** * Creates a client bound to a configured KMS source. diff --git a/common/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java b/common/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java index 7ad5e2acca8..8d37890d906 100644 --- a/common/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java +++ b/common/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java @@ -20,12 +20,13 @@ public class TestFakeKmsClient extends TestKmsClientContract { + private static final String API = "test-kms"; private static final String SOURCE = "test"; private static final String USABLE_KEY = "usable"; private static final String MISSING_KEY = "missing"; private final FakeKmsClient client = - new FakeKmsClient(KmsApi.AWS_KMS, SOURCE).putKey(USABLE_KEY, true, true, true); + new FakeKmsClient(API, SOURCE).putKey(USABLE_KEY, true, true, true); @Override protected KmsClient client() { @@ -34,11 +35,11 @@ protected KmsClient client() { @Override protected KmsReference usableKey() { - return new KmsReference(KmsApi.AWS_KMS, SOURCE, USABLE_KEY); + return new KmsReference(API, SOURCE, USABLE_KEY); } @Override protected KmsReference missingKey() { - return new KmsReference(KmsApi.AWS_KMS, SOURCE, MISSING_KEY); + return new KmsReference(API, SOURCE, MISSING_KEY); } } diff --git a/common/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java b/common/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java index 14fe5d8b4c2..f3b107b2267 100644 --- a/common/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java +++ b/common/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java @@ -23,8 +23,7 @@ public class TestKmsClient { - private static final KmsReference REFERENCE = - new KmsReference(KmsApi.AWS_KMS, "production", "key"); + private static final KmsReference REFERENCE = new KmsReference("test-kms", "production", "key"); @Test void testReturnsProviderProperties() { diff --git a/common/src/test/java/org/apache/gravitino/json/TestJsonUtils.java b/common/src/test/java/org/apache/gravitino/json/TestJsonUtils.java index a0214b99ad3..13124d0341b 100644 --- a/common/src/test/java/org/apache/gravitino/json/TestJsonUtils.java +++ b/common/src/test/java/org/apache/gravitino/json/TestJsonUtils.java @@ -36,7 +36,6 @@ import org.apache.gravitino.dto.rel.partitions.ListPartitionDTO; import org.apache.gravitino.dto.rel.partitions.PartitionDTO; import org.apache.gravitino.dto.rel.partitions.RangePartitionDTO; -import org.apache.gravitino.encryption.kms.KmsApi; import org.apache.gravitino.encryption.kms.KmsReference; import org.apache.gravitino.rel.indexes.Index; import org.apache.gravitino.rel.types.Type; @@ -259,7 +258,7 @@ void testGetLong() throws Exception { @Test void testKmsReferenceAnyFieldMapperSerde() throws JsonProcessingException { KmsReference reference = - new KmsReference(KmsApi.GOOGLE_CLOUD_KMS, "analytics-prod", "projects/p/keys/k"); + new KmsReference("google-cloud-kms", "analytics-prod", "projects/p/keys/k"); String json = JsonUtils.anyFieldMapper().writeValueAsString(reference); @@ -270,6 +269,13 @@ void testKmsReferenceAnyFieldMapperSerde() throws JsonProcessingException { + "\"keyId\":\"projects/p/keys/k\"}"), JsonUtils.anyFieldMapper().readTree(json)); assertEquals(reference, JsonUtils.anyFieldMapper().readValue(json, KmsReference.class)); + assertEquals( + reference, + JsonUtils.anyFieldMapper() + .readValue( + "{\"api\":\" GOOGLE-CLOUD-KMS \",\"source\":\"analytics-prod\"," + + "\"keyId\":\"projects/p/keys/k\"}", + KmsReference.class)); } @Test diff --git a/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java index 183ffb37b98..fc6fb7fc412 100644 --- a/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java +++ b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java @@ -19,23 +19,27 @@ package org.apache.gravitino.encryption.kms; import java.util.HashMap; +import java.util.Locale; import java.util.Map; /** In-memory KMS client for contract and consumer tests. */ public final class FakeKmsClient implements KmsClient { - private final KmsApi api; + private final String api; private final String source; private final Map keys = new HashMap<>(); /** * Creates an empty fake client. * - * @param api KMS API accepted by the client + * @param api canonical KMS API identifier accepted by the client * @param source configured source accepted by the client */ - public FakeKmsClient(KmsApi api, String source) { - this.api = api; + public FakeKmsClient(String api, String source) { + if (api == null || api.trim().isEmpty()) { + throw new IllegalArgumentException("KMS API cannot be blank"); + } + this.api = api.trim().toLowerCase(Locale.ROOT); this.source = source; } @@ -69,7 +73,7 @@ private void requireReference(KmsReference reference) { if (reference == null) { throw new IllegalArgumentException("KMS reference cannot be null"); } - if (reference.api() != api) { + if (!reference.api().equals(api)) { throw new IllegalArgumentException( String.format("Expected KMS API %s but received %s", api, reference.api())); } diff --git a/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java index 855b4bdf6d9..6a5cbb51244 100644 --- a/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java +++ b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java @@ -18,7 +18,6 @@ */ package org.apache.gravitino.encryption.kms; -import java.util.Arrays; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -85,12 +84,8 @@ void testRejectsMismatchedSource() { @Test void testRejectsMismatchedApi() { KmsReference reference = usableKey(); - KmsApi otherApi = - Arrays.stream(KmsApi.values()) - .filter(api -> api != reference.api()) - .findFirst() - .orElseThrow(IllegalStateException::new); - KmsReference mismatched = new KmsReference(otherApi, reference.source(), reference.keyId()); + KmsReference mismatched = + new KmsReference(reference.api() + "-other", reference.source(), reference.keyId()); Assertions.assertThrows( IllegalArgumentException.class, () -> client().getKeyProperties(mismatched)); diff --git a/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java index 4136e5de41e..52eec8f5dc1 100644 --- a/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java +++ b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java @@ -34,9 +34,9 @@ public abstract class TestKmsClientFactoryContract { /** * Returns the API expected from the factory. * - * @return the expected API + * @return the expected API identifier */ - protected abstract KmsApi expectedApi(); + protected abstract String expectedApi(); @Test void testReportsExpectedApi() { From cf4a8c976ad152920c5e879b809964f739ac4e8c Mon Sep 17 00:00:00 2001 From: Nevin Zheng Date: Wed, 22 Jul 2026 13:33:03 -0700 Subject: [PATCH 5/6] refactor(kms): represent missing keys with Optional --- .../gravitino/encryption/kms/KmsClient.java | 10 ++++- .../encryption/kms/KmsKeyProperties.java | 20 +++++----- .../encryption/kms/TestFakeKmsClient.java | 16 +++++++- .../encryption/kms/TestKmsClient.java | 39 +++++++++++++++---- .../encryption/kms/FakeKmsClient.java | 18 +++------ .../encryption/kms/TestKmsClientContract.java | 12 +++--- 6 files changed, 76 insertions(+), 39 deletions(-) diff --git a/common/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java index bcc541f98dd..15b83ab1c43 100644 --- a/common/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java +++ b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java @@ -18,6 +18,7 @@ */ package org.apache.gravitino.encryption.kms; +import java.util.Optional; import org.apache.gravitino.annotation.DeveloperApi; import org.apache.gravitino.exceptions.ConnectionFailedException; @@ -34,12 +35,17 @@ public interface KmsClient extends AutoCloseable { /** * Reads the provider-reported properties of a key. * + *

An empty result means the provider authoritatively reported that the key does not exist. + * Authentication, authorization, timeout, availability, and other indeterminate failures are + * reported as exceptions, never as an empty result. + * * @param reference key to inspect - * @return normalized key properties + * @return normalized key properties, or empty when the key authoritatively does not exist; never + * null * @throws IllegalArgumentException if the reference does not belong to this client * @throws ConnectionFailedException if the provider cannot be queried */ - KmsKeyProperties getKeyProperties(KmsReference reference); + Optional getKeyProperties(KmsReference reference); /** Releases resources owned by this client. */ @Override diff --git a/common/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java index ae1896a5fa8..67dec82f40d 100644 --- a/common/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java +++ b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java @@ -21,9 +21,11 @@ import org.apache.gravitino.annotation.DeveloperApi; /** - * Common properties reported for a KMS key. + * Common properties reported for a successfully located KMS key. * - *

Capabilities describe downstream use; this API does not perform cryptographic operations. + *

Provider-specific lifecycle state may describe the logical key or its selected or current + * version. Capabilities describe key-specific structural support; this API does not perform + * cryptographic operations. */ @DeveloperApi public interface KmsKeyProperties { @@ -36,28 +38,24 @@ public interface KmsKeyProperties { KmsReference reference(); /** - * Returns whether the provider reports that the key exists. + * Returns the provider's normalized lifecycle state for the key. * - * @return whether the key exists - */ - boolean present(); - - /** - * Returns whether the provider reports that the key is enabled. + *

This does not indicate caller authorization, service availability, or that a subsequent + * operation is guaranteed to succeed. * * @return whether the key is enabled */ boolean enabled(); /** - * Returns whether the key supports wrapping data-encryption keys. + * Returns whether the key structurally supports wrapping data-encryption keys. * * @return whether wrapping is supported */ boolean supportsWrapping(); /** - * Returns whether the key supports unwrapping data-encryption keys. + * Returns whether the key structurally supports unwrapping data-encryption keys. * * @return whether unwrapping is supported */ diff --git a/common/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java b/common/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java index 8d37890d906..91b2609b2e1 100644 --- a/common/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java +++ b/common/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java @@ -18,15 +18,29 @@ */ package org.apache.gravitino.encryption.kms; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + public class TestFakeKmsClient extends TestKmsClientContract { private static final String API = "test-kms"; private static final String SOURCE = "test"; private static final String USABLE_KEY = "usable"; + private static final String DISABLED_KEY = "disabled"; private static final String MISSING_KEY = "missing"; private final FakeKmsClient client = - new FakeKmsClient(API, SOURCE).putKey(USABLE_KEY, true, true, true); + new FakeKmsClient(API, SOURCE) + .putKey(USABLE_KEY, true, true, true) + .putKey(DISABLED_KEY, false, true, true); + + @Test + void testReportsDisabledKeyAsPresent() { + KmsKeyProperties properties = + client.getKeyProperties(new KmsReference(API, SOURCE, DISABLED_KEY)).orElseThrow(); + + Assertions.assertFalse(properties.enabled()); + } @Override protected KmsClient client() { diff --git a/common/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java b/common/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java index f3b107b2267..9c2b6b4c085 100644 --- a/common/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java +++ b/common/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java @@ -18,6 +18,8 @@ */ package org.apache.gravitino.encryption.kms; +import java.util.Optional; +import org.apache.gravitino.exceptions.ConnectionFailedException; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -28,20 +30,46 @@ public class TestKmsClient { @Test void testReturnsProviderProperties() { KmsKeyProperties properties = new TestProperties(); - KmsClient client = reference -> properties; + KmsClient client = reference -> Optional.of(properties); - Assertions.assertSame(properties, client.getKeyProperties(REFERENCE)); + Assertions.assertSame(properties, client.getKeyProperties(REFERENCE).orElseThrow()); Assertions.assertTrue(properties.supportsWrapping()); Assertions.assertFalse(properties.supportsUnwrapping()); } @Test void testDefaultClose() { - KmsClient client = reference -> new TestProperties(); + KmsClient client = reference -> Optional.of(new TestProperties()); Assertions.assertDoesNotThrow(client::close); } + @Test + void testPreservesProviderFailureType() { + KmsAuthenticationException authenticationFailure = + new KmsAuthenticationException("authentication failed"); + ConnectionFailedException unavailableFailure = + new ConnectionFailedException("provider unavailable"); + KmsClient authenticationClient = + reference -> { + throw authenticationFailure; + }; + KmsClient unavailableClient = + reference -> { + throw unavailableFailure; + }; + + Assertions.assertSame( + authenticationFailure, + Assertions.assertThrows( + KmsAuthenticationException.class, + () -> authenticationClient.getKeyProperties(REFERENCE))); + Assertions.assertSame( + unavailableFailure, + Assertions.assertThrows( + ConnectionFailedException.class, () -> unavailableClient.getKeyProperties(REFERENCE))); + } + private static final class TestProperties implements KmsKeyProperties { @Override @@ -49,11 +77,6 @@ public KmsReference reference() { return REFERENCE; } - @Override - public boolean present() { - return true; - } - @Override public boolean enabled() { return true; diff --git a/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java index fc6fb7fc412..a55a488c5d8 100644 --- a/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java +++ b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java @@ -21,6 +21,7 @@ import java.util.HashMap; import java.util.Locale; import java.util.Map; +import java.util.Optional; /** In-memory KMS client for contract and consumer tests. */ public final class FakeKmsClient implements KmsClient { @@ -60,13 +61,14 @@ public FakeKmsClient putKey( /** {@inheritDoc} */ @Override - public KmsKeyProperties getKeyProperties(KmsReference reference) { + public Optional getKeyProperties(KmsReference reference) { requireReference(reference); KeyState state = keys.get(reference.keyId()); return state == null - ? new Properties(reference, false, false, false, false) - : new Properties( - reference, true, state.enabled, state.supportsWrapping, state.supportsUnwrapping); + ? Optional.empty() + : Optional.of( + new Properties( + reference, state.enabled, state.supportsWrapping, state.supportsUnwrapping)); } private void requireReference(KmsReference reference) { @@ -100,19 +102,16 @@ private KeyState(boolean enabled, boolean supportsWrapping, boolean supportsUnwr private static final class Properties implements KmsKeyProperties { private final KmsReference reference; - private final boolean present; private final boolean enabled; private final boolean supportsWrapping; private final boolean supportsUnwrapping; private Properties( KmsReference reference, - boolean present, boolean enabled, boolean supportsWrapping, boolean supportsUnwrapping) { this.reference = reference; - this.present = present; this.enabled = enabled; this.supportsWrapping = supportsWrapping; this.supportsUnwrapping = supportsUnwrapping; @@ -123,11 +122,6 @@ public KmsReference reference() { return reference; } - @Override - public boolean present() { - return present; - } - @Override public boolean enabled() { return enabled; diff --git a/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java index 6a5cbb51244..868d64bd62e 100644 --- a/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java +++ b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java @@ -18,6 +18,7 @@ */ package org.apache.gravitino.encryption.kms; +import java.util.Optional; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -48,10 +49,11 @@ public abstract class TestKmsClientContract { @Test void testReportsUsableKeyProperties() { KmsReference reference = usableKey(); - KmsKeyProperties properties = client().getKeyProperties(reference); + Optional result = client().getKeyProperties(reference); + Assertions.assertNotNull(result, "KMS clients must not return null"); + KmsKeyProperties properties = result.orElseThrow(); Assertions.assertEquals(reference, properties.reference()); - Assertions.assertTrue(properties.present()); Assertions.assertTrue(properties.enabled()); Assertions.assertTrue(properties.supportsWrapping()); Assertions.assertTrue(properties.supportsUnwrapping()); @@ -60,10 +62,10 @@ void testReportsUsableKeyProperties() { @Test void testReportsMissingKey() { KmsReference reference = missingKey(); - KmsKeyProperties properties = client().getKeyProperties(reference); + Optional result = client().getKeyProperties(reference); - Assertions.assertEquals(reference, properties.reference()); - Assertions.assertFalse(properties.present()); + Assertions.assertNotNull(result, "KMS clients must not return null"); + Assertions.assertTrue(result.isEmpty()); } @Test From 235980fc537416034c7caea14135c44fc348ab5a Mon Sep 17 00:00:00 2001 From: Nevin Zheng Date: Tue, 21 Jul 2026 15:48:27 -0700 Subject: [PATCH 6/6] feat(kms): configure and route named sources --- .../org/apache/gravitino/GravitinoEnv.java | 37 +- .../encryption/kms/KmsClientRegistry.java | 219 +++++++ .../gravitino/encryption/kms/KmsConfig.java | 147 +++++ .../TestGravitinoEnvKmsClientRegistry.java | 50 ++ .../encryption/kms/TestKmsClientRegistry.java | 533 ++++++++++++++++++ .../encryption/kms/TestKmsConfig.java | 143 +++++ 6 files changed, 1128 insertions(+), 1 deletion(-) create mode 100644 core/src/main/java/org/apache/gravitino/encryption/kms/KmsClientRegistry.java create mode 100644 core/src/main/java/org/apache/gravitino/encryption/kms/KmsConfig.java create mode 100644 core/src/test/java/org/apache/gravitino/TestGravitinoEnvKmsClientRegistry.java create mode 100644 core/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClientRegistry.java create mode 100644 core/src/test/java/org/apache/gravitino/encryption/kms/TestKmsConfig.java diff --git a/core/src/main/java/org/apache/gravitino/GravitinoEnv.java b/core/src/main/java/org/apache/gravitino/GravitinoEnv.java index 11c12911293..56a3cdfcf2c 100644 --- a/core/src/main/java/org/apache/gravitino/GravitinoEnv.java +++ b/core/src/main/java/org/apache/gravitino/GravitinoEnv.java @@ -56,6 +56,7 @@ import org.apache.gravitino.catalog.ViewNormalizeDispatcher; import org.apache.gravitino.catalog.ViewOperationDispatcher; import org.apache.gravitino.credential.CredentialOperationDispatcher; +import org.apache.gravitino.encryption.kms.KmsClientRegistry; import org.apache.gravitino.hook.AccessControlHookDispatcher; import org.apache.gravitino.hook.CatalogHookDispatcher; import org.apache.gravitino.hook.FilesetHookDispatcher; @@ -154,6 +155,8 @@ public class GravitinoEnv { private CredentialOperationDispatcher credentialOperationDispatcher; + private KmsClientRegistry kmsClientRegistry; + private TagDispatcher tagDispatcher; private PolicyDispatcher policyDispatcher; @@ -222,7 +225,13 @@ public void initializeFullComponents(Config config) { FileFetcher.get().initialize(config.get(Configs.BLOCK_UNSAFE_REMOTE_URI)); this.manageFullComponents = true; initBaseComponents(); - initGravitinoServerComponents(); + try { + this.kmsClientRegistry = new KmsClientRegistry(config); + initGravitinoServerComponents(); + } catch (RuntimeException e) { + closeKmsClientRegistry(); + throw e; + } LOG.info("Gravitino full environment is initialized."); } @@ -418,6 +427,18 @@ public CredentialOperationDispatcher credentialOperationDispatcher() { return credentialOperationDispatcher; } + /** + * Get the metadata-only KMS client registry associated with the full Gravitino environment. + * + * @return The KMS client registry instance. + * @throws IllegalStateException if full components have not been initialized + */ + public KmsClientRegistry kmsClientRegistry() { + Preconditions.checkState( + kmsClientRegistry != null, "GravitinoEnv full components are not initialized."); + return kmsClientRegistry; + } + /** * Get the IdGenerator associated with the Gravitino environment. * @@ -586,6 +607,8 @@ public void start() { public void shutdown() { LOG.info("Shutting down Gravitino Environment..."); + closeKmsClientRegistry(); + if (entityStore != null) { try { entityStore.close(); @@ -638,6 +661,18 @@ public void shutdown() { LOG.info("Gravitino Environment is shut down."); } + private void closeKmsClientRegistry() { + KmsClientRegistry registry = kmsClientRegistry; + kmsClientRegistry = null; + if (registry != null) { + try { + registry.close(); + } catch (RuntimeException e) { + LOG.warn("Failed to close KmsClientRegistry.", e); + } + } + } + private void initBaseComponents() { this.metricsSystem = new MetricsSystem(); metricsSystem.register(new JVMMetricsSource()); diff --git a/core/src/main/java/org/apache/gravitino/encryption/kms/KmsClientRegistry.java b/core/src/main/java/org/apache/gravitino/encryption/kms/KmsClientRegistry.java new file mode 100644 index 00000000000..931c0654b83 --- /dev/null +++ b/core/src/main/java/org/apache/gravitino/encryption/kms/KmsClientRegistry.java @@ -0,0 +1,219 @@ +/* + * 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.gravitino.encryption.kms; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Optional; +import java.util.ServiceLoader; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReadWriteLock; +import java.util.concurrent.locks.ReentrantReadWriteLock; +import org.apache.gravitino.Config; + +/** Creates server-private KMS clients and dispatches key inspection by configured source. */ +public final class KmsClientRegistry implements KmsClient { + + private final ReadWriteLock lifecycleLock = new ReentrantReadWriteLock(); + private final Map clients; + private boolean closed; + + /** + * Loads available {@link KmsClientFactory} implementations and creates configured clients. + * + * @param config Gravitino server configuration + * @throws IllegalArgumentException if configuration or factory discovery is invalid + */ + public KmsClientRegistry(Config config) { + this(config, loadFactories()); + } + + KmsClientRegistry(Config config, Iterable factories) { + if (factories == null) { + throw new IllegalArgumentException("KMS client factories cannot be null"); + } + + KmsConfig kmsConfig = new KmsConfig(config); + Map factoriesByApi = indexFactories(factories); + this.clients = createClients(kmsConfig.sources(), factoriesByApi); + } + + /** + * Dispatches the request to the client configured for the reference source. + * + * @param reference key to inspect + * @return normalized key properties, or empty if the provider authoritatively reports that the + * key does not exist + * @throws IllegalArgumentException if the source is unknown or configured for another API + * @throws IllegalStateException if the registry is closed or a provider violates its contract + */ + @Override + public Optional getKeyProperties(KmsReference reference) { + Lock readLock = lifecycleLock.readLock(); + readLock.lock(); + try { + if (closed) { + throw new IllegalStateException("KMS client registry is closed"); + } + if (reference == null) { + throw new IllegalArgumentException("KMS reference cannot be null"); + } + + ConfiguredClient configuredClient = clients.get(reference.source()); + if (configuredClient == null) { + throw new IllegalArgumentException( + String.format("Unknown KMS source '%s'", reference.source())); + } + if (!configuredClient.api.equals(reference.api())) { + throw new IllegalArgumentException( + String.format( + "KMS source '%s' uses API '%s', not '%s'", + reference.source(), configuredClient.api, reference.api())); + } + + Optional properties = configuredClient.client.getKeyProperties(reference); + if (properties == null) { + throw new IllegalStateException( + String.format("KMS client for source '%s' returned null", reference.source())); + } + if (properties.isPresent() && !reference.equals(properties.get().reference())) { + throw new IllegalStateException( + String.format( + "KMS client for source '%s' returned properties for a different reference", + reference.source())); + } + return properties; + } finally { + readLock.unlock(); + } + } + + /** Closes all configured clients. This operation is idempotent. */ + @Override + public void close() { + Lock writeLock = lifecycleLock.writeLock(); + writeLock.lock(); + try { + if (closed) { + return; + } + closed = true; + RuntimeException failure = closeClients(new ArrayList<>(clients.values())); + if (failure != null) { + throw failure; + } + } finally { + writeLock.unlock(); + } + } + + private static Map indexFactories( + Iterable factories) { + Map factoriesByApi = new LinkedHashMap<>(); + for (KmsClientFactory factory : factories) { + if (factory == null) { + throw new IllegalArgumentException("KMS client factory cannot be null"); + } + String api = canonicalizeApi(factory.api()); + KmsClientFactory existing = factoriesByApi.putIfAbsent(api, factory); + if (existing != null) { + throw new IllegalArgumentException( + String.format("Multiple KMS client factories support API '%s'", api)); + } + } + return factoriesByApi; + } + + private static Iterable loadFactories() { + ClassLoader classLoader = Thread.currentThread().getContextClassLoader(); + if (classLoader == null) { + classLoader = KmsClientRegistry.class.getClassLoader(); + } + return ServiceLoader.load(KmsClientFactory.class, classLoader); + } + + private static Map createClients( + Map sourceConfigs, + Map factoriesByApi) { + Map clients = new LinkedHashMap<>(); + try { + sourceConfigs.forEach( + (source, sourceConfig) -> { + KmsClientFactory factory = factoriesByApi.get(sourceConfig.api()); + if (factory == null) { + throw new IllegalArgumentException( + String.format( + "No KMS client factory supports API '%s' for source '%s'", + sourceConfig.api(), source)); + } + KmsClient client = factory.create(source, sourceConfig.properties()); + if (client == null) { + throw new IllegalStateException( + String.format( + "KMS client factory for API '%s' returned null", sourceConfig.api())); + } + clients.put(source, new ConfiguredClient(sourceConfig.api(), client)); + }); + return Collections.unmodifiableMap(clients); + } catch (RuntimeException | Error e) { + RuntimeException closeFailure = closeClients(new ArrayList<>(clients.values())); + if (closeFailure != null) { + e.addSuppressed(closeFailure); + } + throw e; + } + } + + private static RuntimeException closeClients(List clients) { + RuntimeException failure = null; + for (int index = clients.size() - 1; index >= 0; index--) { + try { + clients.get(index).client.close(); + } catch (RuntimeException e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + } + return failure; + } + + private static String canonicalizeApi(String api) { + if (api == null || api.trim().isEmpty()) { + throw new IllegalArgumentException("KMS client factory API cannot be blank"); + } + return api.trim().toLowerCase(Locale.ROOT); + } + + private static final class ConfiguredClient { + private final String api; + private final KmsClient client; + + private ConfiguredClient(String api, KmsClient client) { + this.api = api; + this.client = client; + } + } +} diff --git a/core/src/main/java/org/apache/gravitino/encryption/kms/KmsConfig.java b/core/src/main/java/org/apache/gravitino/encryption/kms/KmsConfig.java new file mode 100644 index 00000000000..8d3a191618d --- /dev/null +++ b/core/src/main/java/org/apache/gravitino/encryption/kms/KmsConfig.java @@ -0,0 +1,147 @@ +/* + * 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.gravitino.encryption.kms; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import org.apache.gravitino.Config; + +final class KmsConfig { + + static final String KMS_CONFIG_PREFIX = "gravitino.kms."; + static final String KMS_SOURCES = KMS_CONFIG_PREFIX + "sources"; + + private static final String SOURCES = "sources"; + private static final String SOURCE_PREFIX = "source."; + private static final String API = "api"; + private static final Pattern SOURCE_NAME_PATTERN = Pattern.compile("[A-Za-z0-9][A-Za-z0-9_-]*"); + private static final Pattern SOURCE_PROPERTY_PATTERN = + Pattern.compile("source\\.([A-Za-z0-9][A-Za-z0-9_-]*)\\.(.+)"); + + private final Map sources; + + KmsConfig(Config config) { + if (config == null) { + throw new KmsConfigurationException("Gravitino configuration cannot be null"); + } + + Map values = config.getConfigsWithPrefix(KMS_CONFIG_PREFIX); + List configuredSources = parseSources(values.get(SOURCES)); + validateKeys(values, configuredSources); + this.sources = parseSourceConfigs(values, configuredSources); + } + + Map sources() { + return sources; + } + + private static List parseSources(String value) { + if (value == null || value.trim().isEmpty()) { + return Collections.emptyList(); + } + + List sources = new ArrayList<>(); + Set uniqueSources = new LinkedHashSet<>(); + for (String item : value.split(",", -1)) { + String source = item.trim(); + if (!SOURCE_NAME_PATTERN.matcher(source).matches()) { + throw new KmsConfigurationException( + "Invalid KMS source name '%s' in %s", source, KMS_SOURCES); + } + if (!uniqueSources.add(source)) { + throw new KmsConfigurationException("Duplicate KMS source '%s' in %s", source, KMS_SOURCES); + } + sources.add(source); + } + return Collections.unmodifiableList(sources); + } + + private static void validateKeys(Map values, List configuredSources) { + Set sourceSet = new LinkedHashSet<>(configuredSources); + for (String key : values.keySet()) { + if (SOURCES.equals(key)) { + continue; + } + + Matcher matcher = SOURCE_PROPERTY_PATTERN.matcher(key); + if (!matcher.matches()) { + throw new KmsConfigurationException( + "Invalid KMS configuration key '%s%s'", KMS_CONFIG_PREFIX, key); + } + if (!sourceSet.contains(matcher.group(1))) { + throw new KmsConfigurationException( + "KMS configuration references unlisted source '%s'", matcher.group(1)); + } + } + } + + private static Map parseSourceConfigs( + Map values, List configuredSources) { + Map sourceConfigs = new LinkedHashMap<>(); + + for (String source : configuredSources) { + String propertyPrefix = SOURCE_PREFIX + source + "."; + String apiKey = propertyPrefix + API; + String apiValue = values.get(apiKey); + if (apiValue == null || apiValue.trim().isEmpty()) { + throw new KmsConfigurationException( + "KMS API property '%s%s' cannot be blank", KMS_CONFIG_PREFIX, apiKey); + } + String api = apiValue.trim().toLowerCase(Locale.ROOT); + + Map properties = new LinkedHashMap<>(); + values.forEach( + (key, value) -> { + if (key.startsWith(propertyPrefix) && !apiKey.equals(key)) { + String sourceProperty = key.substring(propertyPrefix.length()); + properties.put(sourceProperty, value); + } + }); + sourceConfigs.put(source, new SourceConfig(api, properties)); + } + + return Collections.unmodifiableMap(sourceConfigs); + } + + static final class SourceConfig { + private final String api; + private final Map properties; + + private SourceConfig(String api, Map properties) { + this.api = api; + this.properties = Collections.unmodifiableMap(new LinkedHashMap<>(properties)); + } + + String api() { + return api; + } + + Map properties() { + return properties; + } + } +} diff --git a/core/src/test/java/org/apache/gravitino/TestGravitinoEnvKmsClientRegistry.java b/core/src/test/java/org/apache/gravitino/TestGravitinoEnvKmsClientRegistry.java new file mode 100644 index 00000000000..8b9810ec035 --- /dev/null +++ b/core/src/test/java/org/apache/gravitino/TestGravitinoEnvKmsClientRegistry.java @@ -0,0 +1,50 @@ +/* + * 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.gravitino; + +import org.apache.commons.lang3.reflect.FieldUtils; +import org.apache.gravitino.encryption.kms.KmsClientRegistry; +import org.apache.gravitino.encryption.kms.KmsReference; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class TestGravitinoEnvKmsClientRegistry { + + @Test + void testEmptyRegistryIsOptionalAndClosedWithEnvironment() throws IllegalAccessException { + TestGravitinoEnv env = new TestGravitinoEnv(); + Assertions.assertThrows(IllegalStateException.class, env::kmsClientRegistry); + + KmsClientRegistry registry = new KmsClientRegistry(new Config(false) {}); + FieldUtils.writeField(env, "kmsClientRegistry", registry, true); + + Assertions.assertSame(registry, env.kmsClientRegistry()); + KmsReference reference = new KmsReference("test-kms", "missing", "key"); + Assertions.assertThrows( + IllegalArgumentException.class, () -> registry.getKeyProperties(reference)); + + env.shutdown(); + + Assertions.assertThrows(IllegalStateException.class, env::kmsClientRegistry); + Assertions.assertThrows( + IllegalStateException.class, () -> registry.getKeyProperties(reference)); + } + + private static final class TestGravitinoEnv extends GravitinoEnv {} +} diff --git a/core/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClientRegistry.java b/core/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClientRegistry.java new file mode 100644 index 00000000000..fd92a2c91b6 --- /dev/null +++ b/core/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClientRegistry.java @@ -0,0 +1,533 @@ +/* + * 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.gravitino.encryption.kms; + +import java.net.URL; +import java.net.URLClassLoader; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.gravitino.Config; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +public class TestKmsClientRegistry { + + private static final String AWS_API = "aws-kms"; + private static final String GCP_API = "google-cloud-kms"; + private static final String AZURE_API = "azure-key-vault"; + + @Test + void testCreatesAndDispatchesConfiguredClients() { + RecordingFactory awsFactory = new RecordingFactory(" AWS-KMS "); + RecordingFactory gcpFactory = new RecordingFactory(GCP_API); + KmsClientRegistry registry = + new KmsClientRegistry( + config( + "gravitino.kms.sources", "primary,analytics", + "gravitino.kms.source.primary.api", " AWS-KMS ", + "gravitino.kms.source.primary.endpoint.region", "us-west-2", + "gravitino.kms.source.analytics.api", "google-cloud-kms", + "gravitino.kms.source.analytics.endpoint.project", "data-project"), + List.of(awsFactory, gcpFactory)); + + KmsReference awsReference = new KmsReference(AWS_API, "primary", "alias/orders"); + KmsReference gcpReference = + new KmsReference(GCP_API, "analytics", "projects/p/locations/l/keyRings/r/cryptoKeys/k"); + + Assertions.assertEquals( + awsReference, registry.getKeyProperties(awsReference).orElseThrow().reference()); + Assertions.assertEquals( + awsReference, registry.getKeyProperties(awsReference).orElseThrow().reference()); + Assertions.assertEquals( + gcpReference, registry.getKeyProperties(gcpReference).orElseThrow().reference()); + Assertions.assertEquals(Map.of("endpoint.region", "us-west-2"), awsFactory.properties); + Assertions.assertEquals(Map.of("endpoint.project", "data-project"), gcpFactory.properties); + Assertions.assertEquals("alias/orders", awsFactory.providerKeyId); + Assertions.assertEquals( + "projects/p/locations/l/keyRings/r/cryptoKeys/k", gcpFactory.providerKeyId); + Assertions.assertEquals("primary", awsFactory.source); + Assertions.assertEquals("analytics", gcpFactory.source); + Assertions.assertEquals(1, awsFactory.createCount.get()); + Assertions.assertEquals(1, gcpFactory.createCount.get()); + } + + @Test + void testRejectsUnknownSourceAndApiMismatch() { + KmsClientRegistry registry = + new KmsClientRegistry( + config( + "gravitino.kms.sources", "primary", + "gravitino.kms.source.primary.api", "aws-kms"), + List.of(new RecordingFactory(AWS_API))); + + Assertions.assertThrows( + IllegalArgumentException.class, + () -> registry.getKeyProperties(new KmsReference(AWS_API, "other", "key"))); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> registry.getKeyProperties(new KmsReference(GCP_API, "primary", "key"))); + Assertions.assertThrows(IllegalArgumentException.class, () -> registry.getKeyProperties(null)); + } + + @Test + void testCreatesMultipleSourcesForSameApi() { + RecordingFactory factory = new RecordingFactory(AZURE_API); + KmsClientRegistry registry = + new KmsClientRegistry( + config( + "gravitino.kms.sources", "azure-eu,azure-us", + "gravitino.kms.source.azure-eu.api", "azure-key-vault", + "gravitino.kms.source.azure-us.api", "azure-key-vault"), + List.of(factory)); + + Assertions.assertEquals( + "azure-eu", + registry + .getKeyProperties(new KmsReference(AZURE_API, "azure-eu", "primary")) + .orElseThrow() + .reference() + .source()); + Assertions.assertEquals( + "azure-us", + registry + .getKeyProperties(new KmsReference(AZURE_API, "azure-us", "primary")) + .orElseThrow() + .reference() + .source()); + } + + @Test + void testRejectsMissingDuplicateAndInvalidFactories() { + Config awsConfig = + config( + "gravitino.kms.sources", "primary", + "gravitino.kms.source.primary.api", "aws-kms"); + + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsClientRegistry(awsConfig, List.of())); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + new KmsClientRegistry( + awsConfig, + List.of(new RecordingFactory(AWS_API), new RecordingFactory(" AWS-KMS ")))); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> new KmsClientRegistry(awsConfig, List.of(new RecordingFactory(null)))); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> new KmsClientRegistry(awsConfig, List.of(new RecordingFactory(" ")))); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> new KmsClientRegistry(awsConfig, java.util.Arrays.asList((KmsClientFactory) null))); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsClientRegistry(awsConfig, null)); + } + + @Test + void testRejectsUnknownConfiguredApi() { + Config customConfig = + config( + "gravitino.kms.sources", "primary", + "gravitino.kms.source.primary.api", "custom-kms"); + + IllegalArgumentException exception = + Assertions.assertThrows( + IllegalArgumentException.class, + () -> new KmsClientRegistry(customConfig, List.of(new RecordingFactory(AWS_API)))); + Assertions.assertTrue( + exception + .getMessage() + .contains("No KMS client factory supports API 'custom-kms' for source 'primary'")); + } + + @Test + void testPublicConstructorUsesContextClassLoader(@TempDir Path tempDirectory) throws Exception { + Path serviceFile = + tempDirectory.resolve( + "META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory"); + Files.createDirectories(serviceFile.getParent()); + Files.write(serviceFile, ServiceLoadedFactory.class.getName().getBytes(StandardCharsets.UTF_8)); + + ClassLoader originalClassLoader = Thread.currentThread().getContextClassLoader(); + try (URLClassLoader serviceClassLoader = + new URLClassLoader(new URL[] {tempDirectory.toUri().toURL()}, originalClassLoader)) { + Thread.currentThread().setContextClassLoader(serviceClassLoader); + try (KmsClientRegistry registry = + new KmsClientRegistry( + config( + "gravitino.kms.sources", "primary", + "gravitino.kms.source.primary.api", "aws-kms"))) { + KmsReference reference = new KmsReference(AWS_API, "primary", "key"); + Assertions.assertEquals( + reference, registry.getKeyProperties(reference).orElseThrow().reference()); + } + } finally { + Thread.currentThread().setContextClassLoader(originalClassLoader); + } + } + + @Test + void testRejectsInvalidProviderResults() { + KmsReference reference = new KmsReference(AWS_API, "primary", "key"); + Config awsConfig = + config( + "gravitino.kms.sources", "primary", + "gravitino.kms.source.primary.api", "aws-kms"); + + KmsClientFactory nullResultFactory = factory(AWS_API, (source, properties) -> ignored -> null); + KmsClientRegistry nullResultRegistry = + new KmsClientRegistry(awsConfig, List.of(nullResultFactory)); + Assertions.assertThrows( + IllegalStateException.class, () -> nullResultRegistry.getKeyProperties(reference)); + + KmsReference otherReference = new KmsReference(AWS_API, "primary", "other"); + KmsClientFactory wrongReferenceFactory = + factory( + AWS_API, + (source, properties) -> ignored -> Optional.of(new Properties(otherReference))); + KmsClientRegistry wrongReferenceRegistry = + new KmsClientRegistry(awsConfig, List.of(wrongReferenceFactory)); + Assertions.assertThrows( + IllegalStateException.class, () -> wrongReferenceRegistry.getKeyProperties(reference)); + + KmsClientFactory missingKeyFactory = + factory(AWS_API, (source, properties) -> ignored -> Optional.empty()); + KmsClientRegistry missingKeyRegistry = + new KmsClientRegistry(awsConfig, List.of(missingKeyFactory)); + Assertions.assertTrue(missingKeyRegistry.getKeyProperties(reference).isEmpty()); + + KmsClientFactory nullClientFactory = factory(AWS_API, (source, properties) -> null); + Assertions.assertThrows( + IllegalStateException.class, + () -> new KmsClientRegistry(awsConfig, List.of(nullClientFactory))); + } + + @Test + void testClosesClientsInReverseOrderAndIsIdempotent() { + List closeOrder = new ArrayList<>(); + CloseTrackingFactory awsFactory = new CloseTrackingFactory(AWS_API, "aws", closeOrder, null); + CloseTrackingFactory gcpFactory = new CloseTrackingFactory(GCP_API, "gcp", closeOrder, null); + KmsClientRegistry registry = + new KmsClientRegistry( + config( + "gravitino.kms.sources", "primary,analytics", + "gravitino.kms.source.primary.api", "aws-kms", + "gravitino.kms.source.analytics.api", "google-cloud-kms"), + List.of(awsFactory, gcpFactory)); + + registry.close(); + registry.close(); + + Assertions.assertEquals(List.of("gcp", "aws"), closeOrder); + Assertions.assertEquals(1, awsFactory.closeCount.get()); + Assertions.assertEquals(1, gcpFactory.closeCount.get()); + Assertions.assertThrows( + IllegalStateException.class, + () -> registry.getKeyProperties(new KmsReference(AWS_API, "primary", "key"))); + } + + @Test + void testClosesCreatedClientsAfterPartialInitializationFailure() { + CloseTrackingFactory awsFactory = + new CloseTrackingFactory(AWS_API, "aws", new ArrayList<>(), null); + KmsClientFactory failingFactory = + factory( + GCP_API, + (source, properties) -> { + throw new IllegalArgumentException("invalid GCP configuration"); + }); + + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + new KmsClientRegistry( + config( + "gravitino.kms.sources", "primary,analytics", + "gravitino.kms.source.primary.api", "aws-kms", + "gravitino.kms.source.analytics.api", "google-cloud-kms"), + List.of(awsFactory, failingFactory))); + Assertions.assertEquals(1, awsFactory.closeCount.get()); + } + + @Test + void testPreservesInitializationFailureWhenCleanupFails() { + RuntimeException closeFailure = new IllegalStateException("close failed"); + CloseTrackingFactory awsFactory = + new CloseTrackingFactory(AWS_API, "aws", new ArrayList<>(), closeFailure); + IllegalArgumentException creationFailure = + new IllegalArgumentException("invalid GCP configuration"); + KmsClientFactory failingFactory = + factory( + GCP_API, + (source, properties) -> { + throw creationFailure; + }); + + IllegalArgumentException exception = + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + new KmsClientRegistry( + config( + "gravitino.kms.sources", "primary,analytics", + "gravitino.kms.source.primary.api", "aws-kms", + "gravitino.kms.source.analytics.api", "google-cloud-kms"), + List.of(awsFactory, failingFactory))); + + Assertions.assertSame(creationFailure, exception); + Assertions.assertArrayEquals(new Throwable[] {closeFailure}, exception.getSuppressed()); + } + + @Test + void testDispatchesDifferentSourcesConcurrently() throws Exception { + CountDownLatch entered = new CountDownLatch(2); + CountDownLatch release = new CountDownLatch(1); + KmsClientFactory awsFactory = blockingFactory(AWS_API, entered, release); + KmsClientFactory gcpFactory = blockingFactory(GCP_API, entered, release); + KmsClientRegistry registry = + new KmsClientRegistry( + config( + "gravitino.kms.sources", "primary,analytics", + "gravitino.kms.source.primary.api", "aws-kms", + "gravitino.kms.source.analytics.api", "google-cloud-kms"), + List.of(awsFactory, gcpFactory)); + ExecutorService executor = Executors.newFixedThreadPool(2); + + try { + Future> awsResult = + executor.submit( + () -> registry.getKeyProperties(new KmsReference(AWS_API, "primary", "key"))); + Future> gcpResult = + executor.submit( + () -> registry.getKeyProperties(new KmsReference(GCP_API, "analytics", "key"))); + + Assertions.assertTrue(entered.await(5, TimeUnit.SECONDS)); + release.countDown(); + Assertions.assertEquals( + "primary", awsResult.get(5, TimeUnit.SECONDS).orElseThrow().reference().source()); + Assertions.assertEquals( + "analytics", gcpResult.get(5, TimeUnit.SECONDS).orElseThrow().reference().source()); + } finally { + release.countDown(); + executor.shutdownNow(); + registry.close(); + } + } + + @Test + void testAggregatesCloseFailures() { + RuntimeException awsFailure = new IllegalStateException("aws close failed"); + RuntimeException gcpFailure = new IllegalStateException("gcp close failed"); + KmsClientRegistry registry = + new KmsClientRegistry( + config( + "gravitino.kms.sources", "primary,analytics", + "gravitino.kms.source.primary.api", "aws-kms", + "gravitino.kms.source.analytics.api", "google-cloud-kms"), + List.of( + new CloseTrackingFactory(AWS_API, "aws", new ArrayList<>(), awsFailure), + new CloseTrackingFactory(GCP_API, "gcp", new ArrayList<>(), gcpFailure))); + + RuntimeException exception = Assertions.assertThrows(RuntimeException.class, registry::close); + Assertions.assertSame(gcpFailure, exception); + Assertions.assertArrayEquals(new Throwable[] {awsFailure}, exception.getSuppressed()); + } + + private static Config config(String... entries) { + Map properties = new HashMap<>(); + for (int index = 0; index < entries.length; index += 2) { + properties.put(entries[index], entries[index + 1]); + } + return new MapConfig(properties); + } + + private static KmsClientFactory factory(String api, ClientCreator creator) { + return new KmsClientFactory() { + @Override + public String api() { + return api; + } + + @Override + public KmsClient create(String source, Map properties) { + return creator.create(source, properties); + } + }; + } + + private static KmsClientFactory blockingFactory( + String api, CountDownLatch entered, CountDownLatch release) { + return factory( + api, + (source, properties) -> + reference -> { + entered.countDown(); + try { + if (!release.await(5, TimeUnit.SECONDS)) { + throw new IllegalStateException("Timed out waiting to release KMS request"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Interrupted while inspecting KMS key", e); + } + return Optional.of(new Properties(reference)); + }); + } + + private interface ClientCreator { + KmsClient create(String source, Map properties); + } + + private static final class RecordingFactory implements KmsClientFactory { + private final String api; + private String source; + private Map properties; + private String providerKeyId; + private final AtomicInteger createCount = new AtomicInteger(); + + private RecordingFactory(String api) { + this.api = api; + } + + @Override + public String api() { + return api; + } + + @Override + public KmsClient create(String source, Map properties) { + createCount.incrementAndGet(); + this.source = source; + this.properties = properties; + return reference -> { + this.source = reference.source(); + this.providerKeyId = reference.keyId(); + return Optional.of(new Properties(reference)); + }; + } + } + + private static final class CloseTrackingFactory implements KmsClientFactory { + private final String api; + private final String name; + private final List closeOrder; + private final RuntimeException closeFailure; + private final AtomicInteger closeCount = new AtomicInteger(); + + private CloseTrackingFactory( + String api, String name, List closeOrder, RuntimeException closeFailure) { + this.api = api; + this.name = name; + this.closeOrder = closeOrder; + this.closeFailure = closeFailure; + } + + @Override + public String api() { + return api; + } + + @Override + public KmsClient create(String source, Map properties) { + return new KmsClient() { + @Override + public Optional getKeyProperties(KmsReference reference) { + return Optional.of(new Properties(reference)); + } + + @Override + public void close() { + closeCount.incrementAndGet(); + closeOrder.add(name); + if (closeFailure != null) { + throw closeFailure; + } + } + }; + } + } + + private static final class Properties implements KmsKeyProperties { + private final KmsReference reference; + + private Properties(KmsReference reference) { + this.reference = reference; + } + + @Override + public KmsReference reference() { + return reference; + } + + @Override + public boolean enabled() { + return true; + } + + @Override + public boolean supportsWrapping() { + return true; + } + + @Override + public boolean supportsUnwrapping() { + return true; + } + } + + private static final class MapConfig extends Config { + private MapConfig(Map properties) { + super(false); + loadFromMap(properties, key -> true); + } + } + + /** Factory exposed for the context-classloader ServiceLoader test. */ + public static final class ServiceLoadedFactory implements KmsClientFactory { + + /** Creates a test service-loaded factory. */ + public ServiceLoadedFactory() {} + + /** {@inheritDoc} */ + @Override + public String api() { + return AWS_API; + } + + /** {@inheritDoc} */ + @Override + public KmsClient create(String source, Map properties) { + return reference -> Optional.of(new Properties(reference)); + } + } +} diff --git a/core/src/test/java/org/apache/gravitino/encryption/kms/TestKmsConfig.java b/core/src/test/java/org/apache/gravitino/encryption/kms/TestKmsConfig.java new file mode 100644 index 00000000000..b066ef0fd42 --- /dev/null +++ b/core/src/test/java/org/apache/gravitino/encryption/kms/TestKmsConfig.java @@ -0,0 +1,143 @@ +/* + * 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.gravitino.encryption.kms; + +import java.util.HashMap; +import java.util.Map; +import org.apache.gravitino.Config; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class TestKmsConfig { + + private static final String AWS_API = "aws-kms"; + private static final String GCP_API = "google-cloud-kms"; + + @Test + void testParsesSourcesAndProviderProperties() { + KmsConfig config = + parse( + Map.of( + "gravitino.kms.sources", "primary, disaster-recovery", + "gravitino.kms.source.primary.api", " AWS-KMS ", + "gravitino.kms.source.primary.endpoint.region", "us-west-2", + "gravitino.kms.source.primary.credential.method", "default", + "gravitino.kms.source.disaster-recovery.api", "google-cloud-kms", + "gravitino.kms.source.disaster-recovery.endpoint.projectId", "backup-project", + "gravitino.kms.source.disaster-recovery.credential.method", "default")); + + Assertions.assertEquals(2, config.sources().size()); + Assertions.assertEquals(AWS_API, config.sources().get("primary").api()); + Assertions.assertEquals( + Map.of("endpoint.region", "us-west-2", "credential.method", "default"), + config.sources().get("primary").properties()); + Assertions.assertEquals(GCP_API, config.sources().get("disaster-recovery").api()); + Assertions.assertEquals( + Map.of("endpoint.projectId", "backup-project", "credential.method", "default"), + config.sources().get("disaster-recovery").properties()); + Assertions.assertThrows( + UnsupportedOperationException.class, + () -> config.sources().get("primary").properties().put("endpoint.region", "other")); + } + + @Test + void testAllowsNoConfiguredSources() { + Assertions.assertTrue(parse(Map.of()).sources().isEmpty()); + Assertions.assertTrue(parse(Map.of("gravitino.kms.sources", " ")).sources().isEmpty()); + } + + @Test + void testRejectsInvalidOrDuplicateSourceNames() { + assertInvalid(Map.of("gravitino.kms.sources", "primary,"), "Invalid KMS source name"); + assertInvalid(Map.of("gravitino.kms.sources", "bad.name"), "Invalid KMS source name"); + assertInvalid( + Map.of("gravitino.kms.sources", "primary,primary"), "Duplicate KMS source 'primary'"); + } + + @Test + void testRejectsMalformedOrUnlistedSourceProperties() { + assertInvalid(Map.of("gravitino.kms.unexpected", "value"), "Invalid KMS configuration key"); + assertInvalid( + Map.of( + "gravitino.kms.sources", "primary", + "gravitino.kms.source.other.api", "aws-kms"), + "unlisted source 'other'"); + assertInvalid( + Map.of( + "gravitino.kms.sources", "primary", + "gravitino.kms.source.primary.", "value"), + "Invalid KMS configuration key"); + } + + @Test + void testRequiresApi() { + assertInvalid( + Map.of("gravitino.kms.sources", "primary"), + "gravitino.kms.source.primary.api' cannot be blank"); + assertInvalid( + Map.of( + "gravitino.kms.sources", "primary", + "gravitino.kms.source.primary.api", " "), + "gravitino.kms.source.primary.api' cannot be blank"); + + KmsConfig config = + parse( + Map.of( + "gravitino.kms.sources", "primary", + "gravitino.kms.source.primary.api", " Custom-KMS ")); + Assertions.assertEquals("custom-kms", config.sources().get("primary").api()); + } + + @Test + void testAllowsMoreThanOneSourceForAnApi() { + KmsConfig config = + parse( + Map.of( + "gravitino.kms.sources", "primary,secondary", + "gravitino.kms.source.primary.api", "aws-kms", + "gravitino.kms.source.secondary.api", "aws-kms")); + + Assertions.assertEquals(AWS_API, config.sources().get("primary").api()); + Assertions.assertEquals(AWS_API, config.sources().get("secondary").api()); + } + + @Test + void testRejectsNullConfiguration() { + Assertions.assertThrows(KmsConfigurationException.class, () -> new KmsConfig(null)); + } + + private static KmsConfig parse(Map properties) { + return new KmsConfig(new MapConfig(properties)); + } + + private static void assertInvalid(Map properties, String expectedMessage) { + KmsConfigurationException exception = + Assertions.assertThrows(KmsConfigurationException.class, () -> parse(properties)); + Assertions.assertTrue( + exception.getMessage().contains(expectedMessage), + () -> String.format("Expected '%s' in '%s'", expectedMessage, exception.getMessage())); + } + + private static final class MapConfig extends Config { + private MapConfig(Map properties) { + super(false); + loadFromMap(new HashMap<>(properties), key -> true); + } + } +}