javaagent 使用注意

前言

最近做项目，需要实现一个agent，实现运行过程替换字节码，当笔者实现这些功能时发现还是很多注意事项的。而且字节码的替换过程如果类的属性与方法升级了，那么加载就会报错。这种做法的好处是代码无侵入，缺点也很明显，严重依赖特定的jvm版本和中间件等。

javaagent简介

javaagent实际上是JVMTI使用的技术，核心依靠Instrumentation实现。查看这个包，官方文档：java.lang.instrument (Java Platform SE 8 )

其中一句很精髓：Provides services that allow Java programming language agents to instrument programs running on the JVM. The mechanism for instrumentation is modification of the byte-codes of methods. 提供服务，允许Java编程语言代理对JVM上运行的程序进行检测。检测的机制是修改方法的字节码。javaagent有2种实现，一种是jvm参数，一种是动态attach。

实现方式是addTransformer，只要是在addTransformer之前未被加载的类在加载的过程就会被我们自定义的字节码替换，如果已经加载的类需要替换，可以手动retransformClasses，当然也可以redefineClasses，不过就还原来讲，推荐retransformClasses。

准备demo及问题过程

准备字节码替换和demo，先替换一个jdk的类，比如要对File的list进行字节码替换。比如asm javassist等，javassist比较简单，而asm比较常用，比如cglib：https://asm.ow2.io/asm4-guide.pdf

先用javassist试试

ClassPool：CtClass池，可使用classPool.get(类全名)获取CtClass
CtClass：编译时类信息，class文件封装
CtMethod：类中的方法
CtField：类中的属性、变量

写个Controller，触发条件

    @RequestMapping("/file")
    public String[] fileList() {
        File file = new File("/Users/huahua/go");
        return file.list();
    }

agent

    public class Agent {
        private static synchronized void initAgent(String args, Instrumentation inst) {
            System.out.println("agent exec ......");
            inst.addTransformer(new ClassFileTransformer() {
                @Override
                public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                    //字节码修改，替换
                    String refName = className.replace("/", ".");
                    if (MethodFilter.filterClass(refName)) {
                        try {
                            return MethodFilter.getHook(refName).hookMethod(loader, className, classfileBuffer);
                        } catch (NotFoundException | CannotCompileException | IOException e) {
                            throw new RuntimeException(e);
                        }
                    }
                    return classfileBuffer;
                }
            }, true);
//                Class<?> clazz = Class.forName("com.feng.agent.demo.ReTransformDemo");
//                inst.retransformClasses(clazz);
            System.out.println("agent exec end......");
        }

        public static void premain(String args, Instrumentation inst) {
            initAgent(args, inst);
        }

        public static void agentmain(String args, Instrumentation inst) {
            initAgent(args, inst);
        }
    }

hook逻辑

public interface MethodHook {
    byte[] hookMethod(ClassLoader loader, String className, byte[] classfileBuffer) throws NotFoundException, CannotCompileException, IOException;


}

public class FileHook  implements MethodHook{

    public byte[] hookMethod(ClassLoader loader, String className, byte[] classfileBuffer) throws NotFoundException, CannotCompileException, IOException {
        // TODO: 获取ClassPool
        ClassPool classPool = ClassPool.getDefault();
//        CtClass ctClass = classPool.get(className);
        CtClass ctClass = classPool.makeClass(new ByteArrayInputStream(classfileBuffer));
        // TODO: 获取sayHelloFinal方法
        CtMethod ctMethod = ctClass.getMethod("list", "()[Ljava/lang/String;");
        // TODO: 方法前后进行增强
        ctMethod.insertBefore("{ System.out.println(\"start\");}");
        ctMethod.insertAfter("{ System.out.println(\"end\"); }");
        // TODO: CtClass对应的字节码加载到JVM里
//        Class c = ctClass.toClass();
        return ctClass.toBytecode();
    }

}

public class MethodFilter {

    private static Map<String, MethodHook> classMap = new HashMap<>();

    static {
        classMap.put("java.io.File", new FileHook());
    }

    public static boolean filterClass(String classname){
        return classMap.containsKey(classname);
    }

    public static MethodHook getHook(String classname) {
        return classMap.get(classname);
    }
}

问题1 ：前置检查不生效

此时触发第一个注意，已经加载的类必须主动retransformClasses才能生效，否则addTransformer是不会替换类的，addTransformer是前置检查，只有在类载入钱才能执行字节码替换

可以看到实际上类替换未生效，因为File类已经加载了，debug看原因

Arrays.stream(inst.getAllLoadedClasses()).filter((c)->c!=null&&c.getName().startsWith("java.io.File")).collect(Collectors.toList())

如下图，本次替换的File实际上已经加载了，未生效，常用的类还有输入输出流等

解决办法也简单，在addTransformer之后加入retransformClasses即可生效

Class<?>[] classes = inst.getAllLoadedClasses();
    Arrays.stream(classes).filter((c) -> c!=null&& MethodFilter.filterClass(c.getName())).forEach((c)->{
    try {
        inst.retransformClasses(c);
    } catch (UnmodifiableClassException e) {
        throw new RuntimeException(e);
    }
});

