字节码增强技术-Javassist

字节码

什么是字节码，在这里就不在赘述了，网上教程很多。Java 为了能让 Java 程序编译一次到处运行，用 Java 编译器将程序对源代码编译生成固定格式的字节码（.class文件）供 JVM 使用，因此理论上来说，只要符合 JVM 规范的字节码文件，就可以在 JVM 上运行，不同的 JVM 类语言（如Scala、Groovy、Kotlin）编译成字节码都可在 JVM 运行，除此之外，如果你对 JVM 的字节码规范非常了解的话，通过自己按照 JVM 规范自己写也是可以的。

那么什么是字节码增强呢？简单理解就是通过某种手段或者技术修改编译好的字节码，让新生成的字节码能满足我们的定制需求，这里说的需求就有很多了，比如常用的 AOP 底层很多就是使用字节码增强来达到切面拦截，再比如微服务中的链路追踪就使用了字节码增强（仅仅只一些 Java 客户端）来进行埋点标记来记录调用链关系的，所以了解字节码增强对一些框架能有更深入对理解，对问题排查有很大对帮助。

上面说的通过某种手段或者技术到底指哪些呢？我们最常用的 Java Proxy 也是一种增强技术，另外常用的还有 ASM，AspectJ，Javassist等常用的技术，其中ASM在指令层次操作字节码的，需要对JVM的指令有一定的了解，同时众多的指令也很难记住，操作比较高；AspectJ扩展了 Java，定义了一些专门的AOP语法，其中 Spring AOP 就使用了 AspectJ；Javassist 是强调源代码层次操作字节码的框架，操作起来很容易入手。

Javassist

使用Javassist需要使用javassist.jar。

优势：

• 操作简单，容易上手
• 性能高于反射

缺点：

• 性能相比 ASM ，会低一些
• 不支持 continue 和 break 表达式，不支持内部类和匿名类，因此在有些场景是不适合的

Javassist 使用 ClassPool 来操作所有的 Java 类。这个类的工作方式与 JVM 类装载器非常相似，但是有一个重要的区别是它不是将装载的、要执行的类作为应用程序的一部分链接，类池使所装载的类可以通过 Javassist API 作为数据使用。可以使用默认的类池（ClassPool.getDefault()），它是从 JVM 搜索路径中装载的，也可以定义一个搜索您自己的路径列表的类池。甚至可以直接从字节数组或者流中装载二进制类，以及从头开始创建新类。

装载到类池中的类由 CtClass 实例表示。与标准的 Class 类一样， CtClass 提供了检查类数据（如字段和方法）的方法。不过，这只是 CtClass 的部分内容，它还定义了在类中添加新字段、方法和构造函数、以及改变类、父类和接口的方法。

字段、方法和构造函数分别由 CtField、CtMethod 和 CtConstructor 的实例表示。这些类定义了修改由它们所表示的对象的所有方法的方法，包括方法或者构造函数中的实际字节码内容。

Javassist 常用类的说明：

• CtClass（compile-time class）：编译时类信息，它是一个class文件在代码中的抽象表现形式，可以通过一个类的全限定名来获取一个CtClass对象，用来表示这个类文件。
• ClassPool：ClassPool是一张保存CtClass信息的HashTable，key为类的全限定名称，value为类名对应的CtClass对象。当我们需要对某个类进行修改时，就是通过pool.getCtClass(“className”)方法从pool中获取到相应的CtClass。
• CtMethod、CtField：这两个比较好理解，对应的是类中的方法和属性，可以用于定义或者修改一些方法和字段。

Javassist 增强的代码片段是使用字符串来编写的，基本和平时写的 Java 源代码一致，主要的不同是一些是以 $ 开头的标识符，用于表示方法或者构造函数参数、方法返回值等。

比如：

public void method1(String arg1, Object arg2) {
    // 增强代码片段
    {
        System.out.println("入参 1: " + $1); // arg1
        System.out.println("入参 2: " + $2); // arg2
    }
}

入门 Demo

通过一个平时经常需要用的业务日志记录来学习，平时业务日志操作记录基本都是通过 AOP 实现的，这次就使用字节码增强技术来进行实现，对业务代码基本无任何侵入。

首先定义一个业务 service

public class BizService {
    public void bizProcess(Map map) {
        System.out.println("do biz process");
    }
}

编写增强代码片段

public class BizServiceInteceptor {

    public static void preProcess(Map map) {
        System.out.println("preProcess");
        // do log
    }
}

编写测试类

public class JavassistTest {

    public static void main(String[] args) throws NotFoundException, CannotCompileException {
        ClassPool pool = ClassPool.getDefault(); // 获取默认的类池
        CtClass clazz = pool.getOrNull("com.zmc.learning.javassist.BizService");
        if (clazz == null) {
            System.out.println("bizService not found");
            return;
        }
                // 获取需要增强的代码方法
        CtMethod bizProcessMethod = clazz.getDeclaredMethod("bizProcess");
                // 植入增强代码片段
        StringBuffer sb = new StringBuffer();
        sb.append("{");
        sb.append("com.zmc.learning.javassist.BizServiceInteceptor.preProcess($1);"); //获取入参
        sb.append("}");
        bizProcessMethod.insertBefore(sb.toString()); // 是不是有点像 AOP 的 before？
                // 增强后的 class
        clazz.toClass();
        BizService bizService = new BizService();
        bizService.bizProcess(new HashMap());
    }
}

