Java agent 探针技术(2)-JVM 启动后 agentmain 进行类运行时转换

1. 简介

premain 的代理 jar 包需要在 Java 程序启动时指定，并且只能在类加载之前修改字节码，类被加载之后就无能为力了。这种使用方式使其在应用场景上有很大的限制，因为在大部分情况下，我们没有办法在虚拟机启动之时就为其设定代理

为了弥补这个缺点，JDK 1.6 引入了新的 agentmain 用于支持在类加载后再次加载该类，也就是重定义类，在重定义的时候可以修改类。但是这种方式对类的修改有较大的限制，修改后的类要兼容原来的旧类，具体的要求在 Java 官方文档 Instrumentation#retransformClasses()方法介绍 中可以找到： 转换类时禁止添加、删除、重命名成员变量和方法，禁止修改方法的签名，禁止改变类的继承关系

The retransformation may change method bodies, the constant pool and attributes. The retransformation must not add, remove or rename fields or methods, change the signatures of methods, or change inheritance. These restrictions maybe be lifted in future versions. The class file bytes are not checked, verified and installed until after the transformations have been applied, if the resultant bytes are in error this method will throw an exception.

阿里开源的 Java诊断工具 Arthas 就是基于 agentmain ，有兴趣的读者可以前往官方传送门

2. 使用 agentmain 的步骤

与 premain 的使用步骤 类似，agentmain 的使用需要如下几个步骤：

1. 创建一个指定的类作为 Agent-Class ，类中包含 agentmain() 方法，该方法有如下两个声明。JVM 会优先加载方法1，加载成功忽略 2，如果1 没有，则加载 2 方法
   1. public static void agentmain(String agentArgs, Instrumentation inst)：参数 agentArgs 是通过命令行传给 Java agent 的参数， inst 是 Java 的字节码转换工具
   2. public static void agentmain(String agentArgs)
2. 创建 MANIFEST.MF 配置文件，将 Agent-Class指定为包含 agentmain() 方法的类。该配置文件一般会将 Can-Redefine-Classes 和 Can-Retransform-Classes 配置为 true
3. agentmain() 方法的类和 MANIFEST.MF 文件打包成代理 jar 包
4. 使用 java -jar xxx.jar 命令启动一个目标 Java 程序，然后新起一个线程借助 JVMTI（Java Virtual Machine Tool Interface）提供的接口把代理 jar 包 attach 到目标 Java 程序，实现运行时类定义动态转换

在执行第4个步骤后，目标 Java 程序的 JVM 会把代理 jar 包中指定的类卸载，然后重新加载该类，在重新加载的过程中就为类修改提供了机会

3. 使用示例

3.1 创建实现 ClassFileTransformer 接口的类

创建一个 CustomAttachTransformer 类，该类实现了 ClassFileTransformer 接口并重写了 ClassFileTransformer#transform() 方法，主要实现的功能是为sample.attach.TestMain#deal() 方法添加了执行耗时打印，与 premain 使用示例 中的例子基本一致。需要注意的是，在 agentmain 动态修改类定义中不允许为类新增方法，故此处实现方式是在原方法的方法体内插入代码

import javassist.*;

import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.regex.Pattern;

/**
 * 检测方法的执行时间
 */
public class CustomAttachTransformer implements ClassFileTransformer {

    // 被处理的方法列表
    private final static Map<String, List<String>> METHOD_MAP = new ConcurrentHashMap<>();

    private static final String DEFAULT_METHOD = "sample.attach.TestMain.deal";

    private static final String CLASS_REGEX = "^(\\w+\\.)+[\\w]+$";

    private static final Pattern CLASS_PATTERN = Pattern.compile(CLASS_REGEX);

    private CustomAttachTransformer() {
        add(DEFAULT_METHOD);
    }

    public CustomAttachTransformer(String methodString) {
        this();
        if (!CLASS_PATTERN.matcher(methodString).matches()) {
            System.out.println("string:" + methodString + " not a method string");
            return;
        }
        add(methodString);
    }

