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());
+ }
+}
diff --git a/bundles/aws-bundle/build.gradle.kts b/bundles/aws-bundle/build.gradle.kts
index b7ee4927192..268cd47785e 100644
--- a/bundles/aws-bundle/build.gradle.kts
+++ b/bundles/aws-bundle/build.gradle.kts
@@ -29,6 +29,7 @@ dependencies {
implementation(project(":bundles:aws"))
implementation(libs.hadoop3.aws)
implementation(libs.aws.iam)
+ implementation(libs.aws.kms)
implementation(libs.aws.policy)
implementation(libs.aws.sts)
implementation(libs.hadoop3.client.api)
diff --git a/bundles/aws/build.gradle.kts b/bundles/aws/build.gradle.kts
index 7b6d4f1ec12..bf2c933db6c 100644
--- a/bundles/aws/build.gradle.kts
+++ b/bundles/aws/build.gradle.kts
@@ -40,19 +40,32 @@ dependencies {
implementation(libs.guava)
compileOnly(libs.aws.iam)
+ compileOnly(libs.aws.kms)
compileOnly(libs.aws.policy)
compileOnly(libs.aws.sts)
compileOnly(libs.hadoop3.aws)
compileOnly(libs.hadoop3.client.api)
testImplementation(libs.aws.iam)
+ testImplementation(libs.aws.kms)
testImplementation(libs.aws.policy)
testImplementation(libs.aws.sts)
testImplementation(libs.junit.jupiter.api)
testImplementation(libs.junit.jupiter.params)
+ testImplementation(libs.mockito.core)
+ testImplementation(libs.testcontainers)
+ testImplementation(libs.testcontainers.localstack)
+ testImplementation(testFixtures(project(":api")))
testRuntimeOnly(libs.junit.jupiter.engine)
+ testRuntimeOnly(libs.slf4j.jdk14)
}
tasks.compileJava {
dependsOn(":catalogs:catalog-fileset:runtimeJars")
}
+
+tasks.test {
+ if (project.hasProperty("skipITs")) {
+ exclude("**/integration/test/**")
+ }
+}
diff --git a/bundles/aws/src/main/java/org/apache/gravitino/kms/aws/AwsKmsClient.java b/bundles/aws/src/main/java/org/apache/gravitino/kms/aws/AwsKmsClient.java
new file mode 100644
index 00000000000..d0efcfb2b45
--- /dev/null
+++ b/bundles/aws/src/main/java/org/apache/gravitino/kms/aws/AwsKmsClient.java
@@ -0,0 +1,106 @@
+/*
+ * 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.kms.aws;
+
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsAuthenticationException;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.services.kms.model.DescribeKeyRequest;
+import software.amazon.awssdk.services.kms.model.DescribeKeyResponse;
+import software.amazon.awssdk.services.kms.model.InvalidArnException;
+import software.amazon.awssdk.services.kms.model.KeyMetadata;
+import software.amazon.awssdk.services.kms.model.KeyState;
+import software.amazon.awssdk.services.kms.model.KeyUsageType;
+import software.amazon.awssdk.services.kms.model.KmsException;
+import software.amazon.awssdk.services.kms.model.NotFoundException;
+
+final class AwsKmsClient implements KmsClient {
+
+ private final String source;
+ private final software.amazon.awssdk.services.kms.KmsClient delegate;
+
+ AwsKmsClient(String source, software.amazon.awssdk.services.kms.KmsClient delegate) {
+ this.source = source;
+ this.delegate = delegate;
+ }
+
+ @Override
+ public KmsKeyProperties getKeyProperties(KmsReference reference) {
+ validateReference(reference);
+ String providerKeyId = reference.keyId();
+
+ try {
+ DescribeKeyResponse response =
+ delegate.describeKey(DescribeKeyRequest.builder().keyId(providerKeyId).build());
+ if (response == null || response.keyMetadata() == null) {
+ throw new ConnectionFailedException(
+ "AWS KMS returned no metadata for key '%s' from source '%s'", providerKeyId, source);
+ }
+
+ KeyMetadata metadata = response.keyMetadata();
+ boolean enabled =
+ Boolean.TRUE.equals(metadata.enabled()) && metadata.keyState() == KeyState.ENABLED;
+ boolean supportsEncryption = metadata.keyUsage() == KeyUsageType.ENCRYPT_DECRYPT;
+ return new AwsKmsKeyProperties(
+ reference, true, enabled, supportsEncryption, supportsEncryption);
+ } catch (NotFoundException e) {
+ return new AwsKmsKeyProperties(reference, false, false, false, false);
+ } catch (InvalidArnException e) {
+ throw new IllegalArgumentException(
+ String.format("Invalid AWS KMS key ID '%s'", providerKeyId), e);
+ } catch (KmsException e) {
+ if (e.statusCode() == 401 || e.statusCode() == 403) {
+ throw new KmsAuthenticationException(
+ e, "AWS KMS rejected credentials for source '%s'", source);
+ }
+ throw new ConnectionFailedException(
+ e, "Failed to inspect AWS KMS key '%s' from source '%s'", providerKeyId, source);
+ } catch (SdkClientException e) {
+ throw new ConnectionFailedException(
+ e, "Failed to inspect AWS KMS key '%s' from source '%s'", providerKeyId, source);
+ }
+ }
+
+ @Override
+ public void close() {
+ delegate.close();
+ }
+
+ private void validateReference(KmsReference reference) {
+ if (reference == null) {
+ throw new IllegalArgumentException("KMS reference cannot be null");
+ }
+ if (reference.api() != KmsApi.AWS_KMS) {
+ throw new IllegalArgumentException(
+ String.format(
+ "KMS source '%s' requires API '%s', not '%s'",
+ source, KmsApi.AWS_KMS.wireValue(), reference.api().wireValue()));
+ }
+ if (!source.equals(reference.source())) {
+ throw new IllegalArgumentException(
+ String.format(
+ "KMS reference source '%s' does not match AWS KMS source '%s'",
+ reference.source(), source));
+ }
+ }
+}
diff --git a/bundles/aws/src/main/java/org/apache/gravitino/kms/aws/AwsKmsClientFactory.java b/bundles/aws/src/main/java/org/apache/gravitino/kms/aws/AwsKmsClientFactory.java
new file mode 100644
index 00000000000..46c03781272
--- /dev/null
+++ b/bundles/aws/src/main/java/org/apache/gravitino/kms/aws/AwsKmsClientFactory.java
@@ -0,0 +1,149 @@
+/*
+ * 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.kms.aws;
+
+import com.google.common.collect.ImmutableSet;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Map;
+import java.util.Set;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsClientFactory;
+import org.apache.gravitino.encryption.kms.KmsConfigurationException;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.kms.KmsClientBuilder;
+
+/** Creates AWS KMS clients using the AWS SDK default credential provider chain. */
+public final class AwsKmsClientFactory implements KmsClientFactory {
+
+ /** Required AWS region property. */
+ public static final String REGION = "endpoint.region";
+
+ /** Optional HTTP(S) AWS KMS endpoint override property. */
+ public static final String SERVICE_ADDRESS = "endpoint.serviceAddress";
+
+ /** Credential method property. */
+ public static final String CREDENTIAL_METHOD = "credential.method";
+
+ private static final Set SUPPORTED_PROPERTIES =
+ ImmutableSet.of(REGION, SERVICE_ADDRESS, CREDENTIAL_METHOD);
+
+ /** Creates an AWS KMS client factory. */
+ public AwsKmsClientFactory() {}
+
+ /**
+ * Returns the AWS KMS API.
+ *
+ * @return the AWS KMS API
+ */
+ @Override
+ public KmsApi api() {
+ return KmsApi.AWS_KMS;
+ }
+
+ /**
+ * Creates an AWS KMS client for a configured source.
+ *
+ * Credentials are intentionally not accepted as source properties. The AWS SDK resolves
+ * credentials through its default credential provider chain when a request is made.
+ *
+ * @param source logical name of the configured AWS KMS instance
+ * @param properties AWS KMS properties
+ * @return the configured AWS KMS client
+ * @throws IllegalArgumentException if the source or properties are invalid
+ * @throws ConnectionFailedException if the AWS SDK client cannot be initialized
+ */
+ @Override
+ public KmsClient create(String source, Map properties) {
+ validateSource(source);
+ if (properties == null) {
+ throw new KmsConfigurationException("AWS KMS properties cannot be null");
+ }
+ validatePropertyNames(properties);
+
+ String region = requiredProperty(properties, REGION);
+ String credentialMethod = requiredProperty(properties, CREDENTIAL_METHOD);
+ if (!"default".equals(credentialMethod)) {
+ throw new KmsConfigurationException(
+ "Unsupported AWS KMS credential method: %s", credentialMethod);
+ }
+ KmsClientBuilder builder =
+ software.amazon.awssdk.services.kms.KmsClient.builder().region(Region.of(region));
+ if (properties.containsKey(SERVICE_ADDRESS)) {
+ builder.endpointOverride(parseServiceAddress(requiredProperty(properties, SERVICE_ADDRESS)));
+ }
+
+ try {
+ return new AwsKmsClient(source, builder.build());
+ } catch (SdkClientException e) {
+ throw new ConnectionFailedException(
+ e, "Failed to initialize AWS KMS client for source '%s'", source);
+ }
+ }
+
+ private static void validateSource(String source) {
+ if (source == null || source.trim().isEmpty()) {
+ throw new KmsConfigurationException("AWS KMS source cannot be blank");
+ }
+ }
+
+ private static void validatePropertyNames(Map properties) {
+ for (String property : properties.keySet()) {
+ if (!SUPPORTED_PROPERTIES.contains(property)) {
+ throw new KmsConfigurationException("Unsupported AWS KMS property '%s'", property);
+ }
+ }
+ }
+
+ private static String requiredProperty(Map properties, String property) {
+ String value = properties.get(property);
+ if (value == null || value.trim().isEmpty()) {
+ throw new KmsConfigurationException("AWS KMS property '%s' cannot be blank", property);
+ }
+ return value.trim();
+ }
+
+ private static URI parseServiceAddress(String serviceAddress) {
+ URI uri;
+ try {
+ uri = new URI(serviceAddress);
+ } catch (URISyntaxException e) {
+ throw new KmsConfigurationException(
+ e, "Invalid AWS KMS service address '%s'", serviceAddress);
+ }
+
+ boolean http = "http".equalsIgnoreCase(uri.getScheme());
+ boolean https = "https".equalsIgnoreCase(uri.getScheme());
+ boolean rootPath =
+ uri.getPath() == null || uri.getPath().isEmpty() || "/".equals(uri.getPath());
+ if ((!http && !https)
+ || uri.getHost() == null
+ || uri.getUserInfo() != null
+ || !rootPath
+ || uri.getQuery() != null
+ || uri.getFragment() != null) {
+ throw new KmsConfigurationException(
+ "Invalid AWS KMS service address '%s'; expected an HTTP(S) origin", serviceAddress);
+ }
+ return uri;
+ }
+}
diff --git a/bundles/aws/src/main/java/org/apache/gravitino/kms/aws/AwsKmsKeyProperties.java b/bundles/aws/src/main/java/org/apache/gravitino/kms/aws/AwsKmsKeyProperties.java
new file mode 100644
index 00000000000..2442537084a
--- /dev/null
+++ b/bundles/aws/src/main/java/org/apache/gravitino/kms/aws/AwsKmsKeyProperties.java
@@ -0,0 +1,69 @@
+/*
+ * 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.kms.aws;
+
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+
+final class AwsKmsKeyProperties implements KmsKeyProperties {
+
+ private final KmsReference reference;
+ private final boolean present;
+ private final boolean enabled;
+ private final boolean supportsWrapping;
+ private final boolean supportsUnwrapping;
+
+ AwsKmsKeyProperties(
+ 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/bundles/aws/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory b/bundles/aws/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory
new file mode 100644
index 00000000000..9dca5b002b4
--- /dev/null
+++ b/bundles/aws/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory
@@ -0,0 +1,20 @@
+#
+# 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.
+#
+
+org.apache.gravitino.kms.aws.AwsKmsClientFactory
diff --git a/bundles/aws/src/test/java/org/apache/gravitino/kms/aws/TestAwsKmsClient.java b/bundles/aws/src/test/java/org/apache/gravitino/kms/aws/TestAwsKmsClient.java
new file mode 100644
index 00000000000..f4537049787
--- /dev/null
+++ b/bundles/aws/src/test/java/org/apache/gravitino/kms/aws/TestAwsKmsClient.java
@@ -0,0 +1,322 @@
+/*
+ * 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.kms.aws;
+
+import java.util.stream.Stream;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsAuthenticationException;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+import org.apache.gravitino.encryption.kms.TestKmsClientContract;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
+import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.services.kms.model.DependencyTimeoutException;
+import software.amazon.awssdk.services.kms.model.DescribeKeyRequest;
+import software.amazon.awssdk.services.kms.model.DescribeKeyResponse;
+import software.amazon.awssdk.services.kms.model.InvalidArnException;
+import software.amazon.awssdk.services.kms.model.KeyMetadata;
+import software.amazon.awssdk.services.kms.model.KeySpec;
+import software.amazon.awssdk.services.kms.model.KeyState;
+import software.amazon.awssdk.services.kms.model.KeyUsageType;
+import software.amazon.awssdk.services.kms.model.KmsException;
+import software.amazon.awssdk.services.kms.model.KmsInternalException;
+import software.amazon.awssdk.services.kms.model.NotFoundException;
+
+/** Tests AWS KMS key inspection and normalized property mapping. */
+public class TestAwsKmsClient extends TestKmsClientContract {
+
+ private static final String SOURCE = "primary-aws";
+ private static final String USABLE_KEY_ID = "1234abcd-12ab-34cd-56ef-1234567890ab";
+ private static final String MISSING_KEY_ID = "missing-key";
+
+ private software.amazon.awssdk.services.kms.KmsClient delegate;
+ private AwsKmsClient client;
+
+ @BeforeEach
+ void setUp() {
+ delegate = Mockito.mock(software.amazon.awssdk.services.kms.KmsClient.class);
+ client = new AwsKmsClient(SOURCE, delegate);
+ Mockito.when(delegate.describeKey(Mockito.any(DescribeKeyRequest.class)))
+ .thenAnswer(
+ invocation -> {
+ DescribeKeyRequest request = invocation.getArgument(0);
+ if (MISSING_KEY_ID.equals(request.keyId())) {
+ throw NotFoundException.builder().message("Key does not exist").build();
+ }
+ return response(
+ KeyMetadata.builder()
+ .enabled(true)
+ .keyState(KeyState.ENABLED)
+ .keyUsage(KeyUsageType.ENCRYPT_DECRYPT)
+ .build());
+ });
+ }
+
+ @Override
+ protected KmsClient client() {
+ return client;
+ }
+
+ @Override
+ protected KmsReference usableKey() {
+ return reference(USABLE_KEY_ID);
+ }
+
+ @Override
+ protected KmsReference missingKey() {
+ return reference(MISSING_KEY_ID);
+ }
+
+ @ParameterizedTest
+ @ValueSource(
+ strings = {
+ "1234abcd-12ab-34cd-56ef-1234567890ab",
+ "arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab",
+ "alias/application-key",
+ "arn:aws:kms:us-west-2:111122223333:alias/application-key"
+ })
+ void testPassesProviderNativeKeyIdentifierWithoutModification(String keyId) {
+ client.getKeyProperties(reference(keyId));
+
+ ArgumentCaptor requestCaptor =
+ ArgumentCaptor.forClass(DescribeKeyRequest.class);
+ Mockito.verify(delegate).describeKey(requestCaptor.capture());
+ Assertions.assertEquals(keyId, requestCaptor.getValue().keyId());
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = KeySpec.class,
+ names = {"SYMMETRIC_DEFAULT", "RSA_2048"})
+ void testSupportsSymmetricAndAsymmetricEncryptionKeys(KeySpec keySpec) {
+ stubMetadata(
+ KeyMetadata.builder()
+ .enabled(true)
+ .keyState(KeyState.ENABLED)
+ .keyUsage(KeyUsageType.ENCRYPT_DECRYPT)
+ .keySpec(keySpec)
+ .build());
+
+ KmsKeyProperties properties = usableKeyProperties();
+
+ Assertions.assertTrue(properties.enabled());
+ Assertions.assertTrue(properties.supportsWrapping());
+ Assertions.assertTrue(properties.supportsUnwrapping());
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = KeyState.class,
+ names = {
+ "CREATING",
+ "DISABLED",
+ "PENDING_DELETION",
+ "PENDING_IMPORT",
+ "PENDING_REPLICA_DELETION",
+ "UNAVAILABLE",
+ "UPDATING",
+ "UNKNOWN_TO_SDK_VERSION"
+ })
+ void testOnlyEnabledStateIsEnabled(KeyState keyState) {
+ stubMetadata(
+ KeyMetadata.builder()
+ .enabled(true)
+ .keyState(keyState)
+ .keyUsage(KeyUsageType.ENCRYPT_DECRYPT)
+ .build());
+
+ KmsKeyProperties properties = usableKeyProperties();
+
+ Assertions.assertTrue(properties.present());
+ Assertions.assertFalse(properties.enabled());
+ Assertions.assertTrue(properties.supportsWrapping());
+ Assertions.assertTrue(properties.supportsUnwrapping());
+ }
+
+ @Test
+ void testEnabledFlagMustAlsoBeTrue() {
+ stubMetadata(
+ KeyMetadata.builder()
+ .enabled(false)
+ .keyState(KeyState.ENABLED)
+ .keyUsage(KeyUsageType.ENCRYPT_DECRYPT)
+ .build());
+
+ Assertions.assertFalse(usableKeyProperties().enabled());
+ }
+
+ @Test
+ void testMissingStateIsNotEnabled() {
+ stubMetadata(
+ KeyMetadata.builder().enabled(true).keyUsage(KeyUsageType.ENCRYPT_DECRYPT).build());
+
+ Assertions.assertFalse(usableKeyProperties().enabled());
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = KeyUsageType.class,
+ names = {"SIGN_VERIFY", "GENERATE_VERIFY_MAC", "KEY_AGREEMENT", "UNKNOWN_TO_SDK_VERSION"})
+ void testNonEncryptionUsageDoesNotSupportWrapping(KeyUsageType keyUsage) {
+ stubMetadata(
+ KeyMetadata.builder().enabled(true).keyState(KeyState.ENABLED).keyUsage(keyUsage).build());
+
+ KmsKeyProperties properties = usableKeyProperties();
+
+ Assertions.assertTrue(properties.enabled());
+ Assertions.assertFalse(properties.supportsWrapping());
+ Assertions.assertFalse(properties.supportsUnwrapping());
+ }
+
+ @Test
+ void testMissingUsageDoesNotSupportWrapping() {
+ stubMetadata(KeyMetadata.builder().enabled(true).keyState(KeyState.ENABLED).build());
+
+ KmsKeyProperties properties = usableKeyProperties();
+
+ Assertions.assertFalse(properties.supportsWrapping());
+ Assertions.assertFalse(properties.supportsUnwrapping());
+ }
+
+ @Test
+ void testUnknownProviderStateAndUsageAreConservative() {
+ stubMetadata(
+ KeyMetadata.builder()
+ .enabled(true)
+ .keyState("A_NEW_STATE")
+ .keyUsage("A_NEW_USAGE")
+ .build());
+
+ KmsKeyProperties properties = usableKeyProperties();
+
+ Assertions.assertFalse(properties.enabled());
+ Assertions.assertFalse(properties.supportsWrapping());
+ Assertions.assertFalse(properties.supportsUnwrapping());
+ }
+
+ @Test
+ void testMissingKeyReturnsAllFalseProperties() {
+ KmsReference reference = missingKey();
+
+ KmsKeyProperties properties = client.getKeyProperties(reference);
+
+ Assertions.assertEquals(reference, properties.reference());
+ Assertions.assertFalse(properties.present());
+ Assertions.assertFalse(properties.enabled());
+ Assertions.assertFalse(properties.supportsWrapping());
+ Assertions.assertFalse(properties.supportsUnwrapping());
+ }
+
+ @Test
+ void testInvalidArnIsRejectedAsInvalidInput() {
+ InvalidArnException failure = InvalidArnException.builder().message("Invalid ARN").build();
+ Mockito.when(delegate.describeKey(Mockito.any(DescribeKeyRequest.class))).thenThrow(failure);
+
+ IllegalArgumentException thrown =
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () -> client.getKeyProperties(usableKey()));
+
+ Assertions.assertSame(failure, thrown.getCause());
+ }
+
+ @Test
+ void testRejectedCredentialsAreAuthenticationFailures() {
+ KmsException failure =
+ (KmsException) KmsException.builder().message("access denied").statusCode(403).build();
+ Mockito.when(delegate.describeKey(Mockito.any(DescribeKeyRequest.class))).thenThrow(failure);
+
+ KmsAuthenticationException thrown =
+ Assertions.assertThrows(
+ KmsAuthenticationException.class, () -> client.getKeyProperties(usableKey()));
+
+ Assertions.assertSame(failure, thrown.getCause());
+ }
+
+ @ParameterizedTest
+ @MethodSource("connectionFailures")
+ void testProviderAndTransportFailuresAreConnectionFailures(RuntimeException failure) {
+ Mockito.when(delegate.describeKey(Mockito.any(DescribeKeyRequest.class))).thenThrow(failure);
+
+ ConnectionFailedException thrown =
+ Assertions.assertThrows(
+ ConnectionFailedException.class, () -> client.getKeyProperties(usableKey()));
+
+ Assertions.assertSame(failure, thrown.getCause());
+ }
+
+ @Test
+ void testMissingResponseIsConnectionFailure() {
+ Mockito.when(delegate.describeKey(Mockito.any(DescribeKeyRequest.class))).thenReturn(null);
+
+ Assertions.assertThrows(
+ ConnectionFailedException.class, () -> client.getKeyProperties(usableKey()));
+ }
+
+ @Test
+ void testMissingMetadataIsConnectionFailure() {
+ Mockito.when(delegate.describeKey(Mockito.any(DescribeKeyRequest.class)))
+ .thenReturn(DescribeKeyResponse.builder().build());
+
+ Assertions.assertThrows(
+ ConnectionFailedException.class, () -> client.getKeyProperties(usableKey()));
+ }
+
+ @Test
+ void testCloseDelegatesToAwsClient() {
+ client.close();
+
+ Mockito.verify(delegate).close();
+ }
+
+ private static Stream connectionFailures() {
+ return Stream.of(
+ Arguments.of(SdkClientException.create("transport unavailable")),
+ Arguments.of(DependencyTimeoutException.builder().message("dependency timeout").build()),
+ Arguments.of(KmsInternalException.builder().message("internal failure").build()),
+ Arguments.of(KmsException.builder().message("throttled").statusCode(429).build()));
+ }
+
+ private static DescribeKeyResponse response(KeyMetadata metadata) {
+ return DescribeKeyResponse.builder().keyMetadata(metadata).build();
+ }
+
+ private KmsReference reference(String keyId) {
+ return new KmsReference(KmsApi.AWS_KMS, SOURCE, keyId);
+ }
+
+ private KmsKeyProperties usableKeyProperties() {
+ return client.getKeyProperties(usableKey());
+ }
+
+ private void stubMetadata(KeyMetadata metadata) {
+ Mockito.when(delegate.describeKey(Mockito.any(DescribeKeyRequest.class)))
+ .thenReturn(response(metadata));
+ }
+}
diff --git a/bundles/aws/src/test/java/org/apache/gravitino/kms/aws/TestAwsKmsClientFactory.java b/bundles/aws/src/test/java/org/apache/gravitino/kms/aws/TestAwsKmsClientFactory.java
new file mode 100644
index 00000000000..99a314089eb
--- /dev/null
+++ b/bundles/aws/src/test/java/org/apache/gravitino/kms/aws/TestAwsKmsClientFactory.java
@@ -0,0 +1,169 @@
+/*
+ * 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.kms.aws;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.ServiceLoader;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsClientFactory;
+import org.apache.gravitino.encryption.kms.TestKmsClientFactoryContract;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.NullAndEmptySource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+/** Tests AWS KMS factory validation, initialization, and service discovery. */
+public class TestAwsKmsClientFactory extends TestKmsClientFactoryContract {
+
+ private static final String SOURCE = "primary-aws";
+
+ private final AwsKmsClientFactory factory = new AwsKmsClientFactory();
+
+ @Override
+ protected KmsClientFactory factory() {
+ return factory;
+ }
+
+ @Override
+ protected KmsApi expectedApi() {
+ return KmsApi.AWS_KMS;
+ }
+
+ @ParameterizedTest
+ @NullAndEmptySource
+ @ValueSource(strings = {" ", "\t"})
+ void testRejectsBlankSource(String source) {
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () -> factory.create(source, validProperties()));
+ }
+
+ @Test
+ void testRejectsNullProperties() {
+ Assertions.assertThrows(IllegalArgumentException.class, () -> factory.create(SOURCE, null));
+ }
+
+ @Test
+ void testRejectsMissingRegion() {
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> factory.create(SOURCE, Collections.emptyMap()));
+ }
+
+ @ParameterizedTest
+ @NullAndEmptySource
+ @ValueSource(strings = {" ", "\t"})
+ void testRejectsBlankRegion(String region) {
+ Map properties = new HashMap<>();
+ properties.put(AwsKmsClientFactory.REGION, region);
+
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () -> factory.create(SOURCE, properties));
+ }
+
+ @Test
+ void testRejectsUnknownProperty() {
+ Map properties = validProperties();
+ properties.put("accessKeyId", "must-not-be-configured-here");
+
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () -> factory.create(SOURCE, properties));
+ }
+
+ @Test
+ void testRejectsNullPropertyName() {
+ Map properties = validProperties();
+ properties.put(null, "value");
+
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () -> factory.create(SOURCE, properties));
+ }
+
+ @ParameterizedTest
+ @NullAndEmptySource
+ @ValueSource(strings = {" ", "ftp://localhost:4566", "http:///missing-host", "not a uri"})
+ void testRejectsInvalidServiceAddress(String serviceAddress) {
+ Map properties = validProperties();
+ properties.put(AwsKmsClientFactory.SERVICE_ADDRESS, serviceAddress);
+
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () -> factory.create(SOURCE, properties));
+ }
+
+ @ParameterizedTest
+ @ValueSource(
+ strings = {
+ "http://user@localhost:4566",
+ "http://localhost:4566/path",
+ "http://localhost:4566?query=value",
+ "http://localhost:4566#fragment"
+ })
+ void testRejectsServiceAddressThatIsNotAnOrigin(String serviceAddress) {
+ Map properties = validProperties();
+ properties.put(AwsKmsClientFactory.SERVICE_ADDRESS, serviceAddress);
+
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () -> factory.create(SOURCE, properties));
+ }
+
+ @Test
+ void testCreatesClientWithoutResolvingCredentialsOrContactingAws() {
+ try (KmsClient client = factory.create(SOURCE, validProperties())) {
+ Assertions.assertNotNull(client);
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(
+ strings = {
+ "http://localhost:4566",
+ "http://localhost:4566/",
+ "https://kms.us-west-2.amazonaws.com"
+ })
+ void testCreatesClientWithServiceAddress(String serviceAddress) {
+ Map properties = validProperties();
+ properties.put(AwsKmsClientFactory.SERVICE_ADDRESS, serviceAddress);
+
+ try (KmsClient client = factory.create(SOURCE, properties)) {
+ Assertions.assertNotNull(client);
+ }
+ }
+
+ @Test
+ void testFactoryIsDiscoverableWithServiceLoader() {
+ boolean discovered = false;
+ for (KmsClientFactory candidate : ServiceLoader.load(KmsClientFactory.class)) {
+ if (candidate.getClass() == AwsKmsClientFactory.class) {
+ discovered = true;
+ }
+ }
+
+ Assertions.assertTrue(discovered);
+ }
+
+ private static Map validProperties() {
+ Map properties = new HashMap<>();
+ properties.put(AwsKmsClientFactory.REGION, "us-west-2");
+ properties.put(AwsKmsClientFactory.CREDENTIAL_METHOD, "default");
+ return properties;
+ }
+}
diff --git a/bundles/azure-bundle/build.gradle.kts b/bundles/azure-bundle/build.gradle.kts
index 632cc83d539..69d793f331b 100644
--- a/bundles/azure-bundle/build.gradle.kts
+++ b/bundles/azure-bundle/build.gradle.kts
@@ -32,6 +32,7 @@ dependencies {
}
implementation(libs.azure.identity)
+ implementation(libs.azure.security.keyvault.keys)
implementation(libs.azure.storage.file.datalake)
implementation(libs.hadoop3.abs)
implementation(libs.hadoop3.client.api)
diff --git a/bundles/azure/build.gradle.kts b/bundles/azure/build.gradle.kts
index 031c7c5e8cd..4879d620421 100644
--- a/bundles/azure/build.gradle.kts
+++ b/bundles/azure/build.gradle.kts
@@ -40,16 +40,30 @@ dependencies {
implementation(libs.guava)
compileOnly(libs.azure.identity)
+ compileOnly(libs.azure.security.keyvault.keys)
compileOnly(libs.azure.storage.file.datalake)
compileOnly(libs.hadoop3.abs)
compileOnly(libs.hadoop3.client.api)
testImplementation(libs.azure.identity)
+ testImplementation(libs.azure.security.keyvault.keys)
+ testImplementation(testFixtures(project(":api")))
testImplementation(libs.junit.jupiter.api)
testImplementation(libs.junit.jupiter.params)
testRuntimeOnly(libs.junit.jupiter.engine)
+ testRuntimeOnly(libs.slf4j.simple)
}
tasks.compileJava {
dependsOn(":catalogs:catalog-fileset:runtimeJars")
}
+
+tasks.test {
+ if (project.hasProperty("skipITs")) {
+ exclude("**/integration/test/**")
+ } else {
+ testLogging.showStandardStreams = true
+ systemProperty("org.slf4j.simpleLogger.defaultLogLevel", "warn")
+ systemProperty("org.slf4j.simpleLogger.log.com.azure", "off")
+ }
+}
diff --git a/bundles/azure/src/main/java/org/apache/gravitino/encryption/kms/azure/AzureKeyVaultKeyIdentifier.java b/bundles/azure/src/main/java/org/apache/gravitino/encryption/kms/azure/AzureKeyVaultKeyIdentifier.java
new file mode 100644
index 00000000000..1bd6ea12d8e
--- /dev/null
+++ b/bundles/azure/src/main/java/org/apache/gravitino/encryption/kms/azure/AzureKeyVaultKeyIdentifier.java
@@ -0,0 +1,114 @@
+/*
+ * 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.azure;
+
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Locale;
+import java.util.regex.Pattern;
+import javax.annotation.Nullable;
+
+final class AzureKeyVaultKeyIdentifier {
+
+ private static final Pattern KEY_NAME = Pattern.compile("[A-Za-z0-9-]{1,127}");
+ private static final Pattern KEY_VERSION = Pattern.compile("[A-Za-z0-9-]+");
+
+ private final String name;
+ @Nullable private final String version;
+
+ private AzureKeyVaultKeyIdentifier(String name, @Nullable String version) {
+ this.name = name;
+ this.version = version;
+ }
+
+ static AzureKeyVaultKeyIdentifier parse(String keyId, URI vaultUri) {
+ URI keyUri;
+ try {
+ keyUri = new URI(keyId);
+ } catch (URISyntaxException e) {
+ throw new IllegalArgumentException("Azure Key Vault key ID must be a valid URI", e);
+ }
+
+ validateAbsoluteKeyUri(keyUri);
+ if (!sameOrigin(vaultUri, keyUri)) {
+ throw new IllegalArgumentException(
+ String.format("Azure Key Vault key ID must belong to configured vault '%s'", vaultUri));
+ }
+
+ String[] pathSegments = keyUri.getRawPath().split("/", -1);
+ if ((pathSegments.length != 3 && pathSegments.length != 4)
+ || !pathSegments[0].isEmpty()
+ || !"keys".equalsIgnoreCase(pathSegments[1])
+ || !KEY_NAME.matcher(pathSegments[2]).matches()
+ || (pathSegments.length == 4 && !KEY_VERSION.matcher(pathSegments[3]).matches())) {
+ throw new IllegalArgumentException(
+ "Azure Key Vault key ID must have the form https:///keys/[/]");
+ }
+
+ return new AzureKeyVaultKeyIdentifier(
+ pathSegments[2], pathSegments.length == 4 ? pathSegments[3] : null);
+ }
+
+ static URI validateVaultUri(URI vaultUri) {
+ if (vaultUri == null
+ || !vaultUri.isAbsolute()
+ || !"https".equalsIgnoreCase(vaultUri.getScheme())
+ || vaultUri.getHost() == null
+ || vaultUri.getUserInfo() != null
+ || vaultUri.getPort() != -1
+ || vaultUri.getRawQuery() != null
+ || vaultUri.getRawFragment() != null
+ || !(vaultUri.getRawPath().isEmpty() || "/".equals(vaultUri.getRawPath()))) {
+ throw new IllegalArgumentException(
+ "Azure Key Vault URI must be an HTTPS origin without credentials, a port, query, or fragment");
+ }
+ return vaultUri;
+ }
+
+ String name() {
+ return name;
+ }
+
+ @Nullable
+ String version() {
+ return version;
+ }
+
+ private static void validateAbsoluteKeyUri(URI keyUri) {
+ if (!keyUri.isAbsolute()
+ || !"https".equalsIgnoreCase(keyUri.getScheme())
+ || keyUri.getHost() == null
+ || keyUri.getUserInfo() != null
+ || keyUri.getPort() != -1
+ || keyUri.getRawQuery() != null
+ || keyUri.getRawFragment() != null) {
+ throw new IllegalArgumentException(
+ "Azure Key Vault key ID must be an HTTPS URL without credentials, a port, query, or fragment");
+ }
+ }
+
+ private static boolean sameOrigin(URI first, URI second) {
+ return first
+ .getScheme()
+ .toLowerCase(Locale.ROOT)
+ .equals(second.getScheme().toLowerCase(Locale.ROOT))
+ && first.getHost().equalsIgnoreCase(second.getHost())
+ && first.getPort() == second.getPort();
+ }
+}
diff --git a/bundles/azure/src/main/java/org/apache/gravitino/encryption/kms/azure/AzureKeyVaultKmsClient.java b/bundles/azure/src/main/java/org/apache/gravitino/encryption/kms/azure/AzureKeyVaultKmsClient.java
new file mode 100644
index 00000000000..336d854f43e
--- /dev/null
+++ b/bundles/azure/src/main/java/org/apache/gravitino/encryption/kms/azure/AzureKeyVaultKmsClient.java
@@ -0,0 +1,138 @@
+/*
+ * 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.azure;
+
+import com.azure.core.exception.AzureException;
+import com.azure.core.exception.ClientAuthenticationException;
+import com.azure.core.exception.ResourceNotFoundException;
+import com.azure.security.keyvault.keys.KeyClient;
+import com.azure.security.keyvault.keys.models.KeyOperation;
+import com.azure.security.keyvault.keys.models.KeyProperties;
+import com.azure.security.keyvault.keys.models.KeyVaultKey;
+import java.net.URI;
+import java.time.Clock;
+import java.time.Instant;
+import java.time.OffsetDateTime;
+import java.util.List;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsAuthenticationException;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+
+final class AzureKeyVaultKmsClient implements KmsClient {
+
+ private final String source;
+ private final URI vaultUri;
+ private final KeyClient keyClient;
+ private final Clock clock;
+
+ AzureKeyVaultKmsClient(String source, URI vaultUri, KeyClient keyClient) {
+ this(source, vaultUri, keyClient, Clock.systemUTC());
+ }
+
+ AzureKeyVaultKmsClient(String source, URI vaultUri, KeyClient keyClient, Clock clock) {
+ if (source == null || source.trim().isEmpty()) {
+ throw new IllegalArgumentException("Azure Key Vault source cannot be blank");
+ }
+ if (keyClient == null || clock == null) {
+ throw new IllegalArgumentException("Azure Key Vault SDK client and clock cannot be null");
+ }
+
+ this.source = source;
+ this.vaultUri = AzureKeyVaultKeyIdentifier.validateVaultUri(vaultUri);
+ this.keyClient = keyClient;
+ this.clock = clock;
+ }
+
+ @Override
+ public KmsKeyProperties getKeyProperties(KmsReference reference) {
+ validateReference(reference);
+ AzureKeyVaultKeyIdentifier identifier =
+ AzureKeyVaultKeyIdentifier.parse(reference.keyId(), vaultUri);
+
+ try {
+ KeyVaultKey key = keyClient.getKey(identifier.name(), identifier.version());
+ return normalize(reference, key);
+ } catch (ResourceNotFoundException e) {
+ return new AzureKeyVaultKmsKeyProperties(reference, false, false, false, false);
+ } catch (ConnectionFailedException e) {
+ throw e;
+ } catch (ClientAuthenticationException e) {
+ throw new KmsAuthenticationException(
+ e, "Azure Key Vault rejected credentials for source '%s'", source);
+ } catch (AzureException e) {
+ throw connectionFailure(reference, e);
+ } catch (RuntimeException e) {
+ throw connectionFailure(reference, e);
+ }
+ }
+
+ private void validateReference(KmsReference reference) {
+ if (reference == null) {
+ throw new IllegalArgumentException("Azure Key Vault reference cannot be null");
+ }
+ if (reference.api() != KmsApi.AZURE_KEY_VAULT) {
+ throw new IllegalArgumentException(
+ String.format("KMS reference API must be '%s'", KmsApi.AZURE_KEY_VAULT.wireValue()));
+ }
+ if (!source.equals(reference.source())) {
+ throw new IllegalArgumentException(
+ String.format("KMS reference source must be '%s'", source));
+ }
+ }
+
+ private KmsKeyProperties normalize(KmsReference reference, KeyVaultKey key) {
+ if (key == null || key.getProperties() == null) {
+ throw new ConnectionFailedException(
+ "Azure Key Vault returned incomplete properties for key '%s'", reference.keyId());
+ }
+
+ KeyProperties properties = key.getProperties();
+ List operations = key.getKeyOperations();
+ boolean supportsWrapping = operations != null && operations.contains(KeyOperation.WRAP_KEY);
+ boolean supportsUnwrapping = operations != null && operations.contains(KeyOperation.UNWRAP_KEY);
+ return new AzureKeyVaultKmsKeyProperties(
+ reference,
+ true,
+ isEffectivelyEnabled(properties, clock.instant()),
+ supportsWrapping,
+ supportsUnwrapping);
+ }
+
+ private static boolean isEffectivelyEnabled(KeyProperties properties, Instant now) {
+ if (!Boolean.TRUE.equals(properties.isEnabled())) {
+ return false;
+ }
+
+ OffsetDateTime notBefore = properties.getNotBefore();
+ if (notBefore != null && now.isBefore(notBefore.toInstant())) {
+ return false;
+ }
+
+ OffsetDateTime expiresOn = properties.getExpiresOn();
+ return expiresOn == null || now.isBefore(expiresOn.toInstant());
+ }
+
+ private ConnectionFailedException connectionFailure(KmsReference reference, RuntimeException e) {
+ return new ConnectionFailedException(
+ e, "Failed to inspect Azure Key Vault key '%s' for source '%s'", reference.keyId(), source);
+ }
+}
diff --git a/bundles/azure/src/main/java/org/apache/gravitino/encryption/kms/azure/AzureKeyVaultKmsClientFactory.java b/bundles/azure/src/main/java/org/apache/gravitino/encryption/kms/azure/AzureKeyVaultKmsClientFactory.java
new file mode 100644
index 00000000000..959f67d63c3
--- /dev/null
+++ b/bundles/azure/src/main/java/org/apache/gravitino/encryption/kms/azure/AzureKeyVaultKmsClientFactory.java
@@ -0,0 +1,131 @@
+/*
+ * 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.azure;
+
+import com.azure.identity.DefaultAzureCredentialBuilder;
+import com.azure.security.keyvault.keys.KeyClient;
+import com.azure.security.keyvault.keys.KeyClientBuilder;
+import com.google.common.collect.ImmutableSet;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Map;
+import java.util.Set;
+import org.apache.gravitino.annotation.DeveloperApi;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsClientFactory;
+import org.apache.gravitino.encryption.kms.KmsConfigurationException;
+
+/** Creates key-inspection clients for Microsoft Azure Key Vault. */
+@DeveloperApi
+public final class AzureKeyVaultKmsClientFactory implements KmsClientFactory {
+
+ static final String VAULT_URL = "endpoint.vaultUrl";
+ static final String CREDENTIAL_METHOD = "credential.method";
+ static final String MANAGED_IDENTITY_CLIENT_ID = "credential.managedIdentityClientId";
+
+ private static final Set SUPPORTED_PROPERTIES =
+ ImmutableSet.of(VAULT_URL, CREDENTIAL_METHOD, MANAGED_IDENTITY_CLIENT_ID);
+
+ /** Creates an Azure Key Vault KMS client factory. */
+ public AzureKeyVaultKmsClientFactory() {}
+
+ /**
+ * Returns the Azure Key Vault API.
+ *
+ * @return the Azure Key Vault API
+ */
+ @Override
+ public KmsApi api() {
+ return KmsApi.AZURE_KEY_VAULT;
+ }
+
+ /**
+ * Creates a client using Azure's default credential chain.
+ *
+ * The required {@code endpoint.vaultUrl} property is the vault HTTPS origin. The {@code
+ * credential.method} must be {@code default}. The optional {@code
+ * credential.managedIdentityClientId} selects a user-assigned managed identity.
+ *
+ * @param source logical name of the configured Azure Key Vault source
+ * @param properties Azure Key Vault client properties
+ * @return the configured key-inspection client
+ * @throws IllegalArgumentException if the source or properties are invalid
+ */
+ @Override
+ public KmsClient create(String source, Map properties) {
+ validateSource(source);
+ if (properties == null) {
+ throw new KmsConfigurationException("Azure Key Vault properties cannot be null");
+ }
+
+ for (String property : properties.keySet()) {
+ if (!SUPPORTED_PROPERTIES.contains(property)) {
+ throw new KmsConfigurationException("Unsupported Azure Key Vault property '%s'", property);
+ }
+ }
+
+ URI vaultUri = parseVaultUri(requiredProperty(properties, VAULT_URL));
+ String credentialMethod = requiredProperty(properties, CREDENTIAL_METHOD);
+ if (!"default".equals(credentialMethod)) {
+ throw new KmsConfigurationException(
+ "Unsupported Azure Key Vault credential method: %s", credentialMethod);
+ }
+ DefaultAzureCredentialBuilder credentialBuilder = new DefaultAzureCredentialBuilder();
+ String managedIdentityClientId = properties.get(MANAGED_IDENTITY_CLIENT_ID);
+ if (managedIdentityClientId != null) {
+ if (managedIdentityClientId.trim().isEmpty()) {
+ throw new KmsConfigurationException(
+ "Azure Key Vault property '%s' cannot be blank", MANAGED_IDENTITY_CLIENT_ID);
+ }
+ credentialBuilder.managedIdentityClientId(managedIdentityClientId.trim());
+ }
+
+ KeyClient keyClient =
+ new KeyClientBuilder()
+ .vaultUrl(vaultUri.toString())
+ .credential(credentialBuilder.build())
+ .buildClient();
+ return new AzureKeyVaultKmsClient(source, vaultUri, keyClient);
+ }
+
+ private static void validateSource(String source) {
+ if (source == null || source.trim().isEmpty()) {
+ throw new KmsConfigurationException("Azure Key Vault source cannot be blank");
+ }
+ }
+
+ private static String requiredProperty(Map properties, String property) {
+ String value = properties.get(property);
+ if (value == null || value.trim().isEmpty()) {
+ throw new KmsConfigurationException("Azure Key Vault property '%s' is required", property);
+ }
+ return value.trim();
+ }
+
+ private static URI parseVaultUri(String value) {
+ URI vaultUri;
+ try {
+ vaultUri = new URI(value);
+ } catch (URISyntaxException e) {
+ throw new KmsConfigurationException(e, "Azure Key Vault URL must be a valid URI");
+ }
+ return AzureKeyVaultKeyIdentifier.validateVaultUri(vaultUri);
+ }
+}
diff --git a/bundles/azure/src/main/java/org/apache/gravitino/encryption/kms/azure/AzureKeyVaultKmsKeyProperties.java b/bundles/azure/src/main/java/org/apache/gravitino/encryption/kms/azure/AzureKeyVaultKmsKeyProperties.java
new file mode 100644
index 00000000000..91d909da5a4
--- /dev/null
+++ b/bundles/azure/src/main/java/org/apache/gravitino/encryption/kms/azure/AzureKeyVaultKmsKeyProperties.java
@@ -0,0 +1,69 @@
+/*
+ * 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.azure;
+
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+
+final class AzureKeyVaultKmsKeyProperties implements KmsKeyProperties {
+
+ private final KmsReference reference;
+ private final boolean present;
+ private final boolean enabled;
+ private final boolean supportsWrapping;
+ private final boolean supportsUnwrapping;
+
+ AzureKeyVaultKmsKeyProperties(
+ 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/bundles/azure/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory b/bundles/azure/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory
new file mode 100644
index 00000000000..6619c9e5e03
--- /dev/null
+++ b/bundles/azure/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory
@@ -0,0 +1,20 @@
+#
+# 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.
+#
+
+org.apache.gravitino.encryption.kms.azure.AzureKeyVaultKmsClientFactory
diff --git a/bundles/azure/src/test/java/org/apache/gravitino/encryption/kms/azure/TestAzureKeyVaultKmsClient.java b/bundles/azure/src/test/java/org/apache/gravitino/encryption/kms/azure/TestAzureKeyVaultKmsClient.java
new file mode 100644
index 00000000000..709915802f8
--- /dev/null
+++ b/bundles/azure/src/test/java/org/apache/gravitino/encryption/kms/azure/TestAzureKeyVaultKmsClient.java
@@ -0,0 +1,384 @@
+/*
+ * 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.azure;
+
+import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.http.HttpClient;
+import com.azure.core.http.HttpHeaderName;
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.HttpResponse;
+import com.azure.core.http.policy.FixedDelayOptions;
+import com.azure.core.http.policy.RetryOptions;
+import com.azure.security.keyvault.keys.KeyClient;
+import com.azure.security.keyvault.keys.KeyClientBuilder;
+import java.net.URI;
+import java.nio.ByteBuffer;
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
+import java.time.OffsetDateTime;
+import java.time.ZoneOffset;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Function;
+import javax.annotation.Nullable;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+import org.apache.gravitino.encryption.kms.TestKmsClientContract;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+public class TestAzureKeyVaultKmsClient extends TestKmsClientContract {
+
+ private static final String SOURCE = "azure-primary";
+ private static final String VAULT_URL = "https://example.vault.azure.net";
+ private static final String VERSION = "0123456789abcdef0123456789abcdef";
+ private static final Instant NOW = Instant.parse("2030-01-01T00:00:00Z");
+ private static final Clock CLOCK = Clock.fixed(NOW, ZoneOffset.UTC);
+ private static final TokenCredential CREDENTIAL =
+ request ->
+ Mono.just(
+ new AccessToken(
+ "test-token",
+ OffsetDateTime.ofInstant(NOW.plus(Duration.ofDays(1)), ZoneOffset.UTC)));
+
+ private final RecordingHttpClient httpClient = new RecordingHttpClient(this::defaultResponse);
+ private final AzureKeyVaultKmsClient client = createClient(httpClient, CLOCK);
+
+ @Override
+ protected KmsClient client() {
+ return client;
+ }
+
+ @Override
+ protected KmsReference usableKey() {
+ return reference("usable");
+ }
+
+ @Override
+ protected KmsReference missingKey() {
+ return reference("missing");
+ }
+
+ @Test
+ void testRequestsLatestAndPinnedVersionsThroughSdk() {
+ KmsReference latest = reference("usable");
+ KmsReference pinned = reference("usable", VERSION);
+
+ client.getKeyProperties(latest);
+ Assertions.assertEquals("/keys/usable/", httpClient.lastRequest().getUrl().getPath());
+
+ client.getKeyProperties(pinned);
+ Assertions.assertEquals("/keys/usable/" + VERSION, httpClient.lastRequest().getUrl().getPath());
+ Assertions.assertTrue(httpClient.lastRequest().getUrl().getQuery().contains("api-version="));
+ }
+
+ @Test
+ void testPreservesExactRequestedReference() {
+ KmsReference reference = reference("usable", VERSION);
+
+ Assertions.assertSame(reference, client.getKeyProperties(reference).reference());
+ }
+
+ @Test
+ void testNormalizesEffectiveEnabledState() {
+ Assertions.assertTrue(properties(true, null, null, "\"wrapKey\"").enabled());
+ Assertions.assertFalse(properties(false, null, null, "\"wrapKey\"").enabled());
+ Assertions.assertFalse(properties(null, null, null, "\"wrapKey\"").enabled());
+ Assertions.assertFalse(properties(true, NOW.plusSeconds(1), null, "\"wrapKey\"").enabled());
+ Assertions.assertTrue(properties(true, NOW, null, "\"wrapKey\"").enabled());
+ Assertions.assertFalse(properties(true, null, NOW, "\"wrapKey\"").enabled());
+ Assertions.assertTrue(properties(true, null, NOW.plusSeconds(1), "\"wrapKey\"").enabled());
+ }
+
+ @Test
+ void testMapsOnlyNativeWrapCapabilities() {
+ KmsKeyProperties encryptDecrypt = properties(true, null, null, "\"encrypt\",\"decrypt\"");
+ Assertions.assertFalse(encryptDecrypt.supportsWrapping());
+ Assertions.assertFalse(encryptDecrypt.supportsUnwrapping());
+
+ KmsKeyProperties wrap = properties(true, null, null, "\"wrapKey\"");
+ Assertions.assertTrue(wrap.supportsWrapping());
+ Assertions.assertFalse(wrap.supportsUnwrapping());
+
+ KmsKeyProperties unwrap = properties(true, null, null, "\"unwrapKey\"");
+ Assertions.assertFalse(unwrap.supportsWrapping());
+ Assertions.assertTrue(unwrap.supportsUnwrapping());
+ }
+
+ @Test
+ void testMissingKeyHasAllCapabilitiesDisabled() {
+ KmsKeyProperties properties = client.getKeyProperties(missingKey());
+
+ Assertions.assertFalse(properties.present());
+ Assertions.assertFalse(properties.enabled());
+ Assertions.assertFalse(properties.supportsWrapping());
+ Assertions.assertFalse(properties.supportsUnwrapping());
+ }
+
+ @ParameterizedTest
+ @ValueSource(ints = {400, 401, 403, 408, 429, 500, 503})
+ void testNormalizesProviderFailures(int statusCode) {
+ AzureKeyVaultKmsClient failingClient =
+ createClient(
+ new RecordingHttpClient(request -> response(request, statusCode, errorJson())), CLOCK);
+
+ Assertions.assertThrows(
+ ConnectionFailedException.class, () -> failingClient.getKeyProperties(reference("usable")));
+ }
+
+ @Test
+ void testNormalizesTransportFailure() {
+ AzureKeyVaultKmsClient failingClient =
+ createClient(
+ new RecordingHttpClient(
+ request -> {
+ throw new IllegalStateException("transport unavailable");
+ }),
+ CLOCK);
+
+ Assertions.assertThrows(
+ ConnectionFailedException.class, () -> failingClient.getKeyProperties(reference("usable")));
+ }
+
+ @Test
+ void testNormalizesMalformedSuccessfulResponse() {
+ AzureKeyVaultKmsClient malformedClient =
+ createClient(
+ new RecordingHttpClient(request -> response(request, 200, "{\"key\":null}")), CLOCK);
+
+ Assertions.assertThrows(
+ ConnectionFailedException.class,
+ () -> malformedClient.getKeyProperties(reference("usable")));
+ }
+
+ @ParameterizedTest
+ @ValueSource(
+ strings = {
+ "usable",
+ "http://example.vault.azure.net/keys/usable",
+ "https://other.vault.azure.net/keys/usable",
+ "https://example.vault.azure.net/secrets/usable",
+ "https://example.vault.azure.net/keys",
+ "https://example.vault.azure.net/keys/usable/",
+ "https://example.vault.azure.net/keys/usable/version/extra",
+ "https://example.vault.azure.net/keys/usable?version=latest",
+ "https://example.vault.azure.net/keys/usable#fragment",
+ "https://user@example.vault.azure.net/keys/usable",
+ "https://example.vault.azure.net:443/keys/usable",
+ "https://example.vault.azure.net/keys/encoded%2Fname",
+ "https://example.vault.azure.net/keys/not_valid"
+ })
+ void testRejectsMalformedAndCrossVaultKeyIds(String keyId) {
+ int requestCount = httpClient.requestCount();
+
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> client.getKeyProperties(new KmsReference(KmsApi.AZURE_KEY_VAULT, SOURCE, keyId)));
+ Assertions.assertEquals(requestCount, httpClient.requestCount());
+ }
+
+ @Test
+ void testVaultOriginComparisonIsCaseInsensitive() {
+ KmsReference reference =
+ new KmsReference(
+ KmsApi.AZURE_KEY_VAULT, SOURCE, "HTTPS://EXAMPLE.VAULT.AZURE.NET/keys/usable");
+
+ Assertions.assertTrue(client.getKeyProperties(reference).present());
+ }
+
+ @Test
+ void testRejectsInvalidClientConstruction() {
+ KeyClient sdkClient = createSdkClient(new RecordingHttpClient(this::defaultResponse));
+
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> new AzureKeyVaultKmsClient(" ", URI.create(VAULT_URL), sdkClient, CLOCK));
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ new AzureKeyVaultKmsClient(
+ SOURCE, URI.create("http://example.vault.azure.net"), sdkClient, CLOCK));
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ new AzureKeyVaultKmsClient(SOURCE, URI.create(VAULT_URL + "/keys"), sdkClient, CLOCK));
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> new AzureKeyVaultKmsClient(SOURCE, URI.create(VAULT_URL), null, CLOCK));
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> new AzureKeyVaultKmsClient(SOURCE, URI.create(VAULT_URL), sdkClient, null));
+ }
+
+ private KmsKeyProperties properties(
+ @Nullable Boolean enabled,
+ @Nullable Instant notBefore,
+ @Nullable Instant expiresOn,
+ String operations) {
+ RecordingHttpClient transport =
+ new RecordingHttpClient(
+ request ->
+ response(
+ request, 200, keyJson("temporal", enabled, notBefore, expiresOn, operations)));
+ return createClient(transport, CLOCK).getKeyProperties(reference("temporal"));
+ }
+
+ private HttpResponse defaultResponse(HttpRequest request) {
+ if (request.getUrl().getPath().contains("/missing")) {
+ return response(request, 404, errorJson());
+ }
+ return response(request, 200, keyJson("usable", true, null, null, "\"wrapKey\",\"unwrapKey\""));
+ }
+
+ private static AzureKeyVaultKmsClient createClient(HttpClient httpClient, Clock clock) {
+ return new AzureKeyVaultKmsClient(
+ SOURCE, URI.create(VAULT_URL), createSdkClient(httpClient), clock);
+ }
+
+ private static KeyClient createSdkClient(HttpClient httpClient) {
+ return new KeyClientBuilder()
+ .vaultUrl(VAULT_URL)
+ .credential(CREDENTIAL)
+ .httpClient(httpClient)
+ .retryOptions(new RetryOptions(new FixedDelayOptions(0, Duration.ZERO)))
+ .buildClient();
+ }
+
+ private static KmsReference reference(String name) {
+ return new KmsReference(KmsApi.AZURE_KEY_VAULT, SOURCE, VAULT_URL + "/keys/" + name);
+ }
+
+ private static KmsReference reference(String name, String version) {
+ return new KmsReference(
+ KmsApi.AZURE_KEY_VAULT, SOURCE, VAULT_URL + "/keys/" + name + "/" + version);
+ }
+
+ private static String keyJson(
+ String name,
+ @Nullable Boolean enabled,
+ @Nullable Instant notBefore,
+ @Nullable Instant expiresOn,
+ String operations) {
+ String enabledValue = enabled == null ? "null" : enabled.toString();
+ String notBeforeField = notBefore == null ? "" : ",\"nbf\":" + notBefore.getEpochSecond();
+ String expiresOnField = expiresOn == null ? "" : ",\"exp\":" + expiresOn.getEpochSecond();
+ return String.format(
+ "{\"key\":{\"kid\":\"%s/keys/%s/%s\",\"kty\":\"RSA\","
+ + "\"key_ops\":[%s],\"n\":\"AQAB\",\"e\":\"AQAB\"},"
+ + "\"attributes\":{\"enabled\":%s,\"recoveryLevel\":\"Recoverable\"%s%s}}",
+ VAULT_URL, name, VERSION, operations, enabledValue, notBeforeField, expiresOnField);
+ }
+
+ private static String errorJson() {
+ return "{\"error\":{\"code\":\"KeyNotFound\",\"message\":\"not found\"}}";
+ }
+
+ private static HttpResponse response(HttpRequest request, int statusCode, String body) {
+ return new StaticHttpResponse(request, statusCode, body);
+ }
+
+ private static final class RecordingHttpClient implements HttpClient {
+
+ private final Function responseFactory;
+ private final List requests = new ArrayList<>();
+
+ private RecordingHttpClient(Function responseFactory) {
+ this.responseFactory = responseFactory;
+ }
+
+ @Override
+ public Mono send(HttpRequest request) {
+ requests.add(request);
+ return Mono.fromSupplier(() -> responseFactory.apply(request));
+ }
+
+ private HttpRequest lastRequest() {
+ return requests.get(requests.size() - 1);
+ }
+
+ private int requestCount() {
+ return requests.size();
+ }
+ }
+
+ private static final class StaticHttpResponse extends HttpResponse {
+
+ private final int statusCode;
+ private final byte[] body;
+ private final HttpHeaders headers;
+
+ private StaticHttpResponse(HttpRequest request, int statusCode, String body) {
+ super(request);
+ this.statusCode = statusCode;
+ this.body = body.getBytes(StandardCharsets.UTF_8);
+ this.headers =
+ new HttpHeaders()
+ .set(HttpHeaderName.CONTENT_TYPE, "application/json")
+ .set(HttpHeaderName.CONTENT_LENGTH, Integer.toString(this.body.length));
+ }
+
+ @Override
+ public int getStatusCode() {
+ return statusCode;
+ }
+
+ @SuppressWarnings("deprecation")
+ @Override
+ public String getHeaderValue(String name) {
+ return headers.getValue(name);
+ }
+
+ @Override
+ public HttpHeaders getHeaders() {
+ return headers;
+ }
+
+ @Override
+ public Flux getBody() {
+ return Flux.defer(() -> Flux.just(ByteBuffer.wrap(body)));
+ }
+
+ @Override
+ public Mono getBodyAsByteArray() {
+ return Mono.just(body.clone());
+ }
+
+ @Override
+ public Mono getBodyAsString() {
+ return Mono.just(new String(body, StandardCharsets.UTF_8));
+ }
+
+ @Override
+ public Mono getBodyAsString(Charset charset) {
+ return Mono.just(new String(body, charset));
+ }
+ }
+}
diff --git a/bundles/azure/src/test/java/org/apache/gravitino/encryption/kms/azure/TestAzureKeyVaultKmsClientFactory.java b/bundles/azure/src/test/java/org/apache/gravitino/encryption/kms/azure/TestAzureKeyVaultKmsClientFactory.java
new file mode 100644
index 00000000000..bf9db215861
--- /dev/null
+++ b/bundles/azure/src/test/java/org/apache/gravitino/encryption/kms/azure/TestAzureKeyVaultKmsClientFactory.java
@@ -0,0 +1,166 @@
+/*
+ * 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.azure;
+
+import com.google.common.collect.ImmutableMap;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.ServiceLoader;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsClientFactory;
+import org.apache.gravitino.encryption.kms.TestKmsClientFactoryContract;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestAzureKeyVaultKmsClientFactory extends TestKmsClientFactoryContract {
+
+ private static final String VAULT_URL = "https://example.vault.azure.net";
+
+ private final AzureKeyVaultKmsClientFactory factory = new AzureKeyVaultKmsClientFactory();
+
+ @Override
+ protected KmsClientFactory factory() {
+ return factory;
+ }
+
+ @Override
+ protected KmsApi expectedApi() {
+ return KmsApi.AZURE_KEY_VAULT;
+ }
+
+ @Test
+ void testCreatesClientWithDefaultCredentialChain() {
+ KmsClient client =
+ factory.create(
+ "primary",
+ ImmutableMap.of(
+ AzureKeyVaultKmsClientFactory.VAULT_URL,
+ VAULT_URL,
+ AzureKeyVaultKmsClientFactory.CREDENTIAL_METHOD,
+ "default"));
+
+ Assertions.assertInstanceOf(AzureKeyVaultKmsClient.class, client);
+ }
+
+ @Test
+ void testCreatesClientForUserAssignedManagedIdentity() {
+ KmsClient client =
+ factory.create(
+ "primary",
+ ImmutableMap.of(
+ AzureKeyVaultKmsClientFactory.VAULT_URL,
+ VAULT_URL + "/",
+ AzureKeyVaultKmsClientFactory.CREDENTIAL_METHOD,
+ "default",
+ AzureKeyVaultKmsClientFactory.MANAGED_IDENTITY_CLIENT_ID,
+ "client-id"));
+
+ Assertions.assertInstanceOf(AzureKeyVaultKmsClient.class, client);
+ }
+
+ @Test
+ void testLoadsFactoryThroughServiceLoader() {
+ int matchingFactories = 0;
+ for (KmsClientFactory loadedFactory : ServiceLoader.load(KmsClientFactory.class)) {
+ if (loadedFactory instanceof AzureKeyVaultKmsClientFactory) {
+ matchingFactories++;
+ }
+ }
+
+ Assertions.assertEquals(1, matchingFactories);
+ }
+
+ @Test
+ void testRejectsMissingAndUnsupportedProperties() {
+ Assertions.assertThrows(IllegalArgumentException.class, () -> factory.create("primary", null));
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () -> factory.create("primary", ImmutableMap.of()));
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ factory.create(
+ "primary", ImmutableMap.of(AzureKeyVaultKmsClientFactory.VAULT_URL, " ")));
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ factory.create(
+ "primary",
+ ImmutableMap.of(
+ AzureKeyVaultKmsClientFactory.VAULT_URL,
+ VAULT_URL,
+ "clientSecret",
+ "must-not-be-accepted")));
+
+ Map properties = new HashMap<>();
+ properties.put(AzureKeyVaultKmsClientFactory.VAULT_URL, VAULT_URL);
+ properties.put(AzureKeyVaultKmsClientFactory.CREDENTIAL_METHOD, "default");
+ properties.put(null, "value");
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () -> factory.create("primary", properties));
+ }
+
+ @Test
+ void testRejectsInvalidSourcesAndManagedIdentity() {
+ Map properties =
+ ImmutableMap.of(
+ AzureKeyVaultKmsClientFactory.VAULT_URL,
+ VAULT_URL,
+ AzureKeyVaultKmsClientFactory.CREDENTIAL_METHOD,
+ "default");
+
+ Assertions.assertThrows(IllegalArgumentException.class, () -> factory.create(null, properties));
+ Assertions.assertThrows(IllegalArgumentException.class, () -> factory.create(" ", properties));
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ factory.create(
+ "primary",
+ ImmutableMap.of(
+ AzureKeyVaultKmsClientFactory.VAULT_URL,
+ VAULT_URL,
+ AzureKeyVaultKmsClientFactory.CREDENTIAL_METHOD,
+ "default",
+ AzureKeyVaultKmsClientFactory.MANAGED_IDENTITY_CLIENT_ID,
+ " ")));
+ }
+
+ @ParameterizedTest
+ @ValueSource(
+ strings = {
+ "example.vault.azure.net",
+ "http://example.vault.azure.net",
+ "https://",
+ "https://user@example.vault.azure.net",
+ "https://example.vault.azure.net:443",
+ "https://example.vault.azure.net/keys",
+ "https://example.vault.azure.net?tenant=example",
+ "https://example.vault.azure.net#fragment"
+ })
+ void testRejectsInvalidServiceAddresses(String serviceAddress) {
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ factory.create(
+ "primary",
+ ImmutableMap.of(AzureKeyVaultKmsClientFactory.VAULT_URL, serviceAddress)));
+ }
+}
diff --git a/bundles/gcp/build.gradle.kts b/bundles/gcp/build.gradle.kts
index 80658d352a7..b204a062e57 100644
--- a/bundles/gcp/build.gradle.kts
+++ b/bundles/gcp/build.gradle.kts
@@ -37,17 +37,31 @@ dependencies {
}
implementation(libs.commons.lang3)
+ implementation(libs.google.cloud.kms)
implementation(libs.guava)
compileOnly(libs.hadoop3.client.api)
compileOnly(libs.hadoop3.gcs)
compileOnly(libs.google.auth.http)
+ testImplementation(testFixtures(project(":api")))
testImplementation(libs.junit.jupiter.api)
testImplementation(libs.junit.jupiter.params)
+ testImplementation(libs.mockito.inline)
+ testImplementation(libs.slf4j.api)
testRuntimeOnly(libs.junit.jupiter.engine)
+ testRuntimeOnly(libs.slf4j.simple)
}
tasks.compileJava {
dependsOn(":catalogs:catalog-fileset:runtimeJars")
}
+
+tasks.test {
+ val skipITs = project.hasProperty("skipITs")
+ if (skipITs) {
+ exclude("**/integration/test/**")
+ } else {
+ dependsOn(tasks.jar)
+ }
+}
diff --git a/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsClient.java b/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsClient.java
new file mode 100644
index 00000000000..8dfce085b98
--- /dev/null
+++ b/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsClient.java
@@ -0,0 +1,116 @@
+/*
+ * 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.kms.gcp;
+
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+
+final class GoogleCloudKmsClient implements KmsClient {
+
+ private static final Pattern CRYPTO_KEY_NAME =
+ Pattern.compile("projects/([^/]+)/locations/([^/]+)/keyRings/[^/]+/cryptoKeys/[^/]+");
+
+ private final String source;
+ private final String projectId;
+ private final String location;
+ private final GoogleCloudKmsMetadataService metadataService;
+ private final AtomicBoolean closed = new AtomicBoolean();
+
+ GoogleCloudKmsClient(
+ String source,
+ String projectId,
+ String location,
+ GoogleCloudKmsMetadataService metadataService) {
+ if (source == null || source.trim().isEmpty()) {
+ throw new IllegalArgumentException("Google Cloud KMS source cannot be blank");
+ }
+ if (metadataService == null) {
+ throw new IllegalArgumentException("Google Cloud KMS metadata service cannot be null");
+ }
+
+ this.source = source;
+ this.projectId = projectId;
+ this.location = location;
+ this.metadataService = metadataService;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public KmsKeyProperties getKeyProperties(KmsReference reference) {
+ validateReference(reference);
+
+ Optional metadata = metadataService.getKey(reference.keyId());
+ if (metadata == null) {
+ throw new IllegalStateException("Google Cloud KMS metadata service returned null");
+ }
+ if (!metadata.isPresent()) {
+ return GoogleCloudKmsKeyProperties.missing(reference);
+ }
+
+ GoogleCloudKmsKeyMetadata key = metadata.get();
+ return new GoogleCloudKmsKeyProperties(
+ reference,
+ true,
+ key.primaryVersionEnabled(),
+ key.encryptDecryptPurpose(),
+ key.encryptDecryptPurpose());
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public void close() {
+ if (closed.compareAndSet(false, true)) {
+ metadataService.close();
+ }
+ }
+
+ private void validateReference(KmsReference reference) {
+ if (reference == null) {
+ throw new IllegalArgumentException("Google Cloud KMS reference cannot be null");
+ }
+ if (reference.api() != KmsApi.GOOGLE_CLOUD_KMS) {
+ throw new IllegalArgumentException(
+ String.format("KMS API %s does not match Google Cloud KMS", reference.api()));
+ }
+ if (!source.equals(reference.source())) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Google Cloud KMS source %s does not match configured source %s",
+ reference.source(), source));
+ }
+ String providerKeyId = reference.keyId();
+ Matcher matcher = CRYPTO_KEY_NAME.matcher(providerKeyId);
+ if (!matcher.matches()) {
+ throw new IllegalArgumentException(
+ String.format("Invalid Google Cloud KMS CryptoKey resource name: %s", providerKeyId));
+ }
+ if (!projectId.equals(matcher.group(1)) || !location.equals(matcher.group(2))) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Google Cloud KMS key %s is outside configured project %s and location %s",
+ providerKeyId, projectId, location));
+ }
+ }
+}
diff --git a/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsClientFactory.java b/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsClientFactory.java
new file mode 100644
index 00000000000..92e8ea7307c
--- /dev/null
+++ b/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsClientFactory.java
@@ -0,0 +1,124 @@
+/*
+ * 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.kms.gcp;
+
+import com.google.cloud.kms.v1.KeyManagementServiceClient;
+import java.io.IOException;
+import java.util.Map;
+import org.apache.gravitino.annotation.DeveloperApi;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsClientFactory;
+import org.apache.gravitino.encryption.kms.KmsConfigurationException;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+
+/**
+ * Creates metadata-only Google Cloud KMS clients using Application Default Credentials.
+ *
+ * Google Cloud project, location, key ring, and key are supplied in each full CryptoKey resource
+ * name. Credentials are resolved by the Google Cloud client library and are not accepted as source
+ * properties.
+ */
+@DeveloperApi
+public final class GoogleCloudKmsClientFactory implements KmsClientFactory {
+
+ /** Google Cloud project property. */
+ public static final String PROJECT_ID = "endpoint.projectId";
+
+ /** Google Cloud location property. */
+ public static final String LOCATION = "endpoint.location";
+
+ /** Credential method property. */
+ public static final String CREDENTIAL_METHOD = "credential.method";
+
+ private final MetadataServiceFactory metadataServiceFactory;
+
+ /** Creates a factory that uses the default Google Cloud KMS endpoint and credentials. */
+ public GoogleCloudKmsClientFactory() {
+ this(GoogleCloudKmsClientFactory::createDefaultMetadataService);
+ }
+
+ GoogleCloudKmsClientFactory(MetadataServiceFactory metadataServiceFactory) {
+ if (metadataServiceFactory == null) {
+ throw new IllegalArgumentException(
+ "Google Cloud KMS metadata service factory cannot be null");
+ }
+ this.metadataServiceFactory = metadataServiceFactory;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public KmsApi api() {
+ return KmsApi.GOOGLE_CLOUD_KMS;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public KmsClient create(String source, Map properties) {
+ if (source == null || source.trim().isEmpty()) {
+ throw new KmsConfigurationException("Google Cloud KMS source cannot be blank");
+ }
+ if (properties == null) {
+ throw new KmsConfigurationException("Google Cloud KMS properties cannot be null");
+ }
+ for (String property : properties.keySet()) {
+ if (!PROJECT_ID.equals(property)
+ && !LOCATION.equals(property)
+ && !CREDENTIAL_METHOD.equals(property)) {
+ throw new KmsConfigurationException("Unsupported Google Cloud KMS property: %s", property);
+ }
+ }
+
+ String projectId = requiredProperty(properties, PROJECT_ID);
+ String location = requiredProperty(properties, LOCATION);
+ String credentialMethod = requiredProperty(properties, CREDENTIAL_METHOD);
+ if (!"default".equals(credentialMethod)) {
+ throw new KmsConfigurationException(
+ "Unsupported Google Cloud KMS credential method: %s", credentialMethod);
+ }
+
+ GoogleCloudKmsMetadataService metadataService = metadataServiceFactory.create(source);
+ if (metadataService == null) {
+ throw new IllegalStateException("Google Cloud KMS metadata service factory returned null");
+ }
+ return new GoogleCloudKmsClient(source, projectId, location, metadataService);
+ }
+
+ private static String requiredProperty(Map properties, String name) {
+ String value = properties.get(name);
+ if (value == null || value.trim().isEmpty()) {
+ throw new KmsConfigurationException("Google Cloud KMS property %s cannot be blank", name);
+ }
+ return value.trim();
+ }
+
+ private static GoogleCloudKmsMetadataService createDefaultMetadataService(String source) {
+ try {
+ return new GoogleCloudKmsSdkMetadataService(source, KeyManagementServiceClient.create());
+ } catch (IOException | RuntimeException e) {
+ throw new ConnectionFailedException(
+ e, "Google Cloud KMS client initialization failed for source %s", source);
+ }
+ }
+
+ @FunctionalInterface
+ interface MetadataServiceFactory {
+ GoogleCloudKmsMetadataService create(String source);
+ }
+}
diff --git a/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsKeyMetadata.java b/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsKeyMetadata.java
new file mode 100644
index 00000000000..4aad2f5cba2
--- /dev/null
+++ b/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsKeyMetadata.java
@@ -0,0 +1,38 @@
+/*
+ * 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.kms.gcp;
+
+final class GoogleCloudKmsKeyMetadata {
+
+ private final boolean encryptDecryptPurpose;
+ private final boolean primaryVersionEnabled;
+
+ GoogleCloudKmsKeyMetadata(boolean encryptDecryptPurpose, boolean primaryVersionEnabled) {
+ this.encryptDecryptPurpose = encryptDecryptPurpose;
+ this.primaryVersionEnabled = primaryVersionEnabled;
+ }
+
+ boolean encryptDecryptPurpose() {
+ return encryptDecryptPurpose;
+ }
+
+ boolean primaryVersionEnabled() {
+ return primaryVersionEnabled;
+ }
+}
diff --git a/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsKeyProperties.java b/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsKeyProperties.java
new file mode 100644
index 00000000000..2cd352a0126
--- /dev/null
+++ b/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsKeyProperties.java
@@ -0,0 +1,78 @@
+/*
+ * 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.kms.gcp;
+
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+
+final class GoogleCloudKmsKeyProperties implements KmsKeyProperties {
+
+ private final KmsReference reference;
+ private final boolean present;
+ private final boolean enabled;
+ private final boolean supportsWrapping;
+ private final boolean supportsUnwrapping;
+
+ GoogleCloudKmsKeyProperties(
+ 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;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public KmsReference reference() {
+ return reference;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public boolean present() {
+ return present;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public boolean enabled() {
+ return enabled;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public boolean supportsWrapping() {
+ return supportsWrapping;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public boolean supportsUnwrapping() {
+ return supportsUnwrapping;
+ }
+
+ static GoogleCloudKmsKeyProperties missing(KmsReference reference) {
+ return new GoogleCloudKmsKeyProperties(reference, false, false, false, false);
+ }
+}
diff --git a/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsMetadataService.java b/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsMetadataService.java
new file mode 100644
index 00000000000..23c0f374b15
--- /dev/null
+++ b/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsMetadataService.java
@@ -0,0 +1,29 @@
+/*
+ * 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.kms.gcp;
+
+import java.util.Optional;
+
+interface GoogleCloudKmsMetadataService extends AutoCloseable {
+
+ Optional getKey(String resourceName);
+
+ @Override
+ default void close() {}
+}
diff --git a/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsSdkMetadataService.java b/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsSdkMetadataService.java
new file mode 100644
index 00000000000..49164e46ed4
--- /dev/null
+++ b/bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsSdkMetadataService.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.kms.gcp;
+
+import com.google.api.gax.rpc.ApiException;
+import com.google.api.gax.rpc.NotFoundException;
+import com.google.cloud.kms.v1.CryptoKey;
+import com.google.cloud.kms.v1.CryptoKey.CryptoKeyPurpose;
+import com.google.cloud.kms.v1.CryptoKeyVersion.CryptoKeyVersionState;
+import com.google.cloud.kms.v1.KeyManagementServiceClient;
+import java.util.Optional;
+import org.apache.gravitino.encryption.kms.KmsAuthenticationException;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+
+final class GoogleCloudKmsSdkMetadataService implements GoogleCloudKmsMetadataService {
+
+ private final String source;
+ private final KeyManagementServiceClient client;
+
+ GoogleCloudKmsSdkMetadataService(String source, KeyManagementServiceClient client) {
+ if (client == null) {
+ throw new IllegalArgumentException("Google Cloud KMS SDK client cannot be null");
+ }
+ this.source = source;
+ this.client = client;
+ }
+
+ @Override
+ public Optional getKey(String resourceName) {
+ try {
+ CryptoKey key = client.getCryptoKey(resourceName);
+ if (key == null) {
+ throw new ConnectionFailedException(
+ "Google Cloud KMS returned an empty response for source %s", source);
+ }
+ boolean encryptDecryptPurpose = key.getPurpose() == CryptoKeyPurpose.ENCRYPT_DECRYPT;
+ boolean primaryVersionEnabled =
+ key.hasPrimary() && key.getPrimary().getState() == CryptoKeyVersionState.ENABLED;
+ return Optional.of(
+ new GoogleCloudKmsKeyMetadata(encryptDecryptPurpose, primaryVersionEnabled));
+ } catch (NotFoundException e) {
+ return Optional.empty();
+ } catch (ConnectionFailedException e) {
+ throw e;
+ } catch (ApiException e) {
+ switch (e.getStatusCode().getCode()) {
+ case UNAUTHENTICATED:
+ case PERMISSION_DENIED:
+ throw new KmsAuthenticationException(
+ e, "Google Cloud KMS rejected credentials for source %s", source);
+ default:
+ break;
+ }
+ throw unavailable(e);
+ } catch (RuntimeException e) {
+ throw unavailable(e);
+ }
+ }
+
+ @Override
+ public void close() {
+ client.close();
+ }
+
+ private ConnectionFailedException unavailable(RuntimeException cause) {
+ return new ConnectionFailedException(
+ cause, "Google Cloud KMS could not read key properties for source %s", source);
+ }
+}
diff --git a/bundles/gcp/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory b/bundles/gcp/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory
new file mode 100644
index 00000000000..2a8c4e2052e
--- /dev/null
+++ b/bundles/gcp/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory
@@ -0,0 +1,20 @@
+#
+# 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.
+#
+
+org.apache.gravitino.kms.gcp.GoogleCloudKmsClientFactory
diff --git a/bundles/gcp/src/test/java/org/apache/gravitino/kms/gcp/TestGoogleCloudKmsClient.java b/bundles/gcp/src/test/java/org/apache/gravitino/kms/gcp/TestGoogleCloudKmsClient.java
new file mode 100644
index 00000000000..cb6d5f91494
--- /dev/null
+++ b/bundles/gcp/src/test/java/org/apache/gravitino/kms/gcp/TestGoogleCloudKmsClient.java
@@ -0,0 +1,194 @@
+/*
+ * 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.kms.gcp;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+import org.apache.gravitino.encryption.kms.TestKmsClientContract;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestGoogleCloudKmsClient extends TestKmsClientContract {
+
+ private static final String SOURCE = "analytics";
+ private static final String PROJECT_ID = "data-project";
+ private static final String LOCATION = "us-west1";
+ private static final String USABLE_KEY =
+ "projects/data-project/locations/us-west1/keyRings/tables/cryptoKeys/customer";
+ private static final String MISSING_KEY =
+ "projects/data-project/locations/us-west1/keyRings/tables/cryptoKeys/missing";
+
+ private final FakeMetadataService metadataService = new FakeMetadataService();
+
+ private GoogleCloudKmsClient client;
+
+ @BeforeEach
+ void setUp() {
+ metadataService.keys.put(USABLE_KEY, new GoogleCloudKmsKeyMetadata(true, true));
+ client = new GoogleCloudKmsClient(SOURCE, PROJECT_ID, LOCATION, metadataService);
+ }
+
+ @Override
+ protected KmsClient client() {
+ return client;
+ }
+
+ @Override
+ protected KmsReference usableKey() {
+ return reference(USABLE_KEY);
+ }
+
+ @Override
+ protected KmsReference missingKey() {
+ return reference(MISSING_KEY);
+ }
+
+ @Test
+ void readsExactlyOneCryptoKeyResource() {
+ client.getKeyProperties(usableKey());
+
+ assertEquals(1, metadataService.readCount.get());
+ assertEquals(USABLE_KEY, metadataService.lastResourceName.get());
+ }
+
+ @Test
+ void reportsDisabledPrimaryVersion() {
+ metadataService.keys.put(USABLE_KEY, new GoogleCloudKmsKeyMetadata(true, false));
+
+ KmsKeyProperties properties = client.getKeyProperties(usableKey());
+
+ assertTrue(properties.present());
+ assertFalse(properties.enabled());
+ assertTrue(properties.supportsWrapping());
+ assertTrue(properties.supportsUnwrapping());
+ }
+
+ @Test
+ void reportsNonEncryptionPurpose() {
+ metadataService.keys.put(USABLE_KEY, new GoogleCloudKmsKeyMetadata(false, true));
+
+ KmsKeyProperties properties = client.getKeyProperties(usableKey());
+
+ assertTrue(properties.present());
+ assertTrue(properties.enabled());
+ assertFalse(properties.supportsWrapping());
+ assertFalse(properties.supportsUnwrapping());
+ }
+
+ @ParameterizedTest
+ @ValueSource(
+ strings = {
+ "customer",
+ "projects/p/locations/l/keyRings/r/cryptoKeys",
+ "projects/p/locations/l/keyRings/r/cryptoKeys/k/cryptoKeyVersions/1",
+ "projects//locations/l/keyRings/r/cryptoKeys/k",
+ "projects/p/locations//keyRings/r/cryptoKeys/k",
+ "projects/p/locations/l/keyRings//cryptoKeys/k",
+ "projects/p/locations/l/keyRings/r/cryptoKeys/",
+ "https://cloudkms.googleapis.com/v1/projects/p/locations/l/keyRings/r/cryptoKeys/k"
+ })
+ void rejectsInvalidCryptoKeyResourceNames(String keyId) {
+ assertThrows(IllegalArgumentException.class, () -> client.getKeyProperties(reference(keyId)));
+ assertEquals(0, metadataService.readCount.get());
+ }
+
+ @Test
+ void preservesConnectionFailures() {
+ ConnectionFailedException failure =
+ new ConnectionFailedException("Google Cloud KMS is unavailable");
+ metadataService.failure = failure;
+
+ ConnectionFailedException exception =
+ assertThrows(ConnectionFailedException.class, () -> client.getKeyProperties(usableKey()));
+
+ assertEquals(failure, exception);
+ }
+
+ @Test
+ void rejectsBrokenMetadataServiceContract() {
+ metadataService.returnNull = true;
+
+ assertThrows(IllegalStateException.class, () -> client.getKeyProperties(usableKey()));
+ }
+
+ @Test
+ void closesMetadataServiceOnce() {
+ client.close();
+ client.close();
+
+ assertEquals(1, metadataService.closeCount.get());
+ }
+
+ @Test
+ void rejectsInvalidConstruction() {
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> new GoogleCloudKmsClient(" ", PROJECT_ID, LOCATION, metadataService));
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> new GoogleCloudKmsClient(SOURCE, PROJECT_ID, LOCATION, null));
+ }
+
+ private KmsReference reference(String keyId) {
+ return new KmsReference(KmsApi.GOOGLE_CLOUD_KMS, SOURCE, keyId);
+ }
+
+ private static final class FakeMetadataService implements GoogleCloudKmsMetadataService {
+ private final Map keys = new HashMap<>();
+ private final AtomicInteger readCount = new AtomicInteger();
+ private final AtomicInteger closeCount = new AtomicInteger();
+ private final AtomicReference lastResourceName = new AtomicReference<>();
+
+ private ConnectionFailedException failure;
+ private boolean returnNull;
+
+ @Override
+ public Optional getKey(String resourceName) {
+ readCount.incrementAndGet();
+ lastResourceName.set(resourceName);
+ if (failure != null) {
+ throw failure;
+ }
+ if (returnNull) {
+ return null;
+ }
+ return Optional.ofNullable(keys.get(resourceName));
+ }
+
+ @Override
+ public void close() {
+ closeCount.incrementAndGet();
+ }
+ }
+}
diff --git a/bundles/gcp/src/test/java/org/apache/gravitino/kms/gcp/TestGoogleCloudKmsClientFactory.java b/bundles/gcp/src/test/java/org/apache/gravitino/kms/gcp/TestGoogleCloudKmsClientFactory.java
new file mode 100644
index 00000000000..3441a0b4183
--- /dev/null
+++ b/bundles/gcp/src/test/java/org/apache/gravitino/kms/gcp/TestGoogleCloudKmsClientFactory.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.kms.gcp;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.ServiceLoader;
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsClientFactory;
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+import org.apache.gravitino.encryption.kms.TestKmsClientFactoryContract;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestGoogleCloudKmsClientFactory extends TestKmsClientFactoryContract {
+
+ private static final String SOURCE = "analytics";
+ private static final String KEY_NAME =
+ "projects/data-project/locations/us/keyRings/tables/cryptoKeys/customer";
+
+ private final AtomicReference createdSource = new AtomicReference<>();
+ private final GoogleCloudKmsClientFactory factory =
+ new GoogleCloudKmsClientFactory(
+ source -> {
+ createdSource.set(source);
+ return resourceName -> Optional.of(new GoogleCloudKmsKeyMetadata(true, true));
+ });
+
+ @Override
+ protected KmsClientFactory factory() {
+ return factory;
+ }
+
+ @Override
+ protected KmsApi expectedApi() {
+ return KmsApi.GOOGLE_CLOUD_KMS;
+ }
+
+ @Test
+ void createsWorkingClientWithDefaultCredentials() {
+ KmsClient client = factory.create(SOURCE, validProperties());
+ KmsReference reference = new KmsReference(KmsApi.GOOGLE_CLOUD_KMS, SOURCE, KEY_NAME);
+
+ KmsKeyProperties properties = client.getKeyProperties(reference);
+
+ assertEquals(reference, properties.reference());
+ assertTrue(properties.present());
+ assertEquals(SOURCE, createdSource.get());
+ }
+
+ @Test
+ void rejectsMissingOrUnknownConfiguration() {
+ assertThrows(IllegalArgumentException.class, () -> factory.create(SOURCE, null));
+
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> factory.create(SOURCE, Map.of("credentialsFile", "/tmp/credentials.json")));
+ assertTrue(exception.getMessage().contains("credentialsFile"));
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"", " ", "\t"})
+ void rejectsBlankSource(String source) {
+ assertThrows(IllegalArgumentException.class, () -> factory.create(source, validProperties()));
+ }
+
+ @Test
+ void preservesInitializationFailure() {
+ ConnectionFailedException failure =
+ new ConnectionFailedException("Application Default Credentials are unavailable");
+ GoogleCloudKmsClientFactory failingFactory =
+ new GoogleCloudKmsClientFactory(
+ source -> {
+ throw failure;
+ });
+
+ ConnectionFailedException exception =
+ assertThrows(
+ ConnectionFailedException.class,
+ () -> failingFactory.create(SOURCE, validProperties()));
+
+ assertEquals(failure, exception);
+ }
+
+ @Test
+ void rejectsBrokenMetadataServiceFactoryContract() {
+ GoogleCloudKmsClientFactory brokenFactory = new GoogleCloudKmsClientFactory(source -> null);
+
+ assertThrows(
+ IllegalStateException.class, () -> brokenFactory.create(SOURCE, validProperties()));
+ assertThrows(IllegalArgumentException.class, () -> new GoogleCloudKmsClientFactory(null));
+ }
+
+ @Test
+ void serviceLoaderDiscoversFactory() {
+ KmsClientFactory discovered = null;
+ for (KmsClientFactory candidate : ServiceLoader.load(KmsClientFactory.class)) {
+ if (candidate.api() == KmsApi.GOOGLE_CLOUD_KMS) {
+ discovered = candidate;
+ break;
+ }
+ }
+
+ assertInstanceOf(GoogleCloudKmsClientFactory.class, discovered);
+ }
+
+ private static Map validProperties() {
+ Map properties = new HashMap<>();
+ properties.put(GoogleCloudKmsClientFactory.PROJECT_ID, "data-project");
+ properties.put(GoogleCloudKmsClientFactory.LOCATION, "us");
+ properties.put(GoogleCloudKmsClientFactory.CREDENTIAL_METHOD, "default");
+ return properties;
+ }
+}
diff --git a/bundles/gcp/src/test/java/org/apache/gravitino/kms/gcp/TestGoogleCloudKmsSdkMetadataService.java b/bundles/gcp/src/test/java/org/apache/gravitino/kms/gcp/TestGoogleCloudKmsSdkMetadataService.java
new file mode 100644
index 00000000000..4b5ec82e5da
--- /dev/null
+++ b/bundles/gcp/src/test/java/org/apache/gravitino/kms/gcp/TestGoogleCloudKmsSdkMetadataService.java
@@ -0,0 +1,197 @@
+/*
+ * 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.kms.gcp;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import com.google.api.gax.rpc.ApiException;
+import com.google.api.gax.rpc.ApiExceptionFactory;
+import com.google.api.gax.rpc.StatusCode;
+import com.google.cloud.kms.v1.CryptoKey;
+import com.google.cloud.kms.v1.CryptoKey.CryptoKeyPurpose;
+import com.google.cloud.kms.v1.CryptoKeyVersion;
+import com.google.cloud.kms.v1.CryptoKeyVersion.CryptoKeyVersionState;
+import com.google.cloud.kms.v1.KeyManagementServiceClient;
+import java.util.Optional;
+import org.apache.gravitino.encryption.kms.KmsAuthenticationException;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestGoogleCloudKmsSdkMetadataService {
+
+ private static final String SOURCE = "analytics";
+ private static final String KEY_NAME =
+ "projects/data-project/locations/us/keyRings/tables/cryptoKeys/customer";
+
+ private KeyManagementServiceClient sdkClient;
+ private GoogleCloudKmsSdkMetadataService metadataService;
+
+ @BeforeEach
+ void setUp() {
+ sdkClient = mock(KeyManagementServiceClient.class);
+ metadataService = new GoogleCloudKmsSdkMetadataService(SOURCE, sdkClient);
+ }
+
+ @Test
+ void mapsEncryptDecryptKeyWithEnabledPrimary() {
+ when(sdkClient.getCryptoKey(KEY_NAME))
+ .thenReturn(key(CryptoKeyPurpose.ENCRYPT_DECRYPT, CryptoKeyVersionState.ENABLED));
+
+ GoogleCloudKmsKeyMetadata metadata =
+ metadataService.getKey(KEY_NAME).orElseThrow(IllegalStateException::new);
+
+ assertTrue(metadata.encryptDecryptPurpose());
+ assertTrue(metadata.primaryVersionEnabled());
+ verify(sdkClient).getCryptoKey(KEY_NAME);
+ }
+
+ @Test
+ void mapsKeyWithoutPrimaryVersionAsDisabled() {
+ when(sdkClient.getCryptoKey(KEY_NAME))
+ .thenReturn(CryptoKey.newBuilder().setPurpose(CryptoKeyPurpose.ENCRYPT_DECRYPT).build());
+
+ GoogleCloudKmsKeyMetadata metadata =
+ metadataService.getKey(KEY_NAME).orElseThrow(IllegalStateException::new);
+
+ assertTrue(metadata.encryptDecryptPurpose());
+ assertFalse(metadata.primaryVersionEnabled());
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = CryptoKeyVersionState.class,
+ mode = EnumSource.Mode.EXCLUDE,
+ names = {"ENABLED", "UNRECOGNIZED"})
+ void mapsNonEnabledPrimaryVersionsAsDisabled(CryptoKeyVersionState state) {
+ when(sdkClient.getCryptoKey(KEY_NAME)).thenReturn(key(CryptoKeyPurpose.ENCRYPT_DECRYPT, state));
+
+ assertFalse(
+ metadataService
+ .getKey(KEY_NAME)
+ .orElseThrow(IllegalStateException::new)
+ .primaryVersionEnabled());
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = CryptoKeyPurpose.class,
+ mode = EnumSource.Mode.EXCLUDE,
+ names = {"ENCRYPT_DECRYPT", "UNRECOGNIZED"})
+ void mapsOtherPurposesAsNotSupportingEnvelopeEncryption(CryptoKeyPurpose purpose) {
+ when(sdkClient.getCryptoKey(KEY_NAME)).thenReturn(key(purpose, CryptoKeyVersionState.ENABLED));
+
+ assertFalse(
+ metadataService
+ .getKey(KEY_NAME)
+ .orElseThrow(IllegalStateException::new)
+ .encryptDecryptPurpose());
+ }
+
+ @Test
+ void mapsUnrecognizedPurposeAndPrimaryStateAsUnsupportedAndDisabled() {
+ CryptoKey key =
+ CryptoKey.newBuilder()
+ .setPurposeValue(Integer.MAX_VALUE)
+ .setPrimary(CryptoKeyVersion.newBuilder().setStateValue(Integer.MAX_VALUE).build())
+ .build();
+ when(sdkClient.getCryptoKey(KEY_NAME)).thenReturn(key);
+
+ GoogleCloudKmsKeyMetadata metadata =
+ metadataService.getKey(KEY_NAME).orElseThrow(IllegalStateException::new);
+
+ assertFalse(metadata.encryptDecryptPurpose());
+ assertFalse(metadata.primaryVersionEnabled());
+ }
+
+ @Test
+ void mapsNotFoundToEmptyMetadata() {
+ ApiException failure = apiException(StatusCode.Code.NOT_FOUND);
+ when(sdkClient.getCryptoKey(KEY_NAME)).thenThrow(failure);
+
+ assertEquals(Optional.empty(), metadataService.getKey(KEY_NAME));
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = StatusCode.Code.class,
+ mode = EnumSource.Mode.EXCLUDE,
+ names = {"OK", "NOT_FOUND", "UNAUTHENTICATED", "PERMISSION_DENIED"})
+ void mapsProviderFailuresToConnectionFailures(StatusCode.Code code) {
+ ApiException failure = apiException(code);
+ when(sdkClient.getCryptoKey(KEY_NAME)).thenThrow(failure);
+
+ assertThrows(ConnectionFailedException.class, () -> metadataService.getKey(KEY_NAME));
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = StatusCode.Code.class,
+ names = {"UNAUTHENTICATED", "PERMISSION_DENIED"})
+ void mapsRejectedCredentialsToAuthenticationFailures(StatusCode.Code code) {
+ ApiException failure = apiException(code);
+ when(sdkClient.getCryptoKey(KEY_NAME)).thenThrow(failure);
+
+ assertThrows(KmsAuthenticationException.class, () -> metadataService.getKey(KEY_NAME));
+ }
+
+ @Test
+ void mapsEmptyAndUnexpectedResponsesToConnectionFailures() {
+ when(sdkClient.getCryptoKey(KEY_NAME)).thenReturn(null);
+ assertThrows(ConnectionFailedException.class, () -> metadataService.getKey(KEY_NAME));
+
+ when(sdkClient.getCryptoKey(KEY_NAME)).thenThrow(new IllegalStateException("closed transport"));
+ assertThrows(ConnectionFailedException.class, () -> metadataService.getKey(KEY_NAME));
+ }
+
+ @Test
+ void closesSdkClient() {
+ metadataService.close();
+
+ verify(sdkClient).close();
+ }
+
+ @Test
+ void rejectsNullSdkClient() {
+ assertThrows(
+ IllegalArgumentException.class, () -> new GoogleCloudKmsSdkMetadataService(SOURCE, null));
+ }
+
+ private static CryptoKey key(CryptoKeyPurpose purpose, CryptoKeyVersionState state) {
+ return CryptoKey.newBuilder()
+ .setPurpose(purpose)
+ .setPrimary(CryptoKeyVersion.newBuilder().setState(state).build())
+ .build();
+ }
+
+ private static ApiException apiException(StatusCode.Code code) {
+ StatusCode statusCode = mock(StatusCode.class);
+ when(statusCode.getCode()).thenReturn(code);
+ return ApiExceptionFactory.createException(
+ new RuntimeException(code.name()), statusCode, false);
+ }
+}
diff --git a/bundles/transit/README.md b/bundles/transit/README.md
new file mode 100644
index 00000000000..93298f795af
--- /dev/null
+++ b/bundles/transit/README.md
@@ -0,0 +1,74 @@
+
+
+# Transit KMS provider testing
+
+The Transit bundle contains separate OpenBao and HashiCorp Vault implementations of the common KMS
+key-inspection API. The clients only read key metadata; they never encrypt, decrypt, wrap, or unwrap
+data.
+
+## Unit tests
+
+Unit tests use an in-process HTTP server and need no Docker daemon or external credentials.
+
+```shell
+./gradlew :bundles:transit:check -PskipITs
+```
+
+## Docker integration tests
+
+The integration tests start disposable development-mode servers, enable Transit, and create an
+encryption key and a signing-only key. They verify the full factory-to-HTTP metadata path, including
+a missing key. The containers and all key state are removed after the test.
+
+```shell
+./gradlew :bundles:transit:test \
+ -PskipTests \
+ -PskipDockerTests=false
+```
+
+The default immutable image tags are:
+
+| Provider | Image | Optional override |
+| --- | --- | --- |
+| OpenBao | `openbao/openbao:2.6.0` | `GRAVITINO_OPENBAO_DOCKER_IMAGE` |
+| Vault | `hashicorp/vault:2.0.3` | `GRAVITINO_VAULT_DOCKER_IMAGE` |
+
+Pin updates should be made in the test and this README together, then verified on the CI runner
+architecture. The tests use insecure development-mode root tokens only inside disposable
+containers; production credentials are never accepted or required.
+
+## CI configuration
+
+The CI runner needs:
+
+- a supported Docker daemon reachable by Testcontainers;
+- permission to pull the two pinned images; and
+- `-PskipDockerTests=false` on the dedicated Docker-test job.
+
+No CI secrets are required. Ordinary unit-test jobs use `-PskipITs`, so the integration package is
+excluded even if Docker is available.
+
+OpenBao is the primary open-source Transit reference fixture and is published under the Mozilla
+Public License 2.0. Gravitino does not bundle the OpenBao server or its source; the test pulls its
+official image at runtime.
+
+Vault is an interoperability target. Its image may be pulled as an external test fixture, but no
+Vault binary, image, or source is bundled, redistributed, or shipped with Gravitino. Any CI use of
+the image must comply with the image's current HashiCorp terms.
diff --git a/bundles/transit/build.gradle.kts b/bundles/transit/build.gradle.kts
new file mode 100644
index 00000000000..a4664f5ad06
--- /dev/null
+++ b/bundles/transit/build.gradle.kts
@@ -0,0 +1,43 @@
+/*
+ * 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.
+ */
+
+plugins {
+ `maven-publish`
+ id("java")
+}
+
+dependencies {
+ implementation(project(":api")) {
+ exclude("*")
+ }
+ implementation(libs.httpclient5)
+ implementation(libs.jackson.databind)
+
+ testImplementation(testFixtures(project(":api")))
+ testImplementation(libs.junit.jupiter.api)
+ testImplementation(libs.junit.jupiter.params)
+ testImplementation(libs.testcontainers)
+ testRuntimeOnly(libs.junit.jupiter.engine)
+}
+
+tasks.test {
+ if (project.hasProperty("skipITs")) {
+ exclude("**/integration/test/**")
+ }
+}
diff --git a/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/FileTokenSupplier.java b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/FileTokenSupplier.java
new file mode 100644
index 00000000000..1adb7ef69e4
--- /dev/null
+++ b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/FileTokenSupplier.java
@@ -0,0 +1,80 @@
+/*
+ * 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.kms.transit;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+
+final class FileTokenSupplier {
+
+ private static final int MAX_TOKEN_BYTES = 16 * 1024;
+
+ private final String providerName;
+ private final Path tokenFile;
+ private volatile String cachedToken;
+
+ FileTokenSupplier(String providerName, Path tokenFile) {
+ this.providerName = providerName;
+ this.tokenFile = tokenFile;
+ }
+
+ String token() {
+ String token = cachedToken;
+ if (token != null) {
+ return token;
+ }
+ synchronized (this) {
+ if (cachedToken == null) {
+ cachedToken = readToken();
+ }
+ return cachedToken;
+ }
+ }
+
+ synchronized String reload() {
+ cachedToken = readToken();
+ return cachedToken;
+ }
+
+ private String readToken() {
+ byte[] tokenBytes;
+ try {
+ if (Files.size(tokenFile) > MAX_TOKEN_BYTES) {
+ throw new ConnectionFailedException("%s token file exceeds the allowed size", providerName);
+ }
+ tokenBytes = Files.readAllBytes(tokenFile);
+ } catch (ConnectionFailedException e) {
+ throw e;
+ } catch (IOException | RuntimeException e) {
+ throw new ConnectionFailedException(e, "Failed to read the %s token file", providerName);
+ }
+
+ if (tokenBytes.length > MAX_TOKEN_BYTES) {
+ throw new ConnectionFailedException("%s token file exceeds the allowed size", providerName);
+ }
+ String token = new String(tokenBytes, StandardCharsets.UTF_8).trim();
+ if (token.isEmpty() || token.indexOf('\r') >= 0 || token.indexOf('\n') >= 0) {
+ throw new ConnectionFailedException("%s token file is empty or malformed", providerName);
+ }
+ return token;
+ }
+}
diff --git a/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/OpenBaoTransitKmsClient.java b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/OpenBaoTransitKmsClient.java
new file mode 100644
index 00000000000..06b60090065
--- /dev/null
+++ b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/OpenBaoTransitKmsClient.java
@@ -0,0 +1,104 @@
+/*
+ * 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.kms.transit;
+
+import java.net.URI;
+import java.nio.file.Path;
+import java.util.Optional;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+
+final class OpenBaoTransitKmsClient implements KmsClient {
+
+ private final String source;
+ private final TransitApiClient apiClient;
+
+ OpenBaoTransitKmsClient(String source, URI serviceAddress, String transitMount, Path tokenFile) {
+ this.source = source;
+ this.apiClient =
+ new TransitApiClient("OpenBao Transit", serviceAddress, transitMount, tokenFile);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public KmsKeyProperties getKeyProperties(KmsReference reference) {
+ validateReference(reference);
+
+ Optional response =
+ apiClient.readKey(new TransitReadKeyRequest(reference.keyId()));
+ if (!response.isPresent()) {
+ return missingProperties(reference);
+ }
+
+ TransitKeyData data = response.get().data();
+ if (data == null || data.supportsEncryption() == null || data.supportsDecryption() == null) {
+ throw malformedResponse();
+ }
+ if (Boolean.TRUE.equals(data.softDeleted())) {
+ return missingProperties(reference);
+ }
+
+ return new OpenBaoTransitKmsKeyProperties(
+ reference, true, true, data.supportsEncryption(), data.supportsDecryption());
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public void close() {
+ apiClient.close();
+ }
+
+ private void validateReference(KmsReference reference) {
+ if (reference == null) {
+ throw new IllegalArgumentException("OpenBao Transit key reference cannot be null");
+ }
+ if (reference.api() != KmsApi.OPENBAO_TRANSIT) {
+ throw new IllegalArgumentException(
+ String.format("KMS API %s does not match OpenBao Transit", reference.api()));
+ }
+ if (!source.equals(reference.source())) {
+ throw new IllegalArgumentException(
+ String.format(
+ "OpenBao Transit source %s does not match configured source %s",
+ reference.source(), source));
+ }
+ validateKeyId(reference.keyId());
+ }
+
+ private static void validateKeyId(String keyId) {
+ if (keyId == null || keyId.trim().isEmpty()) {
+ throw new IllegalArgumentException("OpenBao Transit key ID cannot be blank");
+ }
+ if (keyId.contains("/") || keyId.contains("\\") || ".".equals(keyId) || "..".equals(keyId)) {
+ throw new IllegalArgumentException(
+ String.format("Invalid OpenBao Transit key ID: %s", keyId));
+ }
+ }
+
+ private static OpenBaoTransitKmsKeyProperties missingProperties(KmsReference reference) {
+ return new OpenBaoTransitKmsKeyProperties(reference, false, false, false, false);
+ }
+
+ private static ConnectionFailedException malformedResponse() {
+ return new ConnectionFailedException("OpenBao Transit returned a malformed response");
+ }
+}
diff --git a/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/OpenBaoTransitKmsClientFactory.java b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/OpenBaoTransitKmsClientFactory.java
new file mode 100644
index 00000000000..171b0c62082
--- /dev/null
+++ b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/OpenBaoTransitKmsClientFactory.java
@@ -0,0 +1,157 @@
+/*
+ * 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.kms.transit;
+
+import java.net.URI;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import org.apache.gravitino.annotation.DeveloperApi;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsClientFactory;
+import org.apache.gravitino.encryption.kms.KmsConfigurationException;
+
+/** Creates clients for the OpenBao Transit secrets engine. */
+@DeveloperApi
+public final class OpenBaoTransitKmsClientFactory implements KmsClientFactory {
+
+ /** Property containing the OpenBao HTTP(S) server base address. */
+ public static final String SERVICE_ADDRESS = "endpoint.address";
+
+ /** Property containing the OpenBao Transit mount path. */
+ public static final String TRANSIT_MOUNT = "endpoint.transitMount";
+
+ /** Property selecting token-file authentication. */
+ public static final String CREDENTIAL_METHOD = "credential.method";
+
+ /** Property containing an absolute path to an OpenBao token file. */
+ public static final String TOKEN_FILE = "credential.path";
+
+ /** Default OpenBao Transit mount path. */
+ public static final String DEFAULT_TRANSIT_MOUNT = "transit";
+
+ private static final Set SUPPORTED_PROPERTIES =
+ Collections.unmodifiableSet(
+ new HashSet<>(
+ Arrays.asList(SERVICE_ADDRESS, TRANSIT_MOUNT, CREDENTIAL_METHOD, TOKEN_FILE)));
+
+ /** {@inheritDoc} */
+ @Override
+ public KmsApi api() {
+ return KmsApi.OPENBAO_TRANSIT;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public KmsClient create(String source, Map properties) {
+ validateSource(source);
+ if (properties == null) {
+ throw new KmsConfigurationException("OpenBao Transit properties cannot be null");
+ }
+ for (String property : properties.keySet()) {
+ if (!SUPPORTED_PROPERTIES.contains(property)) {
+ throw new KmsConfigurationException("Unsupported OpenBao Transit property: %s", property);
+ }
+ }
+
+ URI serviceAddress = parseServiceAddress(requireProperty(properties, SERVICE_ADDRESS));
+ String credentialMethod = requireProperty(properties, CREDENTIAL_METHOD);
+ if (!"token_file".equals(credentialMethod)) {
+ throw new KmsConfigurationException(
+ "Unsupported OpenBao Transit credential method: %s", credentialMethod);
+ }
+ String transitMount = properties.get(TRANSIT_MOUNT);
+ if (transitMount == null) {
+ transitMount = DEFAULT_TRANSIT_MOUNT;
+ } else {
+ transitMount = transitMount.trim();
+ }
+ validateTransitMount(transitMount);
+ Path tokenFile = parseTokenFile(requireProperty(properties, TOKEN_FILE));
+
+ return new OpenBaoTransitKmsClient(source, serviceAddress, transitMount, tokenFile);
+ }
+
+ private static String requireProperty(Map properties, String name) {
+ String value = properties.get(name);
+ if (value == null || value.trim().isEmpty()) {
+ throw new KmsConfigurationException("Missing required OpenBao Transit property: %s", name);
+ }
+ return value.trim();
+ }
+
+ private static URI parseServiceAddress(String value) {
+ URI uri;
+ try {
+ uri = URI.create(value);
+ } catch (IllegalArgumentException e) {
+ throw new KmsConfigurationException(e, "Invalid OpenBao Transit endpoint address");
+ }
+
+ String path = uri.getPath();
+ if (!("http".equalsIgnoreCase(uri.getScheme()) || "https".equalsIgnoreCase(uri.getScheme()))
+ || uri.getHost() == null
+ || uri.getUserInfo() != null
+ || uri.getQuery() != null
+ || uri.getFragment() != null
+ || !(path == null || path.isEmpty() || "/".equals(path))) {
+ throw new KmsConfigurationException(
+ "OpenBao Transit endpoint address must be an HTTP(S) server base address");
+ }
+
+ String address = uri.toString();
+ return URI.create(address.endsWith("/") ? address.substring(0, address.length() - 1) : address);
+ }
+
+ private static Path parseTokenFile(String value) {
+ Path path;
+ try {
+ path = Paths.get(value);
+ } catch (RuntimeException e) {
+ throw new KmsConfigurationException(e, "Invalid OpenBao Transit credential path");
+ }
+ if (!path.isAbsolute()) {
+ throw new KmsConfigurationException(
+ "OpenBao Transit credential path must be an absolute path");
+ }
+ return path;
+ }
+
+ private static void validateSource(String source) {
+ if (source == null || source.trim().isEmpty()) {
+ throw new KmsConfigurationException("OpenBao Transit source cannot be blank");
+ }
+ }
+
+ private static void validateTransitMount(String mount) {
+ if (mount.isEmpty() || mount.startsWith("/") || mount.endsWith("/")) {
+ throw new KmsConfigurationException("Invalid OpenBao Transit mount path");
+ }
+ for (String segment : mount.split("/", -1)) {
+ if (segment.isEmpty() || ".".equals(segment) || "..".equals(segment)) {
+ throw new KmsConfigurationException("Invalid OpenBao Transit mount path");
+ }
+ }
+ }
+}
diff --git a/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/OpenBaoTransitKmsKeyProperties.java b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/OpenBaoTransitKmsKeyProperties.java
new file mode 100644
index 00000000000..94701bf6105
--- /dev/null
+++ b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/OpenBaoTransitKmsKeyProperties.java
@@ -0,0 +1,74 @@
+/*
+ * 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.kms.transit;
+
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+
+final class OpenBaoTransitKmsKeyProperties implements KmsKeyProperties {
+
+ private final KmsReference reference;
+ private final boolean present;
+ private final boolean enabled;
+ private final boolean supportsWrapping;
+ private final boolean supportsUnwrapping;
+
+ OpenBaoTransitKmsKeyProperties(
+ 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;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public KmsReference reference() {
+ return reference;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public boolean present() {
+ return present;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public boolean enabled() {
+ return enabled;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public boolean supportsWrapping() {
+ return supportsWrapping;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public boolean supportsUnwrapping() {
+ return supportsUnwrapping;
+ }
+}
diff --git a/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitApiClient.java b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitApiClient.java
new file mode 100644
index 00000000000..342b5ed16af
--- /dev/null
+++ b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitApiClient.java
@@ -0,0 +1,216 @@
+/*
+ * 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.kms.transit;
+
+import com.fasterxml.jackson.databind.MapperFeature;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.json.JsonMapper;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.net.URI;
+import java.net.URLEncoder;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Path;
+import java.util.Arrays;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import org.apache.gravitino.encryption.kms.KmsAuthenticationException;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+import org.apache.hc.client5.http.classic.methods.HttpGet;
+import org.apache.hc.client5.http.config.ConnectionConfig;
+import org.apache.hc.client5.http.config.RequestConfig;
+import org.apache.hc.client5.http.impl.classic.CloseableHttpClient;
+import org.apache.hc.client5.http.impl.classic.HttpClients;
+import org.apache.hc.client5.http.impl.io.PoolingHttpClientConnectionManagerBuilder;
+import org.apache.hc.core5.http.HttpEntity;
+import org.apache.hc.core5.util.Timeout;
+
+final class TransitApiClient {
+
+ private static final String TOKEN_HEADER = "X-Vault-Token";
+ private static final int MAX_RESPONSE_BYTES = 1024 * 1024;
+ private static final int CONNECT_TIMEOUT_MILLIS = 5_000;
+ private static final int REQUEST_TIMEOUT_MILLIS = 10_000;
+ private static final ObjectMapper OBJECT_MAPPER = createObjectMapper();
+
+ private final String providerName;
+ private final URI serviceAddress;
+ private final String transitMount;
+ private final FileTokenSupplier tokenSupplier;
+ private final CloseableHttpClient httpClient;
+
+ TransitApiClient(String providerName, URI serviceAddress, String transitMount, Path tokenFile) {
+ this(
+ providerName,
+ serviceAddress,
+ transitMount,
+ new FileTokenSupplier(providerName, tokenFile),
+ createHttpClient());
+ }
+
+ TransitApiClient(
+ String providerName,
+ URI serviceAddress,
+ String transitMount,
+ FileTokenSupplier tokenSupplier,
+ CloseableHttpClient httpClient) {
+ this.providerName = providerName;
+ this.serviceAddress = serviceAddress;
+ this.transitMount = transitMount;
+ this.tokenSupplier = tokenSupplier;
+ this.httpClient = httpClient;
+ }
+
+ Optional readKey(TransitReadKeyRequest request) {
+ HttpResult result = execute(request, tokenSupplier.token());
+ if (isAuthenticationFailure(result.statusCode)) {
+ result = execute(request, tokenSupplier.reload());
+ }
+
+ if (result.statusCode == 404) {
+ return Optional.empty();
+ }
+ if (isAuthenticationFailure(result.statusCode)) {
+ throw new KmsAuthenticationException(
+ "%s rejected its configured credentials (HTTP %s)", providerName, result.statusCode);
+ }
+ if (result.statusCode < 200 || result.statusCode >= 300) {
+ throw new ConnectionFailedException(
+ "%s could not inspect key %s (HTTP %s)",
+ providerName, request.keyId(), result.statusCode);
+ }
+
+ try {
+ return Optional.of(OBJECT_MAPPER.readValue(result.body, TransitReadKeyResponse.class));
+ } catch (IOException | RuntimeException e) {
+ throw new ConnectionFailedException(e, "%s returned a malformed response", providerName);
+ }
+ }
+
+ void close() {
+ try {
+ httpClient.close();
+ } catch (IOException e) {
+ throw new ConnectionFailedException(e, "Failed to close the %s HTTP client", providerName);
+ }
+ }
+
+ private HttpResult execute(TransitReadKeyRequest request, String token) {
+ HttpGet httpRequest = new HttpGet(keyMetadataUri(request.keyId()));
+ httpRequest.setHeader("Accept", "application/json");
+ httpRequest.setHeader(TOKEN_HEADER, token);
+ try {
+ return httpClient.execute(
+ httpRequest,
+ response -> new HttpResult(response.getCode(), readResponse(response.getEntity())));
+ } catch (ConnectionFailedException e) {
+ throw e;
+ } catch (IOException | RuntimeException e) {
+ throw new ConnectionFailedException(
+ e, "%s is unavailable while inspecting key %s", providerName, request.keyId());
+ }
+ }
+
+ private URI keyMetadataUri(String keyId) {
+ String encodedMount =
+ Arrays.stream(transitMount.split("/"))
+ .map(TransitApiClient::encodePathSegment)
+ .collect(Collectors.joining("/"));
+ return URI.create(
+ String.format("%s/v1/%s/keys/%s", serviceAddress, encodedMount, encodePathSegment(keyId)));
+ }
+
+ private static byte[] readResponse(HttpEntity entity) throws IOException {
+ if (entity == null) {
+ return new byte[0];
+ }
+ if (entity.getContentLength() > MAX_RESPONSE_BYTES) {
+ throw new ConnectionFailedException("Transit response exceeds the allowed size");
+ }
+
+ try (InputStream responseBody = entity.getContent();
+ ByteArrayOutputStream output = new ByteArrayOutputStream()) {
+ byte[] buffer = new byte[8192];
+ int remaining = MAX_RESPONSE_BYTES + 1;
+ while (remaining > 0) {
+ int read = responseBody.read(buffer, 0, Math.min(buffer.length, remaining));
+ if (read < 0) {
+ break;
+ }
+ output.write(buffer, 0, read);
+ remaining -= read;
+ }
+ if (output.size() > MAX_RESPONSE_BYTES) {
+ throw new ConnectionFailedException("Transit response exceeds the allowed size");
+ }
+ return output.toByteArray();
+ }
+ }
+
+ private static boolean isAuthenticationFailure(int statusCode) {
+ return statusCode == 401 || statusCode == 403;
+ }
+
+ private static String encodePathSegment(String value) {
+ try {
+ return URLEncoder.encode(value, StandardCharsets.UTF_8.name()).replace("+", "%20");
+ } catch (IOException e) {
+ throw new ConnectionFailedException(e, "Failed to encode a Transit key reference");
+ }
+ }
+
+ private static ObjectMapper createObjectMapper() {
+ return JsonMapper.builder().disable(MapperFeature.ALLOW_COERCION_OF_SCALARS).build();
+ }
+
+ private static CloseableHttpClient createHttpClient() {
+ Timeout connectTimeout = Timeout.ofMilliseconds(CONNECT_TIMEOUT_MILLIS);
+ Timeout requestTimeout = Timeout.ofMilliseconds(REQUEST_TIMEOUT_MILLIS);
+ ConnectionConfig connectionConfig =
+ ConnectionConfig.custom()
+ .setConnectTimeout(connectTimeout)
+ .setSocketTimeout(requestTimeout)
+ .build();
+ RequestConfig requestConfig =
+ RequestConfig.custom()
+ .setConnectionRequestTimeout(requestTimeout)
+ .setResponseTimeout(requestTimeout)
+ .build();
+ return HttpClients.custom()
+ .setConnectionManager(
+ PoolingHttpClientConnectionManagerBuilder.create()
+ .setDefaultConnectionConfig(connectionConfig)
+ .build())
+ .setDefaultRequestConfig(requestConfig)
+ .disableAutomaticRetries()
+ .disableRedirectHandling()
+ .build();
+ }
+
+ private static final class HttpResult {
+ private final int statusCode;
+ private final byte[] body;
+
+ private HttpResult(int statusCode, byte[] body) {
+ this.statusCode = statusCode;
+ this.body = body;
+ }
+ }
+}
diff --git a/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitKeyData.java b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitKeyData.java
new file mode 100644
index 00000000000..334222adf4a
--- /dev/null
+++ b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitKeyData.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.kms.transit;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+@JsonIgnoreProperties(ignoreUnknown = true)
+final class TransitKeyData {
+
+ private final Boolean supportsEncryption;
+ private final Boolean supportsDecryption;
+ private final Boolean softDeleted;
+
+ @JsonCreator
+ TransitKeyData(
+ @JsonProperty("supports_encryption") Boolean supportsEncryption,
+ @JsonProperty("supports_decryption") Boolean supportsDecryption,
+ @JsonProperty("soft_deleted") Boolean softDeleted) {
+ this.supportsEncryption = supportsEncryption;
+ this.supportsDecryption = supportsDecryption;
+ this.softDeleted = softDeleted;
+ }
+
+ Boolean supportsEncryption() {
+ return supportsEncryption;
+ }
+
+ Boolean supportsDecryption() {
+ return supportsDecryption;
+ }
+
+ Boolean softDeleted() {
+ return softDeleted;
+ }
+}
diff --git a/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitReadKeyRequest.java b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitReadKeyRequest.java
new file mode 100644
index 00000000000..3b909a1c2ad
--- /dev/null
+++ b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitReadKeyRequest.java
@@ -0,0 +1,35 @@
+/*
+ * 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.kms.transit;
+
+final class TransitReadKeyRequest {
+
+ private final String keyId;
+
+ TransitReadKeyRequest(String keyId) {
+ if (keyId == null || keyId.isEmpty()) {
+ throw new IllegalArgumentException("Transit key ID cannot be empty");
+ }
+ this.keyId = keyId;
+ }
+
+ String keyId() {
+ return keyId;
+ }
+}
diff --git a/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitReadKeyResponse.java b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitReadKeyResponse.java
new file mode 100644
index 00000000000..eed7bd7a337
--- /dev/null
+++ b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitReadKeyResponse.java
@@ -0,0 +1,38 @@
+/*
+ * 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.kms.transit;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+@JsonIgnoreProperties(ignoreUnknown = true)
+final class TransitReadKeyResponse {
+
+ private final TransitKeyData data;
+
+ @JsonCreator
+ TransitReadKeyResponse(@JsonProperty("data") TransitKeyData data) {
+ this.data = data;
+ }
+
+ TransitKeyData data() {
+ return data;
+ }
+}
diff --git a/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/VaultTransitKmsClient.java b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/VaultTransitKmsClient.java
new file mode 100644
index 00000000000..8f504f7cb69
--- /dev/null
+++ b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/VaultTransitKmsClient.java
@@ -0,0 +1,99 @@
+/*
+ * 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.kms.transit;
+
+import java.net.URI;
+import java.nio.file.Path;
+import java.util.Optional;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+
+final class VaultTransitKmsClient implements KmsClient {
+
+ private final String source;
+ private final TransitApiClient apiClient;
+
+ VaultTransitKmsClient(String source, URI serviceAddress, String transitMount, Path tokenFile) {
+ this.source = source;
+ this.apiClient = new TransitApiClient("Vault Transit", serviceAddress, transitMount, tokenFile);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public KmsKeyProperties getKeyProperties(KmsReference reference) {
+ validateReference(reference);
+
+ Optional response =
+ apiClient.readKey(new TransitReadKeyRequest(reference.keyId()));
+ if (!response.isPresent()) {
+ return missingProperties(reference);
+ }
+
+ TransitKeyData data = response.get().data();
+ if (data == null || data.supportsEncryption() == null || data.supportsDecryption() == null) {
+ throw malformedResponse();
+ }
+
+ return new VaultTransitKmsKeyProperties(
+ reference, true, true, data.supportsEncryption(), data.supportsDecryption());
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public void close() {
+ apiClient.close();
+ }
+
+ private void validateReference(KmsReference reference) {
+ if (reference == null) {
+ throw new IllegalArgumentException("Vault Transit key reference cannot be null");
+ }
+ if (reference.api() != KmsApi.VAULT_TRANSIT) {
+ throw new IllegalArgumentException(
+ String.format("KMS API %s does not match Vault Transit", reference.api()));
+ }
+ if (!source.equals(reference.source())) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Vault Transit source %s does not match configured source %s",
+ reference.source(), source));
+ }
+ validateKeyId(reference.keyId());
+ }
+
+ private static void validateKeyId(String keyId) {
+ if (keyId == null || keyId.trim().isEmpty()) {
+ throw new IllegalArgumentException("Vault Transit key ID cannot be blank");
+ }
+ if (keyId.contains("/") || keyId.contains("\\") || ".".equals(keyId) || "..".equals(keyId)) {
+ throw new IllegalArgumentException(String.format("Invalid Vault Transit key ID: %s", keyId));
+ }
+ }
+
+ private static VaultTransitKmsKeyProperties missingProperties(KmsReference reference) {
+ return new VaultTransitKmsKeyProperties(reference, false, false, false, false);
+ }
+
+ private static ConnectionFailedException malformedResponse() {
+ return new ConnectionFailedException("Vault Transit returned a malformed response");
+ }
+}
diff --git a/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/VaultTransitKmsClientFactory.java b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/VaultTransitKmsClientFactory.java
new file mode 100644
index 00000000000..d65a9c7d4df
--- /dev/null
+++ b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/VaultTransitKmsClientFactory.java
@@ -0,0 +1,156 @@
+/*
+ * 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.kms.transit;
+
+import java.net.URI;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import org.apache.gravitino.annotation.DeveloperApi;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsClientFactory;
+import org.apache.gravitino.encryption.kms.KmsConfigurationException;
+
+/** Creates clients for the HashiCorp Vault Transit secrets engine. */
+@DeveloperApi
+public final class VaultTransitKmsClientFactory implements KmsClientFactory {
+
+ /** Property containing the Vault HTTP(S) server base address. */
+ public static final String SERVICE_ADDRESS = "endpoint.address";
+
+ /** Property containing the Vault Transit mount path. */
+ public static final String TRANSIT_MOUNT = "endpoint.transitMount";
+
+ /** Property selecting token-file authentication. */
+ public static final String CREDENTIAL_METHOD = "credential.method";
+
+ /** Property containing an absolute path to a Vault token file. */
+ public static final String TOKEN_FILE = "credential.path";
+
+ /** Default Vault Transit mount path. */
+ public static final String DEFAULT_TRANSIT_MOUNT = "transit";
+
+ private static final Set SUPPORTED_PROPERTIES =
+ Collections.unmodifiableSet(
+ new HashSet<>(
+ Arrays.asList(SERVICE_ADDRESS, TRANSIT_MOUNT, CREDENTIAL_METHOD, TOKEN_FILE)));
+
+ /** {@inheritDoc} */
+ @Override
+ public KmsApi api() {
+ return KmsApi.VAULT_TRANSIT;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public KmsClient create(String source, Map properties) {
+ validateSource(source);
+ if (properties == null) {
+ throw new KmsConfigurationException("Vault Transit properties cannot be null");
+ }
+ for (String property : properties.keySet()) {
+ if (!SUPPORTED_PROPERTIES.contains(property)) {
+ throw new KmsConfigurationException("Unsupported Vault Transit property: %s", property);
+ }
+ }
+
+ URI serviceAddress = parseServiceAddress(requireProperty(properties, SERVICE_ADDRESS));
+ String credentialMethod = requireProperty(properties, CREDENTIAL_METHOD);
+ if (!"token_file".equals(credentialMethod)) {
+ throw new KmsConfigurationException(
+ "Unsupported Vault Transit credential method: %s", credentialMethod);
+ }
+ String transitMount = properties.get(TRANSIT_MOUNT);
+ if (transitMount == null) {
+ transitMount = DEFAULT_TRANSIT_MOUNT;
+ } else {
+ transitMount = transitMount.trim();
+ }
+ validateTransitMount(transitMount);
+ Path tokenFile = parseTokenFile(requireProperty(properties, TOKEN_FILE));
+
+ return new VaultTransitKmsClient(source, serviceAddress, transitMount, tokenFile);
+ }
+
+ private static String requireProperty(Map properties, String name) {
+ String value = properties.get(name);
+ if (value == null || value.trim().isEmpty()) {
+ throw new KmsConfigurationException("Missing required Vault Transit property: %s", name);
+ }
+ return value.trim();
+ }
+
+ private static URI parseServiceAddress(String value) {
+ URI uri;
+ try {
+ uri = URI.create(value);
+ } catch (IllegalArgumentException e) {
+ throw new KmsConfigurationException(e, "Invalid Vault Transit endpoint address");
+ }
+
+ String path = uri.getPath();
+ if (!("http".equalsIgnoreCase(uri.getScheme()) || "https".equalsIgnoreCase(uri.getScheme()))
+ || uri.getHost() == null
+ || uri.getUserInfo() != null
+ || uri.getQuery() != null
+ || uri.getFragment() != null
+ || !(path == null || path.isEmpty() || "/".equals(path))) {
+ throw new KmsConfigurationException(
+ "Vault Transit endpoint address must be an HTTP(S) server base address");
+ }
+
+ String address = uri.toString();
+ return URI.create(address.endsWith("/") ? address.substring(0, address.length() - 1) : address);
+ }
+
+ private static Path parseTokenFile(String value) {
+ Path path;
+ try {
+ path = Paths.get(value);
+ } catch (RuntimeException e) {
+ throw new KmsConfigurationException(e, "Invalid Vault Transit credential path");
+ }
+ if (!path.isAbsolute()) {
+ throw new KmsConfigurationException("Vault Transit credential path must be an absolute path");
+ }
+ return path;
+ }
+
+ private static void validateSource(String source) {
+ if (source == null || source.trim().isEmpty()) {
+ throw new KmsConfigurationException("Vault Transit source cannot be blank");
+ }
+ }
+
+ private static void validateTransitMount(String mount) {
+ if (mount.isEmpty() || mount.startsWith("/") || mount.endsWith("/")) {
+ throw new KmsConfigurationException("Invalid Vault Transit mount path");
+ }
+ for (String segment : mount.split("/", -1)) {
+ if (segment.isEmpty() || ".".equals(segment) || "..".equals(segment)) {
+ throw new KmsConfigurationException("Invalid Vault Transit mount path");
+ }
+ }
+ }
+}
diff --git a/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/VaultTransitKmsKeyProperties.java b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/VaultTransitKmsKeyProperties.java
new file mode 100644
index 00000000000..d8ac48cbc62
--- /dev/null
+++ b/bundles/transit/src/main/java/org/apache/gravitino/kms/transit/VaultTransitKmsKeyProperties.java
@@ -0,0 +1,74 @@
+/*
+ * 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.kms.transit;
+
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+
+final class VaultTransitKmsKeyProperties implements KmsKeyProperties {
+
+ private final KmsReference reference;
+ private final boolean present;
+ private final boolean enabled;
+ private final boolean supportsWrapping;
+ private final boolean supportsUnwrapping;
+
+ VaultTransitKmsKeyProperties(
+ 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;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public KmsReference reference() {
+ return reference;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public boolean present() {
+ return present;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public boolean enabled() {
+ return enabled;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public boolean supportsWrapping() {
+ return supportsWrapping;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public boolean supportsUnwrapping() {
+ return supportsUnwrapping;
+ }
+}
diff --git a/bundles/transit/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory b/bundles/transit/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory
new file mode 100644
index 00000000000..c28f35c1025
--- /dev/null
+++ b/bundles/transit/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory
@@ -0,0 +1,21 @@
+#
+# 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.
+#
+
+org.apache.gravitino.kms.transit.OpenBaoTransitKmsClientFactory
+org.apache.gravitino.kms.transit.VaultTransitKmsClientFactory
diff --git a/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestOpenBaoTransitKmsClient.java b/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestOpenBaoTransitKmsClient.java
new file mode 100644
index 00000000000..17cbca85519
--- /dev/null
+++ b/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestOpenBaoTransitKmsClient.java
@@ -0,0 +1,286 @@
+/*
+ * 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.kms.transit;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import com.sun.net.httpserver.HttpExchange;
+import com.sun.net.httpserver.HttpServer;
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsAuthenticationException;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+import org.apache.gravitino.encryption.kms.TestKmsClientContract;
+import org.apache.gravitino.exceptions.ConnectionFailedException;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestOpenBaoTransitKmsClient extends TestKmsClientContract {
+
+ private static final String SOURCE = "primary";
+ private static final String USABLE_KEY = "customer key";
+ private static final String MISSING_KEY = "missing-key";
+ private static final String USABLE_RESPONSE =
+ "{\"data\":{\"type\":\"aes256-gcm96\",\"supports_encryption\":true,"
+ + "\"supports_decryption\":true,\"soft_deleted\":false}}";
+
+ @TempDir private Path tempDir;
+
+ private final AtomicInteger responseStatus = new AtomicInteger(200);
+ private final AtomicReference responseBody = new AtomicReference<>(USABLE_RESPONSE);
+ private final AtomicReference acceptedToken = new AtomicReference<>();
+ private final List requests = new CopyOnWriteArrayList<>();
+
+ private HttpServer server;
+ private Path tokenFile;
+ private OpenBaoTransitKmsClient client;
+
+ @BeforeEach
+ void startServer() throws IOException {
+ tokenFile = tempDir.resolve("openbao-read-only-token");
+ Files.write(tokenFile, "read-only-token\n".getBytes(StandardCharsets.UTF_8));
+
+ server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+ server.createContext("/", this::respond);
+ server.start();
+ client = createClient("custom/transit", tokenFile);
+ }
+
+ @AfterEach
+ void stopServer() {
+ if (server != null) {
+ server.stop(0);
+ }
+ }
+
+ @Override
+ protected KmsClient client() {
+ return client;
+ }
+
+ @Override
+ protected KmsReference usableKey() {
+ return reference(USABLE_KEY);
+ }
+
+ @Override
+ protected KmsReference missingKey() {
+ return reference(MISSING_KEY);
+ }
+
+ @Test
+ void readsKeyPropertiesUsingMetadataOnlyGet() {
+ client.getKeyProperties(usableKey());
+
+ assertEquals(1, requests.size());
+ RecordedRequest request = requests.get(0);
+ assertEquals("GET", request.method);
+ assertEquals("/v1/custom/transit/keys/customer%20key", request.rawPath);
+ assertEquals("read-only-token", request.token);
+ }
+
+ @Test
+ void reportsSoftDeletedKeyAsMissing() {
+ responseBody.set(
+ "{\"data\":{\"supports_encryption\":true,\"supports_decryption\":true,"
+ + "\"soft_deleted\":true}}");
+
+ KmsKeyProperties properties = client.getKeyProperties(usableKey());
+
+ assertFalse(properties.present());
+ assertFalse(properties.enabled());
+ assertFalse(properties.supportsWrapping());
+ assertFalse(properties.supportsUnwrapping());
+ }
+
+ @Test
+ void treatsMissingSoftDeletedFieldAsEnabledForOlderServers() {
+ responseBody.set("{\"data\":{\"supports_encryption\":true,\"supports_decryption\":true}}");
+
+ assertTrue(client.getKeyProperties(usableKey()).enabled());
+ }
+
+ @Test
+ void mapsTransitOperationSupport() {
+ responseBody.set(
+ "{\"data\":{\"supports_encryption\":false,\"supports_decryption\":true,"
+ + "\"soft_deleted\":false}}");
+
+ KmsKeyProperties properties = client.getKeyProperties(usableKey());
+
+ assertFalse(properties.supportsWrapping());
+ assertTrue(properties.supportsUnwrapping());
+ }
+
+ @ParameterizedTest
+ @ValueSource(ints = {400, 405, 429, 500, 502, 503})
+ void mapsProviderErrorsToConnectionFailures(int statusCode) {
+ responseStatus.set(statusCode);
+
+ assertThrows(ConnectionFailedException.class, () -> client.getKeyProperties(usableKey()));
+ }
+
+ @ParameterizedTest
+ @ValueSource(ints = {401, 403})
+ void mapsRejectedCredentialsToAuthenticationFailures(int statusCode) {
+ responseStatus.set(statusCode);
+
+ assertThrows(KmsAuthenticationException.class, () -> client.getKeyProperties(usableKey()));
+ }
+
+ @ParameterizedTest
+ @ValueSource(
+ strings = {
+ "not-json",
+ "{}",
+ "{\"data\":{}}",
+ "{\"data\":{\"supports_encryption\":\"true\",\"supports_decryption\":true}}",
+ "{\"data\":{\"supports_encryption\":true,\"supports_decryption\":\"true\"}}",
+ "{\"data\":{\"supports_encryption\":true,\"supports_decryption\":true,"
+ + "\"soft_deleted\":\"false\"}}"
+ })
+ void mapsMalformedResponsesToConnectionFailures(String body) {
+ responseBody.set(body);
+
+ assertThrows(ConnectionFailedException.class, () -> client.getKeyProperties(usableKey()));
+ }
+
+ @Test
+ void cachesTokenAcrossSuccessfulRequests() throws IOException {
+ client.getKeyProperties(usableKey());
+ Files.write(tokenFile, "rotated-token".getBytes(StandardCharsets.UTF_8));
+ client.getKeyProperties(usableKey());
+
+ assertEquals("read-only-token", requests.get(0).token);
+ assertEquals("read-only-token", requests.get(1).token);
+ }
+
+ @Test
+ void reloadsTokenAndRetriesOnceAfterAuthenticationFailure() throws IOException {
+ client.getKeyProperties(usableKey());
+ Files.write(tokenFile, "rotated-token".getBytes(StandardCharsets.UTF_8));
+ acceptedToken.set("rotated-token");
+
+ assertTrue(client.getKeyProperties(usableKey()).present());
+
+ assertEquals(3, requests.size());
+ assertEquals("read-only-token", requests.get(1).token);
+ assertEquals("rotated-token", requests.get(2).token);
+ }
+
+ @Test
+ void loadsTokenLazilyOnFirstRequest() throws IOException {
+ Files.write(tokenFile, "before-first-request".getBytes(StandardCharsets.UTF_8));
+
+ client.getKeyProperties(usableKey());
+
+ assertEquals("before-first-request", requests.get(0).token);
+ }
+
+ @Test
+ void closesPooledHttpClient() {
+ client.close();
+
+ assertThrows(ConnectionFailedException.class, () -> client.getKeyProperties(usableKey()));
+ }
+
+ @Test
+ void rejectsInvalidKeyIdWithoutCallingOpenBao() {
+ assertThrows(
+ IllegalArgumentException.class, () -> client.getKeyProperties(reference("nested/key")));
+ assertTrue(requests.isEmpty());
+ }
+
+ @Test
+ void mapsUnreadableTokenAndUnavailableServiceToConnectionFailures() {
+ OpenBaoTransitKmsClient missingTokenClient =
+ createClient("transit", tempDir.resolve("missing-token"));
+ assertThrows(
+ ConnectionFailedException.class, () -> missingTokenClient.getKeyProperties(usableKey()));
+
+ server.stop(0);
+ server = null;
+ assertThrows(ConnectionFailedException.class, () -> client.getKeyProperties(usableKey()));
+ }
+
+ private OpenBaoTransitKmsClient createClient(String transitMount, Path clientTokenFile) {
+ return new OpenBaoTransitKmsClient(
+ SOURCE,
+ URI.create(String.format("http://127.0.0.1:%s", server.getAddress().getPort())),
+ transitMount,
+ clientTokenFile);
+ }
+
+ private KmsReference reference(String keyId) {
+ return new KmsReference(KmsApi.OPENBAO_TRANSIT, SOURCE, keyId);
+ }
+
+ private void respond(HttpExchange exchange) throws IOException {
+ requests.add(
+ new RecordedRequest(
+ exchange.getRequestMethod(),
+ exchange.getRequestURI().getRawPath(),
+ exchange.getRequestHeaders().getFirst("X-Vault-Token")));
+
+ String requestToken = exchange.getRequestHeaders().getFirst("X-Vault-Token");
+ int status;
+ if (exchange.getRequestURI().getRawPath().endsWith("/" + MISSING_KEY)) {
+ status = 404;
+ } else if (acceptedToken.get() != null && !acceptedToken.get().equals(requestToken)) {
+ status = 403;
+ } else {
+ status = responseStatus.get();
+ }
+ byte[] response = responseBody.get().getBytes(StandardCharsets.UTF_8);
+ exchange.getResponseHeaders().set("Content-Type", "application/json");
+ exchange.sendResponseHeaders(status, response.length);
+ exchange.getResponseBody().write(response);
+ exchange.close();
+ }
+
+ private static final class RecordedRequest {
+ private final String method;
+ private final String rawPath;
+ private final String token;
+
+ private RecordedRequest(String method, String rawPath, String token) {
+ this.method = method;
+ this.rawPath = rawPath;
+ this.token = token;
+ }
+ }
+}
diff --git a/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestOpenBaoTransitKmsClientFactory.java b/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestOpenBaoTransitKmsClientFactory.java
new file mode 100644
index 00000000000..d298220b9ba
--- /dev/null
+++ b/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestOpenBaoTransitKmsClientFactory.java
@@ -0,0 +1,173 @@
+/*
+ * 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.kms.transit;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import com.sun.net.httpserver.HttpExchange;
+import com.sun.net.httpserver.HttpServer;
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsClientFactory;
+import org.apache.gravitino.encryption.kms.KmsReference;
+import org.apache.gravitino.encryption.kms.TestKmsClientFactoryContract;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+public class TestOpenBaoTransitKmsClientFactory extends TestKmsClientFactoryContract {
+
+ private static final String SOURCE = "primary";
+
+ @TempDir private Path tempDir;
+
+ private final AtomicReference requestedPath = new AtomicReference<>();
+
+ private HttpServer server;
+ private Path tokenFile;
+
+ @BeforeEach
+ void startServer() throws IOException {
+ tokenFile = tempDir.resolve("openbao-token");
+ Files.write(tokenFile, "read-only-token".getBytes(StandardCharsets.UTF_8));
+
+ server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+ server.createContext("/", this::respond);
+ server.start();
+ }
+
+ @AfterEach
+ void stopServer() {
+ if (server != null) {
+ server.stop(0);
+ }
+ }
+
+ @Override
+ protected KmsClientFactory factory() {
+ return new OpenBaoTransitKmsClientFactory();
+ }
+
+ @Override
+ protected KmsApi expectedApi() {
+ return KmsApi.OPENBAO_TRANSIT;
+ }
+
+ @Test
+ void createsWorkingClientWithDefaultMount() {
+ KmsClient client = factory().create(SOURCE, properties());
+
+ client.getKeyProperties(new KmsReference(KmsApi.OPENBAO_TRANSIT, SOURCE, "customer-key"));
+
+ assertEquals("/v1/transit/keys/customer-key", requestedPath.get());
+ }
+
+ @Test
+ void createsWorkingClientWithCustomMount() {
+ Map properties = properties();
+ properties.put(OpenBaoTransitKmsClientFactory.TRANSIT_MOUNT, "team/transit");
+ KmsClient client = factory().create(SOURCE, properties);
+
+ client.getKeyProperties(new KmsReference(KmsApi.OPENBAO_TRANSIT, SOURCE, "customer-key"));
+
+ assertEquals("/v1/team/transit/keys/customer-key", requestedPath.get());
+ }
+
+ @Test
+ void rejectsMissingRequiredConfiguration() {
+ assertThrows(IllegalArgumentException.class, () -> factory().create(SOURCE, null));
+
+ Map missingAddress = properties();
+ missingAddress.remove(OpenBaoTransitKmsClientFactory.SERVICE_ADDRESS);
+ assertThrows(IllegalArgumentException.class, () -> factory().create(SOURCE, missingAddress));
+
+ Map missingToken = properties();
+ missingToken.remove(OpenBaoTransitKmsClientFactory.TOKEN_FILE);
+ assertThrows(IllegalArgumentException.class, () -> factory().create(SOURCE, missingToken));
+ }
+
+ @Test
+ void rejectsInvalidSourceAndUnknownConfiguration() {
+ assertThrows(IllegalArgumentException.class, () -> factory().create(" ", properties()));
+
+ Map properties = properties();
+ properties.put("token", "secret");
+ assertThrows(IllegalArgumentException.class, () -> factory().create(SOURCE, properties));
+ }
+
+ @Test
+ void rejectsInvalidServiceAddress() {
+ assertInvalidServiceAddress("file:///tmp/openbao");
+ assertInvalidServiceAddress("http://user@localhost");
+ assertInvalidServiceAddress("http://localhost/openbao");
+ assertInvalidServiceAddress("http://localhost?namespace=team");
+ assertInvalidServiceAddress("http://localhost#fragment");
+ }
+
+ @Test
+ void rejectsInvalidMountAndTokenFile() {
+ for (String mount : new String[] {"", "/transit", "transit/", "team//transit", ".", ".."}) {
+ Map properties = properties();
+ properties.put(OpenBaoTransitKmsClientFactory.TRANSIT_MOUNT, mount);
+ assertThrows(IllegalArgumentException.class, () -> factory().create(SOURCE, properties));
+ }
+
+ Map properties = properties();
+ properties.put(OpenBaoTransitKmsClientFactory.TOKEN_FILE, "relative-token");
+ assertThrows(IllegalArgumentException.class, () -> factory().create(SOURCE, properties));
+ }
+
+ private Map properties() {
+ Map properties = new HashMap<>();
+ properties.put(
+ OpenBaoTransitKmsClientFactory.SERVICE_ADDRESS,
+ String.format("http://127.0.0.1:%s", server.getAddress().getPort()));
+ properties.put(OpenBaoTransitKmsClientFactory.TOKEN_FILE, tokenFile.toString());
+ properties.put(OpenBaoTransitKmsClientFactory.CREDENTIAL_METHOD, "token_file");
+ return properties;
+ }
+
+ private void assertInvalidServiceAddress(String address) {
+ Map properties = properties();
+ properties.put(OpenBaoTransitKmsClientFactory.SERVICE_ADDRESS, address);
+ assertThrows(IllegalArgumentException.class, () -> factory().create(SOURCE, properties));
+ }
+
+ private void respond(HttpExchange exchange) throws IOException {
+ requestedPath.set(exchange.getRequestURI().getRawPath());
+ byte[] response =
+ ("{\"data\":{\"supports_encryption\":true,\"supports_decryption\":true,"
+ + "\"soft_deleted\":false}}")
+ .getBytes(StandardCharsets.UTF_8);
+ exchange.getResponseHeaders().set("Content-Type", "application/json");
+ exchange.sendResponseHeaders(200, response.length);
+ exchange.getResponseBody().write(response);
+ exchange.close();
+ }
+}
diff --git a/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestVaultTransitKmsClient.java b/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestVaultTransitKmsClient.java
new file mode 100644
index 00000000000..0c9ba82bfcd
--- /dev/null
+++ b/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestVaultTransitKmsClient.java
@@ -0,0 +1,119 @@
+/*
+ * 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.kms.transit;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import com.sun.net.httpserver.HttpExchange;
+import com.sun.net.httpserver.HttpServer;
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsReference;
+import org.apache.gravitino.encryption.kms.TestKmsClientContract;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+public class TestVaultTransitKmsClient extends TestKmsClientContract {
+
+ private static final String SOURCE = "primary";
+ private static final String USABLE_KEY = "customer-key";
+ private static final String MISSING_KEY = "missing-key";
+
+ @TempDir private Path tempDir;
+
+ private final AtomicReference requestedPath = new AtomicReference<>();
+ private final AtomicReference requestedToken = new AtomicReference<>();
+
+ private HttpServer server;
+ private VaultTransitKmsClient client;
+
+ @BeforeEach
+ void startServer() throws IOException {
+ Path tokenFile = tempDir.resolve("vault-token");
+ Files.write(tokenFile, "read-only-token".getBytes(StandardCharsets.UTF_8));
+
+ server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+ server.createContext("/", this::respond);
+ server.start();
+ client =
+ new VaultTransitKmsClient(
+ SOURCE,
+ URI.create(String.format("http://127.0.0.1:%s", server.getAddress().getPort())),
+ "transit",
+ tokenFile);
+ }
+
+ @AfterEach
+ void stopServer() {
+ if (server != null) {
+ server.stop(0);
+ }
+ }
+
+ @Override
+ protected KmsClient client() {
+ return client;
+ }
+
+ @Override
+ protected KmsReference usableKey() {
+ return reference(USABLE_KEY);
+ }
+
+ @Override
+ protected KmsReference missingKey() {
+ return reference(MISSING_KEY);
+ }
+
+ @Test
+ void readsVaultTransitKeyMetadata() {
+ client.getKeyProperties(usableKey());
+
+ assertEquals("/v1/transit/keys/customer-key", requestedPath.get());
+ assertEquals("read-only-token", requestedToken.get());
+ }
+
+ private KmsReference reference(String keyId) {
+ return new KmsReference(KmsApi.VAULT_TRANSIT, SOURCE, keyId);
+ }
+
+ private void respond(HttpExchange exchange) throws IOException {
+ requestedPath.set(exchange.getRequestURI().getRawPath());
+ requestedToken.set(exchange.getRequestHeaders().getFirst("X-Vault-Token"));
+
+ boolean missing = exchange.getRequestURI().getRawPath().endsWith("/" + MISSING_KEY);
+ byte[] response =
+ ("{\"data\":{\"type\":\"aes256-gcm96\",\"supports_encryption\":true,"
+ + "\"supports_decryption\":true}}")
+ .getBytes(StandardCharsets.UTF_8);
+ exchange.getResponseHeaders().set("Content-Type", "application/json");
+ exchange.sendResponseHeaders(missing ? 404 : 200, response.length);
+ exchange.getResponseBody().write(response);
+ exchange.close();
+ }
+}
diff --git a/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestVaultTransitKmsClientFactory.java b/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestVaultTransitKmsClientFactory.java
new file mode 100644
index 00000000000..c563da7844c
--- /dev/null
+++ b/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestVaultTransitKmsClientFactory.java
@@ -0,0 +1,185 @@
+/*
+ * 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.kms.transit;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import com.sun.net.httpserver.HttpExchange;
+import com.sun.net.httpserver.HttpServer;
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.EnumMap;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.ServiceLoader;
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsClientFactory;
+import org.apache.gravitino.encryption.kms.KmsReference;
+import org.apache.gravitino.encryption.kms.TestKmsClientFactoryContract;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+public class TestVaultTransitKmsClientFactory extends TestKmsClientFactoryContract {
+
+ private static final String SOURCE = "primary";
+
+ @TempDir private Path tempDir;
+
+ private final AtomicReference requestedPath = new AtomicReference<>();
+
+ private HttpServer server;
+ private Path tokenFile;
+
+ @BeforeEach
+ void startServer() throws IOException {
+ tokenFile = tempDir.resolve("vault-token");
+ Files.write(tokenFile, "read-only-token".getBytes(StandardCharsets.UTF_8));
+
+ server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+ server.createContext("/", this::respond);
+ server.start();
+ }
+
+ @AfterEach
+ void stopServer() {
+ if (server != null) {
+ server.stop(0);
+ }
+ }
+
+ @Override
+ protected KmsClientFactory factory() {
+ return new VaultTransitKmsClientFactory();
+ }
+
+ @Override
+ protected KmsApi expectedApi() {
+ return KmsApi.VAULT_TRANSIT;
+ }
+
+ @Test
+ void createsWorkingClientWithDefaultMount() {
+ KmsClient client = factory().create(SOURCE, properties());
+
+ client.getKeyProperties(new KmsReference(KmsApi.VAULT_TRANSIT, SOURCE, "customer-key"));
+
+ assertEquals("/v1/transit/keys/customer-key", requestedPath.get());
+ }
+
+ @Test
+ void createsWorkingClientWithCustomMount() {
+ Map properties = properties();
+ properties.put(VaultTransitKmsClientFactory.TRANSIT_MOUNT, "team/transit");
+ KmsClient client = factory().create(SOURCE, properties);
+
+ client.getKeyProperties(new KmsReference(KmsApi.VAULT_TRANSIT, SOURCE, "customer-key"));
+
+ assertEquals("/v1/team/transit/keys/customer-key", requestedPath.get());
+ }
+
+ @Test
+ void serviceLoaderDiscoversTransitFactories() {
+ Map> factoryClasses = new EnumMap<>(KmsApi.class);
+ for (KmsClientFactory factory : ServiceLoader.load(KmsClientFactory.class)) {
+ factoryClasses.put(factory.api(), factory.getClass());
+ }
+
+ assertEquals(OpenBaoTransitKmsClientFactory.class, factoryClasses.get(KmsApi.OPENBAO_TRANSIT));
+ assertEquals(VaultTransitKmsClientFactory.class, factoryClasses.get(KmsApi.VAULT_TRANSIT));
+ }
+
+ @Test
+ void rejectsMissingRequiredConfiguration() {
+ assertThrows(IllegalArgumentException.class, () -> factory().create(SOURCE, null));
+
+ Map missingAddress = properties();
+ missingAddress.remove(VaultTransitKmsClientFactory.SERVICE_ADDRESS);
+ assertThrows(IllegalArgumentException.class, () -> factory().create(SOURCE, missingAddress));
+
+ Map missingToken = properties();
+ missingToken.remove(VaultTransitKmsClientFactory.TOKEN_FILE);
+ assertThrows(IllegalArgumentException.class, () -> factory().create(SOURCE, missingToken));
+ }
+
+ @Test
+ void rejectsInvalidSourceAndUnknownConfiguration() {
+ assertThrows(IllegalArgumentException.class, () -> factory().create(" ", properties()));
+
+ Map properties = properties();
+ properties.put("token", "secret");
+ assertThrows(IllegalArgumentException.class, () -> factory().create(SOURCE, properties));
+ }
+
+ @Test
+ void rejectsInvalidServiceAddress() {
+ assertInvalidServiceAddress("file:///tmp/vault");
+ assertInvalidServiceAddress("http://user@localhost");
+ assertInvalidServiceAddress("http://localhost/vault");
+ assertInvalidServiceAddress("http://localhost?namespace=team");
+ assertInvalidServiceAddress("http://localhost#fragment");
+ }
+
+ @Test
+ void rejectsInvalidMountAndTokenFile() {
+ for (String mount : new String[] {"", "/transit", "transit/", "team//transit", ".", ".."}) {
+ Map properties = properties();
+ properties.put(VaultTransitKmsClientFactory.TRANSIT_MOUNT, mount);
+ assertThrows(IllegalArgumentException.class, () -> factory().create(SOURCE, properties));
+ }
+
+ Map properties = properties();
+ properties.put(VaultTransitKmsClientFactory.TOKEN_FILE, "relative-token");
+ assertThrows(IllegalArgumentException.class, () -> factory().create(SOURCE, properties));
+ }
+
+ private Map properties() {
+ Map properties = new HashMap<>();
+ properties.put(
+ VaultTransitKmsClientFactory.SERVICE_ADDRESS,
+ String.format("http://127.0.0.1:%s", server.getAddress().getPort()));
+ properties.put(VaultTransitKmsClientFactory.TOKEN_FILE, tokenFile.toString());
+ properties.put(VaultTransitKmsClientFactory.CREDENTIAL_METHOD, "token_file");
+ return properties;
+ }
+
+ private void assertInvalidServiceAddress(String address) {
+ Map properties = properties();
+ properties.put(VaultTransitKmsClientFactory.SERVICE_ADDRESS, address);
+ assertThrows(IllegalArgumentException.class, () -> factory().create(SOURCE, properties));
+ }
+
+ private void respond(HttpExchange exchange) throws IOException {
+ requestedPath.set(exchange.getRequestURI().getRawPath());
+ byte[] response =
+ ("{\"data\":{\"supports_encryption\":true,\"supports_decryption\":true}}")
+ .getBytes(StandardCharsets.UTF_8);
+ exchange.getResponseHeaders().set("Content-Type", "application/json");
+ exchange.sendResponseHeaders(200, response.length);
+ exchange.getResponseBody().write(response);
+ exchange.close();
+ }
+}
diff --git a/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/integration/test/AbstractTransitKmsClientIT.java b/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/integration/test/AbstractTransitKmsClientIT.java
new file mode 100644
index 00000000000..f5db65d6f42
--- /dev/null
+++ b/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/integration/test/AbstractTransitKmsClientIT.java
@@ -0,0 +1,154 @@
+/*
+ * 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.kms.transit.integration.test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.time.Duration;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClient;
+import org.apache.gravitino.encryption.kms.KmsClientFactory;
+import org.apache.gravitino.encryption.kms.KmsKeyProperties;
+import org.apache.gravitino.encryption.kms.KmsReference;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.testcontainers.containers.Container;
+import org.testcontainers.containers.GenericContainer;
+import org.testcontainers.containers.wait.strategy.Wait;
+import org.testcontainers.utility.DockerImageName;
+
+abstract class AbstractTransitKmsClientIT {
+
+ private static final int API_PORT = 8200;
+ private static final String ROOT_TOKEN = "gravitino-kms-test-root-token";
+ private static final String SOURCE = "test";
+ private static final String USABLE_KEY = "usable-key";
+ private static final String SIGNING_KEY = "signing-key";
+
+ private GenericContainer> container;
+ private Path tokenFile;
+
+ @BeforeAll
+ void startBackend() throws Exception {
+ tokenFile = Files.createTempFile("gravitino-transit-token-", ".txt");
+ Files.write(tokenFile, ROOT_TOKEN.getBytes(StandardCharsets.UTF_8));
+
+ container =
+ new GenericContainer<>(DockerImageName.parse(image()))
+ .withCommand(
+ "server",
+ "-dev",
+ String.format("-dev-root-token-id=%s", ROOT_TOKEN),
+ "-dev-listen-address=0.0.0.0:8200")
+ .withEnv(addressEnvironmentVariable(), "http://127.0.0.1:8200")
+ .withEnv(tokenEnvironmentVariable(), ROOT_TOKEN)
+ .withExposedPorts(API_PORT)
+ .waitingFor(
+ Wait.forHttp("/v1/sys/health")
+ .forStatusCode(200)
+ .withStartupTimeout(Duration.ofMinutes(2)));
+ container.start();
+
+ runCommand("secrets", "enable", "transit");
+ runCommand("write", "-f", String.format("transit/keys/%s", USABLE_KEY));
+ runCommand("write", "-f", String.format("transit/keys/%s", SIGNING_KEY), "type=ecdsa-p256");
+ }
+
+ @AfterAll
+ void stopBackend() throws Exception {
+ if (container != null) {
+ container.stop();
+ }
+ if (tokenFile != null) {
+ Files.deleteIfExists(tokenFile);
+ }
+ }
+
+ @Test
+ void inspectsRealTransitKeyMetadata() {
+ try (KmsClient client = factory().create(SOURCE, properties())) {
+ KmsReference usableReference = new KmsReference(api(), SOURCE, USABLE_KEY);
+ KmsKeyProperties usable = client.getKeyProperties(usableReference);
+
+ assertEquals(usableReference, usable.reference());
+ assertTrue(usable.present());
+ assertTrue(usable.enabled());
+ assertTrue(usable.supportsWrapping());
+ assertTrue(usable.supportsUnwrapping());
+
+ KmsKeyProperties signing =
+ client.getKeyProperties(new KmsReference(api(), SOURCE, SIGNING_KEY));
+ assertTrue(signing.present());
+ assertTrue(signing.enabled());
+ assertFalse(signing.supportsWrapping());
+ assertFalse(signing.supportsUnwrapping());
+
+ KmsKeyProperties missing =
+ client.getKeyProperties(new KmsReference(api(), SOURCE, "missing-key"));
+ assertFalse(missing.present());
+ assertFalse(missing.enabled());
+ assertFalse(missing.supportsWrapping());
+ assertFalse(missing.supportsUnwrapping());
+ }
+ }
+
+ protected abstract String image();
+
+ protected abstract String executable();
+
+ protected abstract String addressEnvironmentVariable();
+
+ protected abstract String tokenEnvironmentVariable();
+
+ protected abstract KmsApi api();
+
+ protected abstract KmsClientFactory factory();
+
+ private Map properties() {
+ Map properties = new HashMap<>();
+ properties.put(
+ "endpoint.address",
+ String.format("http://%s:%s", container.getHost(), container.getMappedPort(API_PORT)));
+ properties.put("credential.method", "token_file");
+ properties.put("credential.path", tokenFile.toString());
+ return properties;
+ }
+
+ private void runCommand(String... arguments) throws Exception {
+ String[] command = new String[arguments.length + 1];
+ command[0] = executable();
+ System.arraycopy(arguments, 0, command, 1, arguments.length);
+ Container.ExecResult result = container.execInContainer(command);
+ assertEquals(
+ 0,
+ result.getExitCode(),
+ () ->
+ String.format(
+ "Transit fixture command failed: stdout=%s, stderr=%s",
+ result.getStdout(), result.getStderr()));
+ }
+}
diff --git a/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/integration/test/OpenBaoTransitKmsClientIT.java b/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/integration/test/OpenBaoTransitKmsClientIT.java
new file mode 100644
index 00000000000..06c1871175d
--- /dev/null
+++ b/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/integration/test/OpenBaoTransitKmsClientIT.java
@@ -0,0 +1,66 @@
+/*
+ * 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.kms.transit.integration.test;
+
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClientFactory;
+import org.apache.gravitino.kms.transit.OpenBaoTransitKmsClientFactory;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.TestInstance;
+
+@Tag("gravitino-docker-test")
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+class OpenBaoTransitKmsClientIT extends AbstractTransitKmsClientIT {
+
+ private static final String DEFAULT_IMAGE = "openbao/openbao:2.6.0";
+ private static final String IMAGE_ENVIRONMENT_VARIABLE = "GRAVITINO_OPENBAO_DOCKER_IMAGE";
+
+ @Override
+ protected String image() {
+ String configuredImage = System.getenv(IMAGE_ENVIRONMENT_VARIABLE);
+ return configuredImage == null || configuredImage.trim().isEmpty()
+ ? DEFAULT_IMAGE
+ : configuredImage.trim();
+ }
+
+ @Override
+ protected String executable() {
+ return "bao";
+ }
+
+ @Override
+ protected String addressEnvironmentVariable() {
+ return "BAO_ADDR";
+ }
+
+ @Override
+ protected String tokenEnvironmentVariable() {
+ return "BAO_TOKEN";
+ }
+
+ @Override
+ protected KmsApi api() {
+ return KmsApi.OPENBAO_TRANSIT;
+ }
+
+ @Override
+ protected KmsClientFactory factory() {
+ return new OpenBaoTransitKmsClientFactory();
+ }
+}
diff --git a/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/integration/test/VaultTransitKmsClientIT.java b/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/integration/test/VaultTransitKmsClientIT.java
new file mode 100644
index 00000000000..1e2addac6e3
--- /dev/null
+++ b/bundles/transit/src/test/java/org/apache/gravitino/kms/transit/integration/test/VaultTransitKmsClientIT.java
@@ -0,0 +1,66 @@
+/*
+ * 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.kms.transit.integration.test;
+
+import org.apache.gravitino.encryption.kms.KmsApi;
+import org.apache.gravitino.encryption.kms.KmsClientFactory;
+import org.apache.gravitino.kms.transit.VaultTransitKmsClientFactory;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.TestInstance;
+
+@Tag("gravitino-docker-test")
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+class VaultTransitKmsClientIT extends AbstractTransitKmsClientIT {
+
+ private static final String DEFAULT_IMAGE = "hashicorp/vault:2.0.3";
+ private static final String IMAGE_ENVIRONMENT_VARIABLE = "GRAVITINO_VAULT_DOCKER_IMAGE";
+
+ @Override
+ protected String image() {
+ String configuredImage = System.getenv(IMAGE_ENVIRONMENT_VARIABLE);
+ return configuredImage == null || configuredImage.trim().isEmpty()
+ ? DEFAULT_IMAGE
+ : configuredImage.trim();
+ }
+
+ @Override
+ protected String executable() {
+ return "vault";
+ }
+
+ @Override
+ protected String addressEnvironmentVariable() {
+ return "VAULT_ADDR";
+ }
+
+ @Override
+ protected String tokenEnvironmentVariable() {
+ return "VAULT_TOKEN";
+ }
+
+ @Override
+ protected KmsApi api() {
+ return KmsApi.VAULT_TRANSIT;
+ }
+
+ @Override
+ protected KmsClientFactory factory() {
+ return new VaultTransitKmsClientFactory();
+ }
+}
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..aa121241c5f
--- /dev/null
+++ b/core/src/main/java/org/apache/gravitino/encryption/kms/KmsClientRegistry.java
@@ -0,0 +1,207 @@
+/*
+ * 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.EnumMap;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+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
+ * @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 KmsKeyProperties 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 != reference.api()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "KMS source '%s' uses API '%s', not '%s'",
+ reference.source(), configuredClient.api.wireValue(), reference.api().wireValue()));
+ }
+
+ KmsKeyProperties properties = configuredClient.client.getKeyProperties(reference);
+ if (properties == null || !reference.equals(properties.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 EnumMap<>(KmsApi.class);
+ for (KmsClientFactory factory : factories) {
+ if (factory == null || factory.api() == null) {
+ throw new IllegalArgumentException("KMS client factory and its API cannot be null");
+ }
+ KmsClientFactory existing = factoriesByApi.putIfAbsent(factory.api(), factory);
+ if (existing != null) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Multiple KMS client factories support API '%s'", factory.api().wireValue()));
+ }
+ }
+ 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().wireValue(), 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().wireValue()));
+ }
+ 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 final class ConfiguredClient {
+ private final KmsApi api;
+ private final KmsClient client;
+
+ private ConfiguredClient(KmsApi 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..ebc9baa55a0
--- /dev/null
+++ b/core/src/main/java/org/apache/gravitino/encryption/kms/KmsConfig.java
@@ -0,0 +1,152 @@
+/*
+ * 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.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);
+ }
+ KmsApi api;
+ try {
+ api = KmsApi.fromWireValue(apiValue);
+ } catch (IllegalArgumentException e) {
+ throw new KmsConfigurationException(
+ e, "Invalid KMS API property '%s%s': %s", KMS_CONFIG_PREFIX, apiKey, e.getMessage());
+ }
+
+ 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 KmsApi api;
+ private final Map properties;
+
+ private SourceConfig(KmsApi api, Map properties) {
+ this.api = api;
+ this.properties = Collections.unmodifiableMap(new LinkedHashMap<>(properties));
+ }
+
+ KmsApi 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..538abdea957
--- /dev/null
+++ b/core/src/test/java/org/apache/gravitino/TestGravitinoEnvKmsClientRegistry.java
@@ -0,0 +1,51 @@
+/*
+ * 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.KmsApi;
+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(KmsApi.AWS_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..67450744eaf
--- /dev/null
+++ b/core/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClientRegistry.java
@@ -0,0 +1,503 @@
+/*
+ * 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.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 {
+
+ @Test
+ void testCreatesAndDispatchesConfiguredClients() {
+ RecordingFactory awsFactory = new RecordingFactory(KmsApi.AWS_KMS);
+ RecordingFactory gcpFactory = new RecordingFactory(KmsApi.GOOGLE_CLOUD_KMS);
+ 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(KmsApi.AWS_KMS, "primary", "alias/orders");
+ KmsReference gcpReference =
+ new KmsReference(
+ KmsApi.GOOGLE_CLOUD_KMS, "analytics", "projects/p/locations/l/keyRings/r/cryptoKeys/k");
+
+ Assertions.assertEquals(awsReference, registry.getKeyProperties(awsReference).reference());
+ Assertions.assertEquals(awsReference, registry.getKeyProperties(awsReference).reference());
+ Assertions.assertEquals(gcpReference, registry.getKeyProperties(gcpReference).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(KmsApi.AWS_KMS)));
+
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> registry.getKeyProperties(new KmsReference(KmsApi.AWS_KMS, "other", "key")));
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ registry.getKeyProperties(new KmsReference(KmsApi.GOOGLE_CLOUD_KMS, "primary", "key")));
+ Assertions.assertThrows(IllegalArgumentException.class, () -> registry.getKeyProperties(null));
+ }
+
+ @Test
+ void testCreatesMultipleSourcesForSameApi() {
+ RecordingFactory factory = new RecordingFactory(KmsApi.AZURE_KEY_VAULT);
+ 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(KmsApi.AZURE_KEY_VAULT, "azure-eu", "primary"))
+ .reference()
+ .source());
+ Assertions.assertEquals(
+ "azure-us",
+ registry
+ .getKeyProperties(new KmsReference(KmsApi.AZURE_KEY_VAULT, "azure-us", "primary"))
+ .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(KmsApi.AWS_KMS), new RecordingFactory(KmsApi.AWS_KMS))));
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> new KmsClientRegistry(awsConfig, java.util.Arrays.asList((KmsClientFactory) null)));
+ Assertions.assertThrows(
+ IllegalArgumentException.class, () -> new KmsClientRegistry(awsConfig, null));
+ }
+
+ @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(KmsApi.AWS_KMS, "primary", "key");
+ Assertions.assertEquals(reference, registry.getKeyProperties(reference).reference());
+ }
+ } finally {
+ Thread.currentThread().setContextClassLoader(originalClassLoader);
+ }
+ }
+
+ @Test
+ void testRejectsInvalidProviderResults() {
+ KmsReference reference = new KmsReference(KmsApi.AWS_KMS, "primary", "key");
+ Config awsConfig =
+ config(
+ "gravitino.kms.sources", "primary",
+ "gravitino.kms.source.primary.api", "aws-kms");
+
+ KmsClientFactory nullResultFactory =
+ factory(KmsApi.AWS_KMS, (source, properties) -> ignored -> null);
+ KmsClientRegistry nullResultRegistry =
+ new KmsClientRegistry(awsConfig, List.of(nullResultFactory));
+ Assertions.assertThrows(
+ IllegalStateException.class, () -> nullResultRegistry.getKeyProperties(reference));
+
+ KmsReference otherReference = new KmsReference(KmsApi.AWS_KMS, "primary", "other");
+ KmsClientFactory wrongReferenceFactory =
+ factory(KmsApi.AWS_KMS, (source, properties) -> ignored -> new Properties(otherReference));
+ KmsClientRegistry wrongReferenceRegistry =
+ new KmsClientRegistry(awsConfig, List.of(wrongReferenceFactory));
+ Assertions.assertThrows(
+ IllegalStateException.class, () -> wrongReferenceRegistry.getKeyProperties(reference));
+
+ KmsClientFactory nullClientFactory = factory(KmsApi.AWS_KMS, (source, properties) -> null);
+ Assertions.assertThrows(
+ IllegalStateException.class,
+ () -> new KmsClientRegistry(awsConfig, List.of(nullClientFactory)));
+ }
+
+ @Test
+ void testClosesClientsInReverseOrderAndIsIdempotent() {
+ List closeOrder = new ArrayList<>();
+ CloseTrackingFactory awsFactory =
+ new CloseTrackingFactory(KmsApi.AWS_KMS, "aws", closeOrder, null);
+ CloseTrackingFactory gcpFactory =
+ new CloseTrackingFactory(KmsApi.GOOGLE_CLOUD_KMS, "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(KmsApi.AWS_KMS, "primary", "key")));
+ }
+
+ @Test
+ void testClosesCreatedClientsAfterPartialInitializationFailure() {
+ CloseTrackingFactory awsFactory =
+ new CloseTrackingFactory(KmsApi.AWS_KMS, "aws", new ArrayList<>(), null);
+ KmsClientFactory failingFactory =
+ factory(
+ KmsApi.GOOGLE_CLOUD_KMS,
+ (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(KmsApi.AWS_KMS, "aws", new ArrayList<>(), closeFailure);
+ IllegalArgumentException creationFailure =
+ new IllegalArgumentException("invalid GCP configuration");
+ KmsClientFactory failingFactory =
+ factory(
+ KmsApi.GOOGLE_CLOUD_KMS,
+ (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(KmsApi.AWS_KMS, entered, release);
+ KmsClientFactory gcpFactory = blockingFactory(KmsApi.GOOGLE_CLOUD_KMS, 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(KmsApi.AWS_KMS, "primary", "key")));
+ Future gcpResult =
+ executor.submit(
+ () ->
+ registry.getKeyProperties(
+ new KmsReference(KmsApi.GOOGLE_CLOUD_KMS, "analytics", "key")));
+
+ Assertions.assertTrue(entered.await(5, TimeUnit.SECONDS));
+ release.countDown();
+ Assertions.assertEquals("primary", awsResult.get(5, TimeUnit.SECONDS).reference().source());
+ Assertions.assertEquals("analytics", gcpResult.get(5, TimeUnit.SECONDS).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(KmsApi.AWS_KMS, "aws", new ArrayList<>(), awsFailure),
+ new CloseTrackingFactory(
+ KmsApi.GOOGLE_CLOUD_KMS, "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(KmsApi api, ClientCreator creator) {
+ return new KmsClientFactory() {
+ @Override
+ public KmsApi api() {
+ return api;
+ }
+
+ @Override
+ public KmsClient create(String source, Map properties) {
+ return creator.create(source, properties);
+ }
+ };
+ }
+
+ private static KmsClientFactory blockingFactory(
+ KmsApi 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 new Properties(reference);
+ });
+ }
+
+ private interface ClientCreator {
+ KmsClient create(String source, Map properties);
+ }
+
+ private static final class RecordingFactory implements KmsClientFactory {
+ private final KmsApi api;
+ private String source;
+ private Map properties;
+ private String providerKeyId;
+ private final AtomicInteger createCount = new AtomicInteger();
+
+ private RecordingFactory(KmsApi api) {
+ this.api = api;
+ }
+
+ @Override
+ public KmsApi 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 new Properties(reference);
+ };
+ }
+ }
+
+ private static final class CloseTrackingFactory implements KmsClientFactory {
+ private final KmsApi api;
+ private final String name;
+ private final List closeOrder;
+ private final RuntimeException closeFailure;
+ private final AtomicInteger closeCount = new AtomicInteger();
+
+ private CloseTrackingFactory(
+ KmsApi api, String name, List closeOrder, RuntimeException closeFailure) {
+ this.api = api;
+ this.name = name;
+ this.closeOrder = closeOrder;
+ this.closeFailure = closeFailure;
+ }
+
+ @Override
+ public KmsApi api() {
+ return api;
+ }
+
+ @Override
+ public KmsClient create(String source, Map properties) {
+ return new KmsClient() {
+ @Override
+ public KmsKeyProperties getKeyProperties(KmsReference reference) {
+ return 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 present() {
+ return true;
+ }
+
+ @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 KmsApi api() {
+ return KmsApi.AWS_KMS;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public KmsClient create(String source, Map properties) {
+ return reference -> 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..3aacc64a0dd
--- /dev/null
+++ b/core/src/test/java/org/apache/gravitino/encryption/kms/TestKmsConfig.java
@@ -0,0 +1,134 @@
+/*
+ * 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 {
+
+ @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(KmsApi.AWS_KMS, config.sources().get("primary").api());
+ Assertions.assertEquals(
+ Map.of("endpoint.region", "us-west-2", "credential.method", "default"),
+ config.sources().get("primary").properties());
+ Assertions.assertEquals(
+ KmsApi.GOOGLE_CLOUD_KMS, 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 testRequiresSupportedApi() {
+ 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", "custom"),
+ "Unsupported KMS API 'custom'");
+ }
+
+ @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(KmsApi.AWS_KMS, config.sources().get("primary").api());
+ Assertions.assertEquals(KmsApi.AWS_KMS, 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);
+ }
+ }
+}
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 9222b53ab3d..e200db9c9e9 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -19,6 +19,7 @@
[versions]
awssdk = "2.31.73"
azure-identity = "1.13.1"
+azure-security-keyvault-keys = "4.8.6"
azure-storage-file-datalake = "12.20.0"
reactor-netty-http = "1.2.1"
reactor-netty-core = "1.2.1"
@@ -153,6 +154,7 @@ datanucleus-rdbms = "4.1.19"
datanucleus-jdo = "3.2.0-m3"
hudi = "0.15.0"
google-auth = "1.28.0"
+google-cloud-kms = "2.97.0"
aliyun-credentials = "0.3.12"
aliyun-sdk-oss = "3.10.2"
openlineage = "1.29.0"
@@ -176,6 +178,7 @@ aws-s3 = { group = "software.amazon.awssdk", name = "s3", version.ref = "awssdk"
aws-sts = { group = "software.amazon.awssdk", name = "sts", version.ref = "awssdk" }
aws-kms = { group = "software.amazon.awssdk", name = "kms", version.ref = "awssdk" }
azure-identity = { group = "com.azure", name = "azure-identity", version.ref = "azure-identity"}
+azure-security-keyvault-keys = { group = "com.azure", name = "azure-security-keyvault-keys", version.ref = "azure-security-keyvault-keys"}
azure-storage-file-datalake = { group = "com.azure", name = "azure-storage-file-datalake", version.ref = "azure-storage-file-datalake"}
reactor-netty-http = {group = "io.projectreactor.netty", name = "reactor-netty-http", version.ref = "reactor-netty-http"}
reactor-netty-core = {group = "io.projectreactor.netty", name = "reactor-netty-core", version.ref = "reactor-netty-core"}
@@ -350,6 +353,7 @@ jettison = { group = "org.codehaus.jettison", name = "jettison", version.ref = "
google-auth-http = { group = "com.google.auth", name = "google-auth-library-oauth2-http", version.ref = "google-auth" }
google-auth-credentials = { group = "com.google.auth", name = "google-auth-library-credentials", version.ref = "google-auth" }
+google-cloud-kms = { group = "com.google.cloud", name = "google-cloud-kms", version.ref = "google-cloud-kms" }
aliyun-credentials-sdk = { group='com.aliyun', name='credentials-java', version.ref='aliyun-credentials' }
aliyun-sdk-oss = { module = "com.aliyun.oss:aliyun-sdk-oss", version.ref = "aliyun-sdk-oss" }
diff --git a/settings.gradle.kts b/settings.gradle.kts
index 3f9ec2f06a3..8267eb8d0fa 100644
--- a/settings.gradle.kts
+++ b/settings.gradle.kts
@@ -125,6 +125,7 @@ include(":bundles:gcp", ":bundles:gcp-bundle", ":bundles:iceberg-gcp-bundle")
include(":bundles:aliyun", ":bundles:aliyun-bundle", ":bundles:iceberg-aliyun-bundle")
include(":bundles:tencent", ":bundles:tencent-bundle")
include(":bundles:azure", ":bundles:azure-bundle", ":bundles:iceberg-azure-bundle")
+include(":bundles:transit")
include(":catalogs:hadoop-common")
include(":catalogs:hadoop-auth")
include(":lineage")