博客

Java Agent的实现原理是什么？

  • Post category:Java

“Java Agent的实现原理是什么？”这个问题涉及较多的细节，下面我将从多个方面进行详细讲解，并给出两个示例。

1. Java Agent 简介

Java Agent 是在 JVM 上运行的一种 Java 程序，它能够动态地修改已经加载了的类和字节码，实现在不修改源代码的情况下增强 Java 应用程序的功能。Java Agent 通常用于监控、调试、性能分析、代码替换、代码注入等应用场景。

Java Agent 启动原理可以通过引用 Java 加载器的机制实现，完成了在 Java 程序运行时动态修改字节码的功能。Java Agent 通过 JavaSE 6 引入的 Instrumentation API 进行实现，该 API 提供了一个去修改已存在类和动态添加新类的功能的接口。

2. Java Agent 实现原理

Java Agent 主要实现过程如下：

  • Java Agent Jar 注册到系统 ClassLoader 中；
  • 在 main 函数之前或使用 RuntimeMXBean 获取的代理 ID 信息中；
  • 利用 Instrumentation API 利用 premain 或 agentmain 函数插入（加前缀或加后缀） Java 程序的指定类，获得字节码，存储在文件或数据库中；
  • JVM 动态更新、替换及更新后的类重新加载。

具体步骤：

2.1 编写 Java Agent

Java Agent 一般需要实现一个 premain 或 agentmain 函数，示例代码如下：

public class MyJavaAgent {
    public static void premain(String args, Instrumentation inst) {
        System.out.println("premain");
    }

    public static void agentmain(String args, Instrumentation inst) {
        System.out.println("agentmain");
    }
}

在 premain 和 agentmain 函数执行时，该 Agent 指定在哪个类加载器上注册 agent jar 包。agentmain 可以动态加载并修改已加载类，premain 只能实现在类加载之前的修改，且只能修改某些 Class 的字节码。premain 和 agentmain 只能同时存在一个。

2.2 注册 Java Agent

Java 有两种注册 Agent 的方式：一种是在启动命令行参数中加入 -javaagent 参数，并添加 Agent Jar 包的路径；另一种方法是在 Java 代码中使用 Attach API 动态注册 Agent。

例如，通过 Java 命令启动程序时注册 Agent：

java -javaagent:my-java-agent.jar -jar my-app.jar

通过 Attach API 动态注册 Agent 可以使用以下方式：

VirtualMachine vm = VirtualMachine.attach(pid);
vm.loadAgent(agentJarPath);

2.3 插桩

插桩是指在字节码中插入一些自定义的代码或字节，从而使 Java 程序执行时能够自动执行这些代码。Java Agent 利用 Instrumentation API 修改已经加载的类中的字节码，可以通过添加代码、替换代码等方式对程序进行监控和优化。例如：

public static void main(String[] args) {
    System.out.println("Hello World!");
}

可以插入以下代码：

public static void main(String[] args) {
    System.out.println("My First Agent");
    System.out.println("Hello World!");
}

3. 示例说明

下面给出两个使用示例：

3.1 示例一：动态替换 Class

该示例中，我们使用一个 Java Agent 动态地替换现有的类。

首先，我们创建待替换的类 ReplaceClass：

public class ReplaceClass {
    public void print() {
        System.out.println("ReplaceClass");
    }
}

然后，我们编写 PrintClass：

public class PrintClass {
    public void print() {
        System.out.println("PrintClass");
    }
}

最后，我们编写 MyJavaAgent 类，实现 premain 和 agentmain 函数分别为 ReplaceClass 类替换和还原：

public class MyJavaAgent {
    private static final String targetClass = "ReplaceClass";
    private static final String newClass = "PrintClass";

