Static instrumenter - core (#236)

* static instrumenter core

* Apply suggestions from code review

Co-authored-by: Anuraag Agrawal <[email protected]>

* post review changes

* Apply suggestions from code review

Co-authored-by: Mateusz Rzeszutek <[email protected]>

* code review

* Apply suggestions from code review

Co-authored-by: Trask Stalnaker <[email protected]>

* code review

* code review

* code review

* code review

* code review

* code review

* code review - file cleanup

* forcefuly adding binary file

Co-authored-by: Anuraag Agrawal <[email protected]>
Co-authored-by: Mateusz Rzeszutek <[email protected]>
Co-authored-by: Trask Stalnaker <[email protected]>

Author: 4 people

static-instrumenter/README.md

@@ -6,7 +6,7 @@
 
 Module enhancing OpenTelemetry Java Agent for static instrumentation. The modified agent is capable of instrumenting and saving a new JAR with all relevant instrumentations applied and necessary helper class-code included.
 
-In order to statically instrument a JAR, modified agent needs to be both attached (`-javaagent:`) and run as the main method (`io.opentelemetry.javaagent.StaticInstrumenter` class).
+In order to statically instrument a JAR, modified agent needs to be both attached (`-javaagent:`) and run as the main method (`io.opentelemetry.contrib.staticinstrumenter.Main` class).
 
 ### Gradle plugin
 

static-instrumenter/agent-instrumenter/build.gradle.kts

@@ -3,3 +3,16 @@ plugins {
 }
 
 description = "OpenTelemetry Java Static Instrumentation Agent"
+
+dependencies {
+  implementation("org.slf4j:slf4j-api")
+  runtimeOnly("org.slf4j:slf4j-simple")
+}
+
+tasks {
+  withType<JavaCompile>().configureEach {
+    with(options) {
+      release.set(11)
+    }
+  }
+}

static-instrumenter/agent-instrumenter/src/main/java/io/opentelemetry/contrib/.gitkeep

@@ -0,0 +1,54 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.staticinstrumenter;
+
+class ArchiveEntry {
+
+  private static final ArchiveEntry NOT_CLASS =
+      new ArchiveEntry("", "", /* shouldInstrument= */ false);
+
+  private final String name;
+  private final String path;
+  private final boolean shouldInstrument;
+
+  private ArchiveEntry(String name, String path, boolean shouldInstrument) {
+    this.name = name;
+    this.path = path;
+    this.shouldInstrument = shouldInstrument;
+  }
+
+  static ArchiveEntry fromZipEntryName(String zipEntryName) {
+    if (!isClass(zipEntryName)) {
+      return NOT_CLASS;
+    }
+    String path = zipEntryName.substring(0, zipEntryName.indexOf(".class"));
+    return new ArchiveEntry(className(path), path, !shouldBeSkipped(zipEntryName));
+  }
+
+  private static boolean isClass(String path) {
+    return path.endsWith(".class");
+  }
+
+  private static String className(String path) {
+    return path.replace("/", ".");
+  }
+
+  private static boolean shouldBeSkipped(String zipEntryName) {
+    return zipEntryName.startsWith("io.opentelemetry");
+  }
+
+  String getName() {
+    return name;
+  }
+
+  String getPath() {
+    return path;
+  }
+
+  boolean shouldInstrument() {
+    return shouldInstrument;
+  }
+}

static-instrumenter/agent-instrumenter/src/main/java/io/opentelemetry/contrib/staticinstrumenter/ArchiveEntry.java

@@ -0,0 +1,83 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.staticinstrumenter;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Enumeration;
+import java.util.Map;
+import java.util.jar.JarEntry;
+import java.util.jar.JarFile;
+import java.util.jar.JarOutputStream;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Represents an archive storing classes (JAR, WAR). */
+class ClassArchive {
+
+  interface Factory {
+    ClassArchive createFor(JarFile source, Map<String, byte[]> instrumentedClasses);
+  }
+
+  private static final Logger logger = LoggerFactory.getLogger(ClassArchive.class);
+
+  private final JarFile source;
+  private final Map<String, byte[]> instrumentedClasses;
+
+  ClassArchive(JarFile source, Map<String, byte[]> instrumentedClasses) {
+    this.source = source;
+    this.instrumentedClasses = instrumentedClasses;
+  }
+
+  void copyAllClassesTo(JarOutputStream outJar) throws IOException {
+
+    Enumeration<JarEntry> inEntries = source.entries();
+    while (inEntries.hasMoreElements()) {
+      copyEntry(inEntries.nextElement(), outJar);
+    }
+  }
+
+  private void copyEntry(JarEntry inEntry, JarOutputStream outJar) throws IOException {
+    String inEntryName = inEntry.getName();
+    ZipEntry outEntry =
+        inEntryName.endsWith(".jar") ? new ZipEntry(inEntry) : new ZipEntry(inEntryName);
+
+    try (InputStream entryInputStream = getInputStreamForEntry(inEntry, outEntry)) {
+      outJar.putNextEntry(outEntry);
+      entryInputStream.transferTo(outJar);
+      outJar.closeEntry();
+    } catch (ZipException e) {
+      if (!isEntryDuplicate(e)) {
+        logger.error("Error while creating entry: {}", outEntry.getName(), e);
+        throw e;
+      }
+    }
+  }
+
+  private static boolean isEntryDuplicate(ZipException ze) {
+    return ze.getMessage() != null && ze.getMessage().contains("duplicate");
+  }
+
+  private InputStream getInputStreamForEntry(JarEntry inEntry, ZipEntry outEntry)
+      throws IOException {
+
+    InputStream entryIn = null;
+    ArchiveEntry entry = ArchiveEntry.fromZipEntryName(inEntry.getName());
+    if (entry.shouldInstrument()) {
+      String className = entry.getName();
+      byte[] modified = instrumentedClasses.get(entry.getPath());
+      if (modified != null) {
+        logger.debug("Found instrumented class: " + className);
+        entryIn = new ByteArrayInputStream(modified);
+        outEntry.setSize(modified.length);
+      }
+    }
+    return entryIn == null ? source.getInputStream(inEntry) : entryIn;
+  }
+}

static-instrumenter/agent-instrumenter/src/main/java/io/opentelemetry/contrib/staticinstrumenter/ClassArchive.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.staticinstrumenter;
+
+class CurrentClass {
+
+  private static final ThreadLocal<TransformedClass> currentClass = new ThreadLocal<>();
+
+  private CurrentClass() {}
+
+  static TransformedClass getAndRemove() {
+    TransformedClass tc = currentClass.get();
+    currentClass.remove();
+    return tc;
+  }
+
+  static void set(TransformedClass clazz) {
+    currentClass.set(clazz);
+  }
+}

static-instrumenter/agent-instrumenter/src/main/java/io/opentelemetry/contrib/staticinstrumenter/CurrentClass.java

@@ -0,0 +1,146 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.staticinstrumenter;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.lang.instrument.ClassFileTransformer;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.jar.JarFile;
+import java.util.jar.JarOutputStream;
+import java.util.zip.ZipEntry;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class Main {
+
+  private static final Logger logger = LoggerFactory.getLogger(Main.class);
+
+  private static final Main INSTANCE = new Main(ClassArchive::new);
+
+  private final ClassArchive.Factory classArchiveFactory;
+
+  // key is slashy name, not dotty
+  private final Map<String, byte[]> instrumentedClasses = new ConcurrentHashMap<>();
+  private final Map<String, byte[]> additionalClasses = new ConcurrentHashMap<>();
+
+  public static void main(String[] args) throws Exception {
+
+    if (args.length != 1) {
+      printUsage();
+      return;
+    }
+
+    File outDir = new File(args[0]);
+    if (!outDir.exists()) {
+      outDir.mkdir();
+    }
+
+    String classPath = System.getProperty("java.class.path");
+    logger.debug("Classpath (jars list): " + classPath);
+    String[] jarsList = classPath.split(File.pathSeparator);
+
+    getInstance().saveTransformedJarsTo(jarsList, outDir);
+  }
+
+  @SuppressWarnings("SystemOut")
+  private static void printUsage() {
+    System.out.println(
+        "OpenTelemetry Java Static Instrumenter\n"
+            + "Usage:\njava "
+            + Main.class.getCanonicalName()
+            + " <output directory> (where instrumented archives will be stored)");
+  }
+
+  public static Main getInstance() {
+    return INSTANCE;
+  }
+
+  public static ClassFileTransformer getPreTransformer() {
+    return new PreTransformer();
+  }
+
+  public static ClassFileTransformer getPostTransformer() {
+    return new PostTransformer();
+  }
+
+  // for testing purposes
+  Main(ClassArchive.Factory classArchiveFactory) {
+    this.classArchiveFactory = classArchiveFactory;
+  }
+
+  // FIXME: java 9 / jmod support, proper handling of directories
+  // FIXME: jmod in particular introduces weirdness with adding helpers to the dependencies
+
+  /**
+   * Copies all class archives (JARs, WARs) to outDir. Classes that were instrumented and stored in
+   * instrumentedClasses will get replaced with the new version. All classes added to
+   * additionalClasses will be added to the new archive.
+   *
+   * @param outDir directory where jars will be written
+   * @throws IOException in case of file operation problem
+   */
+  public void saveTransformedJarsTo(String[] jarsList, File outDir) throws IOException {
+
+    for (String pathItem : jarsList) {
+      logger.info("Classpath item processed: " + pathItem);
+      if (isArchive(pathItem)) {
+        saveArchiveTo(new File(pathItem), outDir);
+      }
+    }
+  }
+
+  private static boolean isArchive(String pathItem) {
+    return (pathItem.endsWith(".jar") || pathItem.endsWith(".war"));
+  }
+
+  // FIXME: don't "instrument" our agent jar
+  // FIXME: detect and warn on signed jars (and drop the signing bits)
+  // FIXME: multiple jars with same name
+  private void saveArchiveTo(File inFile, File outDir) throws IOException {
+
+    try (JarFile inJar = new JarFile(inFile);
+        JarOutputStream outJar = jarOutputStreamFor(outDir, inFile.getName())) {
+      ClassArchive inClassArchive = classArchiveFactory.createFor(inJar, instrumentedClasses);
+      inClassArchive.copyAllClassesTo(outJar);
+      injectAdditionalClassesTo(outJar);
+    }
+  }
+
+  private static JarOutputStream jarOutputStreamFor(File outDir, String fileName)
+      throws IOException {
+    File outFile = new File(outDir, fileName);
+    return new JarOutputStream(new FileOutputStream(outFile));
+  }
+
+  // FIXME: only relevant additional classes should be injected
+  private void injectAdditionalClassesTo(JarOutputStream outJar) throws IOException {
+    for (Map.Entry<String, byte[]> entry : additionalClasses.entrySet()) {
+      String className = entry.getKey();
+      byte[] classData = entry.getValue();
+
+      ZipEntry newEntry = new ZipEntry(className);
+      outJar.putNextEntry(newEntry);
+      if (classData != null) {
+        newEntry.setSize(classData.length);
+        outJar.write(classData);
+      }
+      outJar.closeEntry();
+
+      logger.debug("Additional class added: {}", className);
+    }
+  }
+
+  public Map<String, byte[]> getInstrumentedClasses() {
+    return instrumentedClasses;
+  }
+
+  public Map<String, byte[]> getAdditionalClasses() {
+    return additionalClasses;
+  }
+}

static-instrumenter/agent-instrumenter/src/main/java/io/opentelemetry/contrib/staticinstrumenter/Main.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.staticinstrumenter;
+
+import java.lang.instrument.ClassFileTransformer;
+import java.security.ProtectionDomain;
+import java.util.Arrays;
+import javax.annotation.Nullable;
+
+public class PostTransformer implements ClassFileTransformer {
+  @Override
+  @Nullable
+  public byte[] transform(
+      ClassLoader loader,
+      String className,
+      Class<?> classBeingRedefined,
+      ProtectionDomain protectionDomain,
+      byte[] classfileBuffer) {
+
+    TransformedClass pre = CurrentClass.getAndRemove();
+
+    if (pre != null
+        && pre.getName().equals(className)
+        && !Arrays.equals(pre.getClasscode(), classfileBuffer)) {
+      Main.getInstance().getInstrumentedClasses().put(className, classfileBuffer);
+    }
+    return null;
+  }
+}

static-instrumenter/agent-instrumenter/src/main/java/io/opentelemetry/contrib/staticinstrumenter/PostTransformer.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.staticinstrumenter;
+
+import java.lang.instrument.ClassFileTransformer;
+import java.security.ProtectionDomain;
+import javax.annotation.Nullable;
+
+public class PreTransformer implements ClassFileTransformer {
+
+  @Override
+  @Nullable
+  public byte[] transform(
+      ClassLoader loader,
+      String className,
+      Class<?> classBeingRedefined,
+      ProtectionDomain protectionDomain,
+      byte[] classfileBuffer) {
+
+    CurrentClass.set(new TransformedClass(className, classfileBuffer));
+    return null;
+  }
+}