Java Agent使用指南

前言

我们知道，在CMD/SHELL环境中输入java后可以得到java命令的一系列参数，今天重点关注agent相关的东西，如下：

    -agentlib:<libname>[=<选项>]
                  加载本机代理库 <libname>, 例如 -agentlib:hprof
                  另请参阅 -agentlib:jdwp=help 和 -agentlib:hprof=help
    -agentpath:<pathname>[=<选项>]
                  按完整路径名加载本机代理库
    -javaagent:<jarpath>[=<选项>]
                  加载 Java 编程语言代理, 请参阅 java.lang.instrument

我们看到熟悉的-javaagent参数，是不是想到了一些破解软件的方案(狗头保命）。

那我们好好看看java.lang.instrument的说明：

There are two ways to obtain an instance of the Instrumentation interface:
When a JVM is launched in a way that indicates an agent class. In that case an Instrumentation instance is passed to the premain method of the agent class.
When a JVM provides a mechanism to start agents sometime after the JVM is launched. In that case an Instrumentation instance is passed to the agentmain method of the agent code.

重点突出了premain、agentmain两个方法，说明java agent本质上是一个有特殊方法签名的类，通过查阅sun.instrument.InstrumentationImpl#loadClassAndCallAgentmain、sun.instrument.InstrumentationImpl#loadClassAndCallPremain、sun.instrument.InstrumentationImpl#loadClassAndStartAgent的方法体，我们可以容易推断出它需要的签名是这样的：

public static void premain(String agentArgs, Instrumentation inst)
public static void premain(String agentArgs)
public static void agentmain(String agentArgs, Instrumentation inst)
public static void agentmain(String agentArgs)

其中premain是使用-javaagent指定的jar包静态启动时的方法，并且在主类main方法之前运行（如其名pre main），而agentmain则是在VM运行时通过另一种AttachAPI（后续介绍）追加的方法，通过sun.instrument.InstrumentationImpl#loadClassAndStartAgent的逻辑我们知道，包含Instrumentation参数的方法优先级高于不带此参数的。

此时，我们获得了java.lang.instrument.Instrumentation的实例，下面我们就看看如何通过它来达到我们修改字节码的目的。

实现

一般的，我们知道在操作字节码的时候，有ASM、javassist比较流行的框架，这里为了简单演示，我们使用javassist。

看看我的pom文件：

<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
    <modelVersion>4.0.0</modelVersion>
    
    <groupId>com.muzi.byte-code</groupId>
    <artifactId>java-agent</artifactId>
    <name>java-agent</name>
    <description>java agent demo</description>
    <packaging>jar</packaging>

    <properties>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
        <project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
        <maven.compiler.source>8</maven.compiler.source>
        <maven.compiler.target>8</maven.compiler.target>
    </properties>

    <dependencies>
        <!-- 使用javassist修改字节码 -->
        <dependency>
            <groupId>org.javassist</groupId>
            <artifactId>javassist</artifactId>
            <version>3.29.2-GA</version>
        </dependency>
    </dependencies>

    <build>
        <plugins>
            <plugin>
                <groupId>org.apache.maven.plugins</groupId>
                <artifactId>maven-assembly-plugin</artifactId>
                <version>3.1.1</version>
                <executions>
                    <execution>
                        <!--打包时应用配置-->
                        <phase>package</phase>
                        <goals>
                            <!--单个phase中应用配置-->
                            <goal>single</goal>
                        </goals>
                    </execution>
                </executions>
                <configuration>
                    <descriptorRefs>
                        <!--将应用的所有依赖包都打到jar包中。如果依赖了其他jar包，其他的jar包会被解压开，平铺到最终的jar里去，最终输出为jar包-->
                        <descriptorRef>jar-with-dependencies</descriptorRef>
                    </descriptorRefs>
                    <archive>
                        <!-- 设置manifest配置文件-->
                        <manifestEntries>
                            <!--Premain-Class:Agent静态加载时会调用的类全路径名-->
                            <Premain-Class>com.muzi.bytecode.javaagent.MyJavaAgent</Premain-Class>
                            <!--Agent-Class:Agent动态加载时会调用的类全路径名。-->
                            <Agent-Class>com.muzi.bytecode.javaagent.MyJavaAgent</Agent-Class>
                            <!--Can-Redefine-Classes：true表示能重定义此代理所需的类，默认值为 false（可选）-->
                            <Can-Redefine-Classes>true</Can-Redefine-Classes>
                            <!--Can-Retransform-Classes：true 表示能重转换此代理所需的类，默认值为 false （可选）-->
                            <Can-Retransform-Classes>true</Can-Retransform-Classes>
                            <!--Can-Set-Native-Method-Prefix：true表示能设置此代理所需的本机方法前缀，默认值为 false（可选）-->
                            <Can-Set-Native-Method-Prefix>true</Can-Set-Native-Method-Prefix>
                        </manifestEntries>
                    </archive>
                </configuration>
            </plugin>

            <plugin>
                <groupId>org.apache.maven.plugins</groupId>
                <artifactId>maven-compiler-plugin</artifactId>
                <version>3.1</version>
                <configuration>
                    <source>${maven.compiler.source}</source>
                    <target>${maven.compiler.target}</target>
                </configuration>
            </plugin>
        </plugins>
    </build>
</project>

com.muzi.bytecode.javaagent.MyJavaAgent的实现(请仔细看注释内容)

package com.muzi.bytecode.javaagent;

import javassist.CannotCompileException;
import javassist.ClassClassPath;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.Modifier;
import javassist.NotFoundException;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.security.ProtectionDomain;
import java.util.Objects;

public class MyJavaAgent {

    //测试premain的类
    private static final String PREMAIN_NAME = "com.muzi.bytecode.javaagent.test.TestPreMain";
    //测试agentmain的类
    private static final String AGENTMAIN_NAME = "com.muzi.bytecode.javaagent.test.TestAgentMain";


    /**
     * 静态加载。
     * 通过vm参数 -javaagent 指定，会在main方法之前被调用
     */
    public static void premain(String args, Instrumentation instrumentation) {
        System.out.println("--premain[start]--");
        instrumentation.addTransformer(new PrintStartEndTransformer(PREMAIN_NAME), true);
        System.out.println("--premain[end]--");
    }

    /**
     * 动态加载。Java agent指定的agentmain方法，会在main方法之前被调用
     */
    public static void agentmain(String args, Instrumentation instrumentation) {
        System.out.println("--agentmain[start]--");
        instrumentation.addTransformer(new PrintStartEndTransformer(AGENTMAIN_NAME), true);
        Class<?>[] classes = instrumentation.getAllLoadedClasses();
        if (classes != null) {
            for (Class<?> c : classes) {
                if (c.isInterface() || c.isAnnotation() || c.isArray() || c.isEnum()) {
                    continue;
                }
                if (AGENTMAIN_NAME.equals(c.getName())) {
                    try {
                        /*
                         * Instrumentation.retransformClasses()对JVM已经加载的类重新触发类加载。
                         * 使用的就是上面addTransformer的PrintStartEndTransformer，
                         * 可以修改方法体，但是不能变更方法签名、增加和删除方法/类的成员属性
                         */
                        instrumentation.retransformClasses(c);
                    } catch (UnmodifiableClassException e) {
                        System.out.println("-agentmain.error, class: " + c.getName() + ", ex:" + e);
                        e.printStackTrace();
                    }
                }
            }
        }
        System.out.println("--agentmain[end]--");
    }

    static class PrintStartEndTransformer implements ClassFileTransformer {
        private final String thisClassName;
        private final String thisClassPath;

        public PrintStartEndTransformer(String className) {
            this.thisClassName = className;
            this.thisClassPath = className.replaceAll("\\.", "/");
        }

        @Override
        public byte[] transform(ClassLoader loader
                , String className
                , Class<?> classBeingRedefined
                , ProtectionDomain protectionDomain
                , byte[] classfileBuffer) throws IllegalClassFormatException {

            //参数className实际上是字节码里面带斜线的类名
            if (!thisClassPath.equals(className)) {
                //返回null，就不会修改字节码
                return null;
            }
            final ClassPool classPool = ClassPool.getDefault();
            try {
                //千万注意这里，如果是premain，这里是null
                //如果是agentmain，这里已经加载过来，不会是null，必须将其插入ClassPool的classpath
                //否则后面会出现 javassist.NotFoundException
                if (Objects.nonNull(classBeingRedefined)) {
                    classPool.insertClassPath(new ClassClassPath(classBeingRedefined));
                }
                CtClass ctClass = classPool.get(thisClassName);
                CtMethod[] methods = ctClass.getDeclaredMethods();
                for (CtMethod method : methods) {
                    if (Modifier.isNative(method.getModifiers())) {
                        //不处理native方法，因为没有字节码，不过可以通过开启Transformer.isNativeMethodPrefixSupported()
                        // 和 设置Transformer.setNativeMethodPrefix()将Native方法用java代码包裹起来，然后修改新的字节码
                        continue;
                    }
                    //增强字节码：在方法开始结束的位置打印方法信息
                    method.insertBefore("System.out.println(\"Agent ==> "
                            + ctClass.getSimpleName() + "." + method.getName() + " start.\");");
                    method.insertAfter("System.out.println(\"Agent ==> "
                            + ctClass.getSimpleName() + "." + method.getName() + " end.\");", false);
                }
                byte[] bytecode = ctClass.toBytecode();
                //ctClass.detach() 的意思是将内存中曾经被javassist加载过的Class对象移除，
                // 如果下次有需要在内存中找不到会重新走javassist加载
                ctClass.detach();
                return bytecode;
            } catch (NotFoundException e) {
                //遇到这个错，请看上面的classPool.insertClassPath的位置
                e.printStackTrace();
            } catch (CannotCompileException e) {
                e.printStackTrace();
            } catch (IOException e) {
                e.printStackTrace();
            }
            //返回null，不修改字节码
            return null;
        }
    }
}

上面这个代码单独打包，用maven执行package之后，生成java-agent-1.0-SNAPSHOT-jar-with-dependencies.jar，里面有个MANIFEST.MF文件，内容如下：

Manifest-Version: 1.0
Created-By: Apache Maven 3.9.1
Built-By: zimuw
Build-Jdk: 11.0.17
Agent-Class: com.muzi.bytecode.javaagent.MyJavaAgent
Can-Redefine-Classes: true
Can-Retransform-Classes: true
Can-Set-Native-Method-Prefix: true
Premain-Class: com.muzi.bytecode.javaagent.MyJavaAgent
#记得这里必须有空行（删除本行文本）

注意：如果不用maven，要手写MANIFEST.MF文件，记得最后必须留至少一个空行！！！

使用

静态加载

新建一个项目，名称无所谓，直接看代码

package com.muzi.bytecode.javaagent.test;

public class TestPreMain {
    public static void main(String[] args) {
        private_process(1024);
        protected_process(2048);
    }

    private static int private_process(int i) {
        System.out.println("private_process => i=" + i);
        return i;
    }

    static int protected_process(int i) {
        System.out.println("protected_process => i=" + i);
        return i;
    }
}

然后在启动该main方法的vm配置中加上
-javaagent:/你自己的路径/java-agent-1.0-SNAPSHOT-jar-with-dependencies.jar参数，再启动main方法，可以看到输出：

--premain[start]--
--premain[end]--
Agent ==> TestPreMain.main start.
Agent ==> TestPreMain.private_process start.
private_process => i=1024
Agent ==> TestPreMain.private_process end.
Agent ==> TestPreMain.protected_process start.
protected_process => i=2048
Agent ==> TestPreMain.protected_process end.
Agent ==> TestPreMain.main end.

如上，我们的静态增强就成功运行了，接下来你就可以“为所欲为”了！

动态加载

除了静态增加之外，VM还提供了AttachAPI在VM运行时介入。
直接上代码：

TestAgentMain被增强的类的代码

package com.muzi.bytecode.javaagent.test;

public class TestAgentMain {
    /**
     * 无 -javaagent 启动 TestAttachAgent， 随后启动 {@link TestAttachToAgentMain} 观察结果
     *
     * @param args
     */
    public static void main(String[] args) {
        for (int i = 0; i < 100; i++) {
            process(i);
            try {
                //保证能够运行一段时间用来观察
                Thread.sleep(2000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
    }

    public static int process(int i) {
        System.out.println("process => i=" + i);
        return i;
    }
}

执行AttachAPI操作的类

package com.muzi.bytecode.javaagent.test;

import com.sun.tools.attach.AgentInitializationException;
import com.sun.tools.attach.AgentLoadException;
import 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;

public class TestAttachToAgentMain {
    //agent jar包路径
    public static final String AGENT_JAR_PATH = "/你的Agent路径/java-agent-1.0-SNAPSHOT-jar-with-dependencies.jar";


    /**
     * 用来在vm运行时动态增强{@link TestAgentMain}
     *
     * @param args
     * @throws IOException
     * @throws AttachNotSupportedException
     * @throws AgentLoadException
     * @throws AgentInitializationException
     */
    public static void main(String[] args) throws IOException, AttachNotSupportedException, AgentLoadException, AgentInitializationException {
        List<VirtualMachineDescriptor> descriptors = VirtualMachine.list();
        String vmName = "com.muzi.bytecode.javaagent.test.TestAgentMain";
        for (VirtualMachineDescriptor descriptor : descriptors) {
            // 针对指定名称的JVM实例
            if (descriptor.displayName().endsWith(vmName)) {
                System.out.println("将对该进程的vm进行增强：" + vmName + "的vm进程, pid=" + descriptor.id());
                // attach到新JVM
                VirtualMachine vm = VirtualMachine.attach(descriptor);
                // 加载agentmain所在的jar包
                vm.loadAgent(AGENT_JAR_PATH);
                // detach
                vm.detach();
            }
        }
    }
}

运行结果：

process => i=0
process => i=1
process => i=2   
# 在此处运行com.muzi.bytecode.javaagent.test.TestAttachToAgentMain#main
# 就开启了类增强的效果
--agentmain[start]--
--agentmain[end]--
Agent ==> TestAgentMain.process start.
process => i=3
Agent ==> TestAgentMain.process end.
Agent ==> TestAgentMain.process start.
process => i=4
Agent ==> TestAgentMain.process end.
Agent ==> TestAgentMain.process start.
process => i=5
Agent ==> TestAgentMain.process end.
Agent ==> TestAgentMain.process start.
process => i=6
Agent ==> TestAgentMain.process end.
Agent ==> TestAgentMain.process start.
process => i=7
Agent ==> TestAgentMain.process end.
Agent ==> TestAgentMain.process start.
process => i=8
Agent ==> TestAgentMain.process end.

末尾

参考文章

• https://www.cnblogs.com/rickiyang/p/11368932.html
• https://juejin.cn/post/7157684112122183693

在参考文章的基础上解决了调用agentmain时 javassist.NotFoundException的问题，此处调试、验证耽搁了许久。

Java Agent本身很简单，难的还是字节码操作，javassist功能有限，ASM框架更强大。