    public void add(String methodString) {
        String className = methodString.substring(0, methodString.lastIndexOf("."));
        String methodName = methodString.substring(methodString.lastIndexOf(".") + 1);
        List<String> list = METHOD_MAP.computeIfAbsent(className, k -> new ArrayList<>());
        list.add(methodName);
        System.out.println(METHOD_MAP.toString());
    }

    // 重写此方法
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) {
        className = className.replace("/", ".");
        byte[] byteCode = null;
        System.out.println("agent: targetClassName:" + className);
        if (METHOD_MAP.containsKey(className)) {
            // 判断加载的class的包路径是不是需要监控的类
            CtClass ctClass;
            try {
                ClassPool classPool = ClassPool.getDefault();
                // 将要修改的类的classpath加入到ClassPool中，否则找不到该类
                classPool.appendClassPath(new LoaderClassPath(loader));
                ctClass = classPool.get(className);
                for (String methodName : METHOD_MAP.get(className)) {
                    System.out.println("agent: targetMethodName:" + methodName);
                    CtMethod ctMethod = ctClass.getDeclaredMethod(methodName);// 得到这方法实例
                    ctMethod.addLocalVariable("begin", CtClass.longType);
                    ctMethod.addLocalVariable("end", CtClass.longType);
                    ctMethod.insertBefore("begin = System.currentTimeMillis();");
                    ctMethod.insertAfter("end = System.currentTimeMillis();");
                    ctMethod.insertAfter("System.out.println(\"方法" + ctMethod.getName() + "耗时\"+ (end - begin) +\"ms\");");
                }
                byteCode = ctClass.toBytecode();
                // ClassPool中删除该类
                ctClass.detach();
            } catch (Exception e) {
                System.out.println(e.getMessage());
                e.printStackTrace();
            }
        }
        return byteCode;
    }
}

3.2 创建使用 ClassFileTransformer 的 agentmain 类

创建 InstrumentAttach 类，该类需要重点关注的是两个 agentmain() 方法。主要逻辑在两个入参的 agentmain() 方法中，关键步骤如下：

1. 调用 Instrumentation#addTransformer()方法，将自定义的 CustomAttachTransformer 字节码转码器添加到 Instrumentation 中
2. 调用 Instrumentation#retransformClasses()方法，指定需要转化的目标类为 sample.attach.TestMain.class。这一步很关键，运行时类定义动态转换需要指定重新定义的类，否则 JVM 无法处理。另一个类似的方法是 Instrumentation#redefineClasses()，但这个方法是在类加载前使用的，类加载后需使用 Instrumentation#retransformClasses() 方法

import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;

public class InstrumentAttach {

    public static void agentmain(String agentArgs, Instrumentation inst) throws UnmodifiableClassException {
        System.out.println("我是两个参数的 Java Agent agentmain, agentArgs:" + agentArgs);
        inst.addTransformer(new CustomAttachTransformer(agentArgs), true);
        // 指定需要转化的类
        inst.retransformClasses(sample.attach.TestMain.class);
    }

    public static void agentmain(String agentArgs) {
        System.out.println("我是一个参数的 Java Agent agentmain");
    }

    public static void main(String[] args) {
    }
}

3.3 打包代理 jar 包

将包含 agentmain() 方法的类所在模块和 MANIFEST.MF 文件打包成代理 jar 包，IDEA 下打包 jar 包可参考博客 IDEA 打包 jar 包记录，笔者的 MANIFEST.MF 文件内容如下，需要保留最后一行的空行

Manifest-Version: 1.0
Can-Redefine-Classes: true
Agent-Class: sample.attach.InstrumentAttach
Can-Retransform-Classes: true

3.4 打包目标程序 jar 包

目标类 TestMain 如下，逻辑很简单就是在 while 循环中打印字符串，将其和 MANIFEST.MF 文件打包成一个 jar 包，命名为 attachjar.jar

Manifest-Version: 1.0
Main-Class: sample.attach.TestMain

