Stable Config file: target system properties in process_arguments and support template variables in YamlParser

components/cli/src/main/java/datadog/cli/CLIHelper.java

@@ -9,25 +9,30 @@
 import java.nio.file.Paths;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
 
+/** Helper class for retrieving and parsing JVM arguments. */
 public final class CLIHelper {
-  private static final List<String> VM_ARGS = findVmArgs();
+  private static final Map<String, String> VM_ARGS = findVmArgs();
 
-  public static List<String> getVmArgs() {
+  public static Map<String, String> getVmArgs() {
     return VM_ARGS;
   }
 
   @SuppressForbidden
-  private static List<String> findVmArgs() {
+  private static Map<String, String> findVmArgs() {
+    List<String> rawArgs;
+
     // Try ProcFS on Linux
     try {
       if (isLinux()) {
         Path cmdlinePath = Paths.get("/proc/self/cmdline");
         if (Files.exists(cmdlinePath)) {
           try (BufferedReader in = Files.newBufferedReader(cmdlinePath)) {
-            return Arrays.asList(in.readLine().split("\0"));
+            return parseVmArgs(Arrays.asList(in.readLine().split("\0")));
           }
         }
       }
@@ -57,7 +62,8 @@ private static List<String> findVmArgs() {
       }
 
       //noinspection unchecked
-      return (List<String>) vmManagementClass.getMethod("getVmArguments").invoke(vmManagement);
+      rawArgs = (List<String>) vmManagementClass.getMethod("getVmArguments").invoke(vmManagement);
+      return parseVmArgs(rawArgs);
     } catch (final ReflectiveOperationException | UnsatisfiedLinkError ignored) {
       // Ignored exception
     }
@@ -66,20 +72,55 @@ private static List<String> findVmArgs() {
     try {
       final Class<?> VMClass = Class.forName("com.ibm.oti.vm.VM");
       final String[] argArray = (String[]) VMClass.getMethod("getVMArgs").invoke(null);
-      return Arrays.asList(argArray);
+      rawArgs = Arrays.asList(argArray);
+      return parseVmArgs(rawArgs);
     } catch (final ReflectiveOperationException ignored) {
       // Ignored exception
     }
 
     // Fallback to default
     try {
-      return ManagementFactory.getRuntimeMXBean().getInputArguments();
+      rawArgs = ManagementFactory.getRuntimeMXBean().getInputArguments();
+      return parseVmArgs(rawArgs);
     } catch (final Throwable t) {
       // Throws InvocationTargetException on modularized applications
       // with non-opened java.management module
       System.err.println("WARNING: Unable to get VM args using managed beans");
     }
-    return Collections.emptyList();
+    return Collections.emptyMap();
+  }
+
+  /**
+   * Parses a list of VM arguments into a map of key-value pairs. For system properties (-D
+   * arguments), the key includes the -D prefix and the value is everything after the = sign. For
+   * all other VM arguments, the key is the full argument and the value is an empty string.
+   */
+  private static Map<String, String> parseVmArgs(List<String> args) {
+    Map<String, String> result = new HashMap<>();
+
+    // For now, we only support values on system properties (-D arguments)
+    for (String arg : args) {
+      // TODO: Handle other types of VM arguments
+      if (arg.startsWith("-D")) {
+        // Handle system properties (-D arguments)
+        int equalsIndex = arg.indexOf('=');
+
+        if (equalsIndex >= 0) {
+          // Key-value pair
+          String key = arg.substring(0, equalsIndex);
+          String value = arg.substring(equalsIndex + 1);
+          result.put(key, value);
+        } else {
+          // Just a key with no value
+          result.put(arg, "");
+        }
+      } else {
+        // Any other type of VM argument
+        result.put(arg, "");
+      }
+    }
+
+    return result;
   }
 
   private static boolean isLinux() {

components/yaml/build.gradle.kts

@@ -11,4 +11,7 @@ jmh {
 // https://repo1.maven.org/maven2/org/yaml/snakeyaml/2.4/snakeyaml-2.4.pom
 dependencies {
   implementation("org.yaml", "snakeyaml", "2.4")
+  implementation(project(":components:cli"))
+  testImplementation(project(":utils:test-utils"))
+  testImplementation(project(":internal-api"))
 }

components/yaml/src/main/java/datadog/yaml/YamlParser.java

@@ -1,19 +1,99 @@
 package datadog.yaml;
 
-import java.io.FileInputStream;
+import datadog.cli.CLIHelper;
 import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.util.Map;
 import org.yaml.snakeyaml.Yaml;
 
 public class YamlParser {
-  // Supports clazz == null for default yaml parsing
+  private static final Map<String, String> VM_ARGS = CLIHelper.getVmArgs();
+
   public static <T> T parse(String filePath, Class<T> clazz) throws IOException {
     Yaml yaml = new Yaml();
-    try (FileInputStream fis = new FileInputStream(filePath)) {
-      if (clazz == null) {
-        return yaml.load(fis);
-      } else {
-        return yaml.loadAs(fis, clazz);
+    String content = new String(Files.readAllBytes(Paths.get(filePath)));
+    String processedContent = processTemplate(content);
+
+    if (clazz == null) {
+      return yaml.load(processedContent);
+    } else {
+      return yaml.loadAs(processedContent, clazz);
+    }
+  }
+
+  /**
+   * Processes a YAML template by replacing all template variables with their corresponding values.
+   * Template variables are enclosed in double curly braces, e.g. {{variable}}. Returns the
+   * processed content with all template variables resolved.
+   */
+  static String processTemplate(String content) throws IOException {
+    StringBuilder result = new StringBuilder(content.length());
+    String rest = content;
+
+    while (true) {
+      int openIndex = rest.indexOf("{{");
+      if (openIndex == -1) {
+        result.append(rest);
+        break;
+      }
+
+      // Add everything before the template
+      result.append(rest.substring(0, openIndex));
+
+      // Find the closing braces
+      int closeIndex = rest.indexOf("}}", openIndex);
+      if (closeIndex == -1) {
+        throw new IOException("Unterminated template in config");
+      }
+
+      // Extract the template variable
+      String templateVar = rest.substring(openIndex + 2, closeIndex).trim();
+
+      // Process the template variable and get its value
+      String value = processTemplateVar(templateVar);
+
+      // Add the processed value
+      result.append(value);
+
+      // Continue with the rest of the string
+      rest = rest.substring(closeIndex + 2);
+    }
+
+    return result.toString();
+  }
+
+  /**
+   * Processes a template variable by extracting its value from either environment variables or VM
+   * arguments. Template variables should be in the format "{{environment_variables[VAR_NAME]}}" or
+   * "{{process_arguments[-ARG_NAME]}}". Returns "UNDEFINED" if the variable is not found or empty.
+   */
+  static String processTemplateVar(String templateVar) throws IOException {
+    if (templateVar.startsWith("environment_variables['") && templateVar.endsWith("']")) {
+      String envVar =
+          templateVar
+              .substring("environment_variables['".length(), templateVar.length() - 2)
+              .trim();
+      if (envVar.isEmpty()) {
+        throw new IOException("Empty environment variable name in template");
+      }
+      String value = System.getenv(envVar.toUpperCase());
+      if (value == null || value.isEmpty()) {
+        return "UNDEFINED";
+      }
+      return value;
+    } else if (templateVar.startsWith("process_arguments['") && templateVar.endsWith("']")) {
+      String processArg =
+          templateVar.substring("process_arguments['".length(), templateVar.length() - 2).trim();
+      if (processArg.isEmpty()) {
+        throw new IOException("Empty process argument in template");
+      }
+      String value = VM_ARGS.get(processArg);
+      if (value == null || value.isEmpty()) {
+        return "UNDEFINED";
       }
+      return value;
     }
+    return "UNDEFINED";
   }
 }

components/yaml/src/test/groovy/datadog/yaml/YamlParserTest.groovy

@@ -0,0 +1,57 @@
+package datadog.yaml
+
+import datadog.trace.bootstrap.config.provider.stableconfigyaml.StableConfigYaml
+import datadog.trace.test.util.DDSpecification
+import java.nio.file.Path
+import datadog.trace.test.util.FileUtils
+
+class YamlParserTest extends DDSpecification{
+  def "test parse"() {
+    when:
+    String yaml = """
+    apm_configuration_rules:
+      - selectors:
+        - origin: language
+          matches: ["java"]
+          operator: equals
+        configuration:
+          DD_SERVICE: "${templateVar}"
+    """
+    Path filePath = FileUtils.tempFile()
+    try {
+      FileUtils.writeFileRaw(filePath, yaml)
+    } catch (IOException e) {
+      throw new AssertionError("Failed to write to file: ${e.message}")
+    }
+
+    if (envKey != null) {
+      injectEnvConfig(envKey, envVal)
+    }
+    String service
+    try {
+      def data = YamlParser.parse(filePath as String, StableConfigYaml)
+      def configs = data.getApm_configuration_rules().get(0).getConfiguration()
+      service = configs.get("DD_SERVICE").toString()
+    } catch (IOException e) {
+      // parse failed, stable config will be dropped
+      service = null
+    }
+
+    then:
+    service == expect
+
+    where:
+    templateVar | envKey | envVal | expect
+    "{{environment_variables['DD_KEY']}}" | "DD_KEY" | "value" | "value"
+    "{{environment_variables['DD_KEY']}}" | null | null | "UNDEFINED"
+    "{{environment_variables['DD_KEY}}" | "DD_KEY" | "value" | "UNDEFINED"
+    "{{DD_KEY}}" |  "DD_KEY" | "value" | "UNDEFINED"
+    "{{environment_variables['']}}" | null | null | null
+    "{{environment_variables['OTHER_KEY']}}" | "DD_KEY" | "value" | "UNDEFINED"
+    "{{}}" | null | null | "UNDEFINED"
+    "{}" | null | null | "{}"
+    "{{environment_variables[DD_KEY]}}" | "DD_KEY" | "value" | "UNDEFINED"
+    "{{environment_variables['DD_KEY']}" | null | null | null
+    "{{environment_variables['dd_key']}}" | "DD_KEY" | "value" | "value"
+  }
+}

dd-java-agent/src/main/java/datadog/trace/bootstrap/AgentBootstrap.java

@@ -19,6 +19,7 @@
 import java.util.Collections;
 import java.util.IdentityHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.jar.JarFile;
 
@@ -381,8 +382,12 @@ private static List<File> getAgentFilesFromVMArguments() {
       // - On IBM-based JDKs since at least 1.7
       // This prevents custom log managers from working correctly
       // Use reflection to bypass the loading of the class~
-      for (final String argument : CLIHelper.getVmArgs()) {
+      Map<String, String> vmArgs = CLIHelper.getVmArgs();
+      for (final Map.Entry<String, String> entry : vmArgs.entrySet()) {
+        final String argument = entry.getKey();
         if (argument.startsWith(JAVA_AGENT_ARGUMENT)) {
+          // TODO: Modify CLIHelper to parse key-vals on JAVA_AGENT_ARGUMENT arguments, so we don't
+          // have to do the parsing here.
           int index = argument.indexOf('=', JAVA_AGENT_ARGUMENT.length());
           String agentPathname =
               argument.substring(

internal-api/src/main/java/datadog/trace/bootstrap/config/provider/StableConfigParser.java

@@ -9,17 +9,16 @@
 import java.io.IOException;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
-import java.util.Set;
+import java.util.Map;
 import java.util.function.BiPredicate;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 public class StableConfigParser {
   private static final Logger log = LoggerFactory.getLogger(StableConfigParser.class);
 
-  private static final Set<String> VM_ARGS = new HashSet<>(CLIHelper.getVmArgs());
+  private static final Map<String, String> VM_ARGS = CLIHelper.getVmArgs();
 
   /**
    * Parses a configuration file and returns a stable configuration object.
@@ -66,7 +65,9 @@ public static StableConfigSource.StableConfig parse(String filePath) throws IOEx
 
     } catch (IOException e) {
       log.debug(
-          "Stable configuration file either not found or not readable at filepath {}", filePath);
+          "Stable configuration file either not found or not readable at filepath {}. Error: {}",
+          filePath,
+          e.getMessage());
     }
     return StableConfigSource.StableConfig.EMPTY;
   }
@@ -168,7 +169,7 @@ static boolean selectorMatch(String origin, List<String> matches, String operato
       case "process_arguments":
         // For now, always return true if `key` exists in the JVM Args
         // TODO: flesh out the meaning of each operator for process_arguments
-        return VM_ARGS.contains(key);
+        return VM_ARGS.containsKey(key);
       case "tags":
         // TODO: Support this down the line (Must define the source of "tags" first)
         return false;