diff --git a/build.gradle.kts b/build.gradle.kts index f72e12af176e..04fd7ee3a218 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -353,6 +353,7 @@ tasks.register("javaioPreCommit") { dependsOn(":sdks:java:io:jms:build") dependsOn(":sdks:java:io:kafka:build") dependsOn(":sdks:java:io:kafka:upgrade:build") + dependsOn(":sdks:java:io:kafka:file-aware-factories:build") dependsOn(":sdks:java:io:kudu:build") dependsOn(":sdks:java:io:mongodb:build") dependsOn(":sdks:java:io:mqtt:build") diff --git a/sdks/java/extensions/kafka-factories/build.gradle b/sdks/java/extensions/kafka-factories/build.gradle new file mode 100644 index 000000000000..30c5d3fd6642 --- /dev/null +++ b/sdks/java/extensions/kafka-factories/build.gradle @@ -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 { id 'org.apache.beam.module' } +applyJavaNature( + automaticModuleName: 'org.apache.beam.sdk.extensions.kafka.factories', + publish: 'False' +) + +description = "Apache Beam :: SDKs :: Java :: Extensions :: Kafka :: Factories" +ext.summary = "Library to instantiate kafka clients with files from GCS or SecretManager." + +dependencies { + // ------------------------- CORE DEPENDENCIES ------------------------- + implementation project(path: ":sdks:java:core", configuration: "shadow") + provided library.java.kafka_clients + implementation 'com.google.cloud:google-cloud-secretmanager:2.72.0' + implementation library.java.slf4j_api + implementation library.java.vendored_guava_32_1_2_jre + implementation project(path: ":sdks:java:extensions:google-cloud-platform-core") + permitUnusedDeclared project(path: ":sdks:java:extensions:google-cloud-platform-core") + // ------------------------- TEST DEPENDENCIES ------------------------- + testImplementation 'org.apache.kafka:kafka-clients:3.9.0' + testImplementation library.java.junit + testImplementation library.java.mockito_core + testRuntimeOnly library.java.mockito_inline + testImplementation project(path: ":sdks:java:core", configuration: "shadowTest") +} diff --git a/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/FileAwareFactoryFn.java b/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/FileAwareFactoryFn.java new file mode 100644 index 000000000000..a0f15b42382d --- /dev/null +++ b/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/FileAwareFactoryFn.java @@ -0,0 +1,258 @@ +/* + * 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.beam.sdk.extensions.kafka.factories; + +import com.google.cloud.secretmanager.v1.AccessSecretVersionResponse; +import com.google.cloud.secretmanager.v1.SecretManagerServiceClient; +import com.google.cloud.secretmanager.v1.SecretVersionName; +import java.io.File; +import java.io.IOException; +import java.nio.channels.FileChannel; +import java.nio.channels.ReadableByteChannel; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.StandardOpenOption; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import org.apache.beam.sdk.io.FileSystems; +import org.apache.beam.sdk.transforms.SerializableFunction; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * An abstract {@link SerializableFunction} that serves as a base class for factories that need to + * process a configuration map to handle external resources like files and secrets. + * + *

This class is designed to be extended by concrete factory implementations (e.g., for creating + * Kafka consumers). It automates the process of detecting special URI strings within the + * configuration values and transforming them before passing the processed configuration to the + * subclass. + * + *

Supported Patterns:

+ * + * + * + *

Usage:

+ * + *

A subclass must implement the {@link #createObject(Map)} method, which receives the fully + * processed configuration map with all paths localized and secrets resolved. Subclasses can also + * override {@link #downloadAndProcessExtraFiles()} to handle specific preliminary file downloads + * (e.g., a krb5.conf file) before the main configuration processing begins. + * + * @param The type of object this factory creates. + */ +public abstract class FileAwareFactoryFn + implements SerializableFunction, T> { + + public static final String SECRET_VALUE_PREFIX = "secretValue:"; + public static final String DIRECTORY_PREFIX = "/tmp"; + private static final Pattern PATH_PATTERN = + Pattern.compile("([a-zA-Z0-9]+://[^\"]+)|(secretValue:[^\"]+)|(secretFile:[^\"]+)"); + + private static final Map secretCache = new ConcurrentHashMap<>(); + + private final String factoryType; + private static final Logger LOG = LoggerFactory.getLogger(FileAwareFactoryFn.class); + + public FileAwareFactoryFn(String factoryType) { + Preconditions.checkNotNull(factoryType); + this.factoryType = factoryType; + } + + protected abstract T createObject(Map config); + + @Override + public T apply(Map config) { + if (config == null) { + return createObject(config); + } + + Map processedConfig = new HashMap<>(config); + + String key = ""; + Object value = null; + try { + downloadAndProcessExtraFiles(); + + for (Map.Entry e : config.entrySet()) { + try { + key = e.getKey(); + value = e.getValue(); + if (value instanceof String) { + String originalValue = (String) value; + Matcher matcher = PATH_PATTERN.matcher(originalValue); + StringBuffer sb = new StringBuffer(); + + while (matcher.find()) { + String externalPath = matcher.group(1); + String secretValue = matcher.group(2); + String secretFile = matcher.group(3); + + if (externalPath != null) { + try { + String tmpPath = replacePathWithLocal(externalPath); + String localPath = downloadExternalFile(externalPath, tmpPath); + matcher.appendReplacement(sb, Matcher.quoteReplacement(localPath)); + LOG.info("Downloaded {} to {}", externalPath, localPath); + } catch (IOException io) { + throw new IOException("Failed to download file : " + externalPath, io); + } + } else if (secretValue != null) { + try { + String secretId = secretValue.substring(SECRET_VALUE_PREFIX.length()); + String processedSecret = + processSecret(originalValue, secretId, getSecretWithCache(secretId)); + + matcher.appendReplacement(sb, Matcher.quoteReplacement(processedSecret)); + } catch (IllegalArgumentException ia) { + throw new IllegalArgumentException("Failed to get secret.", ia); + } + } else if (secretFile != null) { + throw new UnsupportedOperationException("Not yet implemented."); + } + } + matcher.appendTail(sb); + String processedValue = sb.toString(); + processedConfig.put(key, processedValue); + } + } catch (IOException ex) { + throw new RuntimeException("Failed trying to process value for key " + key + ".", ex); + } + } + } catch (IOException e) { + throw new RuntimeException("Failed trying to process extra files.", e); + } + + return createObject(processedConfig); + } + + /** + * A function to download files from their specified external storage path and copy them to the + * provided local filepath. The local filepath is provided by the replacePathWithLocal. + * + * @param externalFilePath + * @param outputFileString + * @return + * @throws IOException + */ + protected static synchronized String downloadExternalFile( + String externalFilePath, String outputFileString) throws IOException { + // create the file only if it doesn't exist + if (new File(outputFileString).exists()) { + return outputFileString; + } + Path outputFilePath = Paths.get(outputFileString); + Path parentDir = outputFilePath.getParent(); + if (parentDir != null) { + Files.createDirectories(parentDir); + } + LOG.info("Staging external file [{}] to [{}]", externalFilePath, outputFileString); + Set options = new HashSet<>(2); + options.add(StandardOpenOption.CREATE); + options.add(StandardOpenOption.WRITE); + + // Copy the external file into a local file and will throw an I/O exception in case file not + // found. + try (ReadableByteChannel readerChannel = + FileSystems.open(FileSystems.matchSingleFileSpec(externalFilePath).resourceId())) { + try (FileChannel writeChannel = FileChannel.open(outputFilePath, options)) { + writeChannel.transferFrom(readerChannel, 0, Long.MAX_VALUE); + } + } + return outputFileString; + } + + protected byte[] getSecretWithCache(String secretId) { + return secretCache.computeIfAbsent(secretId, this::getSecret); + } + + /** + * A helper method to create a new string with the external paths replaced with their local path + * and subdirectory based on the factory type in the /tmp directory. For example, the kerberos + * factory type will replace the file paths with /tmp/kerberos/file.path + * + * @param externalPath + * @return a string with all instances of external paths converted to the local paths where the + * files sit. + */ + private String replacePathWithLocal(String externalPath) throws IOException { + String externalBucketPrefixIdentifier = "://"; + int externalBucketPrefixIndex = externalPath.lastIndexOf(externalBucketPrefixIdentifier); + if (externalBucketPrefixIndex == -1) { + // if we don't find a known bucket prefix then we will error early. + throw new RuntimeException( + "The provided external bucket could not be matched to a known source."); + } + + int prefixLength = externalBucketPrefixIndex + externalBucketPrefixIdentifier.length(); + return DIRECTORY_PREFIX + "/" + factoryType + "/" + externalPath.substring(prefixLength); + } + + /** + * @throws IOException A hook for subclasses to download and process specific files before the + * main configuration is handled. For example, the kerberos factory can use this to download a + * krb5.conf and set a system property. + */ + protected void downloadAndProcessExtraFiles() throws IOException { + // Default implementation should do nothing. + } + + protected String getBaseDirectory() { + return DIRECTORY_PREFIX; + } + + protected byte[] getSecret(String secretVersion) { + SecretVersionName secretVersionName; + if (SecretVersionName.isParsableFrom(secretVersion)) { + secretVersionName = SecretVersionName.parse(secretVersion); + } else { + throw new IllegalArgumentException( + "Provided Secret must be in the form" + + " projects/{project}/secrets/{secret}/versions/{secret_version}"); + } + try (SecretManagerServiceClient client = SecretManagerServiceClient.create()) { + AccessSecretVersionResponse response = client.accessSecretVersion(secretVersionName); + return response.getPayload().getData().toByteArray(); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + protected String processSecret(String originalValue, String secretId, byte[] secretValue) { + // By Default, this will return the secret value directly. This function can be overridden by + // derived classes. + return new String(secretValue, StandardCharsets.UTF_8); + } +} diff --git a/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/KerberosConsumerFactoryFn.java b/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/KerberosConsumerFactoryFn.java new file mode 100644 index 000000000000..409904b667f9 --- /dev/null +++ b/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/KerberosConsumerFactoryFn.java @@ -0,0 +1,150 @@ +/* + * 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.beam.sdk.extensions.kafka.factories; + +import java.io.File; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.attribute.PosixFilePermission; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import javax.security.auth.login.Configuration; +import org.apache.kafka.clients.consumer.Consumer; +import org.apache.kafka.clients.consumer.KafkaConsumer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class KerberosConsumerFactoryFn extends FileAwareFactoryFn> { + private static final String LOCAL_FACTORY_TYPE = "kerberos"; + private String krb5ConfigPath = ""; + private static volatile String localKrb5ConfPath = ""; + + private static final Object lock = new Object(); + + // Standard Kafka property for SASL JAAS configuration + private static final String JAAS_CONFIG_PROPERTY = "sasl.jaas.config"; + private static final String KEYTAB_SECRET_PREFIX = "keyTab=\"secretValue:"; + private static final Pattern KEYTAB_SECRET_PATTERN = + Pattern.compile("(keyTab=\"secretValue:[^\"]+)"); + + private static final Logger LOG = LoggerFactory.getLogger(KerberosConsumerFactoryFn.class); + + public KerberosConsumerFactoryFn(String krb5ConfigPath) { + super("kerberos"); + this.krb5ConfigPath = krb5ConfigPath; + } + + @Override + protected Consumer createObject(Map config) { + // This will be called after the config map processing has occurred. Therefore, we know that the + // property will have had it's value replaced with a local directory. + // We don't need to worry about the external bucket prefix in this case. + try { + String jaasConfig = (String) config.get(JAAS_CONFIG_PROPERTY); + String localKeytabPath = ""; + if (jaasConfig != null && !jaasConfig.isEmpty()) { + localKeytabPath = + jaasConfig.substring( + jaasConfig.indexOf("keyTab=\"") + 8, jaasConfig.lastIndexOf("\" principal")); + } + + // Set the permissions on the file to be as strict as possible for security reasons. The + // keytab contains sensitive information and should be as locked down as possible. + Path path = Paths.get(localKeytabPath); + Set perms = new HashSet<>(); + perms.add(PosixFilePermission.OWNER_READ); + Files.setPosixFilePermissions(path, perms); + } catch (IOException e) { + throw new RuntimeException( + "Could not access keytab file. Make sure that the sasl.jaas.config config property " + + "is set correctly.", + e); + } + return new KafkaConsumer<>(config); + } + + @Override + protected void downloadAndProcessExtraFiles() throws IOException { + synchronized (lock) { + // we only want a new krb5 file if there is not already one present. + if (localKrb5ConfPath.isEmpty()) { + if (this.krb5ConfigPath != null && !this.krb5ConfigPath.isEmpty()) { + String localPath = + super.getBaseDirectory() + "/" + LOCAL_FACTORY_TYPE + "/" + "krb5.conf"; + localKrb5ConfPath = downloadExternalFile(this.krb5ConfigPath, localPath); + + System.setProperty("java.security.krb5.conf", localKrb5ConfPath); + Configuration.getConfiguration().refresh(); + } + } + } + } + + @Override + protected String processSecret(String originalValue, String secretId, byte[] secretValue) + throws RuntimeException { + Matcher matcher = KEYTAB_SECRET_PATTERN.matcher(originalValue); + String localFileString = ""; + while (matcher.find()) { + String currentSecretId = matcher.group(1); + if (currentSecretId == null || currentSecretId.isEmpty()) { + throw new RuntimeException( + "Error matching values. Secret was discovered but its value is null"); + } + currentSecretId = currentSecretId.substring(KEYTAB_SECRET_PREFIX.length()); + if (!currentSecretId.equals(secretId)) { + // A sasl.jaas.config can contain multiple keytabs in one string. Therefore, we must assume + // that there can + // also be multiple keytab secrets in the same string. If the currently matched secret does + // not equal + // the secret that we are processing (passed in via secretId) then we do not want to create + // a keytab file and overwrite it. + continue; + } + String filename = "kafka-client-" + UUID.randomUUID().toString() + ".keytab"; + + localFileString = super.getBaseDirectory() + "/" + LOCAL_FACTORY_TYPE + "/" + filename; + Path localFilePath = Paths.get(localFileString); + Path parentDir = localFilePath.getParent(); + try { + if (parentDir != null) { + Files.createDirectories(parentDir); + } + Files.write(localFilePath, secretValue); + if (!new File(localFileString).canRead()) { + LOG.warn("The file is not readable"); + } + LOG.info("Successfully wrote file to path: {}", localFilePath); + } catch (IOException e) { + throw new RuntimeException("Unable to create the keytab file for the provided secret."); + } + } + // if no localFile was created, then we can assume that the secret is meant to be kept as a + // value. + return localFileString.isEmpty() + ? new String(secretValue, StandardCharsets.UTF_8) + : localFileString; + } +} diff --git a/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/package-info.java b/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/package-info.java new file mode 100644 index 000000000000..da12c8203a64 --- /dev/null +++ b/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/package-info.java @@ -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. + */ + +/** ConsumerFactoryFns for file paths that exist in GCS or Google SecretManager. */ +package org.apache.beam.sdk.extensions.kafka.factories; diff --git a/sdks/java/extensions/kafka-factories/src/test/java/org/apache/beam/sdk/extensions/kafka/factories/FileAwareFactoryFnTest.java b/sdks/java/extensions/kafka-factories/src/test/java/org/apache/beam/sdk/extensions/kafka/factories/FileAwareFactoryFnTest.java new file mode 100644 index 000000000000..0ad096e856dc --- /dev/null +++ b/sdks/java/extensions/kafka-factories/src/test/java/org/apache/beam/sdk/extensions/kafka/factories/FileAwareFactoryFnTest.java @@ -0,0 +1,268 @@ +/* + * 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.beam.sdk.extensions.kafka.factories; + +import java.io.ByteArrayInputStream; +import java.io.File; +import java.io.IOException; +import java.nio.channels.Channels; +import java.nio.channels.FileChannel; +import java.nio.channels.ReadableByteChannel; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import org.apache.beam.sdk.io.FileSystems; +import org.apache.beam.sdk.io.fs.MatchResult; +import org.apache.beam.sdk.io.fs.ResourceId; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.ArgumentMatchers; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +@RunWith(JUnit4.class) +public class FileAwareFactoryFnTest { + + @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder(); + + private TestFactoryFn factory; + private String baseDir; + private static final String TEST_FACTORY_TYPE = "test-factory"; + + // A concrete implementation for testing the abstract FileAwareFactoryFn + static class TestFactoryFn extends FileAwareFactoryFn { + public TestFactoryFn() { + super(TEST_FACTORY_TYPE); + } + + @Override + protected Object createObject(Map config) { + // Return the processed config for easy assertion + return config; + } + } + + @Before + public void setup() throws IOException { + baseDir = "/tmp/" + TEST_FACTORY_TYPE; + factory = Mockito.spy(new TestFactoryFn()); + Mockito.doReturn(baseDir).when(factory).getBaseDirectory(); + } + + @Test + public void testHappyPathReplacesExternalPath() { + // Arrange + String gcsPath = "gs://test-bucket/config-file.json"; + String expectedLocalPath = + FileAwareFactoryFn.DIRECTORY_PREFIX + + "/" + + TEST_FACTORY_TYPE + + "/test-bucket/config-file.json"; + Map config = new HashMap<>(); + config.put("config.file.path", gcsPath); + + // Act & Assert + // Use try-with-resources to manage the scope of the static mock on FileSystems + try (MockedStatic mockedFileSystems = Mockito.mockStatic(FileSystems.class)) { + // 1. Mock the underlying static FileSystems calls to avoid real network I/O + MatchResult.Metadata metadata = Mockito.mock(MatchResult.Metadata.class); + ResourceId resourceId = Mockito.mock(ResourceId.class); + Mockito.when(metadata.resourceId()).thenReturn(resourceId); + mockedFileSystems.when(() -> FileSystems.matchSingleFileSpec(gcsPath)).thenReturn(metadata); + + // 2. Mock 'open' to return a channel with no data, simulating a successful download + ReadableByteChannel channel = Channels.newChannel(new ByteArrayInputStream(new byte[0])); + mockedFileSystems.when(() -> FileSystems.open(resourceId)).thenReturn(channel); + + // Act + Map processedConfig = (Map) factory.apply(config); + + // Assert + Assert.assertEquals(expectedLocalPath, processedConfig.get("config.file.path")); + Assert.assertTrue( + "Local file should have been created", new File(expectedLocalPath).exists()); + } + } + + @Test + public void testApplyFailurePathThrowsRuntimeExceptionOnDownloadFailure() { + // Arrange + String gcsPath = "gs://test-bucket/failing-file.txt"; + Map config = new HashMap<>(); + config.put("critical.file", gcsPath); + + // Mock the static FileSystems.matchSingleFileSpec to throw an exception + try (MockedStatic mockedFileSystems = Mockito.mockStatic(FileSystems.class)) { + mockedFileSystems + .when(() -> FileSystems.matchSingleFileSpec(gcsPath)) + .thenThrow(new IOException("GCS file not found")); + // Act & Assert + RuntimeException exception = + Assert.assertThrows(RuntimeException.class, () -> factory.apply(config)); + Assert.assertTrue(exception.getMessage().contains("Failed trying to process value")); + Assert.assertTrue(exception.getCause() instanceof IOException); + Assert.assertTrue(exception.getCause().getMessage().contains("Failed to download file")); + } + } + + @Test + public void testApplyHappyPathIgnoresNonExternalValues() { + // Arrange + Map config = new HashMap<>(); + config.put("some.string", "/local/path/file.txt"); + config.put("some.number", 42); + config.put("some.boolean", false); + + // Act + Map processedConfig = (Map) factory.apply(config); + + // Assert + Assert.assertEquals(config, processedConfig); + } + + @Test + public void testApplyEdgeCaseMultipleExternalPathsInSingleValue() { + // Arrange + String gcsPath1 = "gs://bucket/keytab.keytab"; + String gcsPath2 = "gs://bucket/trust.jks"; + String originalValue = + "jaas_config keyTab=\"" + gcsPath1 + "\" trustStore=\"" + gcsPath2 + "\""; + + String expectedLocalPath1 = + FileAwareFactoryFn.DIRECTORY_PREFIX + "/" + TEST_FACTORY_TYPE + "/bucket/keytab.keytab"; + String expectedLocalPath2 = + FileAwareFactoryFn.DIRECTORY_PREFIX + "/" + TEST_FACTORY_TYPE + "/bucket/trust.jks"; + String expectedProcessedValue = + "jaas_config keyTab=\"" + + expectedLocalPath1 + + "\" trustStore=\"" + + expectedLocalPath2 + + "\""; + + Map config = new HashMap<>(); + config.put("jaas.config", originalValue); + + try (MockedStatic mockedFileSystems = Mockito.mockStatic(FileSystems.class)) { + // Mock GCS calls for both paths + mockSuccessfulDownload(mockedFileSystems, gcsPath1); + mockSuccessfulDownload(mockedFileSystems, gcsPath2); + + // Act + Map processedConfig = (Map) factory.apply(config); + + // Assert + Assert.assertEquals(expectedProcessedValue, processedConfig.get("jaas.config")); + } + } + + @Test + public void testApplyEdgeCaseLocalFileWriteFails() throws IOException { + // Arrange + String gcsPath = "gs://test-bucket/some-file.txt"; + Map config = new HashMap<>(); + config.put("a.file", gcsPath); + + // Mock GCS part to succeed + try (MockedStatic mockedFileSystems = Mockito.mockStatic(FileSystems.class); + MockedStatic mockedFileChannel = Mockito.mockStatic(FileChannel.class)) { + mockSuccessfulDownload(mockedFileSystems, gcsPath); + + // Mock the local file writing part to fail + mockedFileChannel + .when( + () -> + FileChannel.open( + ArgumentMatchers.any(Path.class), ArgumentMatchers.any(Set.class))) + .thenThrow(new IOException("Permission denied")); + + // Act & Assert + RuntimeException exception = + Assert.assertThrows(RuntimeException.class, () -> factory.apply(config)); + Assert.assertTrue(exception.getMessage().contains("Failed trying to process value")); + Assert.assertTrue(exception.getCause() instanceof IOException); + // Check that the root cause is our "Permission denied" mock + Assert.assertTrue(exception.getCause().getCause().getMessage().contains("Permission denied")); + } + } + + @Test + public void testApplyHappyPathResolvesSecretValue() { + // Arrange + String secretVersion = "secretValue:projects/p/secrets/s/versions/v"; + String secretVersionParsed = "projects/p/secrets/s/versions/v"; + String secretValue = "my-secret-password"; + String originalValue = "password=" + secretVersion; + String expectedProcessedValue = "password=" + secretValue; + + Map config = new HashMap<>(); + config.put("db.password", originalValue); + + TestFactoryFn factoryWithMockedSecret = + new TestFactoryFn() { + @Override + public byte[] getSecret(String secretIdentifier) { + // Assert that the correct identifier is passed + Assert.assertEquals(secretVersionParsed, secretIdentifier); + // Return a predictable, hardcoded value for the test + return secretValue.getBytes(StandardCharsets.UTF_8); + } + }; + + // Act + @SuppressWarnings("unchecked") + Map processedConfig = + (Map) factoryWithMockedSecret.apply(config); + + // Assert + Assert.assertEquals(expectedProcessedValue, processedConfig.get("db.password")); + } + + @Test + public void testApplyFailurePathThrowsExceptionForInvalidSecretFormat() { + // Arrange + String invalidSecret = "secretValue:not-a-valid-secret-path"; + Map config = new HashMap<>(); + config.put("db.password", "password=" + invalidSecret); + + // Act & Assert + RuntimeException ex = Assert.assertThrows(RuntimeException.class, () -> factory.apply(config)); + Assert.assertEquals(IllegalArgumentException.class, ex.getCause().getClass()); + } + + // Helper method to reduce boilerplate in mocking successful GCS downloads + private void mockSuccessfulDownload(MockedStatic mockedFileSystems, String gcsPath) { + MatchResult.Metadata metadata = Mockito.mock(MatchResult.Metadata.class); + ResourceId resourceId = Mockito.mock(ResourceId.class); + Mockito.when(metadata.resourceId()).thenReturn(resourceId); + mockedFileSystems + .when(() -> FileSystems.matchSingleFileSpec(ArgumentMatchers.eq(gcsPath))) + .thenReturn(metadata); + + ReadableByteChannel channel = Channels.newChannel(new ByteArrayInputStream(new byte[0])); + mockedFileSystems + .when(() -> FileSystems.open(ArgumentMatchers.eq(resourceId))) + .thenReturn(channel); + } +} diff --git a/sdks/java/extensions/kafka-factories/src/test/java/org/apache/beam/sdk/extensions/kafka/factories/KerberosConsumerFactoryFnTest.java b/sdks/java/extensions/kafka-factories/src/test/java/org/apache/beam/sdk/extensions/kafka/factories/KerberosConsumerFactoryFnTest.java new file mode 100644 index 000000000000..503b2f8f10c0 --- /dev/null +++ b/sdks/java/extensions/kafka-factories/src/test/java/org/apache/beam/sdk/extensions/kafka/factories/KerberosConsumerFactoryFnTest.java @@ -0,0 +1,242 @@ +/* + * 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.beam.sdk.extensions.kafka.factories; + +import static org.mockito.Mockito.spy; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Comparator; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import java.util.stream.Stream; +import javax.security.auth.login.Configuration; +import org.apache.kafka.clients.consumer.KafkaConsumer; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.ArgumentCaptor; +import org.mockito.ArgumentMatchers; +import org.mockito.MockedConstruction; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +@RunWith(JUnit4.class) +public class KerberosConsumerFactoryFnTest { + + private KerberosConsumerFactoryFn factory; + private String originalKrb5Conf; + private static final String KRB5_GCS_PATH = "gs://sec-bucket/kerberos/krb5.conf"; + private static final String KRB5_S3_PATH = "s3://sec-bucket/kerberos/krb5.conf"; + private static final String LOCAL_FACTORY_TYPE = "kerberos"; + + @Before + public void setup() { + try { + java.lang.reflect.Field field = + KerberosConsumerFactoryFn.class.getDeclaredField("localKrb5ConfPath"); + field.setAccessible(true); + field.set(null, ""); + } catch (Exception e) { + throw new RuntimeException(e); + } + originalKrb5Conf = System.getProperty("java.security.krb5.conf"); + } + + @After + public void tearDown() throws IOException { + // Clean up system property to avoid affecting other tests + if (originalKrb5Conf != null) { + System.setProperty("java.security.krb5.conf", originalKrb5Conf); + } else { + System.clearProperty("java.security.krb5.conf"); + } + + // Clean up the directory created outside of the JUnit TemporaryFolder rule. + Path pathToDelete = Paths.get(FileAwareFactoryFn.DIRECTORY_PREFIX, LOCAL_FACTORY_TYPE); + if (Files.exists(pathToDelete)) { + try (Stream walk = Files.walk(pathToDelete)) { + walk.sorted(Comparator.reverseOrder()).map(Path::toFile).forEach(File::delete); + } + } + } + + @Test + @SuppressWarnings("rawtypes") + public void testHappyGcsPath() { + String keytabGcsPath = "gs://sec-bucket/keytabs/my.keytab"; + String expectedKrb5LocalPath = "/tmp/kerberos/sec-bucket/kerberos/krb5.conf"; + String expectedKeytabLocalPath = "/tmp/kerberos/sec-bucket/keytabs/my.keytab"; + + Map config = new HashMap<>(); + config.put( + "sasl.jaas.config", + "com.sun.security.auth.module.Krb5LoginModule required keyTab=\"" + + keytabGcsPath + + "\" principal=\"user@REALM\";"); + + factory = spy(new KerberosConsumerFactoryFn(KRB5_GCS_PATH)); + // This mock prevents the spy from calling the real createObject method, + // which would otherwise crash. + Mockito.doReturn(null).when(factory).createObject(ArgumentMatchers.anyMap()); + + try (MockedStatic mockedStaticFactory = + Mockito.mockStatic(FileAwareFactoryFn.class, Mockito.CALLS_REAL_METHODS); + MockedStatic mockedConfiguration = Mockito.mockStatic(Configuration.class); + MockedStatic mockedFiles = Mockito.mockStatic(Files.class); + MockedConstruction mockedConsumer = + Mockito.mockConstruction(KafkaConsumer.class)) { + + Assert.assertNotNull(mockedConsumer); + // Mock the static downloadExternalFile method to prevent any GCS interaction + mockedStaticFactory + .when( + () -> + FileAwareFactoryFn.downloadExternalFile( + ArgumentMatchers.eq(KRB5_GCS_PATH), ArgumentMatchers.anyString())) + .thenReturn(expectedKrb5LocalPath); + mockedStaticFactory + .when( + () -> + FileAwareFactoryFn.downloadExternalFile( + ArgumentMatchers.eq(keytabGcsPath), ArgumentMatchers.anyString())) + .thenReturn(expectedKeytabLocalPath); + + Configuration mockConf = Mockito.mock(Configuration.class); + mockedConfiguration.when(Configuration::getConfiguration).thenReturn(mockConf); + mockedFiles + .when( + () -> + Files.setPosixFilePermissions( + ArgumentMatchers.any(Path.class), ArgumentMatchers.any(Set.class))) + .thenReturn(null); + mockedFiles + .when(() -> Files.createDirectories(ArgumentMatchers.any(Path.class))) + .thenReturn(null); + + // Act + factory.apply(config); + + // Assert + // 1. Verify that the krb5.conf system property was set correctly. + Assert.assertEquals(expectedKrb5LocalPath, System.getProperty("java.security.krb5.conf")); + + // 2. Capture the config passed to createObject and verify the keytab path was replaced. + ArgumentCaptor> configCaptor = ArgumentCaptor.forClass(Map.class); + Mockito.verify(factory).createObject(configCaptor.capture()); + Map capturedConfig = configCaptor.getValue(); + String processedJaasConfig = (String) capturedConfig.get("sasl.jaas.config"); + Assert.assertTrue(processedJaasConfig.contains("keyTab=\"" + expectedKeytabLocalPath + "\"")); + + // 3. Verify that the JAAS configuration was refreshed. + Mockito.verify(mockConf).refresh(); + } + } + + @Test + @SuppressWarnings("rawtypes") + public void testHappyS3Path() { + String keytabPath = "s3://sec-bucket/keytabs/my.keytab"; + String expectedKrb5LocalPath = "/tmp/kerberos/sec-bucket/kerberos/krb5.conf"; + String expectedKeytabLocalPath = "/tmp/kerberos/sec-bucket/keytabs/my.keytab"; + + Map config = new HashMap<>(); + config.put( + "sasl.jaas.config", + "com.sun.security.auth.module.Krb5LoginModule required keyTab=\"" + + keytabPath + + "\" principal=\"user@REALM\";"); + factory = spy(new KerberosConsumerFactoryFn(KRB5_S3_PATH)); + // This mock prevents the spy from calling the real createObject method, + // which would otherwise crash. + Mockito.doReturn(null).when(factory).createObject(ArgumentMatchers.anyMap()); + + try (MockedStatic mockedStaticFactory = + Mockito.mockStatic(FileAwareFactoryFn.class, Mockito.CALLS_REAL_METHODS); + MockedStatic mockedConfiguration = Mockito.mockStatic(Configuration.class); + MockedStatic mockedFiles = Mockito.mockStatic(Files.class); + MockedConstruction mockedConsumer = + Mockito.mockConstruction(KafkaConsumer.class)) { + + Assert.assertNotNull(mockedConsumer); + // Mock the static downloadExternalFile method to prevent any interaction + mockedStaticFactory + .when( + () -> + FileAwareFactoryFn.downloadExternalFile( + ArgumentMatchers.eq(KRB5_S3_PATH), ArgumentMatchers.anyString())) + .thenReturn(expectedKrb5LocalPath); + mockedStaticFactory + .when( + () -> + FileAwareFactoryFn.downloadExternalFile( + ArgumentMatchers.eq(keytabPath), ArgumentMatchers.anyString())) + .thenReturn(expectedKeytabLocalPath); + + Configuration mockConf = Mockito.mock(Configuration.class); + mockedConfiguration.when(Configuration::getConfiguration).thenReturn(mockConf); + mockedFiles + .when( + () -> + Files.setPosixFilePermissions( + ArgumentMatchers.any(Path.class), ArgumentMatchers.any(Set.class))) + .thenReturn(null); + mockedFiles + .when(() -> Files.createDirectories(ArgumentMatchers.any(Path.class))) + .thenReturn(null); + + // Act + factory.apply(config); + + // Assert + // 1. Verify that the krb5.conf system property was set correctly. + Assert.assertEquals(expectedKrb5LocalPath, System.getProperty("java.security.krb5.conf")); + + // 2. Capture the config passed to createObject and verify the keytab path was replaced. + ArgumentCaptor> configCaptor = ArgumentCaptor.forClass(Map.class); + Mockito.verify(factory).createObject(configCaptor.capture()); + Map capturedConfig = configCaptor.getValue(); + String processedJaasConfig = (String) capturedConfig.get("sasl.jaas.config"); + Assert.assertTrue(processedJaasConfig.contains("keyTab=\"" + expectedKeytabLocalPath + "\"")); + + // 3. Verify that the JAAS configuration was refreshed. + Mockito.verify(mockConf).refresh(); + } + } + + @Test + public void testInvalidKrb5ConfPathThrowsException() { + // Arrange + String invalidPath = "not-a-gcs-path"; // This path is missing the "gs://" prefix + factory = new KerberosConsumerFactoryFn(invalidPath); + Map config = new HashMap<>(); + + // Act & Assert + RuntimeException ex = Assert.assertThrows(RuntimeException.class, () -> factory.apply(config)); + + Assert.assertTrue(ex.getMessage().contains("Failed trying to process extra files")); + Assert.assertTrue(ex.getCause() instanceof IOException); + } +} diff --git a/settings.gradle.kts b/settings.gradle.kts index c867e7ae2314..d887c8ad9d41 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -186,6 +186,7 @@ include(":sdks:java:extensions:kryo") include(":sdks:java:extensions:google-cloud-platform-core") include(":sdks:java:extensions:jackson") include(":sdks:java:extensions:join-library") +include(":sdks:java:extensions:kafka-factories") include(":sdks:java:extensions:ml") include(":sdks:java:extensions:ordered") include(":sdks:java:extensions:protobuf")