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..93f8ec90d01
--- /dev/null
+++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java
@@ -0,0 +1,112 @@
+/*
+ * 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.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;
+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 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 identifier
+ * @param source configured KMS client-instance name
+ * @param keyId provider-native key identifier
+ * @throws IllegalArgumentException if any argument is null or blank
+ */
+ @JsonCreator
+ public KmsReference(
+ @JsonProperty("api") String api,
+ @JsonProperty("source") String source,
+ @JsonProperty("keyId") String keyId) {
+ 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.trim().toLowerCase(Locale.ROOT);
+ this.source = source;
+ this.keyId = keyId;
+ }
+
+ /**
+ * Returns the explicitly selected KMS API identifier.
+ *
+ * @return the canonical KMS API identifier
+ */
+ public String 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.equals(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/TestKmsReference.java b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsReference.java
new file mode 100644
index 00000000000..32b6b45b062
--- /dev/null
+++ b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsReference.java
@@ -0,0 +1,75 @@
+/*
+ * 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 testStoresCanonicalApiAndPreservesProviderKey() {
+ KmsReference reference = new KmsReference(" AWS-KMS ", "production", " alias/Customer-Key ");
+
+ Assertions.assertEquals("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("", "production", "key"));
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () -> new KmsReference(" ", "production", "key"));
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () -> new KmsReference("aws-kms", null, "key"));
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () -> new KmsReference("aws-kms", "production", null));
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () -> new KmsReference("aws-kms", "", "key"));
+ Assertions.assertThrows(
+ 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("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());
+ 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/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/common/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java
new file mode 100644
index 00000000000..d199c7cbd63
--- /dev/null
+++ b/common/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/common/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java
new file mode 100644
index 00000000000..15b83ab1c43
--- /dev/null
+++ b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java
@@ -0,0 +1,53 @@
+/*
+ * 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.Optional;
+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.
+ *
+ *
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, 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
+ */
+ Optional getKeyProperties(KmsReference reference);
+
+ /** Releases resources owned by this client. */
+ @Override
+ default void close() {}
+}
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
new file mode 100644
index 00000000000..f533d4c2253
--- /dev/null
+++ b/common/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 stable, canonical KMS API identifier implemented by this factory.
+ *
+ * @return the supported KMS API identifier
+ */
+ String 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/common/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java
new file mode 100644
index 00000000000..ec1613666a6
--- /dev/null
+++ b/common/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/common/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java
new file mode 100644
index 00000000000..67dec82f40d
--- /dev/null
+++ b/common/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java
@@ -0,0 +1,63 @@
+/*
+ * 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 successfully located KMS key.
+ *
+ * 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 {
+
+ /**
+ * Returns the requested key reference.
+ *
+ * @return the key reference
+ */
+ KmsReference reference();
+
+ /**
+ * Returns the provider's normalized lifecycle state for the key.
+ *
+ *
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 structurally supports wrapping data-encryption keys.
+ *
+ * @return whether wrapping is supported
+ */
+ boolean supportsWrapping();
+
+ /**
+ * Returns whether the key structurally supports unwrapping data-encryption keys.
+ *
+ * @return whether unwrapping is supported
+ */
+ boolean supportsUnwrapping();
+}
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
new file mode 100644
index 00000000000..91b2609b2e1
--- /dev/null
+++ b/common/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java
@@ -0,0 +1,59 @@
+/*
+ * 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 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)
+ .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() {
+ return client;
+ }
+
+ @Override
+ protected KmsReference usableKey() {
+ return new KmsReference(API, SOURCE, USABLE_KEY);
+ }
+
+ @Override
+ protected KmsReference missingKey() {
+ 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
new file mode 100644
index 00000000000..9c2b6b4c085
--- /dev/null
+++ b/common/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java
@@ -0,0 +1,95 @@
+/*
+ * 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.Optional;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class TestKmsClient {
+
+ private static final KmsReference REFERENCE = new KmsReference("test-kms", "production", "key");
+
+ @Test
+ void testReturnsProviderProperties() {
+ KmsKeyProperties properties = new TestProperties();
+ KmsClient client = reference -> Optional.of(properties);
+
+ Assertions.assertSame(properties, client.getKeyProperties(REFERENCE).orElseThrow());
+ Assertions.assertTrue(properties.supportsWrapping());
+ Assertions.assertFalse(properties.supportsUnwrapping());
+ }
+
+ @Test
+ void testDefaultClose() {
+ 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
+ public KmsReference reference() {
+ return REFERENCE;
+ }
+
+ @Override
+ public boolean enabled() {
+ return true;
+ }
+
+ @Override
+ public boolean supportsWrapping() {
+ return true;
+ }
+
+ @Override
+ public boolean supportsUnwrapping() {
+ return false;
+ }
+ }
+}
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..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,6 +36,7 @@
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.KmsReference;
import org.apache.gravitino.rel.indexes.Index;
import org.apache.gravitino.rel.types.Type;
import org.apache.gravitino.rel.types.Types;
@@ -254,6 +255,29 @@ void testGetLong() throws Exception {
assertEquals(1L, result);
}
+ @Test
+ void testKmsReferenceAnyFieldMapperSerde() throws JsonProcessingException {
+ KmsReference reference =
+ new KmsReference("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));
+ assertEquals(
+ reference,
+ JsonUtils.anyFieldMapper()
+ .readValue(
+ "{\"api\":\" GOOGLE-CLOUD-KMS \",\"source\":\"analytics-prod\","
+ + "\"keyId\":\"projects/p/keys/k\"}",
+ KmsReference.class));
+ }
+
@Test
void testPartitionDTOSerde() throws JsonProcessingException {
String[] field1 = {"dt"};
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
new file mode 100644
index 00000000000..a55a488c5d8
--- /dev/null
+++ b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java
@@ -0,0 +1,140 @@
+/*
+ * 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.Locale;
+import java.util.Map;
+import java.util.Optional;
+
+/** In-memory KMS client for contract and consumer tests. */
+public final class FakeKmsClient implements KmsClient {
+
+ private final String api;
+ private final String source;
+ private final Map keys = new HashMap<>();
+
+ /**
+ * Creates an empty fake client.
+ *
+ * @param api canonical KMS API identifier accepted by the client
+ * @param source configured source accepted by the client
+ */
+ 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;
+ }
+
+ /**
+ * 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 Optional getKeyProperties(KmsReference reference) {
+ requireReference(reference);
+ KeyState state = keys.get(reference.keyId());
+ return state == null
+ ? Optional.empty()
+ : Optional.of(
+ new Properties(
+ reference, 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().equals(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 enabled;
+ private final boolean supportsWrapping;
+ private final boolean supportsUnwrapping;
+
+ private Properties(
+ KmsReference reference,
+ boolean enabled,
+ boolean supportsWrapping,
+ boolean supportsUnwrapping) {
+ this.reference = reference;
+ this.enabled = enabled;
+ this.supportsWrapping = supportsWrapping;
+ this.supportsUnwrapping = supportsUnwrapping;
+ }
+
+ @Override
+ public KmsReference reference() {
+ return reference;
+ }
+
+ @Override
+ public boolean enabled() {
+ return enabled;
+ }
+
+ @Override
+ public boolean supportsWrapping() {
+ return supportsWrapping;
+ }
+
+ @Override
+ public boolean supportsUnwrapping() {
+ return supportsUnwrapping;
+ }
+ }
+}
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
new file mode 100644
index 00000000000..868d64bd62e
--- /dev/null
+++ b/common/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java
@@ -0,0 +1,95 @@
+/*
+ * 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.Optional;
+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();
+ 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.enabled());
+ Assertions.assertTrue(properties.supportsWrapping());
+ Assertions.assertTrue(properties.supportsUnwrapping());
+ }
+
+ @Test
+ void testReportsMissingKey() {
+ KmsReference reference = missingKey();
+ Optional result = client().getKeyProperties(reference);
+
+ Assertions.assertNotNull(result, "KMS clients must not return null");
+ Assertions.assertTrue(result.isEmpty());
+ }
+
+ @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();
+ 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
new file mode 100644
index 00000000000..52eec8f5dc1
--- /dev/null
+++ b/common/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 identifier
+ */
+ protected abstract String expectedApi();
+
+ @Test
+ void testReportsExpectedApi() {
+ Assertions.assertEquals(expectedApi(), factory().api());
+ }
+}
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);
+ }
+ }
+}