From 42832e77f5b450896419c7c08b847e77bf2fc16d Mon Sep 17 00:00:00 2001
From: Andrew Kent <andrew.kent@datadoghq.com>
Date: Thu, 7 Dec 2017 14:44:34 -0800
Subject: [PATCH] Add agent's classloader to default instrumentation location
 strategy

The transformer needs the agent's classloader to resolve advice
classes at transform time. Also stop using reflection to register classloads.
---
 .../com/datadoghq/agent/TracingAgent.java     | 18 ++++++++--------
 .../src/main/java/dd/trace/DDAdvice.java      | 21 +++++++++++++++++++
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/dd-java-agent/src/main/java/com/datadoghq/agent/TracingAgent.java b/dd-java-agent/src/main/java/com/datadoghq/agent/TracingAgent.java
index dd0866e5f6..5ced172430 100644
--- a/dd-java-agent/src/main/java/com/datadoghq/agent/TracingAgent.java
+++ b/dd-java-agent/src/main/java/com/datadoghq/agent/TracingAgent.java
@@ -25,7 +25,6 @@ import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
 
 import dd.trace.Instrumenter;
 import java.lang.instrument.Instrumentation;
-import java.lang.reflect.Method;
 import java.util.Collections;
 import java.util.ServiceLoader;
 import java.util.Set;
@@ -40,7 +39,12 @@ import net.bytebuddy.utility.JavaModule;
 @Slf4j
 public class TracingAgent {
 
-  public static void premain(final String agentArgs, final Instrumentation inst) throws Exception {
+  /** Return the classloader the core agent is running on. */
+  public static ClassLoader getAgentClassLoader() {
+    return TracingAgent.class.getClassLoader();
+  }
+
+  public static void premain(String agentArgs, final Instrumentation inst) throws Exception {
     log.debug("Using premain for loading {}", TracingAgent.class.getSimpleName());
     addByteBuddy(inst);
     AgentRulesManager.initialize();
@@ -111,7 +115,7 @@ public class TracingAgent {
         final JavaModule module,
         final boolean loaded,
         final DynamicType dynamicType) {
-      log.debug("Transformed {}", typeDescription);
+      log.debug("Transformed {} -- {}", typeDescription, classLoader);
 
       if (classLoader == null) {
         return;
@@ -123,13 +127,9 @@ public class TracingAgent {
         initializedClassloaders.add(classLoader);
 
         try {
-          final Class<?> rulesManager =
-              Class.forName("com.datadoghq.agent.InstrumentationRulesManager", true, classLoader);
-          final Method registerClassLoad =
-              rulesManager.getDeclaredMethod("registerClassLoad", Object.class);
-          registerClassLoad.invoke(null, classLoader);
+          InstrumentationRulesManager.registerClassLoad(classLoader);
         } catch (final Throwable e) {
-          log.info("ClassLoad Registration for target " + classLoader, e);
+          log.error("Failed ClassLoad Registration for target " + classLoader, e);
         }
       }
     }
diff --git a/dd-java-agent/tooling/src/main/java/dd/trace/DDAdvice.java b/dd-java-agent/tooling/src/main/java/dd/trace/DDAdvice.java
index b8503c8568..bcb43d2e11 100644
--- a/dd-java-agent/tooling/src/main/java/dd/trace/DDAdvice.java
+++ b/dd-java-agent/tooling/src/main/java/dd/trace/DDAdvice.java
@@ -1,11 +1,32 @@
 package dd.trace;
 
+import java.lang.reflect.Method;
+import lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.dynamic.ClassFileLocator;
 
 /** A bytebuddy advice builder with default DataDog settings. */
+@Slf4j
 public class DDAdvice extends AgentBuilder.Transformer.ForAdvice {
+  private static ClassLoader AGENT_CLASSLOADER;
+
+  static {
+    try {
+      Class<?> agentClass =
+          DDAdvice.class.getClassLoader().loadClass("com.datadoghq.agent.TracingAgent");
+      Method getAgentClassloaderMethod = agentClass.getMethod("getAgentClassLoader");
+      AGENT_CLASSLOADER = (ClassLoader) getAgentClassloaderMethod.invoke(null);
+    } catch (Throwable t) {
+      log.error("Unable to locate agent classloader. Falling back to System Classloader");
+      AGENT_CLASSLOADER = ClassLoader.getSystemClassLoader();
+    }
+  }
+
   public static AgentBuilder.Transformer.ForAdvice create() {
     return new DDAdvice()
+        .with(
+            new AgentBuilder.LocationStrategy.Simple(
+                ClassFileLocator.ForClassLoader.of(AGENT_CLASSLOADER)))
         .withExceptionHandler(ExceptionHandlers.defaultExceptionHandler());
   }