Support generation of glue bytecode resources as well as glue java files

buildSrc/src/main/kotlin/instrument-glue.gradle.kts

@@ -4,44 +4,47 @@ plugins {
   id("java-common")
 }
 
-val generatedGlueDir = layout.buildDirectory.dir("generated/glue/java/").get()
+val generatedGlueResources = layout.buildDirectory.dir("generated/glue/resources/").get()
+val generatedGlueJava = layout.buildDirectory.dir("generated/glue/java/").get()
 sourceSets {
   create("glue") {
-    output.dir(generatedGlueDir)
+    output.dir(generatedGlueResources)
+    output.dir(generatedGlueJava)
   }
   main {
-    java.srcDir(generatedGlueDir)
+    resources.srcDir(generatedGlueResources)
+    java.srcDir(generatedGlueJava)
   }
 }
 
 val glueImplementation by configurations
 dependencies {
   glueImplementation(libs.asm)
+  glueImplementation(libs.spotbugs.annotations)
   glueImplementation(project(":utils"))
 }
 
 tasks.register<JavaExec>("generateGlue") {
   val glue: String by project.extra
 
-  val javaFile = generatedGlueDir.file("datadog/instrument/glue/${glue}.java")
+  // generate glue files under a consistent packaging location
+  val resourcesDir = generatedGlueResources.dir("datadog/instrument/glue/")
+  val javaDir = generatedGlueJava.dir("datadog/instrument/glue/")
 
   group = "Build"
   description = "Generate ${glue}"
   mainClass = "datadog.instrument.glue.${glue}Generator"
   classpath = sourceSets["glue"].runtimeClasspath
-  args = listOf(javaFile.toString())
-  outputs.files(javaFile)
+  args = listOf(resourcesDir.toString(), javaDir.toString())
+  outputs.dirs(resourcesDir, javaDir)
 }
 
+tasks.processResources {
+  dependsOn(tasks.named("generateGlue"))
+}
 tasks.compileJava {
   dependsOn(tasks.named("generateGlue"))
 }
 tasks.named("sourcesJar") {
   dependsOn(tasks.named("generateGlue"))
 }
-
-tasks.jar {
-  // glue classes only contain large string constants that get inlined into other classes
-  // - the inlining means we can safely drop these classes from the final jar to save space
-  excludes.add("datadog/instrument/glue/**")
-}

class-inject/build.gradle.kts

