From b675a5c859fecc2052b524c872fc1c1d411efa1f Mon Sep 17 00:00:00 2001 From: Nevin Zheng Date: Tue, 21 Jul 2026 15:48:15 -0700 Subject: [PATCH 1/5] feat(kms): define key inspection API --- api/build.gradle.kts | 3 + .../gravitino/encryption/kms/KmsApi.java | 85 +++++++++++ .../kms/KmsAuthenticationException.java | 50 ++++++ .../gravitino/encryption/kms/KmsClient.java | 47 ++++++ .../encryption/kms/KmsClientFactory.java | 49 ++++++ .../kms/KmsConfigurationException.java | 49 ++++++ .../encryption/kms/KmsKeyProperties.java | 65 ++++++++ .../encryption/kms/KmsReference.java | 105 +++++++++++++ .../encryption/kms/TestFakeKmsClient.java | 44 ++++++ .../gravitino/encryption/kms/TestKmsApi.java | 47 ++++++ .../encryption/kms/TestKmsClient.java | 73 +++++++++ .../encryption/kms/TestKmsReference.java | 71 +++++++++ .../encryption/kms/FakeKmsClient.java | 142 ++++++++++++++++++ .../encryption/kms/TestKmsClientContract.java | 98 ++++++++++++ .../kms/TestKmsClientFactoryContract.java | 45 ++++++ 15 files changed, 973 insertions(+) create mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java create mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java create mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java create mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java create mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java create mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java create mode 100644 api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java create mode 100644 api/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java create mode 100644 api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsApi.java create mode 100644 api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java create mode 100644 api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsReference.java create mode 100644 api/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java create mode 100644 api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java create mode 100644 api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java diff --git a/api/build.gradle.kts b/api/build.gradle.kts index f0fe3ba5ee9..616703cb1eb 100644 --- a/api/build.gradle.kts +++ b/api/build.gradle.kts @@ -18,6 +18,7 @@ */ plugins { `maven-publish` + `java-test-fixtures` id("java") id("idea") } @@ -32,6 +33,8 @@ dependencies { testImplementation(libs.junit.jupiter.api) testImplementation(libs.junit.jupiter.params) testRuntimeOnly(libs.junit.jupiter.engine) + + testFixturesApi(libs.junit.jupiter.api) } tasks.build { diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java new file mode 100644 index 00000000000..e32e17d201f --- /dev/null +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsApi.java @@ -0,0 +1,85 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.encryption.kms; + +import java.util.Arrays; +import java.util.Locale; +import java.util.stream.Collectors; +import org.apache.gravitino.annotation.DeveloperApi; + +/** The KMS product API that defines key identifier and operation semantics. */ +@DeveloperApi +public enum KmsApi { + /** Amazon Web Services Key Management Service. */ + AWS_KMS("aws-kms"), + + /** Google Cloud Key Management Service. */ + GOOGLE_CLOUD_KMS("google-cloud-kms"), + + /** Microsoft Azure Key Vault. */ + AZURE_KEY_VAULT("azure-key-vault"), + + /** The OpenBao Transit secrets engine API. */ + OPENBAO_TRANSIT("openbao-transit"), + + /** The HashiCorp Vault Transit secrets engine API. */ + VAULT_TRANSIT("vault-transit"); + + private final String wireValue; + + KmsApi(String wireValue) { + this.wireValue = wireValue; + } + + /** + * Returns the stable configuration value for this API. + * + * @return the configuration value + */ + public String wireValue() { + return wireValue; + } + + /** + * Parses a configured API value. + * + * @param value configured API value + * @return the matching API + * @throws IllegalArgumentException if the value is null, blank, or unsupported + */ + public static KmsApi fromWireValue(String value) { + if (value == null || value.trim().isEmpty()) { + throw new IllegalArgumentException("KMS API cannot be blank"); + } + + String normalized = value.trim().toLowerCase(Locale.ROOT); + return Arrays.stream(values()) + .filter(api -> api.wireValue.equals(normalized)) + .findFirst() + .orElseThrow( + () -> + new IllegalArgumentException( + String.format( + "Unsupported KMS API '%s'; supported values are %s", + value, + Arrays.stream(values()) + .map(KmsApi::wireValue) + .collect(Collectors.joining(", "))))); + } +} diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java new file mode 100644 index 00000000000..d199c7cbd63 --- /dev/null +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsAuthenticationException.java @@ -0,0 +1,50 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.encryption.kms; + +import com.google.errorprone.annotations.FormatMethod; +import com.google.errorprone.annotations.FormatString; +import org.apache.gravitino.exceptions.ConnectionFailedException; + +/** Indicates that a KMS backend rejected or could not resolve configured credentials. */ +public class KmsAuthenticationException extends ConnectionFailedException { + + /** + * Creates an authentication exception with the specified cause and detail message. + * + * @param cause the cause + * @param message the detail message + * @param args the arguments to the message + */ + @FormatMethod + public KmsAuthenticationException(Throwable cause, @FormatString String message, Object... args) { + super(cause, message, args); + } + + /** + * Creates an authentication exception with the specified detail message. + * + * @param message the detail message + * @param args the arguments to the message + */ + @FormatMethod + public KmsAuthenticationException(@FormatString String message, Object... args) { + super(message, args); + } +} diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java new file mode 100644 index 00000000000..bcc541f98dd --- /dev/null +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsClient.java @@ -0,0 +1,47 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.encryption.kms; + +import org.apache.gravitino.annotation.DeveloperApi; +import org.apache.gravitino.exceptions.ConnectionFailedException; + +/** + * Inspects keys managed by a configured KMS. + * + *

This is a server-side operation client, not a credential-vending API. Provider credentials + * authenticate calls made by the client and must never be returned to callers. This client does not + * perform cryptographic operations. + */ +@DeveloperApi +public interface KmsClient extends AutoCloseable { + + /** + * Reads the provider-reported properties of a key. + * + * @param reference key to inspect + * @return normalized key properties + * @throws IllegalArgumentException if the reference does not belong to this client + * @throws ConnectionFailedException if the provider cannot be queried + */ + KmsKeyProperties getKeyProperties(KmsReference reference); + + /** Releases resources owned by this client. */ + @Override + default void close() {} +} diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java new file mode 100644 index 00000000000..3cfc4898492 --- /dev/null +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsClientFactory.java @@ -0,0 +1,49 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.encryption.kms; + +import java.util.Map; +import org.apache.gravitino.annotation.DeveloperApi; +import org.apache.gravitino.exceptions.ConnectionFailedException; + +/** Creates server-side KMS clients for one supported KMS API. */ +@DeveloperApi +public interface KmsClientFactory { + + /** + * Returns the KMS API implemented by this factory. + * + * @return the supported KMS API + */ + KmsApi api(); + + /** + * Creates a client bound to a configured KMS source. + * + *

Provider credentials are private implementation details of the returned client. They must + * not be exposed as Gravitino credentials or key properties. + * + * @param source logical name of the configured KMS instance + * @param properties provider-specific configuration + * @return the configured client + * @throws IllegalArgumentException if the source or configuration is invalid + * @throws ConnectionFailedException if required external initialization fails + */ + KmsClient create(String source, Map properties); +} diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java new file mode 100644 index 00000000000..ec1613666a6 --- /dev/null +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsConfigurationException.java @@ -0,0 +1,49 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.encryption.kms; + +import com.google.errorprone.annotations.FormatMethod; +import com.google.errorprone.annotations.FormatString; + +/** Indicates invalid KMS configuration detected during initialization. */ +public class KmsConfigurationException extends IllegalArgumentException { + + /** + * Creates a configuration exception with the specified detail message. + * + * @param message the detail message + * @param args the arguments to the message + */ + @FormatMethod + public KmsConfigurationException(@FormatString String message, Object... args) { + super(args.length == 0 ? message : String.format(message, args)); + } + + /** + * Creates a configuration exception with the specified cause and detail message. + * + * @param cause the cause + * @param message the detail message + * @param args the arguments to the message + */ + @FormatMethod + public KmsConfigurationException(Throwable cause, @FormatString String message, Object... args) { + super(args.length == 0 ? message : String.format(message, args), cause); + } +} diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java new file mode 100644 index 00000000000..ae1896a5fa8 --- /dev/null +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsKeyProperties.java @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.encryption.kms; + +import org.apache.gravitino.annotation.DeveloperApi; + +/** + * Common properties reported for a KMS key. + * + *

Capabilities describe downstream use; this API does not perform cryptographic operations. + */ +@DeveloperApi +public interface KmsKeyProperties { + + /** + * Returns the requested key reference. + * + * @return the key reference + */ + KmsReference reference(); + + /** + * Returns whether the provider reports that the key exists. + * + * @return whether the key exists + */ + boolean present(); + + /** + * Returns whether the provider reports that the key is enabled. + * + * @return whether the key is enabled + */ + boolean enabled(); + + /** + * Returns whether the key supports wrapping data-encryption keys. + * + * @return whether wrapping is supported + */ + boolean supportsWrapping(); + + /** + * Returns whether the key supports unwrapping data-encryption keys. + * + * @return whether unwrapping is supported + */ + boolean supportsUnwrapping(); +} diff --git a/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java new file mode 100644 index 00000000000..9350e8844f5 --- /dev/null +++ b/api/src/main/java/org/apache/gravitino/encryption/kms/KmsReference.java @@ -0,0 +1,105 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.encryption.kms; + +import com.google.common.base.Preconditions; +import java.util.Objects; +import javax.annotation.Nullable; +import org.apache.commons.lang3.StringUtils; +import org.apache.gravitino.annotation.DeveloperApi; + +/** + * Identifies a key owned by a configured KMS source. + * + *

Contains no credentials or key material. + */ +@DeveloperApi +public final class KmsReference { + + private final KmsApi api; + private final String source; + private final String keyId; + + /** + * Creates a structurally valid key reference without contacting the provider. + * + * @param api explicitly selected KMS API + * @param source configured KMS client-instance name + * @param keyId provider-native key identifier + * @throws IllegalArgumentException if the API is null or either string is null or blank + */ + public KmsReference(KmsApi api, String source, String keyId) { + Preconditions.checkArgument(api != null, "KMS API cannot be null"); + Preconditions.checkArgument(StringUtils.isNotBlank(source), "KMS source cannot be blank"); + Preconditions.checkArgument(StringUtils.isNotBlank(keyId), "KMS key ID cannot be blank"); + + this.api = api; + this.source = source; + this.keyId = keyId; + } + + /** + * Returns the explicitly selected KMS API. + * + * @return the KMS API + */ + public KmsApi api() { + return api; + } + + /** + * Returns the configured KMS client-instance name. + * + * @return the source name + */ + public String source() { + return source; + } + + /** + * Returns the provider-native key identifier. + * + * @return the key identifier + */ + public String keyId() { + return keyId; + } + + @Override + public boolean equals(@Nullable Object other) { + if (this == other) { + return true; + } + if (!(other instanceof KmsReference)) { + return false; + } + KmsReference that = (KmsReference) other; + return api == that.api && source.equals(that.source) && keyId.equals(that.keyId); + } + + @Override + public int hashCode() { + return Objects.hash(api, source, keyId); + } + + @Override + public String toString() { + return String.format("KmsReference{api=%s, source='%s', keyId='%s'}", api, source, keyId); + } +} diff --git a/api/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java b/api/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java new file mode 100644 index 00000000000..7ad5e2acca8 --- /dev/null +++ b/api/src/test/java/org/apache/gravitino/encryption/kms/TestFakeKmsClient.java @@ -0,0 +1,44 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.encryption.kms; + +public class TestFakeKmsClient extends TestKmsClientContract { + + private static final String SOURCE = "test"; + private static final String USABLE_KEY = "usable"; + private static final String MISSING_KEY = "missing"; + + private final FakeKmsClient client = + new FakeKmsClient(KmsApi.AWS_KMS, SOURCE).putKey(USABLE_KEY, true, true, true); + + @Override + protected KmsClient client() { + return client; + } + + @Override + protected KmsReference usableKey() { + return new KmsReference(KmsApi.AWS_KMS, SOURCE, USABLE_KEY); + } + + @Override + protected KmsReference missingKey() { + return new KmsReference(KmsApi.AWS_KMS, SOURCE, MISSING_KEY); + } +} diff --git a/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsApi.java b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsApi.java new file mode 100644 index 00000000000..cae91fb0303 --- /dev/null +++ b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsApi.java @@ -0,0 +1,47 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.encryption.kms; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class TestKmsApi { + + @Test + void testParsesStableWireValues() { + Assertions.assertEquals(KmsApi.AWS_KMS, KmsApi.fromWireValue("aws-kms")); + Assertions.assertEquals(KmsApi.GOOGLE_CLOUD_KMS, KmsApi.fromWireValue(" GOOGLE-CLOUD-KMS ")); + Assertions.assertEquals("azure-key-vault", KmsApi.AZURE_KEY_VAULT.wireValue()); + Assertions.assertEquals( + KmsApi.OPENBAO_TRANSIT, KmsApi.fromWireValue(KmsApi.OPENBAO_TRANSIT.wireValue())); + Assertions.assertEquals(KmsApi.VAULT_TRANSIT, KmsApi.fromWireValue("vault-transit")); + } + + @Test + void testRejectsUnknownValues() { + Assertions.assertThrows(IllegalArgumentException.class, () -> KmsApi.fromWireValue(null)); + Assertions.assertThrows(IllegalArgumentException.class, () -> KmsApi.fromWireValue(" ")); + + IllegalArgumentException exception = + Assertions.assertThrows( + IllegalArgumentException.class, () -> KmsApi.fromWireValue("custom")); + Assertions.assertTrue(exception.getMessage().contains("Unsupported KMS API 'custom'")); + Assertions.assertTrue(exception.getMessage().contains("openbao-transit")); + } +} diff --git a/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java new file mode 100644 index 00000000000..14fe5d8b4c2 --- /dev/null +++ b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClient.java @@ -0,0 +1,73 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.encryption.kms; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class TestKmsClient { + + private static final KmsReference REFERENCE = + new KmsReference(KmsApi.AWS_KMS, "production", "key"); + + @Test + void testReturnsProviderProperties() { + KmsKeyProperties properties = new TestProperties(); + KmsClient client = reference -> properties; + + Assertions.assertSame(properties, client.getKeyProperties(REFERENCE)); + Assertions.assertTrue(properties.supportsWrapping()); + Assertions.assertFalse(properties.supportsUnwrapping()); + } + + @Test + void testDefaultClose() { + KmsClient client = reference -> new TestProperties(); + + Assertions.assertDoesNotThrow(client::close); + } + + private static final class TestProperties implements KmsKeyProperties { + + @Override + public KmsReference reference() { + return REFERENCE; + } + + @Override + public boolean present() { + return true; + } + + @Override + public boolean enabled() { + return true; + } + + @Override + public boolean supportsWrapping() { + return true; + } + + @Override + public boolean supportsUnwrapping() { + return false; + } + } +} diff --git a/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsReference.java b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsReference.java new file mode 100644 index 00000000000..fe6b566faa0 --- /dev/null +++ b/api/src/test/java/org/apache/gravitino/encryption/kms/TestKmsReference.java @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.encryption.kms; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class TestKmsReference { + + @Test + void testStoresReferenceExactly() { + KmsReference reference = new KmsReference(KmsApi.AWS_KMS, "production", " alias/Customer-Key "); + + Assertions.assertEquals(KmsApi.AWS_KMS, reference.api()); + Assertions.assertEquals("production", reference.source()); + Assertions.assertEquals(" alias/Customer-Key ", reference.keyId()); + } + + @Test + void testRejectsMissingFields() { + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference(null, "production", "key")); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, null, "key")); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, "production", null)); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, "", "key")); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, " ", "key")); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, "production", "")); + Assertions.assertThrows( + IllegalArgumentException.class, () -> new KmsReference(KmsApi.AWS_KMS, "production", " ")); + } + + @Test + void testValueSemantics() { + KmsReference first = new KmsReference(KmsApi.AWS_KMS, "production", "key"); + KmsReference same = new KmsReference(KmsApi.AWS_KMS, "production", "key"); + KmsReference differentApi = new KmsReference(KmsApi.GOOGLE_CLOUD_KMS, "production", "key"); + KmsReference differentSource = new KmsReference(KmsApi.AWS_KMS, "recovery", "key"); + KmsReference differentKey = new KmsReference(KmsApi.AWS_KMS, "production", "another-key"); + + Assertions.assertEquals(first, same); + Assertions.assertEquals(first.hashCode(), same.hashCode()); + Assertions.assertNotEquals(first, differentApi); + Assertions.assertNotEquals(first, differentSource); + Assertions.assertNotEquals(first, differentKey); + Assertions.assertNotEquals(first, null); + Assertions.assertNotEquals(first, "key"); + Assertions.assertEquals( + "KmsReference{api=AWS_KMS, source='production', keyId='key'}", first.toString()); + } +} diff --git a/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java b/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java new file mode 100644 index 00000000000..183ffb37b98 --- /dev/null +++ b/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/FakeKmsClient.java @@ -0,0 +1,142 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.encryption.kms; + +import java.util.HashMap; +import java.util.Map; + +/** In-memory KMS client for contract and consumer tests. */ +public final class FakeKmsClient implements KmsClient { + + private final KmsApi api; + private final String source; + private final Map keys = new HashMap<>(); + + /** + * Creates an empty fake client. + * + * @param api KMS API accepted by the client + * @param source configured source accepted by the client + */ + public FakeKmsClient(KmsApi api, String source) { + this.api = api; + this.source = source; + } + + /** + * Adds or replaces a key reported by this client. + * + * @param keyId provider-native key identifier + * @param enabled whether the key is enabled + * @param supportsWrapping whether the key supports wrapping + * @param supportsUnwrapping whether the key supports unwrapping + * @return this client + */ + public FakeKmsClient putKey( + String keyId, boolean enabled, boolean supportsWrapping, boolean supportsUnwrapping) { + keys.put(keyId, new KeyState(enabled, supportsWrapping, supportsUnwrapping)); + return this; + } + + /** {@inheritDoc} */ + @Override + public KmsKeyProperties getKeyProperties(KmsReference reference) { + requireReference(reference); + KeyState state = keys.get(reference.keyId()); + return state == null + ? new Properties(reference, false, false, false, false) + : new Properties( + reference, true, state.enabled, state.supportsWrapping, state.supportsUnwrapping); + } + + private void requireReference(KmsReference reference) { + if (reference == null) { + throw new IllegalArgumentException("KMS reference cannot be null"); + } + if (reference.api() != api) { + throw new IllegalArgumentException( + String.format("Expected KMS API %s but received %s", api, reference.api())); + } + if (!source.equals(reference.source())) { + throw new IllegalArgumentException( + String.format( + "KMS source %s does not match configured source %s", reference.source(), source)); + } + } + + private static final class KeyState { + + private final boolean enabled; + private final boolean supportsWrapping; + private final boolean supportsUnwrapping; + + private KeyState(boolean enabled, boolean supportsWrapping, boolean supportsUnwrapping) { + this.enabled = enabled; + this.supportsWrapping = supportsWrapping; + this.supportsUnwrapping = supportsUnwrapping; + } + } + + private static final class Properties implements KmsKeyProperties { + + private final KmsReference reference; + private final boolean present; + private final boolean enabled; + private final boolean supportsWrapping; + private final boolean supportsUnwrapping; + + private Properties( + KmsReference reference, + boolean present, + boolean enabled, + boolean supportsWrapping, + boolean supportsUnwrapping) { + this.reference = reference; + this.present = present; + this.enabled = enabled; + this.supportsWrapping = supportsWrapping; + this.supportsUnwrapping = supportsUnwrapping; + } + + @Override + public KmsReference reference() { + return reference; + } + + @Override + public boolean present() { + return present; + } + + @Override + public boolean enabled() { + return enabled; + } + + @Override + public boolean supportsWrapping() { + return supportsWrapping; + } + + @Override + public boolean supportsUnwrapping() { + return supportsUnwrapping; + } + } +} diff --git a/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java b/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java new file mode 100644 index 00000000000..855b4bdf6d9 --- /dev/null +++ b/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientContract.java @@ -0,0 +1,98 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.encryption.kms; + +import java.util.Arrays; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** Common inspection contract for implemented KMS clients. */ +public abstract class TestKmsClientContract { + + /** + * Returns the client under test. + * + * @return the client + */ + protected abstract KmsClient client(); + + /** + * Returns a reference to an enabled key that supports wrapping and unwrapping. + * + * @return the usable key reference + */ + protected abstract KmsReference usableKey(); + + /** + * Returns a reference to a key that does not exist. + * + * @return the missing key reference + */ + protected abstract KmsReference missingKey(); + + @Test + void testReportsUsableKeyProperties() { + KmsReference reference = usableKey(); + KmsKeyProperties properties = client().getKeyProperties(reference); + + Assertions.assertEquals(reference, properties.reference()); + Assertions.assertTrue(properties.present()); + Assertions.assertTrue(properties.enabled()); + Assertions.assertTrue(properties.supportsWrapping()); + Assertions.assertTrue(properties.supportsUnwrapping()); + } + + @Test + void testReportsMissingKey() { + KmsReference reference = missingKey(); + KmsKeyProperties properties = client().getKeyProperties(reference); + + Assertions.assertEquals(reference, properties.reference()); + Assertions.assertFalse(properties.present()); + } + + @Test + void testRejectsNullReference() { + Assertions.assertThrows(IllegalArgumentException.class, () -> client().getKeyProperties(null)); + } + + @Test + void testRejectsMismatchedSource() { + KmsReference reference = usableKey(); + KmsReference mismatched = + new KmsReference(reference.api(), reference.source() + "-other", reference.keyId()); + + Assertions.assertThrows( + IllegalArgumentException.class, () -> client().getKeyProperties(mismatched)); + } + + @Test + void testRejectsMismatchedApi() { + KmsReference reference = usableKey(); + KmsApi otherApi = + Arrays.stream(KmsApi.values()) + .filter(api -> api != reference.api()) + .findFirst() + .orElseThrow(IllegalStateException::new); + KmsReference mismatched = new KmsReference(otherApi, reference.source(), reference.keyId()); + + Assertions.assertThrows( + IllegalArgumentException.class, () -> client().getKeyProperties(mismatched)); + } +} diff --git a/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java b/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java new file mode 100644 index 00000000000..4136e5de41e --- /dev/null +++ b/api/src/testFixtures/java/org/apache/gravitino/encryption/kms/TestKmsClientFactoryContract.java @@ -0,0 +1,45 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.encryption.kms; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** Common contract for KMS client factories. */ +public abstract class TestKmsClientFactoryContract { + + /** + * Returns the factory under test. + * + * @return the factory + */ + protected abstract KmsClientFactory factory(); + + /** + * Returns the API expected from the factory. + * + * @return the expected API + */ + protected abstract KmsApi expectedApi(); + + @Test + void testReportsExpectedApi() { + Assertions.assertEquals(expectedApi(), factory().api()); + } +} From f070523c4ece2fe5f27501b3b5014b9a92cb4e7d Mon Sep 17 00:00:00 2001 From: Nevin Zheng Date: Tue, 21 Jul 2026 15:48:27 -0700 Subject: [PATCH 2/5] feat(kms): configure and route named sources --- .../org/apache/gravitino/GravitinoEnv.java | 37 +- .../encryption/kms/KmsClientRegistry.java | 207 +++++++ .../gravitino/encryption/kms/KmsConfig.java | 152 ++++++ .../TestGravitinoEnvKmsClientRegistry.java | 51 ++ .../encryption/kms/TestKmsClientRegistry.java | 503 ++++++++++++++++++ .../encryption/kms/TestKmsConfig.java | 134 +++++ 6 files changed, 1083 insertions(+), 1 deletion(-) create mode 100644 core/src/main/java/org/apache/gravitino/encryption/kms/KmsClientRegistry.java create mode 100644 core/src/main/java/org/apache/gravitino/encryption/kms/KmsConfig.java create mode 100644 core/src/test/java/org/apache/gravitino/TestGravitinoEnvKmsClientRegistry.java create mode 100644 core/src/test/java/org/apache/gravitino/encryption/kms/TestKmsClientRegistry.java create mode 100644 core/src/test/java/org/apache/gravitino/encryption/kms/TestKmsConfig.java diff --git a/core/src/main/java/org/apache/gravitino/GravitinoEnv.java b/core/src/main/java/org/apache/gravitino/GravitinoEnv.java index 11c12911293..56a3cdfcf2c 100644 --- a/core/src/main/java/org/apache/gravitino/GravitinoEnv.java +++ b/core/src/main/java/org/apache/gravitino/GravitinoEnv.java @@ -56,6 +56,7 @@ import org.apache.gravitino.catalog.ViewNormalizeDispatcher; import org.apache.gravitino.catalog.ViewOperationDispatcher; import org.apache.gravitino.credential.CredentialOperationDispatcher; +import org.apache.gravitino.encryption.kms.KmsClientRegistry; import org.apache.gravitino.hook.AccessControlHookDispatcher; import org.apache.gravitino.hook.CatalogHookDispatcher; import org.apache.gravitino.hook.FilesetHookDispatcher; @@ -154,6 +155,8 @@ public class GravitinoEnv { private CredentialOperationDispatcher credentialOperationDispatcher; + private KmsClientRegistry kmsClientRegistry; + private TagDispatcher tagDispatcher; private PolicyDispatcher policyDispatcher; @@ -222,7 +225,13 @@ public void initializeFullComponents(Config config) { FileFetcher.get().initialize(config.get(Configs.BLOCK_UNSAFE_REMOTE_URI)); this.manageFullComponents = true; initBaseComponents(); - initGravitinoServerComponents(); + try { + this.kmsClientRegistry = new KmsClientRegistry(config); + initGravitinoServerComponents(); + } catch (RuntimeException e) { + closeKmsClientRegistry(); + throw e; + } LOG.info("Gravitino full environment is initialized."); } @@ -418,6 +427,18 @@ public CredentialOperationDispatcher credentialOperationDispatcher() { return credentialOperationDispatcher; } + /** + * Get the metadata-only KMS client registry associated with the full Gravitino environment. + * + * @return The KMS client registry instance. + * @throws IllegalStateException if full components have not been initialized + */ + public KmsClientRegistry kmsClientRegistry() { + Preconditions.checkState( + kmsClientRegistry != null, "GravitinoEnv full components are not initialized."); + return kmsClientRegistry; + } + /** * Get the IdGenerator associated with the Gravitino environment. * @@ -586,6 +607,8 @@ public void start() { public void shutdown() { LOG.info("Shutting down Gravitino Environment..."); + closeKmsClientRegistry(); + if (entityStore != null) { try { entityStore.close(); @@ -638,6 +661,18 @@ public void shutdown() { LOG.info("Gravitino Environment is shut down."); } + private void closeKmsClientRegistry() { + KmsClientRegistry registry = kmsClientRegistry; + kmsClientRegistry = null; + if (registry != null) { + try { + registry.close(); + } catch (RuntimeException e) { + LOG.warn("Failed to close KmsClientRegistry.", e); + } + } + } + private void initBaseComponents() { this.metricsSystem = new MetricsSystem(); metricsSystem.register(new JVMMetricsSource()); diff --git a/core/src/main/java/org/apache/gravitino/encryption/kms/KmsClientRegistry.java b/core/src/main/java/org/apache/gravitino/encryption/kms/KmsClientRegistry.java new file mode 100644 index 00000000000..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); + } + } +} From 198b42978cd1d0c8818317c15ff3bab701f277d9 Mon Sep 17 00:00:00 2001 From: Nevin Zheng Date: Tue, 21 Jul 2026 15:48:40 -0700 Subject: [PATCH 3/5] feat(kms): add Transit key inspection clients --- bundles/transit/build.gradle.kts | 43 +++ .../kms/transit/FileTokenSupplier.java | 80 +++++ .../kms/transit/OpenBaoTransitKmsClient.java | 104 +++++++ .../OpenBaoTransitKmsClientFactory.java | 157 ++++++++++ .../OpenBaoTransitKmsKeyProperties.java | 74 +++++ .../kms/transit/TransitApiClient.java | 216 +++++++++++++ .../gravitino/kms/transit/TransitKeyData.java | 53 ++++ .../kms/transit/TransitReadKeyRequest.java | 35 +++ .../kms/transit/TransitReadKeyResponse.java | 38 +++ .../kms/transit/VaultTransitKmsClient.java | 99 ++++++ .../transit/VaultTransitKmsClientFactory.java | 156 ++++++++++ .../transit/VaultTransitKmsKeyProperties.java | 74 +++++ ....gravitino.encryption.kms.KmsClientFactory | 21 ++ .../transit/TestOpenBaoTransitKmsClient.java | 286 ++++++++++++++++++ .../TestOpenBaoTransitKmsClientFactory.java | 173 +++++++++++ .../transit/TestVaultTransitKmsClient.java | 119 ++++++++ .../TestVaultTransitKmsClientFactory.java | 185 +++++++++++ settings.gradle.kts | 1 + 18 files changed, 1914 insertions(+) create mode 100644 bundles/transit/build.gradle.kts create mode 100644 bundles/transit/src/main/java/org/apache/gravitino/kms/transit/FileTokenSupplier.java create mode 100644 bundles/transit/src/main/java/org/apache/gravitino/kms/transit/OpenBaoTransitKmsClient.java create mode 100644 bundles/transit/src/main/java/org/apache/gravitino/kms/transit/OpenBaoTransitKmsClientFactory.java create mode 100644 bundles/transit/src/main/java/org/apache/gravitino/kms/transit/OpenBaoTransitKmsKeyProperties.java create mode 100644 bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitApiClient.java create mode 100644 bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitKeyData.java create mode 100644 bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitReadKeyRequest.java create mode 100644 bundles/transit/src/main/java/org/apache/gravitino/kms/transit/TransitReadKeyResponse.java create mode 100644 bundles/transit/src/main/java/org/apache/gravitino/kms/transit/VaultTransitKmsClient.java create mode 100644 bundles/transit/src/main/java/org/apache/gravitino/kms/transit/VaultTransitKmsClientFactory.java create mode 100644 bundles/transit/src/main/java/org/apache/gravitino/kms/transit/VaultTransitKmsKeyProperties.java create mode 100644 bundles/transit/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory create mode 100644 bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestOpenBaoTransitKmsClient.java create mode 100644 bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestOpenBaoTransitKmsClientFactory.java create mode 100644 bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestVaultTransitKmsClient.java create mode 100644 bundles/transit/src/test/java/org/apache/gravitino/kms/transit/TestVaultTransitKmsClientFactory.java 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/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") From 6476e5f8d1cf1a32fefa80dfd479c840c8667230 Mon Sep 17 00:00:00 2001 From: Nevin Zheng Date: Tue, 21 Jul 2026 15:48:49 -0700 Subject: [PATCH 4/5] feat(aws): add KMS key inspection --- bundles/aws-bundle/build.gradle.kts | 1 + bundles/aws/build.gradle.kts | 13 + .../gravitino/kms/aws/AwsKmsClient.java | 106 ++++++ .../kms/aws/AwsKmsClientFactory.java | 149 ++++++++ .../kms/aws/AwsKmsKeyProperties.java | 69 ++++ ....gravitino.encryption.kms.KmsClientFactory | 20 ++ .../gravitino/kms/aws/TestAwsKmsClient.java | 322 ++++++++++++++++++ .../kms/aws/TestAwsKmsClientFactory.java | 169 +++++++++ 8 files changed, 849 insertions(+) create mode 100644 bundles/aws/src/main/java/org/apache/gravitino/kms/aws/AwsKmsClient.java create mode 100644 bundles/aws/src/main/java/org/apache/gravitino/kms/aws/AwsKmsClientFactory.java create mode 100644 bundles/aws/src/main/java/org/apache/gravitino/kms/aws/AwsKmsKeyProperties.java create mode 100644 bundles/aws/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory create mode 100644 bundles/aws/src/test/java/org/apache/gravitino/kms/aws/TestAwsKmsClient.java create mode 100644 bundles/aws/src/test/java/org/apache/gravitino/kms/aws/TestAwsKmsClientFactory.java 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; + } +} From 218034b56a5995d3a0c00c03b9bb7fe2a1a43c95 Mon Sep 17 00:00:00 2001 From: Nevin Zheng Date: Tue, 21 Jul 2026 15:48:59 -0700 Subject: [PATCH 5/5] feat(gcp): add Cloud KMS key inspection --- bundles/gcp/build.gradle.kts | 14 ++ .../kms/gcp/GoogleCloudKmsClient.java | 116 +++++++++++ .../kms/gcp/GoogleCloudKmsClientFactory.java | 124 +++++++++++ .../kms/gcp/GoogleCloudKmsKeyMetadata.java | 38 ++++ .../kms/gcp/GoogleCloudKmsKeyProperties.java | 78 +++++++ .../gcp/GoogleCloudKmsMetadataService.java | 29 +++ .../gcp/GoogleCloudKmsSdkMetadataService.java | 85 ++++++++ ....gravitino.encryption.kms.KmsClientFactory | 20 ++ .../kms/gcp/TestGoogleCloudKmsClient.java | 194 +++++++++++++++++ .../gcp/TestGoogleCloudKmsClientFactory.java | 142 +++++++++++++ .../TestGoogleCloudKmsSdkMetadataService.java | 197 ++++++++++++++++++ gradle/libs.versions.toml | 2 + 12 files changed, 1039 insertions(+) create mode 100644 bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsClient.java create mode 100644 bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsClientFactory.java create mode 100644 bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsKeyMetadata.java create mode 100644 bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsKeyProperties.java create mode 100644 bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsMetadataService.java create mode 100644 bundles/gcp/src/main/java/org/apache/gravitino/kms/gcp/GoogleCloudKmsSdkMetadataService.java create mode 100644 bundles/gcp/src/main/resources/META-INF/services/org.apache.gravitino.encryption.kms.KmsClientFactory create mode 100644 bundles/gcp/src/test/java/org/apache/gravitino/kms/gcp/TestGoogleCloudKmsClient.java create mode 100644 bundles/gcp/src/test/java/org/apache/gravitino/kms/gcp/TestGoogleCloudKmsClientFactory.java create mode 100644 bundles/gcp/src/test/java/org/apache/gravitino/kms/gcp/TestGoogleCloudKmsSdkMetadataService.java 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/gradle/libs.versions.toml b/gradle/libs.versions.toml index 9222b53ab3d..c314d99eff5 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -153,6 +153,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" @@ -350,6 +351,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" }