Skip to content

Commit

Permalink
feat(bootstrap): Add lib-injection multiple JVM agent guardrails
Browse files Browse the repository at this point in the history
  • Loading branch information
@PerfectSlayer
PerfectSlayer committed Jun 4, 2024
1 parent fccff88 commit 8940a4e
Show file tree
Hide file tree
Showing 2 changed files with 127 additions and 57 deletions.
142 changes: 85 additions & 57 deletions dd-java-agent/src/main/java/datadog/trace/bootstrap/AgentBootstrap.java
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
package datadog.trace.bootstrap;

import static java.nio.charset.StandardCharsets.UTF_8;

import de.thetaphi.forbiddenapis.SuppressForbidden;
import java.io.BufferedReader;
import java.io.File;
Expand All @@ -13,16 +15,14 @@
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.security.CodeSource;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Set;
import java.util.jar.JarFile;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
* Entry point for initializing the agent.
Expand All @@ -48,6 +48,9 @@
public final class AgentBootstrap {
private static final Class<?> thisClass = AgentBootstrap.class;
private static final int MAX_EXCEPTION_CHAIN_LENGTH = 99;
private static final String JAVA_AGENT_ARGUMENT = "-javaagent:";
static final String LIB_INJECTION_ENABLED_FLAG = "DD_INJECTION_ENABLED";
static final String LIB_INJECTION_FORCE_FLAG = "DD_INJECT_FORCE";

private static boolean initialized = false;

Expand All @@ -59,9 +62,14 @@ public static void agentmain(final String agentArgs, final Instrumentation inst)
if (alreadyInitialized() || lessThanJava8() || isJdkTool()) {
return;
}
// Prevent incompatibility issues with other JVM agents when setup using lib-injection
List<File> agentFiles = getAgentFilesFromVMArguments();
if (shouldAbortDueToOtherJavaAgents(agentFiles)) {
return;
}

try {
final URL agentJarURL = installAgentJar(inst);
final URL agentJarURL = installAgentJar(inst, agentFiles);
final Class<?> agentClass = Class.forName("datadog.trace.bootstrap.Agent", true, null);
if (agentClass.getClassLoader() != null) {
throw new IllegalStateException("DD Java Agent NOT added to bootstrap classpath.");
Expand Down Expand Up @@ -93,6 +101,16 @@ static boolean exceptionCauseChainContains(Throwable ex, String exClassName) {
return false;
}

private static boolean alreadyInitialized() {
if (initialized) {
System.out.println(
"Warning: dd-java-agent is being initialized more than once. Please check that you are defining -javaagent:dd-java-agent.jar only once.");
return true;
}
initialized = true;
return false;
}

private static boolean lessThanJava8() {
return lessThanJava8(System.getProperty("java.version"), System.out);
}
Expand All @@ -119,16 +137,6 @@ static boolean lessThanJava8(String version, PrintStream output) {
return false;
}

private static boolean alreadyInitialized() {
if (initialized) {
System.out.println(
"Warning: dd-java-agent is being initialized more than once. Please check that you are defining -javaagent:dd-java-agent.jar only once.");
return true;
}
initialized = true;
return false;
}

private static boolean isJdkTool() {
String moduleMain = System.getProperty("jdk.module.main");
if (null != moduleMain && !moduleMain.isEmpty() && moduleMain.charAt(0) == 'j') {
Expand Down Expand Up @@ -195,14 +203,43 @@ static int parseJavaMajorVersion(String version) {
return major;
}

static boolean shouldAbortDueToOtherJavaAgents(List<File> agentFiles) {
// Considered enabled if defined
boolean libInjectionEnabled = System.getenv(LIB_INJECTION_ENABLED_FLAG) != null;
// Considered forced if set to TRUE|True|1
String libInjectionForceFlag = System.getenv(LIB_INJECTION_FORCE_FLAG);
boolean libInjectionForced =
libInjectionForceFlag != null
&& (libInjectionForceFlag.equals("TRUE")
|| libInjectionForceFlag.equals("True")
|| libInjectionForceFlag.equals("1"));

// Simply considering having multiple agents
if (libInjectionEnabled && agentFiles.size() > 1) {
if (libInjectionForced) {
System.out.println(
"Warning: multiple JVM agents detected."
+ "Loading multiple APM/Tracing agent is not a recommended or supported configuration.");
} else {
System.out.println(
"Warning: multiple JVM agents detected."
+ "Loading multiple APM/Tracing agent is not a recommended or supported configuration."
+ "Please set the DD_INJECT_FORCE configuration to TRUE to load Datadog APM/Tracing agent.");
return true;
}
}
return false;
}

public static void main(final String[] args) {
if (lessThanJava8()) {
return;
}
AgentJar.main(args);
}

private static synchronized URL installAgentJar(final Instrumentation inst)
private static synchronized URL installAgentJar(
final Instrumentation inst, final List<File> agentFiles)
throws IOException, URISyntaxException {
// First try Code Source
final CodeSource codeSource = thisClass.getProtectionDomain().getCodeSource();
Expand All @@ -219,7 +256,7 @@ private static synchronized URL installAgentJar(final Instrumentation inst)
}

System.out.println("Could not get bootstrap jar from code source, using -javaagent arg");
File javaagentFile = getAgentFileFromJavaagentArg(inst);
File javaagentFile = getAgentFileFromJavaagentArg(agentFiles);
if (javaagentFile != null) {
URL ddJavaAgentJarURL = javaagentFile.toURI().toURL();
return appendAgentToBootstrapClassLoaderSearch(inst, ddJavaAgentJarURL, javaagentFile);
Expand All @@ -244,51 +281,43 @@ private static URL appendAgentToBootstrapClassLoaderSearch(
return ddJavaAgentJarURL;
}

private static File getAgentFileFromJavaagentArg(Instrumentation inst) throws IOException {
URL ddJavaAgentJarURL;
private static File getAgentFileFromJavaagentArg(List<File> agentFiles) {
if (agentFiles.isEmpty()) {
System.out.println("Could not get bootstrap jar from -javaagent arg: no argument specified");
return null;
} else if (agentFiles.size() > 1) {
System.out.println(
"Could not get bootstrap jar from -javaagent arg: multiple javaagents specified");
return null;
} else {
return agentFiles.get(0);
}
}

private static List<File> getAgentFilesFromVMArguments() {
final List<File> agentFiles = new ArrayList<>();
// 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 = getVMArgumentsThroughReflection();

String agentArgument = null;
for (final String arg : arguments) {
if (arg.startsWith("-javaagent")) {
if (agentArgument == null) {
agentArgument = arg;
// Use reflection to bypass the loading of the class~
for (final String argument : getVMArgumentsThroughReflection()) {
if (argument.startsWith(JAVA_AGENT_ARGUMENT)) {
int index = argument.indexOf('=', JAVA_AGENT_ARGUMENT.length());
String agentPathname =
argument.substring(
JAVA_AGENT_ARGUMENT.length(), index == -1 ? argument.length() : index);
File agentFile = new File(agentPathname);
if (agentFile.exists() && agentFile.isFile()) {
agentFiles.add(agentFile);
} else {
System.out.println(
"Could not get bootstrap jar from -javaagent arg: multiple javaagents specified");
return null;
"Could not get bootstrap jar from -javaagent arg: unable to find javaagent file: "
+ agentFile);
}
}
}

if (agentArgument == null) {
System.out.println("Could not get bootstrap jar from -javaagent arg: no argument specified");
return null;
}

// argument is of the form -javaagent:/path/to/dd-java-agent.jar=optionalargumentstring
final Matcher matcher = Pattern.compile("-javaagent:([^=]+).*").matcher(agentArgument);

if (!matcher.matches()) {
System.out.println(
"Could not get bootstrap jar from -javaagent arg: unable to parse javaagent parameter: "
+ agentArgument);
return null;
}

final File javaagentFile = new File(matcher.group(1));
if (!(javaagentFile.exists() || javaagentFile.isFile())) {
System.out.println(
"Could not get bootstrap jar from -javaagent arg: unable to find javaagent file: "
+ javaagentFile);
return null;
}
return javaagentFile;
return agentFiles;
}

@SuppressForbidden
Expand Down Expand Up @@ -345,23 +374,22 @@ private static List<String> getVMArgumentsThroughReflection() {
} 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");
"WARNING: Unable to get VM args through reflection. A custom java.util.logging.LogManager may not work correctly");

return ManagementFactory.getRuntimeMXBean().getInputArguments();
}
}
}

private static boolean checkJarManifestMainClassIsThis(final URL jarUrl) throws IOException {
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();
try (final BufferedReader reader =
new BufferedReader(
new InputStreamReader(manifestUrl.openStream(), StandardCharsets.UTF_8))) {
new BufferedReader(new InputStreamReader(manifestUrl.openStream(), UTF_8))) {
String line;
while ((line = reader.readLine()) != null) {
if (line.equals(mainClassLine)) {
return true;
return;
}
}
}
Expand Down
42 changes: 42 additions & 0 deletions dd-java-agent/src/test/groovy/datadog/trace/bootstrap/AgentBootstrapTest.groovy
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,17 @@
package datadog.trace.bootstrap

import org.junit.Rule
import org.junit.contrib.java.lang.system.EnvironmentVariables
import spock.lang.Specification

import static datadog.trace.bootstrap.AgentBootstrap.LIB_INJECTION_ENABLED_FLAG
import static datadog.trace.bootstrap.AgentBootstrap.LIB_INJECTION_FORCE_FLAG
import static datadog.trace.bootstrap.AgentBootstrap.shouldAbortDueToOtherJavaAgents

class AgentBootstrapTest extends Specification {
@Rule
public final EnvironmentVariables environmentVariables = new EnvironmentVariables()

def 'parse java.version strings'() {
when:
def major = AgentBootstrap.parseJavaMajorVersion(version)
Expand Down Expand Up @@ -138,4 +147,37 @@ class AgentBootstrapTest extends Specification {
then:
!causeChainContainsException
}

def 'proceed when not setup through lib-injection even with multiple agents'() {
setup:
def agentFiles = [new File("agent1.jar"), new File("agent2.jar")]

expect:
!shouldAbortDueToOtherJavaAgents(agentFiles)
}

def 'abort when setup through lib-injection and having multiples agents'() {
setup:
this.environmentVariables.set(LIB_INJECTION_ENABLED_FLAG, "TRUE")
def agentFiles = [new File("agent1.jar"), new File("agent2.jar")]

expect:
shouldAbortDueToOtherJavaAgents(agentFiles)

cleanup:
this.environmentVariables.clear(LIB_INJECTION_ENABLED_FLAG)
}

def 'proceed when setup through forced lib-injection even with multiple agents'() {
setup:
this.environmentVariables.set(LIB_INJECTION_ENABLED_FLAG, "TRUE")
this.environmentVariables.set(LIB_INJECTION_FORCE_FLAG, "TRUE")
def agentFiles = [new File("agent1.jar"), new File("agent2.jar")]

expect:
!shouldAbortDueToOtherJavaAgents(agentFiles)

cleanup:
this.environmentVariables.clear(LIB_INJECTION_ENABLED_FLAG, LIB_INJECTION_FORCE_FLAG)
}
}

0 comments on commit 8940a4e

Please sign in to comment.