如何使用ASM配合Gradle在编译期修改javac生成的Java字节码

核心逻辑 #

编译期修改class不需要用到agent，核心思路是在javac任务执行完成、class文件生成后，插入自定义任务读取所有class文件，通过ASM修改字节码后覆盖原文件即可在编译期完成插桩。

代码存放位置 #

ASM插桩逻辑需要放到独立的Gradle插件模块中，结构如下：

• 根项目下新建独立插件模块，命名为asm-plugin
• 该模块单独配置ASM和Gradle插件开发依赖
• 插件开发完成后 apply 到你的业务模块，即可自动介入项目编译流程

最简实现示例 #

1. 插件模块配置 #

asm-plugin/build.gradle 配置如下：

plugins {
    id 'java-gradle-plugin'
    id 'groovy'
}

dependencies {
    implementation gradleApi()
    implementation localGroovy()
    // 引入ASM依赖
    implementation 'org.ow2.asm:asm:9.5'
    implementation 'org.ow2.asm:asm-commons:9.5'
}

gradlePlugin {
    plugins {
        asmInject {
            id = 'com.example.asm-inject'
            implementationClass = 'com.example.AsmInjectPlugin'
        }
    }
}

2. 插件入口实现 #

插件入口负责hook javac任务，插入插桩逻辑：

package com.example;

import org.gradle.api.Plugin;
import org.gradle.api.Project;
import org.gradle.api.tasks.compile.JavaCompile;
import java.io.*;
import org.objectweb.asm.*;

public class AsmInjectPlugin implements Plugin<Project> {
    @Override
    public void apply(Project project) {
        project.getTasks().withType(JavaCompile.class, compileTask -> {
            // javac执行完成后执行插桩
            compileTask.doLast(task -> {
                File classOutputDir = compileTask.getDestinationDir();
                // 遍历所有class文件修改
                modifyClassFiles(classOutputDir);
            });
        });
    }

    private void modifyClassFiles(File dir) {
        for (File file : dir.listFiles()) {
            if (file.isDirectory()) {
                modifyClassFiles(file);
            } else if (file.getName().endsWith(".class") 
                && !file.getName().matches("(R|BuildConfig).*class")) {
                try {
                    FileInputStream fis = new FileInputStream(file);
                    ClassReader cr = new ClassReader(fis);
                    ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_FRAMES);
                    // 自定义ClassVisitor处理字节码
                    cr.accept(new CustomClassVisitor(Opcodes.ASM9, cw), 0);
                    // 写回修改后的字节码
                    FileOutputStream fos = new FileOutputStream(file);
                    fos.write(cw.toByteArray());
                    fos.close();
                    fis.close();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        }
    }
}

3. ASM字节码修改逻辑 #

示例实现给所有普通方法开头插入打印逻辑：

package com.example;

import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;

public class CustomClassVisitor extends ClassVisitor {
    public CustomClassVisitor(int api, ClassVisitor cv) {
        super(api, cv);
    }

    @Override
    public MethodVisitor visitMethod(int access, String name, String desc, String sign, String[] exceptions) {
        MethodVisitor mv = super.visitMethod(access, name, desc, sign, exceptions);
        // 跳过构造方法和静态代码块
        if ("<init>".equals(name) || "<clinit>".equals(name)) {
            return mv;
        }
        return new MethodVisitor(Opcodes.ASM9, mv) {
            @Override
            public void visitCode() {
                // 插入System.out.println("hello ASM")
                visitFieldInsn(Opcodes.GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
                visitLdcInsn("hello ASM");
                visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V", false);
                super.visitCode();
            }
        };
    }
}

4. 主项目引用插件 #

根目录build.gradle添加插件引用：

plugins {
    id 'com.example.asm-inject' version '1.0-SNAPSHOT' apply false
}

业务模块的build.gradle应用插件：

plugins {
    id 'java'
    id 'com.example.asm-inject'
}

验证方法 #

执行build任务后，反编译业务模块生成的class文件，即可看到所有普通方法开头已经插入了指定的打印逻辑。

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