public class TestMain {

    public static void main(String[] args) throws InterruptedException {
        int count = 0;
        do {
            deal(count);
            Thread.sleep(2000);
            count++;
        } while (count <= 50);
    }

    public  static void deal(Integer count) {
        System.out.println("deal handle:" + count);
    }
}

3.5 编写 attach 处理程序

创建一个 AttachThread 类用于模拟 attach 操作，其关键点如下：

1. 获取当前程序启动时的本机上的 JVM 集合，类似 jps 命令
2. 遍历JVM 集合，找到目标 JVM 并对其进行 attach 链接
3. 当前JVM 链接上目标JVM，调用 loadAgent() 方法加载代理 jar 包且传递参数，最后 detach 目标JVM

---

VirtualMachine#attach() 方法会链接到目标 JVM 上，并返回一个目标 VirtualMachine 实例。通过该实例的 VirtualMachine#loadAgent() 方法可以将代理 jar 包中的 ClassFileTransformer 转换器注册到目标 JVM 中，当目标 JVM 中的类重新加载时就会触发转换器的 ClassFileTransformer#transform() 方法完成动态转换

Attach API 有 2 个主要的类：

1. VirtualMachine
   代表一个 Java 虚拟机，也就是程序需要监控的目标虚拟机，提供了 Attach 动作(在目标 JVM 上面附加一个代理) 和 Detach 动作等等
2. VirtualMachineDescriptor
   描述虚拟机的容器类，配合 VirtualMachine 完成各种功能

mport com.sun.tools.attach.AttachNotSupportedException;
import com.sun.tools.attach.VirtualMachine;
import com.sun.tools.attach.VirtualMachineDescriptor;

import java.io.IOException;
import java.util.List;
import java.util.Objects;

public class AttachThread extends Thread {
    /**
     * 记录程序启动时的 VM 集合
     */
    private final List<VirtualMachineDescriptor> listBefore;
    /**
     * 要加载的agent.jar
     */
    private final String jar;

    private AttachThread(String attachJar, List<VirtualMachineDescriptor> vms) {
        listBefore = vms;
        jar = attachJar;
    }

    @Override
    public void run() {
        VirtualMachine vm;
        List<VirtualMachineDescriptor> latestList;
        int count = 0;

        try {
            while (true) {
                latestList = VirtualMachine.list();
                vm = hasTargetVm(latestList);
                if (vm == null) {
                    System.out.println("没有目标 jvm 程序，请手动指定java pid");
                    try {
                        vm = VirtualMachine.attach("46358");
                    } catch (AttachNotSupportedException e) {
                        //System.out.println("拒绝访问 Disconnected from the target VM");
                    }
                }

                Thread.sleep(1000);
                System.out.println(count++);
                if (Objects.nonNull(vm) || count >= 100) {
                    break;
                }
            }
            Objects.requireNonNull(vm).loadAgent(jar, "hello");
            vm.detach();
        } catch (Exception e) {
            System.out.println("异常：" + e);
        }
    }

    /**
     * 判断是否有目标 JVM 程序运行
     */
    private VirtualMachine hasTargetVm(List<VirtualMachineDescriptor> listAfter) throws IOException, AttachNotSupportedException {
        for (VirtualMachineDescriptor vmd : listAfter) {
            if (vmd.displayName().endsWith("TestMain"))
                return VirtualMachine.attach(vmd);
        }
        return null;
    }

    public static void main(String[] args) {
        new AttachThread("/Users/xxxxxx/workspace/demo/out/artifacts/attach/attach.jar", VirtualMachine.list()).start();
    }
}

3.6 测试

1. 使用命令 java -cp /Users/xxxxxx/workspace/demo/out/artifacts/attachjar/attachjar.jar sample.attach.TestMain 启动目标程序 jar 包，本例中为 attachjar.jar
2. 运行AttachThread#main() 方法，将代理 jar 包 attach.jar 通过 attach 加载到目标 JVM 中，完成运行时动态修改 sample.attach.TestMain 类定义