Java agent原理浅析+实战demo

在工作中有使用sandbox 修改类中的方法体、返回等，现通过demo学习，了解其底层实现的原理。

agent在Java生态中使用广泛，例如arthas、trace记录等框架等，都是通过agent实现。

Java中的agent分为两类：

1. premain agent： 在Jvm启动前，将agent随jvm一起启动生效。
2. attach agent：在Jvm启动后，通过将agent attach到指定pid的jvm上生效。

不论是修改代码逻辑、修改函数返回，本质上来说都修改是加载类的字节码文件，而这两种agent分类均是通过jvm提供的扩展点在类加载前增加逻辑来修改字节码，达到修改实际生效的字节码文件的效果。

区别在于：

premain agent由于是在jvm启动之前生效的，所以可以在类的初次加载前完成扩展逻辑，在类的初次加载时对类的字节码文件进行修改。

而attach agent生效时，类都已经加载完毕，所以此时需要触发想要修改的类 重新进行加载，进而修改其字节码文件。

Premain agent

premain，在main函数之前启动的意思。

通常是通过jvm启动参数指定agent

例如

-javaagent:/Users/xxx/IdeaProjects/AgentDemo/target/AgentDemo-1.0-SNAPSHOT-jar-with-dependencies.jar

执行入口

我们知道，当我们将Java代码打包成jar包时，需要指定入口main方法，main方法即为整体的触发点。

Premain agent也类似，需要指定一个premain agent的入口方法，不论是通过何种打包方式，最后在jar包的MANIFEST.INF文件中，通过

Premain-Class: org.example.PreMainAgent

来指定该premain agent的入口类，同时默认调用以下两个方法

//优先级更高
public static void premain(String agentArgs, Instrumentation inst)
    
public static void premain(String agentArgs)

也就是说，当工程中存在一个这样的类，打包成存在premain入口的jar包。

即可通过测试的java程序中加上对应jvm启动参数，实现premain agent随jvm启动

package org.example;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;

import javassist.*;


public class PreMainAgent {

    public static void premain(String agentArgs, Instrumentation inst) {
        System.out.println("PreMainAgent premain enter.args=" + agentArgs);
    }
}

如何修改字节码文件

当进入触发入口的之后，通过入参中的Instrumentation来对字节码进行修改。

Instrumentation为jvm提供的 允许开发者在Java程序运行时检查和修改应用程序的行为 的工具类，在agent中主要是使用其对类的字节码文件进行修改。

void addTransformer(ClassFileTransformer transformer, boolean canRetransform);

Instrumentation中存在Transformer的概念，transformer是在类加载前可以对类进行修改、检查的扩展点。

当我们调用 inst.addTransformer(new DemoTransFormer(), true);添加了一个transformer后，该transformer就会在类加载前生效。

例如下面这个transformer，即可打印出所有加载类的名字。

static class DemoTransFormer implements ClassFileTransformer {
        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
            System.out.println("PreMainAgent start transform. className=" + className);
            return null;
        }
    }

而其返回的是一个byte[]数组，即类加载后的字节码文件，也就是，如果想要对指定的类进行字节码文件修改，只需要在这修改，然后将修改后的字节码文件返回即可。

下面为示例，该transformer对“org/example/simple/Hello”进行transform，在这个类的sayHello方法的最后，插入了一个打印hello world的语句。然后将修改后的字节码文件返回。举例使用的是修改字节码文件使用的Javaassist库。

static class DemoTransFormer implements ClassFileTransformer {


        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
            System.out.println("PreMainAgent start transform. className=" + className);
            if ("org/example/simple/Hello".equals(className)) {
                try {
                    // 从ClassPool获得CtClass对象
                    final ClassPool classPool = ClassPool.getDefault();
                    final CtClass clazz = classPool.get("org.example.simple.Hello");
                    CtMethod sayHello = clazz.getDeclaredMethod("sayHello");
                    String methodBody = "System.out.println("hello world!");";
                    sayHello.insertAfter(methodBody);

                    // 返回字节码，并且detachCtClass对象
                    byte[] byteCode = clazz.toBytecode();
                    //detach的意思是将内存中曾经被javassist加载过的Date对象移除，如果下次有需要在内存中找不到会重新走javassist加载
                    clazz.detach();
                    return byteCode;
                } catch (Throwable ex) {
                    System.out.println(ex.getClass().getCanonicalName());
                    ex.printStackTrace();
                }
            }
            return null;
        }
    }

