浅谈Java Agent

Posted on 2023-05-18 Symbols count in article: 9.2k Reading time ≈ 8 mins.

在JDK5以后，官方提供了一种新特性 Java Agent，也叫Java 探针技术，它可以帮助我们构建一个和主程序独立的代理程序，通过代理程序可以在不修改主程序代码的情况下对主程序进行增强，如实现性能监测、日志记录、热加载等，很多著名的软件/工具都使用了这个技术，如arthas、skywalking等

下面通过示例来学习如何使用Java Agent

关于Java Agent的一些简要说明

通过-javaagent:agent.jar参数加载agent程序，可以多次使用以加载多个agent
Java Agent有两个方法，一个是premain(String [, Instrumentation])，一个是agentmain(String[, Instrumentation])，Instrumentation参数可选
premain方法会在主程序JVM启动前执行，因此可以在此方法中进行一些初始化工作；agentmain方法会在主程序JVM启动后执行，但是不能使用-javaagent参数来加载，而是要使用VirtualMachine的loadAgent方法来进行加载
Instrumentation接口提供了在程序运行期间对程序进行动态调整的能力，比如修改字节码、替换class

开始动手

前提：

创建两个maven项目，一个是主程序，一个是agent

主程序

Main.java

public class Main {
    public static void main(String[] args) {
        System.out.println("hello");
    }
}

pom.xml

<build>
    <plugins>
        <plugin>
            <groupId>org.apache.maven.plugins</groupId>
            <artifactId>maven-jar-plugin</artifactId>
            <configuration>
                <archive>
                    <manifest>
                        <mainClass>com.codecho.demo.Main</mainClass>
                    </manifest>
                </archive>
            </configuration>
        </plugin>
    </plugins>
</build>

agent程序

LogAgent.java

public class LogAgent {
    
    // 主程序启动前，代理先执行，如果代理抛出异常，主程序无法正常启动
    // 参数 Instrumentation 可选
    public static void premain(String args, Instrumentation instrumentation) {
        System.out.println("LogAgent#premain executed, args: " + args);
    }
    
}

pom.xml（这里的配置是为了生成MANIFEST.MF文件，里面有premain class的信息）

<build>
    <plugins>
        <plugin>
            <groupId>org.apache.maven.plugins</groupId>
            <artifactId>maven-jar-plugin</artifactId>
            <configuration>
                <archive>
                    <manifestEntries>
                        <Premain-Class>com.codecho.agent.LogAgent</Premain-Class>
                    </manifestEntries>
                </archive>
            </configuration>
        </plugin>
    </plugins>
</build>

测试premain

使用maven package得到主程序和agent的jar包

通过-javaagent参数加载agent（这里我把两个jar包放到同一目录了）

# 不指定参数
java -javaagent:log-agent-1.0-SNAPSHOT.jar -jar agent-demo-1.0-SNAPSHOT.jar

# 指定参数
java -javaagent:log-agent-1.0-SNAPSHOT.jar=param1,param2,param3 -jar agent-demo-1.0-SNAPSHOT.jar

测试agentmain

修改agent程序的LogAgent.java和pom.xml

LogAgent.java

添加agentmain方法

public class LogAgent {
    
    // 主程序启动前，代理先执行，如果代理抛出异常，主程序无法正常启动
    // 参数 Instrumentation 可选
    public static void premain(String args, Instrumentation instrumentation) {
        System.out.println("LogAgent#premain executed, args: " + args);
    }

    // 主程序启动后，代理后执行
    // 参数 Instrumentation 可选
    public static void agentmain(String args, Instrumentation instrumentation) {
        System.out.println("LogAgent#agentmain executed, args: " + args);
    }
    
}

pom.xml

将Premain-Class改为Agent-Class

<build>
    <plugins>		
        <plugin>
            <groupId>org.apache.maven.plugins</groupId>
            <artifactId>maven-jar-plugin</artifactId>
            <configuration>
                <archive>
                    <manifestEntries>
                        <Agent-Class>com.codecho.agent.LogAgent</Agent-Class>
                    </manifestEntries>
                </archive>
            </configuration>
        </plugin>
    </plugins>
</build>

修改主程序的Main.java和pom.xml，新增Agent.java

Main.java

因为agentmain方法是在主程序JVM启动后执行，因此这里使用输入流保证main方法不会马上结束

public class Main {
    
    public static void main(String[] args) throws IOException {
        System.in.read();
    }
    
}

pom.xml

<!--引入jdk的tools-->
<dependencies>
    <dependency>
        <groupId>com.sun</groupId>
        <artifactId>tools</artifactId>
        <version>1.8</version>
        <scope>system</scope>
        <systemPath>${JAVA_HOME}\lib\tools.jar</systemPath>
    </dependency>
</dependencies>

Attach.java

public class Attach {

    public static void main(String[] args) {
        VirtualMachine vm = null;
        try {
            // 这里的24816是Main主程序的进程号，运行Main.java后使用jps命令查看pid
            vm = VirtualMachine.attach("24816");
            vm.loadAgent("D:/IdeaWorkspace/log-agent/target/log-agent-1.0-SNAPSHOT.jar", "hello");
        } catch (AttachNotSupportedException | IOException | AgentLoadException | AgentInitializationException e) {
            throw new RuntimeException(e);
        }
    }
    
}

