From a697881efbcceb8a3c957d1473ac38cb4f4d9610 Mon Sep 17 00:00:00 2001
From: Andrew Kent <andrew.kent@datadoghq.com>
Date: Thu, 11 Oct 2018 19:49:40 -0700
Subject: [PATCH] Assert no bootstrap classes in test fields or method
 signatures

---
 .../datadog/trace/agent/test/SpockRunner.java | 40 ++++++++++++++++---
 .../test/groovy/AgentTestRunnerTest.groovy    |  9 +----
 2 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/dd-java-agent/testing/src/main/java/datadog/trace/agent/test/SpockRunner.java b/dd-java-agent/testing/src/main/java/datadog/trace/agent/test/SpockRunner.java
index f7653e1d8b..57ce0c5487 100644
--- a/dd-java-agent/testing/src/main/java/datadog/trace/agent/test/SpockRunner.java
+++ b/dd-java-agent/testing/src/main/java/datadog/trace/agent/test/SpockRunner.java
@@ -4,6 +4,7 @@ import com.google.common.reflect.ClassPath;
 import java.io.File;
 import java.io.IOException;
 import java.lang.reflect.Field;
+import java.lang.reflect.Method;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.Set;
@@ -66,6 +67,7 @@ public class SpockRunner extends Sputnik {
       throws InitializationError, NoSuchFieldException, SecurityException, IllegalArgumentException,
           IllegalAccessException {
     super(shadowTestClass(clazz));
+    assertNoBootstrapClassesInTestClass(clazz);
     // access the classloader created in shadowTestClass above
     Field clazzField = Sputnik.class.getDeclaredField("clazz");
     try {
@@ -77,6 +79,37 @@ public class SpockRunner extends Sputnik {
     }
   }
 
+  private static void assertNoBootstrapClassesInTestClass(Class<?> testClass) {
+    for (final Field field : testClass.getDeclaredFields()) {
+      assertNotBootstrapClass(testClass, field.getType());
+    }
+
+    for (final Method method : testClass.getDeclaredMethods()) {
+      assertNotBootstrapClass(testClass, method.getReturnType());
+      for (final Class paramType : method.getParameterTypes()) {
+        assertNotBootstrapClass(testClass, paramType);
+      }
+    }
+  }
+
+  private static void assertNotBootstrapClass(Class<?> testClass, Class<?> clazz) {
+    if ((!clazz.isPrimitive()) && isBootstrapClass(clazz.getName())) {
+      throw new IllegalStateException(
+          testClass.getName()
+              + ": Bootstrap classes are not allowed in test class field or method signatures. Offending class: "
+              + clazz.getName());
+    }
+  }
+
+  private static boolean isBootstrapClass(String className) {
+    for (int i = 0; i < TEST_BOOTSTRAP_PREFIXES.length; ++i) {
+      if (className.startsWith(TEST_BOOTSTRAP_PREFIXES[i])) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   // Shadow the test class with bytes loaded by InstrumentationClassLoader
   private static Class<?> shadowTestClass(final Class<?> clazz) {
     try {
@@ -117,11 +150,8 @@ public class SpockRunner extends Sputnik {
     Set<String> bootstrapClasses = new HashSet<String>();
     for (ClassPath.ClassInfo info : TestUtils.getTestClasspath().getAllClasses()) {
       // if info starts with bootstrap prefix: add to bootstrap jar
-      for (int i = 0; i < TEST_BOOTSTRAP_PREFIXES.length; ++i) {
-        if (info.getName().startsWith(TEST_BOOTSTRAP_PREFIXES[i])) {
-          bootstrapClasses.add(info.getResourceName());
-          break;
-        }
+      if (isBootstrapClass(info.getName())) {
+        bootstrapClasses.add(info.getResourceName());
       }
     }
     return new File(
diff --git a/dd-java-agent/testing/src/test/groovy/AgentTestRunnerTest.groovy b/dd-java-agent/testing/src/test/groovy/AgentTestRunnerTest.groovy
index cad6839c25..33b68e49fa 100644
--- a/dd-java-agent/testing/src/test/groovy/AgentTestRunnerTest.groovy
+++ b/dd-java-agent/testing/src/test/groovy/AgentTestRunnerTest.groovy
@@ -13,10 +13,6 @@ class AgentTestRunnerTest extends AgentTestRunner {
   private static final ClassLoader BOOTSTRAP_CLASSLOADER = null
   private static final ClassLoader OT_LOADER
   private static final boolean AGENT_INSTALLED_IN_CLINIT
-  // having opentracing class in test field should not cause problems
-  private static final Tracer A_TRACER = null
-  // having dd tracer api class in test field should not cause problems
-  private static final datadog.trace.api.Tracer DD_API_TRACER = null
 
   @Shared
   private Class sharedSpanClass
@@ -52,10 +48,9 @@ class AgentTestRunnerTest extends AgentTestRunner {
     }
 
     expect:
-    // a shared OT class should cause no trouble
+    // shared OT classes should cause no trouble
     sharedSpanClass.getClassLoader() == BOOTSTRAP_CLASSLOADER
-    A_TRACER == null
-    DD_API_TRACER  == null
+    Tracer.getClassLoader() == BOOTSTRAP_CLASSLOADER
     !AGENT_INSTALLED_IN_CLINIT
     getTestTracer() == TestUtils.getUnderlyingGlobalTracer()
     getAgentTransformer() != null