Getting the -javaagent parameter is necessary in some cases

dd-java-agent/dd-java-agent.gradle

@@ -121,7 +121,6 @@ dependencies {
 tasks.withType(Test).configureEach {
   jvmArgs "-Ddd.service.name=java-agent-tests"
   jvmArgs "-Ddd.writer.type=LoggingWriter"
-  jvmArgs "-Dagent.jar.to.forward=${shadowJar.archivePath}"
   // Multi-threaded logging seems to be causing deadlocks with Gradle's log capture.
 //  jvmArgs "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=debug"
 //  jvmArgs "-Dorg.slf4j.simpleLogger.defaultLogLevel=debug"

dd-java-agent/src/main/java/datadog/trace/agent/TracingAgent.java

@@ -5,14 +5,20 @@ import java.io.File;
 import java.io.IOException;
 import java.io.InputStreamReader;
 import java.lang.instrument.Instrumentation;
+import java.lang.management.ManagementFactory;
 import java.lang.reflect.Constructor;
+import java.lang.reflect.Field;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.net.URISyntaxException;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
 import java.security.CodeSource;
+import java.util.Arrays;
+import java.util.List;
 import java.util.jar.JarFile;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 /** Entry point for initializing the agent. */
 public class TracingAgent {
@@ -151,16 +157,103 @@ public class TracingAgent {
 
   private static synchronized URL installBootstrapJar(
       final Instrumentation inst, final boolean usingCustomLogManager)
-      throws IOException, URISyntaxException, ClassNotFoundException {
+      throws IOException, URISyntaxException {
+    URL bootstrapURL = null;
+
+    // First try Code Source
     final CodeSource codeSource = TracingAgent.class.getProtectionDomain().getCodeSource();
 
     if (codeSource != null) {
-      final URL bootstrapURL = codeSource.getLocation();
+      bootstrapURL = codeSource.getLocation();
+      final File bootstrapFile = new File(bootstrapURL.toURI());
 
-      inst.appendToBootstrapClassLoaderSearch(new JarFile(new File(bootstrapURL.toURI())));
-      return bootstrapURL;
+      if (!bootstrapFile.isDirectory()) {
+        inst.appendToBootstrapClassLoaderSearch(new JarFile(bootstrapFile));
+        return bootstrapURL;
+      }
+    }
+
+    System.out.println("Could not get bootstrap jar from code source, using -javaagent arg");
+
+    // ManagementFactory indirectly references java.util.logging.LogManager
+    // - On Oracle-based JDKs after 1.8
+    // - 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
+    final List<String> arguments;
+    if (usingCustomLogManager) {
+      System.out.println("Custom log manager detected: getting vm args through reflection");
+      arguments = getVMArgumentsThroughReflection();
     } else {
-      throw new RuntimeException("Unable to install bootstrap jar.  Code source not found");
+      arguments = ManagementFactory.getRuntimeMXBean().getInputArguments();
+    }
+
+    String agentArgument = null;
+    for (final String arg : arguments) {
+      if (arg.startsWith("-javaagent")) {
+        if (agentArgument == null) {
+          agentArgument = arg;
+        } else {
+          throw new RuntimeException(
+              "Multiple javaagents specified and code source unavailable, not installing tracing agent");
+        }
+      }
+    }
+
+    if (agentArgument == null) {
+      throw new RuntimeException(
+          "Could not find javaagent parameter and code source unavailable, not installing tracing agent");
+    }
+
+    // argument is of the form -javaagent:/path/to/dd-java-agent.jar=optionalargumentstring
+    final Matcher matcher = Pattern.compile("-javaagent:([^=]+).*").matcher(agentArgument);
+
+    if (!matcher.matches()) {
+      throw new RuntimeException("Unable to parse javaagent parameter: " + agentArgument);
+    }
+
+    bootstrapURL = new URL("file:" + matcher.group(1));
+    inst.appendToBootstrapClassLoaderSearch(new JarFile(new File(bootstrapURL.toURI())));
+
+    return bootstrapURL;
+  }
+
+  private static List<String> getVMArgumentsThroughReflection() {
+    try {
+      // Try Oracle-based
+      final Class managementFactoryHelperClass =
+          TracingAgent.class.getClassLoader().loadClass("sun.management.ManagementFactoryHelper");
+
+      final Class vmManagementClass =
+          TracingAgent.class.getClassLoader().loadClass("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);
+      }
+
+      return (List<String>) vmManagementClass.getMethod("getVmArguments").invoke(vmManagement);
+
+    } catch (final ReflectiveOperationException e) {
+      try { // Try IBM-based.
+        final Class VMClass = TracingAgent.class.getClassLoader().loadClass("com.ibm.oti.vm.VM");
+        final String[] argArray = (String[]) VMClass.getMethod("getVMArgs").invoke(null);
+        return Arrays.asList(argArray);
+      } catch (final ReflectiveOperationException e1) {
+        // Fallback to default
+        System.out.println(
+            "WARNING: Unable to get VM args through reflection.  A custom java.util.logging.LogManager may not work correctly");
+
+        return ManagementFactory.getRuntimeMXBean().getInputArguments();
+      }
     }
   }
 

dd-java-agent/src/test/groovy/datadog/trace/agent/AgentLoadedIntoBootstrapTest.groovy

@@ -5,8 +5,7 @@ import jvmbootstraptest.AgentLoadedChecker
 import spock.lang.Specification
 
 class AgentLoadedIntoBootstrapTest extends Specification {
-
-
+  
   def "Agent loads in when separate jvm is launched"() {
     expect:
     IntegrationTestUtils.runOnSeparateJvm(AgentLoadedChecker.getName()

dd-java-agent/src/test/java/datadog/trace/agent/test/IntegrationTestUtils.java

@@ -10,6 +10,8 @@ import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
+import java.lang.management.ManagementFactory;
+import java.lang.management.RuntimeMXBean;
 import java.lang.reflect.Field;
 import java.net.URL;
 import java.util.ArrayList;
@@ -175,6 +177,17 @@ public class IntegrationTestUtils {
     return (String[]) f.get(null);
   }
 
+  private static String getAgentArgument() {
+    final RuntimeMXBean runtimeMxBean = ManagementFactory.getRuntimeMXBean();
+    for (final String arg : runtimeMxBean.getInputArguments()) {
+      if (arg.startsWith("-javaagent")) {
+        return arg;
+      }
+    }
+
+    throw new RuntimeException("Agent jar not found");
+  }
+
   /**
    * On a separate JVM, run the main method for a given class.
    *
@@ -192,18 +205,11 @@ public class IntegrationTestUtils {
       throws Exception {
 
     final String separator = System.getProperty("file.separator");
-    final String agentJar = System.getProperty("agent.jar.to.forward");
-
-    if (agentJar == null) {
-      throw new RuntimeException("agent.jar.to.forward property must be set");
-    }
-
-    final String classpath =
-        agentJar + System.getProperty("path.separator") + System.getProperty("java.class.path");
+    final String classpath = System.getProperty("java.class.path");
     final String path = System.getProperty("java.home") + separator + "bin" + separator + "java";
 
     final List<String> vmArgsList = new ArrayList<>(Arrays.asList(jvmArgs));
-    vmArgsList.add("-javaagent:" + agentJar);
+    vmArgsList.add(getAgentArgument());
 
     final List<String> commands = new ArrayList<>();
     commands.add(path);
@@ -212,7 +218,6 @@ public class IntegrationTestUtils {
     commands.add(classpath);
     commands.add(mainClassName);
     commands.addAll(Arrays.asList(mainMethodArgs));
-
     final ProcessBuilder processBuilder = new ProcessBuilder(commands.toArray(new String[0]));
     processBuilder.environment().putAll(envVars);
     final Process process = processBuilder.start();