测试加入代码后果然生效

问题2：jdk的类替换问题

笔者这里使用的jdk自带的system.out，如果我自己写一个类呢，实际情况很常见。

public class FileCheck {

    public void checkFilePath(File file){
        if (file.getAbsolutePath().startsWith("/Users")) {
            System.out.println("user dir");
        }
        System.out.println("File start " + file.getPath());
    }
}

ctMethod.insertBefore("{ FileCheck.checkFilePath(this);}");

会触发

javassist.CannotCompileException: [source error] no such class: FileCheck

因为修改的是JDK的类，但是JDK的类是bootstrap加载的，那么我们自己写的类呢

bootstrap的classloader是没办法加载AppClassloader的类的

所以需要 appendToBootstrapClassLoaderSearch，把我们写的类放进jdk的搜索范围，为此修改插桩技术，因为需要静态方法才好插桩，当然也可以用非静态方法，用反射插桩。

public class FileCheck {

    public static void checkFilePath(File file){
        if (file.getAbsolutePath().startsWith("/Users")) {
            System.out.println("user dir");
        }
        System.out.println("File start " + file.getPath());
    }
}



    public byte[] hookMethod(ClassLoader loader, String className, byte[] classfileBuffer) throws NotFoundException, CannotCompileException, IOException {
        // TODO: 获取ClassPool
        ClassPool classPool = ClassPool.getDefault();
//        CtClass ctClass = classPool.get(className);
        CtClass ctClass = classPool.makeClass(new ByteArrayInputStream(classfileBuffer));
        // TODO: 获取sayHelloFinal方法
        CtMethod ctMethod = ctClass.getMethod("list", "()[Ljava/lang/String;");
        // TODO: 方法前后进行增强
        ctMethod.insertBefore("{com.feng.agent.FileCheck.checkFilePath($0);}");
        ctMethod.insertAfter("{System.out.println(\"end\");}");
        // TODO: CtClass对应的字节码加载到JVM里
//        Class c = ctClass.toClass();
        return ctClass.toBytecode();
    }

修改后执行正常

问题3：classloader的问题

经过上面的处理，虽然jdk的类可以替换了，但是是通过把agent的jar加到appendToBootstrapClassLoaderSearch搜索解决的，但是BootstrapClassLoader类加载器并不会加载一些额外的类，就会造成多次使用多次加载的现象。示例如下

public class CheckStatus {

    private static Map<String, Boolean> statusMap = new HashMap<>();

    public static void initStatus(){
        statusMap.put("FILE_STATUS", true);
    }

    public static Boolean getStatus(String statusKey){
        if (!statusMap.containsKey(statusKey)) return false;
        return statusMap.get(statusKey);
    }
}

然后通过agent初始化

然后在字节码替换的地方加入

public class FileCheck {

    public static void checkFilePath(File file){
        if (file.getAbsolutePath().startsWith("/Users")) {
            System.out.println("user dir");
            System.out.println("CheckStatus: " + CheckStatus.getStatus("FILE_STATUS"));
        }
        System.out.println("File start " + file.getPath());
    }
}

执行后发现CheckStatus的值是false

原因也很简单，因为appendToBootstrapClassLoaderSearch前的载入classloader是APPclassloader，但是appendToBootstrapClassLoaderSearch后使用的bootstrapclassloader，所以只要颠倒顺序即可解决

实际上应该把jdk替换和非jdk的区分，因为代码复用的情况，但是有时候又不能严格区分，此时就会有矛盾的处理，因为双亲委派和依赖加载，所以很多时候是自定义classloader，把agent的核心jar用自定义classloader反射执行。但是涉及jdk相关的类需要使用jdk原有逻辑加载

public class AgentClassloader extends URLClassLoader {
    public AgentClassloader(URL[] urls) {
        super(urls, ClassLoader.getSystemClassLoader().getParent());
    }

    @Override
    protected synchronized Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
        final Class<?> loadedClass = findLoadedClass(name);
        if (loadedClass != null) {
            return loadedClass;
        }

        // 优先从parent（SystemClassLoader）里加载系统类，避免抛出ClassNotFoundException
        if (name != null && (name.startsWith("sun.") || name.startsWith("java."))) {
            return super.loadClass(name, resolve);
        }
        try {
            Class<?> aClass = findClass(name);
            if (resolve) {
                resolveClass(aClass);
            }
            return aClass;
        } catch (Exception e) {
            // ignore
        }
        return super.loadClass(name, resolve);
    }
}

总结

实际上agent本身的技术很简单，但是涉及类加载就复杂多了，类有classloader，线程有classloader，而且线程的classloader和类的可以不一样，当子classloader加载可以去parent里面去找，但是parent不能向下查找，此时就只能自己加载。

另外agent的原理是类加载前执行替换，那么一些jdk的类就会出现替换失败，且jdk的类是bootstrapclassloader加载的，所以经常容易处理不好，加载异常，需要把jdk替换的相关类加入bootstrap查找，而且appclassloader或者自定义加载的bootstrap还会重复加载。