最后完整的premain agent类如下

package org.example;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;

import javassist.*;


public class PreMainAgent {


    public static void premain(String agentArgs, Instrumentation inst) {
        System.out.println("PreMainAgent premain enter.args=" + agentArgs);
        inst.addTransformer(new DemoTransFormer(), true);
    }

    static class DemoTransFormer implements ClassFileTransformer {


        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
            if (className.startsWith("java") || className.startsWith("sun")
                    || className.startsWith("com/intellij") || className.startsWith("jdk")) {
                return null;
            }
            System.out.println("PreMainAgent start transform. className=" + className);
            if ("org/example/simple/Hello".equals(className)) {
                try {
                    // 从ClassPool获得CtClass对象
                    final ClassPool classPool = ClassPool.getDefault();
                    final CtClass clazz = classPool.get("org.example.simple.Hello");
                    CtMethod sayHello = clazz.getDeclaredMethod("sayHello");
                    String methodBody = "System.out.println("hello world!");";
                    sayHello.insertAfter(methodBody);

                    // 返回字节码，并且detachCtClass对象
                    byte[] byteCode = clazz.toBytecode();
                    //detach的意思是将内存中曾经被javassist加载过的Date对象移除，如果下次有需要在内存中找不到会重新走javassist加载
                    clazz.detach();
                    return byteCode;
                } catch (Throwable ex) {
                    System.out.println(ex.getClass().getCanonicalName());
                    ex.printStackTrace();
                }
            }
            return null;
        }
    }

}

测试使用的java代码如下：

package org.example;

import org.example.simple.Hello;

/**
 * Hello world!
 *
 */
public class App 
{
    public static void main( String[] args ) throws InterruptedException {
        System.out.println( "entry main." );
        while (true) {
            Hello.sayHello();
            Thread.sleep(1000);
        }
    }
}

package org.example.simple;

public class Hello {

    public static void sayHello(){
        System.out.println("hello");
    }
}

启动输出如下：

PreMainAgent premain enter.args=null
//先加载main方法所在的类
PreMainAgent start transform. className=org/example/App
//执行main方法
entry main.
//加载Hello类 触发transform
PreMainAgent start transform. className=org/example/simple/Hello
//transform后，增加打印hello world的语句，并循环执行
hello
hello world!
hello
hello world!
hello
hello world!

premain agent的流程图大致如下

Attach agent

attach agent是在jvm已经启动之后，再附着到指定的java程序中，他需要解决和premain agent相同的两个问题：

1. 执行入口
2. 如何修改字节码文件

执行入口

premain agent随着用户本身的jvm启动，因此从用户的视角来看的话，只启动一个有特殊jvm参数的程序。

而attach agent不同，本身已经有一个java程序了，然后需要再启动一个java程序，后启动的java程序需要连接到之前的java程序中，将agent的jar包发送过去并load，从而完成attach。

整个过程的流程图如下：

其中attach的入口与premain不同，为

//优先级更高
public static void agentmain (String agentArgs, Instrumentation inst)

public static void agentmain (String agentArgs)

package org.example;

import javassist.NotFoundException;

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

public class AttachAgent {

    public static void agentmain(String agentArgs, Instrumentation instrumentation) throws NotFoundException, ClassNotFoundException, UnmodifiableClassException, InterruptedException {
        System.out.println("entry AttachAgent main.");
    }
}

同样需要在jar包的MANIFEST.INF文件中，通过

Agent-Class: org.example.AttachAgent

指定attach的jar包的agent入口即可

连接并load attach agent的代码

package org.example;

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

import java.util.List;
import java.util.Scanner;


public class AttachAgentMain {

    public static void main(String[] args) {
        //获取当前系统中所有 运行中的 虚拟机
        System.out.println("Attach test agent start.");
        System.out.println("Please select the jvm to attch.");
        List<VirtualMachineDescriptor> list = VirtualMachine.list();
        System.out.println("------------------------------------------");
        for (int i = 0; i < list.size(); i++) {
            VirtualMachineDescriptor vmd = list.get(i);
            System.out.println(i + " : " + vmd.displayName());
        }
        System.out.println("------------------------------------------");
        try {
        Scanner scanner = new Scanner(System.in);
        //选择并连接指定的虚拟机
        int vmSelect = scanner.nextInt();
            VirtualMachineDescriptor vmd = list.get(vmSelect);
            System.out.println("vmd = "+vmd.displayName());
            System.out.println("vmd pid = "+vmd.id());
            VirtualMachine virtualMachine = VirtualMachine.attach(vmd.id());
            System.out.println("virtualMachine ="+virtualMachine.toString());
            //load attach agent
            virtualMachine.loadAgent("/Users/xxx/IdeaProjects/AgentDemo/target/AgentDemo-1.0-SNAPSHOT-jar-with-dependencies.jar");
            //结束
            virtualMachine.detach();
        }catch (Throwable e){
            e.printStackTrace();
        }
    }
}

