Java Agent与字节码注入

概述

Java Agent是一种在JVM启动时或运行时修改字节码的机制。本教程介绍Java Agent的使用。

1. premain方式

import java.lang.instrument.Instrumentation;

public class MyAgent {
    public static void premain(String agentArgs, Instrumentation inst) {
        System.out.println("Agent启动: " + agentArgs);
        
        inst.addTransformer((className, classfileBuffer) -> {
            if (className.contains("MyClass")) {
                System.out.println("修改类: " + className);
                // 返回修改后的字节码
                return modifyClass(classfileBuffer);
            }
            return null;
        });
    }
    
    private static byte[] modifyClass(byte[] classfileBuffer) {
        // 使用ASM或Javassist修改字节码
        return classfileBuffer;
    }
}

# MANIFEST.MF
Premain-Class: MyAgent

2. agentmain方式

import java.lang.instrument.Instrumentation;
import com.sun.tools.attach.VirtualMachine;

public class RuntimeAgent {
    public static void agentmain(String agentArgs, Instrumentation inst) {
        System.out.println("运行时Agent启动");
        
        inst.addTransformer((className, classfileBuffer) -> {
            // 修改字节码
            return null;
        });
    }
    
    public static void attachAgent(String pid) {
        try {
            VirtualMachine vm = VirtualMachine.attach(pid);
            vm.loadAgent("agent.jar", "agentArgs");
            vm.detach();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}

3. 实际应用示例

AOP实现

public class AOPAgent {
    public static void premain(String agentArgs, Instrumentation inst) {
        inst.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className,
                                    Class<?> classBeingRedefined,
                                    ProtectionDomain protectionDomain,
                                    byte[] classfileBuffer) {
                if (className.equals("com/example/MyClass")) {
                    return addLogging(classfileBuffer);
                }
                return null;
            }
        });
    }
    
    private static byte[] addLogging(byte[] classfileBuffer) {
        // 使用ASM添加日志
        ClassReader cr = new ClassReader(classfileBuffer);
        ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_FRAMES);
        
        ClassVisitor cv = new ClassVisitor(ASM9, cw) {
            @Override
            public MethodVisitor visitMethod(int access, String name, String descriptor,
                                           String signature, String[] exceptions) {
                MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
                return new MethodVisitor(ASM9, mv) {
                    @Override
                    public void visitCode() {
                        super.visitCode();
                        mv.visitFieldInsn(GETSTATIC, "java/lang/System", "out",
                                        "Ljava/io/PrintStream;");
                        mv.visitLdcInsn("方法被调用: " + name);
                        mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream",
                                         "println", "(Ljava/lang/String;)V", false);
                    }
                };
            }
        };
        
        cr.accept(cv, 0);
        return cw.toByteArray();
    }
}

性能监控

public class PerformanceAgent {
    public static void premain(String agentArgs, Instrumentation inst) {
        inst.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className,
                                    Class<?> classBeingRedefined,
                                    ProtectionDomain protectionDomain,
                                    byte[] classfileBuffer) {
                if (className.startsWith("com/example/service")) {
                    return addTiming(classfileBuffer);
                }
                return null;
            }
        });
    }
    
    private static byte[] addTiming(byte[] classfileBuffer) {
        // 添加方法执行时间统计
        return classfileBuffer;
    }
}

4. 最佳实践

1. 选择合适的注入方式：premain或agentmain
2. 处理异常情况：确保Agent的稳定性
3. 性能考虑：避免过度字节码操作
4. 安全考虑：确保Agent的安全性
5. 测试验证：确保Agent正确工作

总结

Java Agent是一种强大的字节码注入机制。掌握Java Agent的使用，可以实现AOP、性能监控等高级功能。