Skip to content

Commit

Permalink
Load file config YAML using core schema, ensure that env var substiut… (
Browse files Browse the repository at this point in the history
  • Loading branch information
jack-berg committed May 20, 2024
1 parent 7953048 commit c71c4d9
Show file tree
Hide file tree
Showing 2 changed files with 114 additions and 36 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,15 @@
import java.util.regex.Pattern;
import org.snakeyaml.engine.v2.api.Load;
import org.snakeyaml.engine.v2.api.LoadSettings;
import org.snakeyaml.engine.v2.common.ScalarStyle;
import org.snakeyaml.engine.v2.constructor.StandardConstructor;
import org.snakeyaml.engine.v2.exceptions.ConstructorException;
import org.snakeyaml.engine.v2.exceptions.YamlEngineException;
import org.snakeyaml.engine.v2.nodes.MappingNode;
import org.yaml.snakeyaml.Yaml;
import org.snakeyaml.engine.v2.nodes.Node;
import org.snakeyaml.engine.v2.nodes.NodeTuple;
import org.snakeyaml.engine.v2.nodes.ScalarNode;
import org.snakeyaml.engine.v2.schema.CoreSchema;

/**
* Configure {@link OpenTelemetrySdk} from YAML configuration files conforming to the schema in <a
Expand Down Expand Up @@ -127,7 +133,7 @@ static OpenTelemetryConfiguration parse(

// Visible for testing
static Object loadYaml(InputStream inputStream, Map<String, String> environmentVariables) {
LoadSettings settings = LoadSettings.builder().build();
LoadSettings settings = LoadSettings.builder().setSchema(new CoreSchema()).build();
Load yaml = new Load(settings, new EnvSubstitutionConstructor(settings, environmentVariables));
return yaml.loadFromInputStream(inputStream);
}
Expand All @@ -145,52 +151,94 @@ static Object loadYaml(InputStream inputStream, Map<String, String> environmentV
*/
private static final class EnvSubstitutionConstructor extends StandardConstructor {

// Yaml is not thread safe but this instance is always used on the same thread
private final Yaml yaml = new Yaml();
// Load is not thread safe but this instance is always used on the same thread
private final Load load;
private final Map<String, String> environmentVariables;

private EnvSubstitutionConstructor(
LoadSettings loadSettings, Map<String, String> environmentVariables) {
super(loadSettings);
load = new Load(loadSettings);
this.environmentVariables = environmentVariables;
}

/**
* Implementation is same as {@link
* org.snakeyaml.engine.v2.constructor.BaseConstructor#constructMapping(MappingNode)} except we
* override the resolution of values with our custom {@link #constructValueObject(Node)}, which
* performs environment variable substitution.
*/
@Override
@SuppressWarnings({"ReturnValueIgnored", "CatchingUnchecked"})
protected Map<Object, Object> constructMapping(MappingNode node) {
// First call the super to construct mapping from MappingNode as usual
Map<Object, Object> result = super.constructMapping(node);

// Iterate through the map entries, and:
// 1. Identify entries which are scalar strings eligible for environment variable substitution
// 2. Apply environment variable substitution
// 3. Re-parse substituted value so it has correct type (i.e. yaml.load(newVal))
for (Map.Entry<Object, Object> entry : result.entrySet()) {
Object value = entry.getValue();
if (!(value instanceof String)) {
continue;
Map<Object, Object> mapping = settings.getDefaultMap().apply(node.getValue().size());
List<NodeTuple> nodeValue = node.getValue();
for (NodeTuple tuple : nodeValue) {
Node keyNode = tuple.getKeyNode();
Object key = constructObject(keyNode);
if (key != null) {
try {
key.hashCode(); // check circular dependencies
} catch (Exception e) {
throw new ConstructorException(
"while constructing a mapping",
node.getStartMark(),
"found unacceptable key " + key,
tuple.getKeyNode().getStartMark(),
e);
}
}

String val = (String) value;
Matcher matcher = ENV_VARIABLE_REFERENCE.matcher(val);
if (!matcher.find()) {
continue;
Node valueNode = tuple.getValueNode();
Object value = constructValueObject(valueNode);
if (keyNode.isRecursive()) {
if (settings.getAllowRecursiveKeys()) {
postponeMapFilling(mapping, key, value);
} else {
throw new YamlEngineException(
"Recursive key for mapping is detected but it is not configured to be allowed.");
}
} else {
mapping.put(key, value);
}
}

int offset = 0;
StringBuilder newVal = new StringBuilder();
do {
MatchResult matchResult = matcher.toMatchResult();
String replacement = environmentVariables.getOrDefault(matcher.group(1), "");
newVal.append(val, offset, matchResult.start()).append(replacement);
offset = matchResult.end();
} while (matcher.find());
if (offset != val.length()) {
newVal.append(val, offset, val.length());
}
entry.setValue(yaml.load(newVal.toString()));
return mapping;
}

private Object constructValueObject(Node node) {
Object value = constructObject(node);
if (!(node instanceof ScalarNode)) {
return value;
}
if (!(value instanceof String)) {
return value;
}

String val = (String) value;
Matcher matcher = ENV_VARIABLE_REFERENCE.matcher(val);
if (!matcher.find()) {
return value;
}

return result;
int offset = 0;
StringBuilder newVal = new StringBuilder();
ScalarStyle scalarStyle = ((ScalarNode) node).getScalarStyle();
do {
MatchResult matchResult = matcher.toMatchResult();
String replacement = environmentVariables.getOrDefault(matcher.group(1), "");
newVal.append(val, offset, matchResult.start()).append(replacement);
offset = matchResult.end();
} while (matcher.find());
if (offset != val.length()) {
newVal.append(val, offset, val.length());
}
// If the value was double quoted, retain the double quotes so we don't change a value
// intended to be a string to a different type after environment variable substitution
if (scalarStyle == ScalarStyle.DOUBLE_QUOTED) {
newVal.insert(0, "\"");
newVal.append("\"");
}
return load.loadFromString(newVal.toString());
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -387,15 +387,36 @@ void parse_nullBoxedPrimitivesParsedToNull() {
new Sampler().withTraceIdRatioBased(new TraceIdRatioBased()))));
}

@ParameterizedTest
@MethodSource("coreSchemaValuesArgs")
void coreSchemaValues(String rawYaml, Object expectedYamlResult) {
Object yaml =
FileConfiguration.loadYaml(
new ByteArrayInputStream(rawYaml.getBytes(StandardCharsets.UTF_8)),
Collections.emptyMap());
assertThat(yaml).isEqualTo(expectedYamlResult);
}

@SuppressWarnings("unchecked")
private static java.util.stream.Stream<Arguments> coreSchemaValuesArgs() {
return java.util.stream.Stream.of(
Arguments.of("key1: 0o123\n", mapOf(entry("key1", 83))),
Arguments.of("key1: 0123\n", mapOf(entry("key1", 123))),
Arguments.of("key1: 0xdeadbeef\n", mapOf(entry("key1", 3735928559L))),
Arguments.of("key1: \"0xdeadbeef\"\n", mapOf(entry("key1", "0xdeadbeef"))));
}

@ParameterizedTest
@MethodSource("envVarSubstitutionArgs")
void envSubstituteAndLoadYaml(String rawYaml, Object expectedYamlResult) {
Map<String, String> environmentVariables = new HashMap<>();
environmentVariables.put("STR_1", "value1");
environmentVariables.put("STR_2", "value2");
environmentVariables.put("EMPTY_STR", "");
environmentVariables.put("BOOL", "true");
environmentVariables.put("INT", "1");
environmentVariables.put("FLOAT", "1.1");
environmentVariables.put("HEX", "0xdeadbeef");

Object yaml =
FileConfiguration.loadYaml(
Expand All @@ -412,6 +433,7 @@ private static java.util.stream.Stream<Arguments> envVarSubstitutionArgs() {
Arguments.of("key1: ${BOOL}\n", mapOf(entry("key1", true))),
Arguments.of("key1: ${INT}\n", mapOf(entry("key1", 1))),
Arguments.of("key1: ${FLOAT}\n", mapOf(entry("key1", 1.1))),
Arguments.of("key1: ${HEX}\n", mapOf(entry("key1", 3735928559L))),
Arguments.of(
"key1: ${STR_1}\n" + "key2: value2\n",
mapOf(entry("key1", "value1"), entry("key2", "value2"))),
Expand All @@ -421,7 +443,8 @@ private static java.util.stream.Stream<Arguments> envVarSubstitutionArgs() {
// Multiple environment variables referenced
Arguments.of("key1: ${STR_1}${STR_2}\n", mapOf(entry("key1", "value1value2"))),
Arguments.of("key1: ${STR_1} ${STR_2}\n", mapOf(entry("key1", "value1 value2"))),
// Undefined environment variable
// Undefined / empty environment variable
Arguments.of("key1: ${EMPTY_STR}\n", mapOf(entry("key1", null))),
Arguments.of("key1: ${STR_3}\n", mapOf(entry("key1", null))),
Arguments.of("key1: ${STR_1} ${STR_3}\n", mapOf(entry("key1", "value1"))),
// Environment variable keys must match pattern: [a-zA-Z_]+[a-zA-Z0-9_]*
Expand All @@ -432,7 +455,14 @@ private static java.util.stream.Stream<Arguments> envVarSubstitutionArgs() {
"key1:\n ${STR_1}: value1\n",
mapOf(entry("key1", mapOf(entry("${STR_1}", "value1"))))),
Arguments.of(
"key1:\n - ${STR_1}\n", mapOf(entry("key1", Collections.singletonList("${STR_1}")))));
"key1:\n - ${STR_1}\n", mapOf(entry("key1", Collections.singletonList("${STR_1}")))),
// Quoted environment variables
Arguments.of("key1: \"${HEX}\"\n", mapOf(entry("key1", "0xdeadbeef"))),
Arguments.of("key1: \"${STR_1}\"\n", mapOf(entry("key1", "value1"))),
Arguments.of("key1: \"${EMPTY_STR}\"\n", mapOf(entry("key1", ""))),
Arguments.of("key1: \"${BOOL}\"\n", mapOf(entry("key1", "true"))),
Arguments.of("key1: \"${INT}\"\n", mapOf(entry("key1", "1"))),
Arguments.of("key1: \"${FLOAT}\"\n", mapOf(entry("key1", "1.1"))));
}

private static <K, V> Map.Entry<K, V> entry(K key, @Nullable V value) {
Expand Down Expand Up @@ -461,7 +491,7 @@ void read_WithEnvironmentVariables() {
+ " - batch:\n"
+ " exporter:\n"
+ " otlp:\n"
+ " endpoint: \"${UNSET_ENV_VAR}\"\n";
+ " endpoint: ${UNSET_ENV_VAR}\n";
Map<String, String> envVars = new HashMap<>();
envVars.put("OTEL_EXPORTER_OTLP_ENDPOINT", "http://collector:4317");
OpenTelemetryConfiguration model =
Expand Down

0 comments on commit c71c4d9

Please sign in to comment.