代码也很简单，需要注意的是，BizServiceInteceptor 是静态方法，同时入参需要和原方法一致。这个 demo 是不是很 easy，看的出来 Javassist 在源代码层面上操作字节码对程序员还是很友好的。

Demo 2 增强 Http Client

需求：需求也比较简单，在跨系统进行 Http 调用的时候，需要记录请求的来源和调用的链路，如果是你你会怎么做呢。

正常硬编码的方式

// 基于 Apache httpclient 进行 http 调用
// 调用方
{
    httpRequest.setHeader("source", "sys1");
}

// 处理方
{
    Header header = httpRequest.getHeaders("source")[0];
    String sourceSystem = header.getValue();
}

这样也能比较容易的实现链路的记录，但是这样的方式明显不适合。为什么呢？第一这段代码其实跟业务关系不大，每个业务方都需要编写 http 调用的时候加上这么一段前置的逻辑，仔细想想，如果有10个系统通过 http 相互调用呢。第二如果 A 系统 set header 的 key 为 source，那下游也得知道你 set 的 key 值，同时下游再调用下游的的时候，如果它set 的 key 不是 source 呢，那是不是就不统一了，这时你可能会说各个系统协调好，统一不就好了，是的，这样确实可以，那为什么不使用类似 AOP 一样的技术进行拦截统一加呢，这样又不会对代码有侵入性。

使用Javassist处理

http client 依赖

            org.apache.httpcomponents
            httpclient
            4.5.8

增强的代码片段

public class HttpClientInterceptor {

    public static void intercept(HttpRequest httpRequest) {
        httpRequest.setHeader("source", "test");
    }
}

增强的工具类

public class HttpClientInstrumentation {

    private static final String ENHANCE_CLASS = "org.apache.http.impl.client.InternalHttpClient"; // 增强的 client
    private static final String ENHANCE_METHOD = "doExecute"; // 增强的方法

    public static void enhance() throws NotFoundException, CannotCompileException {
        ClassPool classPool = ClassPool.getDefault();
        CtClass ctClass = classPool.getOrNull(ENHANCE_CLASS);
        if (ctClass == null) {
            System.out.println("http client not found");
            return;
        }
        CtMethod doExecuteMethod = ctClass.getDeclaredMethod(ENHANCE_METHOD);
        String sb = "{" +
                "com.zmc.learning.javassist.HttpClientInterceptor.intercept" + "($2);" + // 获取入参 HttpRequest
                "}";
        doExecuteMethod.insertBefore(sb); // 植入代码片段
        ctClass.toClass();
    }
}

测试类

public class HttpClientTest {

    public static void main(String[] args) throws NotFoundException, CannotCompileException {
        HttpClientInstrumentation.enhance();
        HttpGet httpGet = sendGet();
        Header header = httpGet.getHeaders("source")[0];
        System.out.println(header.getValue());
    }