先运行Main主程序，通过jps查看主程序pid，再运行Attach程序

测试Instrumentation#addTransformer方法(作用在premain方法)

修改agent程序的LogAgent.java和pom.xml

LogAgent.java

修改premain方法

public class LogAgent {
    
    public static void premain(String args, Instrumentation instrumentation) {
        System.out.println("LogAgent#premain executed, args: " + args);

        instrumentation.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
                // 当前类为com/codecho/demo/TrafficService时，才替换class
                if (!className.equals("com/codecho/demo/TrafficService")) {
                    return classfileBuffer;
                }

                System.out.println("replace class: " + className);
                ByteArrayOutputStream bos;
                // 类加载前替换新的class
                try (FileInputStream fis = new FileInputStream("D:/IdeaWorkspace/agent-demo/target/classes/com/codecho/demo/TrafficService.class")) {
                    bos = new ByteArrayOutputStream();
                    byte[] buffer = new byte[fis.available()];
                    fis.read(buffer);
                    bos.write(buffer, 0, buffer.length);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }

                return bos.toByteArray();
            }
        });
    }
    
}

pom.xml

<build>
    <plugins>
        <plugin>
            <groupId>org.apache.maven.plugins</groupId>
            <artifactId>maven-jar-plugin</artifactId>
            <configuration>
                <archive>
                    <manifestEntries>
                        <Premain-Class>com.codecho.agent.LogAgent</Premain-Class>
                    </manifestEntries>
                </archive>
            </configuration>
        </plugin>
    </plugins>
</build>

修改主程序的Main.java，新增TrafficService.java

TrafficService

public class TrafficService {
    
    public String transport() {
        return "take the train";
    }
    
}

Main.java

public class Main {
    
    public static void main(String[] args) {
        String transport = new TrafficService().transport();
        System.out.println("transport: " + transport);
    }
    
}

单独运行Main.java，查看输出

1 2	PS D:\IdeaWorkspace\agent-demo\target> java -jar agent-demo-1.0-SNAPSHOT.jar transport: take the train

修改TrafficService.java并编译

public class TrafficService {
    
    public String transport() {
        return "take the plane";
    }
    
}

通过-javaagent参数加载agent，查看输出

测试Instrumentation#retransformClasses方法(作用在agentmain方法)

修改agent程序LogAgent.java和pom.xml

LogAgent.java

public class LogAgent {

    // 主程序启动前，代理先执行，如果代理抛出异常，主程序无法正常启动
    // 参数 Instrumentation 可选
    public static void premain(String args, Instrumentation instrumentation) {
        System.out.println("LogAgent#premain executed, args: " + args);
    }

    // 主程序启动后，代理后执行
    // 参数 Instrumentation 可选
    public static void agentmain(String args, Instrumentation instrumentation) throws ClassNotFoundException, UnmodifiableClassException {
        System.out.println("LogAgent#agentmain executed, args: " + args);

        instrumentation.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
                // 当前类为com/codecho/demo/TrafficService时，才替换class
                if (!className.equals("com/codecho/demo/TrafficService")) {
                    return classfileBuffer;
                }

                ByteArrayOutputStream bos;
                // 类加载前替换新的class
                try (FileInputStream fis = new FileInputStream("D:/IdeaWorkspace/agent-demo/target/classes/com/codecho/demo/TrafficService.class")) {
                    bos = new ByteArrayOutputStream();
                    byte[] buffer = new byte[fis.available()];
                    fis.read(buffer);
                    bos.write(buffer, 0, buffer.length);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }

                return bos.toByteArray();
            }
        }, true);

        instrumentation.retransformClasses(Class.forName("com.codecho.demo.TrafficService", false, ClassLoader.getSystemClassLoader()));
    }

}

pom.xml

<build>
    <plugins>
        <plugin>
            <groupId>org.apache.maven.plugins</groupId>
            <artifactId>maven-jar-plugin</artifactId>
            <configuration>
                <archive>
                    <manifestEntries>
                        <Agent-Class>com.codecho.agent.LogAgent</Agent-Class>
                        <!--不加此配置会导致load失败-->
                        <!--Agent JAR loaded but agent failed to initialize-->
                        <Can-Retransform-Classes>true</Can-Retransform-Classes>
                    </manifestEntries>
                </archive>
            </configuration>
        </plugin>
    </plugins>
</build>

修改主程序的Main.java

Main.java

通过死循环判断class是否重新加载

public class Main {
    
    public static void main(String[] args) {
        for (;;) {
            String transport = new TrafficService().transport();
            System.out.println("transport: " + transport);
            try {
                TimeUnit.SECONDS.sleep(5);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
    }
    
}

TrafficService.java

public class TrafficService {
    
    public String transport() {
        return "take the bus";
    }
    
}

先运行Main主程序，再修改TrafficService的transport方法，通过jps查看主程序pid，再运行Attach程序

可以看到TrafficService#transport的输出从bus变为subway，表明TrafficService类确实重新加载了

总结

Java Agent是JDK提供的一种技术，可以对Java程序进行增强、监测、分析等
Java Agent主要有premain和agentmain两个入口方法，Instrumentation接口可以让我们在程序运行期间动态地更改字节码、替换class等
可以通过Javassist或ASM等工具来方便地操作字节码