@@ -6,3 +6,9 @@ plugins {
 // class-inject generates glue bytecode at build-time to provide access to Unsafe.defineClass
 // we inline this glue as a string constant, then install it at runtime using instrumentation
 extra["glue"] = "DefineClassGlue"
+
+tasks.jar {
+  // DefineClassGlue only contains large string constants that get inlined into ClassInjector
+  // - the inlining means we can safely drop DefineClassGlue from the final jar to save space
+  excludes.add("datadog/instrument/glue/DefineClassGlue.class")
+}

class-inject/src/glue/java/datadog/instrument/glue/DefineClassGlueGenerator.java

@@ -11,10 +11,11 @@
 import static org.objectweb.asm.ClassWriter.COMPUTE_FRAMES;
 import static org.objectweb.asm.Opcodes.*;
 
-import java.io.File;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.BiFunction;
@@ -36,7 +37,7 @@
  */
 final class DefineClassGlueGenerator {
   // our glue must be located inside the java.lang namespace for accessibility reasons
-  private static final String DEFINECLASS_GLUE_CLASS = "java/lang/$Datadog$DefineClass$Glue$";
+  private static final String DEFINECLASSGLUE_CLASS = "java/lang/$Datadog$DefineClassGlue";
 
   private static final String OBJECT_CLASS = "java/lang/Object";
   private static final String CLASS_CLASS = "java/lang/Class";
@@ -94,7 +95,7 @@ private static byte[] generateBytecode(String unsafeNamespace) {
     cw.visit(
         V1_8,
         ACC_PUBLIC | ACC_FINAL,
-        DEFINECLASS_GLUE_CLASS,
+        DEFINECLASSGLUE_CLASS,
         null,
         OBJECT_CLASS,
         new String[] {BIFUNCTION_CLASS});
@@ -108,7 +109,7 @@ private static byte[] generateBytecode(String unsafeNamespace) {
     mv = cw.visitMethod(ACC_STATIC, "<clinit>", "()V", null, null);
     mv.visitCode();
     mv.visitMethodInsn(INVOKESTATIC, unsafeClass, "getUnsafe", "()" + unsafeDescriptor, false);
-    mv.visitFieldInsn(PUTSTATIC, DEFINECLASS_GLUE_CLASS, "UNSAFE", unsafeDescriptor);
+    mv.visitFieldInsn(PUTSTATIC, DEFINECLASSGLUE_CLASS, "UNSAFE", unsafeDescriptor);
     mv.visitInsn(RETURN);
     mv.visitMaxs(1, 0);
     mv.visitEnd();
@@ -291,7 +292,7 @@ private static byte[] generateBytecode(String unsafeNamespace) {
 
     // load Unsafe field into local variable for performance
     mv.visitLabel(setupUnsafeDefiner);
-    mv.visitFieldInsn(GETSTATIC, DEFINECLASS_GLUE_CLASS, "UNSAFE", unsafeDescriptor);
+    mv.visitFieldInsn(GETSTATIC, DEFINECLASSGLUE_CLASS, "UNSAFE", unsafeDescriptor);
     mv.visitVarInsn(ASTORE, unsafeInstance);
 
     // check if we've defined all the given bytecode
@@ -351,22 +352,21 @@ private static byte[] generateBytecode(String unsafeNamespace) {
   }
 
   public static void main(String[] args) throws IOException {
-    if (args.length < 1 || !args[0].endsWith(".java")) {
-      throw new IllegalArgumentException("Expected: java-file");
+    if (args.length != 2) {
+      throw new IllegalArgumentException("Expected: resources-dir java-dir");
     }
-    File file = new File(args[0]);
-    String name = file.getName();
+    Path defineClassGlue = Paths.get(args[1], "DefineClassGlue.java");
     List<String> lines = new ArrayList<>();
-    classHeader(lines, name.substring(0, name.length() - 5));
+    classHeader(lines, "DefineClassGlue");
     lines.add("  /** Glue Id */");
-    lines.add("  String ID = \"" + DEFINECLASS_GLUE_CLASS.replace('/', '.') + "\";");
+    lines.add("  String ID = \"" + DEFINECLASSGLUE_CLASS.replace('/', '.') + "\";");
     lines.add("  /** Packed Java 8 bytecode */");
     lines.add("  String V8 =");
     packBytecode(lines, generateBytecode("sun/misc"));
     lines.add("  /** Packed Java 9+ bytecode */");
     lines.add("  String V9 =");
     packBytecode(lines, generateBytecode("jdk/internal/misc"));
     lines.add("}");
-    Files.write(file.toPath(), lines, StandardCharsets.UTF_8);
+    Files.write(defineClassGlue, lines, StandardCharsets.UTF_8);
   }
 }

utils/src/main/java/datadog/instrument/utils/Glue.java

@@ -8,7 +8,10 @@
 
 import static java.nio.charset.StandardCharsets.UTF_16BE;
 
+import java.io.ByteArrayOutputStream;
+import java.io.InputStream;
 import java.util.List;
+import java.util.MissingResourceException;
 
 /** Methods for packing glue bytecode into strings that can be stored in the constant pool. */
 public final class Glue {
@@ -69,4 +72,38 @@ public static void packBytecode(List<String> lines, byte[] bytecode) {
   public static byte[] unpackBytecode(String bytecode) {
     return bytecode.getBytes(UTF_16BE);
   }
+
+  private static final String GLUE_RESOURCE_PREFIX = "/datadog/instrument/glue/";
+
+  private static final int BUFFER_SIZE = 8192;
+
+  /**
+   * Loads glue bytecode with the given name from the given host.
+   *
+   * @param host the host class used to load the glue resource
+   * @param glueName the glue resource containing the bytecode
+   * @return the glue bytecode
+   * @throws MissingResourceException if the bytecode cannot be read
+   */
+  @SuppressWarnings({"Since15", "DataFlowIssue"})
+  public static byte[] loadBytecode(Class<?> host, String glueName) {
+    String glueResource = GLUE_RESOURCE_PREFIX + glueName;
+    try (InputStream is = host.getResourceAsStream(glueResource)) {
+      if (JVM.atLeastJava(9)) {
+        return is.readAllBytes();
+      } else {
+        try (ByteArrayOutputStream os = new ByteArrayOutputStream()) {
+          int bytesRead;
+          byte[] buf = new byte[BUFFER_SIZE];
+          while ((bytesRead = is.read(buf, 0, BUFFER_SIZE)) != -1) {
+            os.write(buf, 0, bytesRead);
+          }
+          return os.toByteArray();
+        }
+      }
+    } catch (Throwable e) {
+      String detail = "Cannot load " + glueResource + " from " + host + ": " + e;
+      throw new MissingResourceException(detail, host.getName(), glueResource);
+    }
+  }
 }

utils/src/test/java/datadog/instrument/utils/GlueTest.java

@@ -12,6 +12,7 @@
 import java.net.URLClassLoader;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.MissingResourceException;
 import javax.tools.JavaCompiler;
 import javax.tools.SimpleJavaFileObject;
 import javax.tools.StandardJavaFileManager;
@@ -21,17 +22,17 @@
 
 class GlueTest {
 
+  private static final byte[] exampleBytecode = {
+    -54, -2, -70, -66, 0, 0, 0, 52, 0, 10, 10, 0, 2, 0, 3, 7, 0, 4, 12, 0, 5, 0, 6, 1, 0, 16, 106,
+    97, 118, 97, 47, 108, 97, 110, 103, 47, 79, 98, 106, 101, 99, 116, 1, 0, 6, 60, 105, 110, 105,
+    116, 62, 1, 0, 3, 40, 41, 86, 7, 0, 8, 1, 0, 7, 69, 120, 97, 109, 112, 108, 101, 1, 0, 4, 67,
+    111, 100, 101, 0, 32, 0, 7, 0, 2, 0, 0, 0, 0, 0, 1, 0, 0, 0, 5, 0, 6, 0, 1, 0, 9, 0, 0, 0, 17,
+    0, 1, 0, 1, 0, 0, 0, 5, 42, -73, 0, 1, -79, 0, 0, 0, 0, 0, 0
+  };
+
   @Test
   void glueGeneration(@TempDir File classesDir) throws Exception {
 
-    byte[] exampleBytecode = {
-      -54, -2, -70, -66, 0, 0, 0, 52, 0, 10, 10, 0, 2, 0, 3, 7, 0, 4, 12, 0, 5, 0, 6, 1, 0, 16, 106,
-      97, 118, 97, 47, 108, 97, 110, 103, 47, 79, 98, 106, 101, 99, 116, 1, 0, 6, 60, 105, 110, 105,
-      116, 62, 1, 0, 3, 40, 41, 86, 7, 0, 8, 1, 0, 7, 69, 120, 97, 109, 112, 108, 101, 1, 0, 4, 67,
-      111, 100, 101, 0, 32, 0, 7, 0, 2, 0, 0, 0, 0, 0, 1, 0, 0, 0, 5, 0, 6, 0, 1, 0, 9, 0, 0, 0, 17,
-      0, 1, 0, 1, 0, 0, 0, 5, 42, -73, 0, 1, -79, 0, 0, 0, 0, 0, 0
-    };
-
     JavaCompiler compiler = getSystemJavaCompiler();
     StandardJavaFileManager fileManager = compiler.getStandardFileManager(null, null, null);
     fileManager.setLocation(StandardLocation.CLASS_OUTPUT, singletonList(classesDir));
@@ -62,7 +63,7 @@ void glueGeneration(@TempDir File classesDir) throws Exception {
   @Test
   void paddingRequired() {
 
-    byte[] exampleBytecode = {
+    byte[] exampleBytecodeNeedsPadding = {
       -54, -2, -70, -66, 0, 0, 0, 52, 0, 10, 10, 0, 2, 0, 3, 7, 0, 4, 12, 0, 5, 0, 6, 1, 0, 16, 106,
       97, 118, 97, 47, 108, 97, 110, 103, 47, 79, 98, 106, 101, 99, 116, 1, 0, 6, 60, 105, 110, 105,
       116, 62, 1, 0, 3, 40, 41, 86, 7, 0, 8, 1, 0, 2, 69, 103, 1, 0, 4, 67, 111, 100, 101, 0, 32, 0,
@@ -72,10 +73,22 @@ void paddingRequired() {
 
     assertThrows(
         IllegalStateException.class,
-        () -> Glue.packBytecode(new ArrayList<>(), exampleBytecode),
+        () -> Glue.packBytecode(new ArrayList<>(), exampleBytecodeNeedsPadding),
         "Bytecode length is not even; requires padding");
   }
 
+  @Test
+  void glueResource() {
+    byte[] loadedBytecode = Glue.loadBytecode(GlueTest.class, "test.glue");
+    assertArrayEquals(exampleBytecode, loadedBytecode);
+  }
+
+  @Test
+  void missingGlueResource() {
+    assertThrows(
+        MissingResourceException.class, () -> Glue.loadBytecode(GlueTest.class, "missing.glue"));
+  }
+
   static class InMemoryJavaFile extends SimpleJavaFileObject {
     private final List<String> lines;