premain

package org.example;

public class MyAgent {
  public static void premain(String args) {
    System.out.println("premain");
  }
}

• Java 虚拟机所能识别的premain方法接收的是字符串类型的参数，而并非类似于main方法的字符串数组
• 需要将其打包成 jar 包，并在其中的 MANIFEST.MF 配置文件中，指定所谓的Premain-class，才能以java agent运行该premain方法

# 注意第一条命令会向manifest.txt文件写入两行数据，其中包括一行空行
$ echo 'Premain-Class: org.example.MyAgent
' > manifest.txt
$ jar cvmf manifest.txt myagent.jar org/
$ java -javaagent:myagent.jar HelloWorld
premain
Hello, World

• 也可以通过 Attach API 远程加载

import java.io.IOException;

import com.sun.tools.attach.*;

public class AttachTest {
  public static void main(String[] args)
      throws AttachNotSupportedException, IOException, AgentLoadException, AgentInitializationException {
    if (args.length <= 1) {
      System.out.println("Usage: java AttachTest <PID> /PATH/TO/AGENT.jar");
      return;
    }
    VirtualMachine vm = VirtualMachine.attach(args[0]);
    vm.loadAgent(args[1]);
  }
}

• Attach API 远程加载的 Java agent 不会再先于main方法执
• 运行的也不是premain方法，而是名为agentmain的方法

public class MyAgent { 
  public static void agentmain(String args) {
    System.out.println("agentmain");
  }
}

• 需要更新 jar 包中的 manifest 文件，使其包含Agent-Class的配置

$ echo 'Agent-Class: org.example.MyAgent
' > manifest.txt
$ jar cvmf manifest.txt myagent.jar org/
$ java HelloWorld
Hello, World
$ jps
$ java AttachTest <pid> myagent.jar
agentmain
// 最后一句输出来自于运行HelloWorld的Java进程

• Java 虚拟机并不限制 Java agent 的数量。你可以在 java 命令后附上多个-javaagent参数，或者远程 attach 多个 Java agent，Java 虚拟机会按照定义顺序，或者 attach 的顺序逐个执行这些 Java agent
• Java agent 还提供了一套 instrumentation 机制，允许应用程序拦截类加载事件，并且更改该类的字节码

字节码注入

package org.example;

import java.lang.instrument.*;
import java.security.ProtectionDomain;

public class MyAgent {
  public static void premain(String args, Instrumentation instrumentation) {
    instrumentation.addTransformer(new MyTransformer());
  }

  static class MyTransformer implements ClassFileTransformer {
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
        ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
      System.out.printf("Loaded %s: 0x%X%X%X%X\n", className, classfileBuffer[0], classfileBuffer[1],
          classfileBuffer[2], classfileBuffer[3]);
      return null;
    }
  }
}

• 可以通过Instrumentation来注册类加载事件的拦截器。该拦截器需要实现ClassFileTransformer接口，并重写其中的transform方法
• transform方法接收一个 byte 数组类型的参数，它代表的是正在被加载的类的字节码
• transform方法将返回一个 byte 数组，代表更新过后的类的字节码
  • 当方法返回之后，Java 虚拟机会使用所返回的 byte 数组，来完成接下来的类加载工作
  • 如果transform方法返回 null 或者抛出异常，那么 Java 虚拟机将使用原来的 byte 数组完成类加载工作
• 基于这一类加载事件的拦截功能，我们可以实现字节码注入（bytecode instrumentation），往正在被加载的类中插入额外的字节码
• Java agent 还提供了另外两个功能redefine和retransform
  • 针对的是已加载的类，并要求用户传入所要redefine或者retransform的类实例
• redefine指的是舍弃原本的字节码，并替换成由用户提供的 byte 数组。该功能比较危险，一般用于修复出错了的字节码
• retransform则将针对所传入的类，重新调用所有已注册的ClassFileTransformer的transform方法
  • 在执行premain或者agentmain方法前，Java 虚拟机早已加载了不少类，而这些类的加载事件并没有被拦截，因此也没有被注入。使用retransform功能可以注入这些已加载但未注入的类
  • 在定义了多个 Java agent，多个注入的情况下，我们可能需要移除其中的部分注入。当调用Instrumentation.removeTransformer去除某个注入类后，我们可以调用retransform功能，重新从原始 byte 数组开始进行注入
• Java agent 的这些功能都是通过 JVMTI agent，也就是 C agent 来实现的
  • JVMTI 是一个事件驱动的工具实现接口，通常，我们会在 C agent 加载后的入口方法Agent_OnLoad处注册各个事件的钩子（hook）方法
  • 当 Java 虚拟机触发了这些事件时，便会调用对应的钩子方法
  • 可以为 JVMTI 中的ClassFileLoadHook事件设置钩子，从而在 C 层面拦截所有的类加载事件

