Skip to content

Commit d642870

Browse files
fozzie15fozzie15
committed
Add tests for file aware factory fn
1 parent 7d03e90 commit d642870

4 files changed

Lines changed: 825 additions & 0 deletions

File tree

‎sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/FileAwareFactoryFn.java‎

Lines changed: 239 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,239 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.sdk.io.kafka;
19+
20+
/*
21+
* Licensed to the Apache Software Foundation (ASF) under one
22+
* or more contributor license agreements. See the NOTICE file
23+
* distributed with this work for additional information
24+
* regarding copyright ownership. The ASF licenses this file
25+
* to you under the Apache License, Version 2.0 (the
26+
* "License"); you may not use this file except in compliance
27+
* with the License. You may obtain a copy of the License at
28+
*
29+
* http://www.apache.org/licenses/LICENSE-2.0
30+
*
31+
* Unless required by applicable law or agreed to in writing, software
32+
* distributed under the License is distributed on an "AS IS" BASIS,
33+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
34+
* See the License for the specific language governing permissions and
35+
* limitations under the License.
36+
*/
37+
38+
import com.google.cloud.secretmanager.v1.AccessSecretVersionResponse;
39+
import com.google.cloud.secretmanager.v1.SecretManagerServiceClient;
40+
import com.google.cloud.secretmanager.v1.SecretVersionName;
41+
import com.google.common.base.Preconditions;
42+
import java.io.File;
43+
import java.io.IOException;
44+
import java.nio.channels.FileChannel;
45+
import java.nio.channels.ReadableByteChannel;
46+
import java.nio.charset.StandardCharsets;
47+
import java.nio.file.Files;
48+
import java.nio.file.Path;
49+
import java.nio.file.Paths;
50+
import java.nio.file.StandardOpenOption;
51+
import java.util.HashMap;
52+
import java.util.HashSet;
53+
import java.util.Map;
54+
import java.util.Set;
55+
import java.util.concurrent.ConcurrentHashMap;
56+
import java.util.regex.Matcher;
57+
import java.util.regex.Pattern;
58+
import org.apache.beam.sdk.io.FileSystems;
59+
import org.apache.beam.sdk.transforms.SerializableFunction;
60+
import org.slf4j.Logger;
61+
import org.slf4j.LoggerFactory;
62+
63+
public abstract class FileAwareFactoryFn<T>
64+
implements SerializableFunction<Map<String, Object>, T> {
65+
66+
public static final String GCS_PATH_PREFIX = "gs://";
67+
public static final String SECRET_VALUE_PREFIX = "secretValue:";
68+
public static final String DIRECTORY_PREFIX = "/tmp";
69+
private static final Pattern PATH_PATTERN =
70+
Pattern.compile("(gs://[^\"]+)|(secretValue:[^\"]+)|(secretFile:[^\"]+)");
71+
72+
private static final Map<String, byte[]> secretCache = new ConcurrentHashMap<>();
73+
74+
private final String factoryType;
75+
private static final Logger LOG = LoggerFactory.getLogger(FileAwareFactoryFn.class);
76+
77+
public FileAwareFactoryFn(String factoryType) {
78+
Preconditions.checkNotNull(factoryType);
79+
this.factoryType = factoryType;
80+
}
81+
82+
protected abstract T createObject(Map<String, Object> config);
83+
84+
@Override
85+
public T apply(Map<String, Object> config) {
86+
if (config == null) {
87+
return createObject(config);
88+
}
89+
90+
Map<String, Object> processedConfig = new HashMap<>(config);
91+
92+
String key = "";
93+
Object value = null;
94+
try {
95+
downloadAndProcessExtraFiles();
96+
97+
for (Map.Entry<String, Object> e : config.entrySet()) {
98+
try {
99+
key = e.getKey();
100+
value = e.getValue();
101+
if (value instanceof String) {
102+
String originalValue = (String) value;
103+
Matcher matcher = PATH_PATTERN.matcher(originalValue);
104+
StringBuffer sb = new StringBuffer();
105+
106+
while (matcher.find()) {
107+
String gcsPath = matcher.group(1);
108+
String secretValue = matcher.group(2);
109+
String secretFile = matcher.group(3);
110+
111+
if (gcsPath != null) {
112+
try {
113+
String tmpPath = replacePathWithLocal(gcsPath);
114+
String localPath = downloadGcsFile(gcsPath, tmpPath);
115+
matcher.appendReplacement(sb, Matcher.quoteReplacement(localPath));
116+
LOG.info("Downloaded {} to {}", gcsPath, localPath);
117+
} catch (IOException io) {
118+
throw new IOException("Failed to download file : " + gcsPath, io);
119+
}
120+
} else if (secretValue != null) {
121+
try {
122+
String secretId = secretValue.substring(SECRET_VALUE_PREFIX.length());
123+
String processedSecret =
124+
processSecret(originalValue, secretId, getSecretWithCache(secretId));
125+
126+
matcher.appendReplacement(sb, Matcher.quoteReplacement(processedSecret));
127+
} catch (IllegalArgumentException ia) {
128+
throw new IllegalArgumentException("Failed to get secret.", ia);
129+
}
130+
} else if (secretFile != null) {
131+
throw new UnsupportedOperationException("Not yet implemented.");
132+
}
133+
}
134+
matcher.appendTail(sb);
135+
String processedValue = sb.toString();
136+
processedConfig.put(key, processedValue);
137+
}
138+
} catch (IOException ex) {
139+
throw new RuntimeException(
140+
"Failed trying to process value " + value + " for key " + key + ".", ex);
141+
}
142+
}
143+
} catch (IOException e) {
144+
throw new RuntimeException("Failed trying to process extra files.", e);
145+
}
146+
147+
LOG.info("ProcessedConfig: {}", processedConfig);
148+
return createObject(processedConfig);
149+
}
150+
151+
/**
152+
* A function to download files from their specified gcs path and copy them to the provided local
153+
* filepath. The local filepath is provided by the replacePathWithLocal.
154+
*
155+
* @param gcsFilePath
156+
* @param outputFileString
157+
* @return
158+
* @throws IOException
159+
*/
160+
protected static synchronized String downloadGcsFile(String gcsFilePath, String outputFileString)
161+
throws IOException {
162+
// create the file only if it doesn't exist
163+
if (!new File(outputFileString).exists()) {
164+
Path outputFilePath = Paths.get(outputFileString);
165+
Path parentDir = outputFilePath.getParent();
166+
if (parentDir != null) {
167+
Files.createDirectories(parentDir);
168+
}
169+
170+
LOG.info("Staging GCS file [{}] to [{}]", gcsFilePath, outputFileString);
171+
Set<StandardOpenOption> options = new HashSet<>(2);
172+
options.add(StandardOpenOption.CREATE);
173+
options.add(StandardOpenOption.WRITE);
174+
175+
// Copy the GCS file into a local file and will throw an I/O exception in case file not found.
176+
try (ReadableByteChannel readerChannel =
177+
FileSystems.open(FileSystems.matchSingleFileSpec(gcsFilePath).resourceId())) {
178+
try (FileChannel writeChannel = FileChannel.open(outputFilePath, options)) {
179+
writeChannel.transferFrom(readerChannel, 0, Long.MAX_VALUE);
180+
}
181+
}
182+
}
183+
return outputFileString;
184+
}
185+
186+
protected byte[] getSecretWithCache(String secretId) {
187+
return secretCache.computeIfAbsent(secretId, this::getSecret);
188+
}
189+
190+
/**
191+
* A helper method to create a new string with the gcs paths replaced with their local path and
192+
* subdirectory based on the factory type in the /tmp directory. For example, the kerberos factory
193+
* type will replace the file paths with /tmp/kerberos/file.path
194+
*
195+
* @param gcsPath
196+
* @return a string with all instances of GCS paths converted to the local paths where the files
197+
* sit.
198+
*/
199+
private String replacePathWithLocal(String gcsPath) throws IOException {
200+
return DIRECTORY_PREFIX + "/" + factoryType + "/" + gcsPath.substring(GCS_PATH_PREFIX.length());
201+
}
202+
203+
/**
204+
* @throws IOException A hook for subclasses to download and process specific files before the
205+
* main configuration is handled. For example, the kerberos factory can use this to download a
206+
* krb5.conf and set a system property.
207+
*/
208+
protected void downloadAndProcessExtraFiles() throws IOException {
209+
// Default implementation should do nothing.
210+
}
211+
212+
protected String getBaseDirectory() {
213+
return DIRECTORY_PREFIX;
214+
}
215+
216+
protected byte[] getSecret(String secretVersion) {
217+
SecretVersionName secretVersionName;
218+
if (SecretVersionName.isParsableFrom(secretVersion)) {
219+
secretVersionName = SecretVersionName.parse(secretVersion);
220+
} else {
221+
throw new IllegalArgumentException(
222+
"Provided Secret must be in the form"
223+
+ " projects/{project}/secrets/{secret}/versions/{secret_version}");
224+
}
225+
try (SecretManagerServiceClient client = SecretManagerServiceClient.create()) {
226+
AccessSecretVersionResponse response = client.accessSecretVersion(secretVersionName);
227+
return response.getPayload().getData().toByteArray();
228+
} catch (IOException e) {
229+
throw new RuntimeException(e);
230+
}
231+
}
232+
233+
protected String processSecret(String originalValue, String secretId, byte[] secretValue)
234+
throws RuntimeException {
235+
// By Default, this will return the secret value directly. This function can be overridden by
236+
// derived classes.
237+
return new String(secretValue, StandardCharsets.UTF_8);
238+
}
239+
}