    public static void replaceClass(String className, Instrumentation inst) throws Exception {
        inst.addTransformer(new ClassFileTransformer() {
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                if (!className.equals(targetClass)) {
                    return classfileBuffer;
                }

                try {
                    return Files.readAllBytes(Paths.get("PrintClass.class"));
                } catch (IOException e) {
                    e.printStackTrace();
                }

                return classfileBuffer;
            }
        });
        Class<?>[] classes = inst.getAllLoadedClasses();
        for (Class<?> clazz : classes) {
            if (clazz.getName().equals(targetClass)) {
                inst.redefineClasses(new ClassDefinition(clazz, inst.getTransformer(targetClass).transform(clazz.getClassLoader(), clazz.getName(), clazz.getProtectionDomain(), clazz.getInterfaces(), clazz.getDeclaredFields())));
            }
        }
    }

    public static void restoreClass(String className, Instrumentation inst) throws ClassNotFoundException, UnmodifiableClassException {
        Class<?> clazz = Class.forName(className);
        // 此处代码可以去掉
        byte[] classBytes;
        try {
            classBytes = Files.readAllBytes(Paths.get("ReplaceClass.class"));
        } catch (IOException e) {
            e.printStackTrace();
            return;
        }
        inst.redefineClasses(new ClassDefinition(clazz, classBytes));
    }

    public static void premain(String args, Instrumentation inst) throws Exception {
        replaceClass(targetClass, inst);
    }
    // 如果需要动态更新类可以使用agentmain
    public static void agentmain(String args, Instrumentation inst) throws Exception {
        restoreClass(targetClass, inst);
        replaceClass(newClass, inst);
    }
}

3.2 示例二：动态添加 Class

该示例中，我们使用一个 Java Agent 动态地添加新类。

首先，我们创建 ClassAdder 类：

public class ClassAdder {
    public void print() {
        System.out.println("ClassAdder");
    }
}

然后，我们编写 MyJavaAgent 类，实现 premain 函数向已经加载的类 MyMainClass 中添加 ClassAdder 类：

public class MyJavaAgent {
    private static final String targetClass = "MyMainClass";
    private static final String newClass = "ClassAdder";

    public static void addClass(String newClassName, Instrumentation inst) {
        inst.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                if (className.equals(targetClass.replace(".", "/"))) {
                    try {
                        ClassPool cp = ClassPool.getDefault();
                        CtClass cc = cp.get(newClassName);
                        for (CtMethod method : cc.getDeclaredMethods()) {
                            method.insertBefore("System.out.println(\"add class: " + newClassName + "\");");
                        }

                        byte[] byteCode = cc.toBytecode();
                        cc.detach();
                        return byteCode;
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }
                return classfileBuffer;
            }
        });
    }

    public static void premain(String args, Instrumentation inst) throws Exception {
        addClass(newClass, inst);

        Class[] allLoadedClasses = inst.getAllLoadedClasses();
        for (Class clazz : allLoadedClasses) {
            if (clazz.getName().equals(targetClass)) {
                inst.redefineClasses(new ClassDefinition(clazz, getClassfileBytes(clazz)));
                break;
            }
        }
    }

    private static byte[] getClassfileBytes(Class clazz) throws IOException {
        String clazzAsPath = clazz.getName().replace('.', '/').concat(".class");
        URL classUrl = clazz.getClassLoader().getResource(clazzAsPath);
        byte[] bytes = IOUtils.toByteArray(classUrl);
        return bytes;
    }
}

然后在 MyMainClass 中新增加代码验证：

public class MyMainClass {
  public static void main(String[] args) throws Exception {
        System.out.println("MyMainClass");
        new ClassAdder().print();
    }
}

最后在启动时添加 Java Agent：

java -javaagent:my-java-agent.jar -cp . MyMainClass

4. 总结

Java Agent 是应用程序级别的代码修改工具，它可以在运行时实现类级别的代码修改而不需要预编译或修改代码。该工具可用于检测内存、收集性能数据、修改类以匹配特定平台、避免性能问题等。本文介绍了如何使用 Java Agent 实现类的替换和动态添加，同时也介绍了 Java Agent 的基本原理和实现方式。

Tags: