如何在Java Agent中插装方法并调用自定义回调函数

步骤1：定义自定义回调类 #

在Agent项目中创建一个回调类，提供静态方法供插装后的目标方法调用，方便统一处理逻辑：

public class AgentCallbacks {
    // 方法执行前的回调逻辑
    public static void beforeMethod(String className, String methodName, Object[] args) {
        // 这里可以实现日志、监控、参数校验等自定义逻辑
        System.out.printf("前置回调：类[%s] 方法[%s]，参数：%s%n", className, methodName, args);
    }

    // 方法执行后的回调逻辑
    public static void afterMethod(String className, String methodName, Object result) {
        // 这里可以实现结果统计、异常捕获等自定义逻辑
        System.out.printf("后置回调：类[%s] 方法[%s]，返回值：%s%n", className, methodName, result);
    }
}

步骤2：确保回调类能被目标类加载器访问 #

由于目标类的类加载器（如应用类加载器）默认无法加载Agent内部的类，需要将Agent的JAR包添加到Bootstrap类加载器的搜索路径。在Agent的premain方法中通过Instrumentation实现：

public class MyAgent {
    public static void premain(String agentArgs, Instrumentation inst) {
        // 将Agent自身JAR包加入Bootstrap类加载器路径
        try {
            JarFile agentJar = new JarFile(new File(
                MyAgent.class.getProtectionDomain().getCodeSource().getLocation().toURI()
            ));
            inst.appendToBootstrapClassLoaderSearch(agentJar);
        } catch (URISyntaxException | IOException e) {
            e.printStackTrace();
        }

        // 注册类转换器
        inst.addTransformer(new MyClassTransformer());
    }
}

步骤3：修改Transform方法调用回调函数 #

在你的transform方法中，替换原有的日志代码，改为调用自定义回调类的静态方法，利用Javassist内置变量传递方法的关键信息：

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

    byte[] byteCode = classfileBuffer;
    
    try {
        ClassPool classPool = scopedClassPoolFactory.create(loader, rootPool,
                    ScopedClassPoolRepositoryImpl.getInstance());
        CtClass ctClass = classPool.makeClass(new ByteArrayInputStream(classfileBuffer));
        
        // 将类名格式从"/分隔"转为"."分隔的全限定类名
        String fullClassName = className.replace("/", ".");
                          
        CtMethod[] methods = ctClass.getDeclaredMethods();

        for (CtMethod method : methods) {
            Annotation annotation = getAnnotation(method);
            if (annotation != null) {
                log.info("Going to Transform the class " + fullClassName);
                
                // 插入前置回调代码
                String beforeCode = String.format(
                    "com.example.agent.AgentCallbacks.beforeMethod(\"%s\", \"%s\", $args);",
                    fullClassName, method.getName()
                );
                method.insertBefore(beforeCode);

                // 插入后置回调代码，$_代表方法返回值（void方法时为null）
                String afterCode = String.format(
                    "com.example.agent.AgentCallbacks.afterMethod(\"%s\", \"%s\", $_);",
                    fullClassName, method.getName()
                );
                method.insertAfter(afterCode);
            }
        }
        byteCode = ctClass.toBytecode();
        ctClass.detach();
    } catch (Throwable ex) {
        log.log(Level.SEVERE, "Error in transforming the class: " + className, ex);
    }
    return byteCode;
}

关键注意事项 #

• 全限定类名：调用回调方法时必须使用完整的包名+类名，确保JVM能准确定位到回调类。
• Javassist内置变量：
  • $args：当前方法的参数数组（Object[]类型）
  • $_：当前方法的返回值，void方法时该值为null
  • $0：当前对象实例（仅非静态方法可用）
• 类加载隔离：如果不将Agent类加入Bootstrap类加载器，目标类的类加载器会因找不到回调类抛出ClassNotFoundException。
• 异常防护：建议在回调方法内部添加异常捕获逻辑，避免回调代码的异常影响目标方法的正常执行。

内容的提问来源于stack exchange，提问作者Proy