‎sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KerberosConsumerFactoryFn.java‎

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,158 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.sdk.io.kafka;
19+
20+
import java.io.File;
21+
import java.io.IOException;
22+
import java.nio.charset.StandardCharsets;
23+
import java.nio.file.Files;
24+
import java.nio.file.Path;
25+
import java.nio.file.Paths;
26+
import java.nio.file.attribute.PosixFilePermission;
27+
import java.util.HashSet;
28+
import java.util.Map;
29+
import java.util.Set;
30+
import java.util.UUID;
31+
import java.util.regex.Matcher;
32+
import java.util.regex.Pattern;
33+
import javax.security.auth.login.Configuration;
34+
import org.apache.kafka.clients.consumer.Consumer;
35+
import org.apache.kafka.clients.consumer.KafkaConsumer;
36+
import org.slf4j.Logger;
37+
import org.slf4j.LoggerFactory;
38+
39+
public class KerberosConsumerFactoryFn extends FileAwareFactoryFn<Consumer<byte[], byte[]>> {
40+
private static final String LOCAL_FACTORY_TYPE = "kerberos";
41+
private String krb5ConfigGcsPath = "";
42+
private static volatile String localKrb5ConfPath = "";
43+
44+
private static final Object lock = new Object();
45+
46+
// Standard Kafka property for SASL JAAS configuration
47+
private static final String JAAS_CONFIG_PROPERTY = "sasl.jaas.config";
48+
private static final String KEYTAB_SECRET_PREFIX = "keyTab=\"secretValue:";
49+
private static final Pattern KEYTAB_SECRET_PATTERN =
50+
Pattern.compile("(keyTab=\"secretValue:[^\"]+)");
51+
52+
private static final Logger LOG = LoggerFactory.getLogger(KerberosConsumerFactoryFn.class);
53+
54+
public KerberosConsumerFactoryFn(String krb5ConfigGcsPath) {
55+
super("kerberos");
56+
this.krb5ConfigGcsPath = krb5ConfigGcsPath;
57+
}
58+
59+
@Override
60+
protected Consumer<byte[], byte[]> createObject(Map<String, Object> config) {
61+
// This will be called after the config map processing has occurred. Therefore, we know that the
62+
// property will have
63+
// had it's value replaced with a local directory. We don't need to worry about the GCS prefix
64+
// in this case.
65+
LOG.info("config when creating the objects: {}", config);
66+
try {
67+
String jaasConfig = (String) config.get(JAAS_CONFIG_PROPERTY);
68+
String localKeytabPath = "";
69+
if (jaasConfig != null && !jaasConfig.isEmpty()) {
70+
localKeytabPath =
71+
jaasConfig.substring(
72+
jaasConfig.indexOf("keyTab=\"") + 8, jaasConfig.lastIndexOf("\" principal"));
73+
}
74+
75+
// Set the permissions on the file to be as strict as possible for security reasons. The
76+
// keytab contains
77+
// sensitive information and should be as locked down as possible.
78+
Path path = Paths.get(localKeytabPath);
79+
Set<PosixFilePermission> perms = new HashSet<>();
80+
perms.add(PosixFilePermission.OWNER_READ);
81+
Files.setPosixFilePermissions(path, perms);
82+
} catch (IOException e) {
83+
throw new RuntimeException(
84+
"Could not access keytab file. Make sure that the sasl.jaas.config config property "
85+
+ "is set correctly.",
86+
e);
87+
}
88+
return new KafkaConsumer<>(config);
89+
}
90+
91+
@Override
92+
protected void downloadAndProcessExtraFiles() throws IOException {
93+
synchronized (lock) {
94+
// we only want a new krb5 file if there is not already one present.
95+
if (localKrb5ConfPath.isEmpty()) {
96+
if (this.krb5ConfigGcsPath != null && !this.krb5ConfigGcsPath.isEmpty()) {
97+
String localPath =
98+
super.getBaseDirectory() + "/" + LOCAL_FACTORY_TYPE + "/" + "krb5.conf";
99+
localKrb5ConfPath = downloadGcsFile(this.krb5ConfigGcsPath, localPath);
100+
101+
System.setProperty("java.security.krb5.conf", localKrb5ConfPath);
102+
Configuration.getConfiguration().refresh();
103+
LOG.info(
104+
"Successfully set and refreshed java.security.krb5.conf to {}", localKrb5ConfPath);
105+
}
106+
}
107+
}
108+
}
109+
110+
@Override
111+
protected String processSecret(String originalValue, String secretId, byte[] secretValue)
112+
throws RuntimeException {
113+
Matcher matcher = KEYTAB_SECRET_PATTERN.matcher(originalValue);
114+
String localFileString = "";
115+
while (matcher.find()) {
116+
String currentSecretId = matcher.group(1);
117+
if (currentSecretId == null || currentSecretId.isEmpty()) {
118+
throw new RuntimeException(
119+
"Error matching values. Secret was discovered but its value is null");
120+
}
121+
currentSecretId = currentSecretId.substring(KEYTAB_SECRET_PREFIX.length());
122+
LOG.info("currentSecretId: {} and secretId: {}", currentSecretId, secretId);
123+
if (!currentSecretId.equals(secretId)) {
124+
// A sasl.jaas.config can contain multiple keytabs in one string. Therefore, we must assume
125+
// that there can
126+
// also be multiple keytab secrets in the same string. If the currently matched secret does
127+
// not equal
128+
// the secret that we are processing (passed in via secretId) then we do not want to create
129+
// a keytab file and
130+
// overwrite it.
131+
continue;
132+
}
133+
String filename = "kafka-client-" + UUID.randomUUID().toString() + ".keytab";
134+
135+
localFileString = super.getBaseDirectory() + "/" + LOCAL_FACTORY_TYPE + "/" + filename;
136+
Path localFilePath = Paths.get(localFileString);
137+
Path parentDir = localFilePath.getParent();
138+
try {
139+
if (parentDir != null) {
140+
Files.createDirectories(parentDir);
141+
}
142+
Files.write(localFilePath, secretValue);
143+
if (!new File(localFileString).canRead()) {
144+
LOG.info("The file is not readable");
145+
}
146+
LOG.info("Successfully wrote file to path: {}", localFilePath);
147+
} catch (IOException e) {
148+
throw new RuntimeException("Unable to create the keytab file for the provided secret.");
149+
}
150+
}
151+
// if no localFile was created, then we can assume that the secret is meant to be kept as a
152+
// value.
153+
LOG.info("LocalFilestring: {}", localFileString);
154+
return localFileString.isEmpty()
155+
? new String(secretValue, StandardCharsets.UTF_8)
156+
: localFileString;
157+
}
158+
}

0 commit comments

Comments
 (0)