如何修改字节码

修改字节码的方式和premain agent相同，都是通过transformer生效来修改。

但是transformer只会在加载类的时候生效，而attach agent生效时，jvm都已经运行了，此时类应该都已经加载完了。

所以此时需要比premain agent多一个重新transformer的操作，会重新拉取类本身的定义，然后按照现在定义的的transformer重新transform一遍，这样才可以做到在运行起来之后仍然修改类的字节码文件。

示例如下：

package org.example;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.NotFoundException;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.security.ProtectionDomain;

public class AttachAgent {

    public static void agentmain(String agentArgs, Instrumentation instrumentation) throws NotFoundException, ClassNotFoundException, UnmodifiableClassException, InterruptedException {
        DemoTransFormer demoTransFormer = new DemoTransFormer();
        instrumentation.addTransformer( demoTransFormer, true);
        Class<?> claz = Class.forName("org.example.simple.Hello");
        try {
            System.out.println("start to retransform claz="+claz.getName());
            //重新transform指定的类
            instrumentation.retransformClasses(claz);
        }catch (Throwable e){
            e.printStackTrace();
        }
    }

    static class DemoTransFormer implements ClassFileTransformer {


        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
            System.out.println("AttachAgent start transform. className=" + className);
            if ("org/example/simple/Hello".equals(className)) {
                try {
                    // 从ClassPool获得CtClass对象
                    final ClassPool classPool = ClassPool.getDefault();
                    System.out.println("classPool="+classPool);
                    final CtClass clazz = classPool.get("org.example.simple.Hello");
                    System.out.println("clazz="+clazz);
                    CtMethod sayHello = clazz.getDeclaredMethod("sayHello");
                    String methodBody = "System.out.println("hello world after attach agent transform!");";
                    sayHello.insertAfter(methodBody);

                    // 返回字节码，并且detachCtClass对象
                    byte[] byteCode = clazz.toBytecode();
                    //detach的意思是将内存中曾经被javassist加载过的Date对象移除，如果下次有需要在内存中找不到会重新走javassist加载
                    clazz.detach();
                    return byteCode;
                } catch (Throwable ex) {
                    System.out.println(ex.getClass().getCanonicalName());
                    ex.printStackTrace();
                }
            }
            return null;
        }
    }
}

打包agent范例

使用maven-assembly-plugin打包插件

值得注意的是：

1. 打包的时候最好把依赖一起打包成一个jar包，否则在启动agent执行逻辑后会报java.lang.NoClassDefFoundError

<build>
  <plugins>
    <plugin>
      <groupId>org.apache.maven.plugins</groupId>
      <artifactId>maven-assembly-plugin</artifactId>
      <version>2.2-beta-5</version>
      <configuration>
        <descriptorRefs>
          <descriptorRef>jar-with-dependencies</descriptorRef>
        </descriptorRefs>
        <archive>
          <!--自动添加META-INF/MANIFEST.MF -->
          <manifest>
            <addClasspath>true</addClasspath>
          </manifest>
          <manifestEntries>
            <!--premain agent 入口设置 -->
            <Premain-Class>org.example.PreMainAgent</Premain-Class>
            <!--attach agent 入口设置 -->
            <Agent-Class>org.example.AttachAgent</Agent-Class>
            <Can-Redefine-Classes>true</Can-Redefine-Classes>
            <Can-Retransform-Classes>true</Can-Retransform-Classes>
          </manifestEntries>
        </archive>
      </configuration>
      <executions>
        <execution>
          <goals>
            <goal>assembly</goal>
          </goals>
          <phase>package</phase>
        </execution>
      </executions>
    </plugin>
  </plugins>
</build>

参考文章：

1. www.cnblogs.com/rickiyang/p…
2. lotabout.me/2024/Java-A…
3. www.cnblogs.com/qisi/p/java…