基于字节码注入的 profiler

• 我们可以利用字节码注入来实现代码覆盖工具
• 定义一个运行时类，并在某一程序行为的周围，注入对该运行时类中方法的调用，以表示该程序行为正要发生或者已经发生

package org.example;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

public class MyProfiler {
  public static ConcurrentHashMap<Class<?>, AtomicInteger> data = new ConcurrentHashMap<>();

  public static void fireAllocationEvent(Class<?> klass) {
    data.computeIfAbsent(klass, kls -> new AtomicInteger())
        .incrementAndGet();
  }

  public static void dump() {
    data.forEach((kls, counter) -> {
      System.err.printf("%s: %d\n", kls.getName(), counter.get());
    });
  }

  static {
    Runtime.getRuntime().addShutdownHook(new Thread(MyProfiler::dump));
  }
}

• 上面这段代码便是一个运行时类。该类维护了一个HashMap，用来统计每个类所新建实例的数目。当程序退出时，我们将逐个打印出每个类的名字，以及其新建实例的数目
• 在 Java agent 中，我们可以截获正在加载的类，并且在每条new字节码之后插入对fireAllocationEvent方法的调用，以表示当前正在新建某个类的实例

package org.example;

import java.lang.instrument.*;
import java.security.ProtectionDomain;

import org.objectweb.asm.*;
import org.objectweb.asm.tree.*;

public class MyAgent {

  public static void premain(String args, Instrumentation instrumentation) {
    instrumentation.addTransformer(new MyTransformer());
  }

  static class MyTransformer implements ClassFileTransformer, Opcodes {
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
        ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
      if (className.startsWith("java")    ||
          className.startsWith("javax")   || 
          className.startsWith("jdk")     ||
          className.startsWith("sun")     ||
          className.startsWith("com/sun") ||
          className.startsWith("org/example")) {
        // Skip JDK classes and profiler classes
        return null;
      }

      ClassReader cr = new ClassReader(classfileBuffer);
      ClassNode classNode = new ClassNode(ASM7);
      cr.accept(classNode, ClassReader.SKIP_FRAMES);

      for (MethodNode methodNode : classNode.methods) {
        for (AbstractInsnNode node : methodNode.instructions.toArray()) {
          if (node.getOpcode() == NEW) {
            TypeInsnNode typeInsnNode = (TypeInsnNode) node;

            InsnList instrumentation = new InsnList();
            instrumentation.add(new LdcInsnNode(Type.getObjectType(typeInsnNode.desc)));
            instrumentation.add(new MethodInsnNode(INVOKESTATIC, "org/example/MyProfiler", "fireAllocationEvent",
                "(Ljava/lang/Class;)V", false));

            methodNode.instructions.insert(node, instrumentation);
          }
        }
      }

      ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
      classNode.accept(cw);
      return cw.toByteArray();
    }
  }
}

• 不得不排除对 JDK 类以及该运行时类的注入。这是因为，对这些类的注入很可能造成死循环调用，并最终抛出StackOverflowException异常
• 解决方案：设置一个线程私有的标识位，用以区分应用代码的上下文以及注入代码的上下文
  • 当即将执行注入代码时，我们将根据标识位判断是否已经位于注入代码的上下文之中
  • 如果不是，则设置标识位并正常执行注入代码
  • 如果是，则直接返回，不再执行注入代码
• 字节码注入的另一个技术难点则是命名空间
  • 比如：不少应用程序都依赖于字节码工程库 ASM。当我们的注入逻辑依赖于 ASM 时，便有可能出现注入使用最新版本的 ASM，而应用程序使用较低版本的 ASM 的问题
  • JDK 本身也使用了 ASM 库，如用来生成 Lambda 表达式的适配器类。JDK 的做法是重命名整个 ASM 库，为所有类的包名添加jdk.internal前缀
  • 开发者可能需要借助自定义类加载器来隔离命名空间
• 观察者效应（observer effect）对所收集的数据造成的影响
• 当使用字节码注入开发 profiler 时，需要辩证地看待所收集的数据。它仅能表示在被注入的情况下程序的执行状态，而非没有注入情况下的程序执行状态

面向方面编程

• 面向方面编程的核心理念是定义切入点（pointcut）以及通知（advice）。程序控制流中所有匹配该切入点的连接点（joinpoint）都将执行这段通知代码
• 面向方面编程的其中一种实现方式便是字节码注入，比如AspectJ