Provide access to VM arguments via shared CLIHelper

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

@@ -2,22 +2,20 @@
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 
+import datadog.trace.bootstrap.config.provider.VMArgsCache;
 import de.thetaphi.forbiddenapis.SuppressForbidden;
 import java.io.BufferedReader;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStreamReader;
 import java.io.PrintStream;
 import java.lang.instrument.Instrumentation;
-import java.lang.management.ManagementFactory;
-import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.URL;
 import java.security.CodeSource;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.IdentityHashMap;
 import java.util.List;
@@ -383,7 +381,7 @@ 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 : getVMArgumentsThroughReflection()) {
+      for (final String argument : VMArgsCache.ARGS.getJvmArgs()) {
         if (argument.startsWith(JAVA_AGENT_ARGUMENT)) {
           int index = argument.indexOf('=', JAVA_AGENT_ARGUMENT.length());
           String agentPathname =
@@ -424,57 +422,6 @@ private static File getAgentFileUsingClassLoaderLookup() throws URISyntaxExcepti
     return javaagentFile;
   }
 
-  @SuppressForbidden
-  private static List<String> getVMArgumentsThroughReflection() {
-    // Try Oracle-based
-    // IBM Semeru Runtime 1.8.0_345-b01 will throw UnsatisfiedLinkError here.
-    try {
-      final Class<?> managementFactoryHelperClass =
-          Class.forName("sun.management.ManagementFactoryHelper");
-
-      final Class<?> vmManagementClass = Class.forName("sun.management.VMManagement");
-
-      Object vmManagement;
-
-      try {
-        vmManagement =
-            managementFactoryHelperClass.getDeclaredMethod("getVMManagement").invoke(null);
-      } catch (final NoSuchMethodException e) {
-        // Older vm before getVMManagement() existed
-        final Field field = managementFactoryHelperClass.getDeclaredField("jvm");
-        field.setAccessible(true);
-        vmManagement = field.get(null);
-        field.setAccessible(false);
-      }
-
-      //noinspection unchecked
-      return (List<String>) vmManagementClass.getMethod("getVmArguments").invoke(vmManagement);
-    } catch (final ReflectiveOperationException | UnsatisfiedLinkError ignored) {
-      // Ignored exception
-    }
-
-    // Try IBM-based.
-    try {
-      final Class<?> VMClass = Class.forName("com.ibm.oti.vm.VM");
-      final String[] argArray = (String[]) VMClass.getMethod("getVMArgs").invoke(null);
-      return Arrays.asList(argArray);
-    } catch (final ReflectiveOperationException ignored) {
-      // Ignored exception
-    }
-
-    // Fallback to default
-    try {
-      System.err.println(
-          "WARNING: Unable to get VM args through reflection. A custom java.util.logging.LogManager may not work correctly");
-      return ManagementFactory.getRuntimeMXBean().getInputArguments();
-    } 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();
-  }
-
   private static void checkJarManifestMainClassIsThis(final URL jarUrl) throws IOException {
     final URL manifestUrl = new URL("jar:" + jarUrl + "!/META-INF/MANIFEST.MF");
     final String mainClassLine = "Main-Class: " + thisClass.getCanonicalName();

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

@@ -0,0 +1,122 @@
+package datadog.trace.bootstrap.config.provider;
+
+import java.io.BufferedReader;
+import java.io.FileReader;
+import java.io.IOException;
+import java.lang.management.ManagementFactory;
+import java.lang.reflect.Field;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+
+/** VMArgsCache stores JVM Arguments applied to the current process */
+public class VMArgsCache {
+  public static final VMArgsCache ARGS = new VMArgsCache(initJvmArgs());
+
+  private final HashSet<String> args;
+
+  public VMArgsCache(List<String> args) {
+    this.args = new HashSet<>(args);
+  }
+
+  public HashSet<String> getJvmArgs() {
+    return this.args;
+  }
+
+  public boolean contains(String argument) {
+    return this.args.contains(argument);
+  }
+
+  private static List<String> initJvmArgs() {
+    // If linux OS, use procfs
+    // TODO: equals, or contains?
+    if (System.getProperty("os.name").equalsIgnoreCase("linux")) {
+      // Get the current process PID from /proc/self/status
+      try {
+        String pid = getPidFromProcStatus();
+        if (pid != null) {
+          // Get the JVM arguments from /proc/[pid]/cmdline
+          return getJvmArgsFromProcCmdline(pid);
+        }
+      } catch (IOException e) {
+        // ignore exception, try other methods
+      }
+    }
+    // Try Oracle-based
+    // IBM Semeru Runtime 1.8.0_345-b01 will throw UnsatisfiedLinkError here.
+    try {
+      final Class<?> managementFactoryHelperClass =
+          Class.forName("sun.management.ManagementFactoryHelper");
+
+      final Class<?> vmManagementClass = Class.forName("sun.management.VMManagement");
+
+      Object vmManagement;
+
+      try {
+        vmManagement =
+            managementFactoryHelperClass.getDeclaredMethod("getVMManagement").invoke(null);
+      } catch (final NoSuchMethodException e) {
+        // Older vm before getVMManagement() existed
+        final Field field = managementFactoryHelperClass.getDeclaredField("jvm");
+        field.setAccessible(true);
+        vmManagement = field.get(null);
+        field.setAccessible(false);
+      }
+
+      //noinspection unchecked
+      return (List<String>) vmManagementClass.getMethod("getVmArguments").invoke(vmManagement);
+    } catch (final ReflectiveOperationException | UnsatisfiedLinkError ignored) {
+      // Ignored exception
+    }
+
+    // Try IBM-based.
+    try {
+      final Class<?> VMClass = Class.forName("com.ibm.oti.vm.VM");
+      final String[] argArray = (String[]) VMClass.getMethod("getVMArgs").invoke(null);
+      return Arrays.asList(argArray);
+    } catch (final ReflectiveOperationException ignored) {
+      // Ignored exception
+    }
+
+    // Fallback to default
+    try {
+      System.err.println(
+          "WARNING: Unable to get VM args through reflection. A custom java.util.logging.LogManager may not work correctly");
+      return ManagementFactory.getRuntimeMXBean().getInputArguments();
+    } 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();
+  }
+
+  // Helper methods for getting process information from linux proc dir
+  private static String getPidFromProcStatus() throws IOException {
+    String pid = null;
+    // Read /proc/self/status to find the current process's PID
+    try (BufferedReader pidReader = new BufferedReader(new FileReader("/proc/self/status"))) {
+      String line;
+      while ((line = pidReader.readLine()) != null) {
+        if (line.startsWith("Pid:")) {
+          pid = line.split(":")[1].trim();
+          break;
+        }
+      }
+    }
+    return pid;
+  }
+
+  private static List<String> getJvmArgsFromProcCmdline(String pid) throws IOException {
+    // Read /proc/[pid]/cmdline to get JVM arguments
+    BufferedReader argsReader = new BufferedReader(new FileReader("/proc/" + pid + "/cmdline"));
+    String cmdLine = argsReader.readLine();
+    if (cmdLine != null) {
+      // Return JVM arguments as a list of strings split by null characters
+      return Arrays.asList(cmdLine.split("\0"));
+    } else {
+      return null;
+    }
+  }
+}