Load file config YAML using core schema, ensure that env var substiut… (#6436)

jack-berg · web-flow · commit c71c4d983cf2 · 2024-05-20T14:39:41.000-05:00
diff --git a/sdk-extensions/incubator/src/main/java/io/opentelemetry/sdk/extension/incubator/fileconfig/FileConfiguration.java b/sdk-extensions/incubator/src/main/java/io/opentelemetry/sdk/extension/incubator/fileconfig/FileConfiguration.java
@@ -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
@@ -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);
   }
@@ -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());
     }
   }
 }
diff --git a/sdk-extensions/incubator/src/test/java/io/opentelemetry/sdk/extension/incubator/fileconfig/FileConfigurationParseTest.java b/sdk-extensions/incubator/src/test/java/io/opentelemetry/sdk/extension/incubator/fileconfig/FileConfigurationParseTest.java
@@ -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(
@@ -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"))),
@@ -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_]*
@@ -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) {
@@ -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 =