    private static HttpGet sendGet() {
        //创建默认的httpClient实例
        CloseableHttpClient httpClient = HttpClients.createDefault();
        HttpGet get = null;

        try {
            //用get方法发送http请求
            get = new HttpGet("http://www.baidu.com/");
            System.out.println("执行get请求:...." + get.getURI());
            CloseableHttpResponse httpResponse = null;

//            get.setHeader("source", "abc");

            //发送get请求
            httpResponse = httpClient.execute(get);
            try {
                //response实体
                HttpEntity entity = httpResponse.getEntity();
                if (null != entity) {
                    System.out.println("响应状态码:" + httpResponse.getStatusLine());
                    System.out.println("-------------------------------------------------");
                    System.out.println("响应内容:" + EntityUtils.toString(entity));
                    System.out.println("-------------------------------------------------");
                }
            } finally {
                httpResponse.close();
            }
        } catch (Exception ignore) {
            ;
        } finally {
            try {
                if (httpClient != null) {
                    httpClient.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }

        return get;
    }
}

这样一个 http client 的增强就好了，在业务方需要使用 http client 调用的系统中只需要在启动的时候调用 HttpClientInstrumentation.enhance 方法即可，对业务方来说是透明的，需要注意的是为什么增强的目标方法是 doExecute 呢？因为这个方法是 http client 最底层的方法，上层的方法最后都会到这个方法来发起请求，因此拦截这个方法即可，另外，这个 demo 中只拦截了 InternalHttpClient 类，如果需要全面一点的话可能还需要拦截 MinimalHttpClient,AbstractHttpClient 等，需要注意的是这个仅仅对 http client 同步的客户端进行了增强。

有了这个 demo 是不是对其他对跨系统调用的客户端也可以实现呢，redis，kafka，nsq，grpc？大致的思路其实是一致的。说到这里再补一句，其实链路追踪的框架调用链这块大致的思路也是这样的，通过增强各个中间件的客户端来把系统的调用链串起来。

Demo 3 实现代理

参照Jdk Java Proxy 编写一个动态代理的工具，Java Proxy 相关的请自行了解。

编写ProxyFactory

public class ProxyFactory {

    public static Object newProxyInstance(ClassLoader classLoader, Class interfaceClass, InvocationHandler h) throws Throwable {
        ClassPool pool = ClassPool.getDefault();
        
        // 1.创建代理类 ProxyClass
        CtClass proxyClass = pool.makeClass("ProxyClass");

        //  2.给代理类添加字段：private InvocationHandler handler;
        CtClass handlerCc = pool.get(InvocationHandler.class.getName());
        CtField handlerField = new CtField(handlerCc, "handler", proxyClass); // CtField(CtClass fieldType, String fieldName, CtClass addToThisClass)
        handlerField.setModifiers(AccessFlag.PRIVATE);
        proxyClass.addField(handlerField);

        // 3.添加构造函数：public NewProxyClass(InvocationHandler handler) { this.handler = handler; } 
        CtConstructor ctConstructor = new CtConstructor(new CtClass[] { handlerCc }, proxyClass);
        ctConstructor.setBody("$0.handler = $1;"); // $0代表this, $1代表构造函数的第1个参数
        proxyClass.addConstructor(ctConstructor);
        
        // 4.为代理类添加相应接口方法及实现 
        CtClass interfaceCc = pool.get(interfaceClass.getName());
        
        // 4.1 为代理类添加接口：public class ProxyClass implements IHello
        proxyClass.addInterface(interfaceCc);
        
        // 4.2 为代理类添加相应方法及实现
        CtMethod[] ctMethods = interfaceCc.getDeclaredMethods();
        for (CtMethod ctMethod : ctMethods) {
            String methodFieldName = ctMethod.getName(); // 新的方法名

            // 4.2.1 为代理类添加反射方法字段
            // 如：private static Method method1 = Class.forName("com.zmc.learning.javassist.IHello").getDeclaredMethod("sayHello", new Class[] { Integer.TYPE });
            // 构造反射字段声明及赋值语句
            String classParamsStr = "new Class[0]"; // 方法的多个参数类型以英文逗号分隔
            if (ctMethod.getParameterTypes().length > 0) { // getParameterTypes获取方法参数类型列表
                for (CtClass clazz : ctMethod.getParameterTypes()) {
                    classParamsStr = (("new Class[0]".equals(classParamsStr)) ? clazz.getName() : classParamsStr + "," + clazz.getName()) + ".class";
                }
                classParamsStr = "new Class[] {" + classParamsStr + "}";
            }
            String methodFieldTpl = "private static java.lang.reflect.Method %s=Class.forName(\"%s\").getDeclaredMethod(\"%s\", %s);";
            String methodFieldBody = String.format(methodFieldTpl, ctMethod.getName(), interfaceClass.getName(), ctMethod.getName(), classParamsStr);
            // 为代理类添加反射方法字段. CtField.make(String sourceCodeText, CtClass addToThisClass)
            CtField methodField = CtField.make(methodFieldBody, proxyClass);
            proxyClass.addField(methodField);

            // 4.2.2 为方法添加方法体
            // 构造方法体. this.handler.invoke(this, 反射字段名, 方法参数列表); 
            String methodBody = "$0.handler.invoke($0, " + methodFieldName + ", $args)";
            // 如果方法有返回类型，则需要转换为相应类型后返回，因为invoke方法的返回类型为Object
            if (CtPrimitiveType.voidType != ctMethod.getReturnType()) {
                // 对8个基本类型进行转型
                // 例如：((Integer)this.handler.invoke(this, this.m2, new Object[] { paramString, new Boolean(paramBoolean), paramObject })).intValue();
                if (ctMethod.getReturnType() instanceof CtPrimitiveType) {
                    CtPrimitiveType ctPrimitiveType = (CtPrimitiveType) ctMethod.getReturnType();
                    methodBody = "return ((" + ctPrimitiveType.getWrapperName() + ") " + methodBody + ")." + ctPrimitiveType.getGetMethodName() + "()";
                } else { // 对于非基本类型直接转型即可
                    methodBody = "return (" + ctMethod.getReturnType().getName() + ") " + methodBody;
                }
            }
            methodBody += ";";
            // 为代理类添加方法. CtMethod(CtClass returnType, String methodName, CtClass[] parameterTypes, CtClass addToThisClass)
            CtMethod newMethod = new CtMethod(ctMethod.getReturnType(), ctMethod.getName(),
                    ctMethod.getParameterTypes(), proxyClass);
            newMethod.setBody(methodBody);
            proxyClass.addMethod(newMethod);
        }
        
        // 5.生成代理实例. 将入参InvocationHandler handler设置到代理类的InvocationHandler handler变量
        Class newClass = proxyClass.toClass(classLoader, null);
        return newClass.getConstructor(InvocationHandler.class).newInstance(h);
    }
}

ProxyFactory 和 Jdk Java Proxy 方法签名都一致，使用的方式也是一致的，需要注意的是对8个基本类型进行转型。