[备份] 如何实现 Java 的内联汇编

[备份] 如何实现Java的内联汇编

介绍

Java是否可以在代码中嵌入ASM代码，类似C++的__asm__呢。通过一些摸索，在老外的代码中学习到了一种方式。下方这段代码，仅初始化了Runtime和Process对象，当执行完内联汇编之后，将会执行Runtime.exec方法并使Process p指向返回值，或者说得到返回值。本文就是解释如何实现下方这一段代码内容，虽然暂时想不到合适的使用场景


xxxxxxxxxx
public class TestExec {
    public static void main(String[] args) {
        Runtime r = Runtime.getRuntime();
        Process p = null;
        // 内联汇编
        __asm__(asm -> asm.INIT()
                .ALOAD(1)
                .LDC("open /System/Applications/Calculator.app")
                .INVOKEVIRTUAL(Runtime.class, "exec",
                        MethodType.methodType(Process.class, String.class))
                .ASTORE(2)
        );
        // 成功执行命令并得到返回值
        if (p != null) {
            System.out.println("not null");
        }
    }
}

动态调用

将Java代码编译为字节码之后，会遇到一个问题，我们编写的__asm__方法和其中的表达式是无法直接变成我们希望的字节码的，而是这样的字节码：


xxxxxxxxxx
INVOKEDYNAMIC ...; [
    java/lang/invoke/LambdaMetafactory...
    (Ljava/lang/Object;)V, 
    org/sec/asm/test/TestExec.lambda$main$0...
    ...
]
INVOKESTATIC org/sec/asm/core/ASM.__asm__...

在当前类中会产生一个新的方法


xxxxxxxxxx
private static synthetic lambda$main$0(Lorg/sec/asm/core/ASMOpcodes;)

这个方法中是我们表达式的字节码，可以看到我们希望的ALOAD 1实际上是ICONST_1压栈常量1后INVOKE某个类的ALOAD方法。我们需要做的事情是将这一段替换为实际上期望的ALOAD 1指令，然后写入到原来的类中


xxxxxxxxxx
ICONST_1 ...
INVOKEINTERFACE org/sec/asm/core/ASMOpcodes.ALOAD...

收集

第一步要做的是Collect收集所有的__asm__方法，在我源码的CollectVisitor中核心代码如下，可以大致浏览下，我将逐步介绍


xxxxxxxxxx
@Override
public MethodVisitor visitMethod(...) {
    if (Modifier.isStatic(access) && Constants.BLOCK_TYPE_DESC.equals(descriptor)) {
        ClassWriter writer = new ClassWriter(0);
        String className = String.format("%s$%s$%s", this.name, name, System.currentTimeMillis());
        writer.visit(version, Opcodes.ACC_PUBLIC, className,
                null, "java/lang/Object", null);
        MethodVisitor proxy = writer.visitMethod(access, name, descriptor, signature, exceptions);
        writers.put(new MethodRef(name, descriptor), writer);
        return proxy;
    }
    return null;
}

这里的BLOCK_TYPE_DESC代码如下：某一个返回值为void且入参类型是ASMOpcodes的方法。注释中的方法正好符合这个条件，这也是JVM根据Lambda表达式生成的新方法：返回值为空且参数是ASMOpcodes接口（自定义的核心接口）


xxxxxxxxxx
// lambda$main$0(Lorg/sec/asm/core/ASMOpcodes;)
MethodType BLOCK_TYPE = MethodType.methodType(Void.TYPE, ASMOpcodes.class);

参数类型是ASMOpcodes因为__asm__方法泛型中已经指定


xxxxxxxxxx
static void __asm__(Consumer<ASMOpcodes> ignored) {
    __asm__internal__();
}

回到上文，接下来我创建了一个根据当前类名和方法名生成的唯一类名，然后将当前表达式生成的新方法（当前visit的方法）写入了这个唯一类中，交给后续处理


xxxxxxxxxx
String className = String.format("%s$%s$%s", this.name, name, System.currentTimeMillis());
        writer.visit(version, Opcodes.ACC_PUBLIC, className,
                null, "java/lang/Object", null);
MethodVisitor proxy = writer.visitMethod(access, name, descriptor, signature, exceptions);

写好的ClassWriter对象保存到某一个Map中，注意这个Map记录了每一个__asm__内联汇编方法对应的一个唯一类ClassWriter对象，包含了一个新生成的方法，后续我们将会对这个对象做操作


xxxxxxxxxx
writers.put(new MethodRef(name, descriptor), writer);

最后我们return proxy没有特殊意义，这里没有对字节码进行修改。如果当前方法不是新生成方法而是普通方法，返回null表示不感兴趣不是目标

核心

第一处收集完毕接下来是核心部分，我们重新对字节码内容进行分析，修改字节码。之前保存的Map对象有用处，在Map中保存的Key包含了所有内联汇编生成的新方法，如果分析到达了新生成的方法处直接返回，我们不需要对新生成方法内的具体代码做操作


xxxxxxxxxx
public MethodVisitor visitMethod(...) {
    if (methods.containsKey(new MethodRef(name, descriptor))) {
        return null;
    }
    return new CoreAdapter(Opcodes.ASM9,
            super.visitMethod(access, name, descriptor, signature, exceptions)
            loader, methods,this);
}

进入到CoreAdapter方法，处理其他所有的方法。注意如果普通方法中遇到了__asm__的INVOKE指令则直接返回，不处理这种情况，换句话来说是删除这条指令。而这个指令对比文章开头，是与INVOKEDYNAMIC配合的指令。遇到其他的方法调用正常继续分析即可


xxxxxxxxxx
@Override
public void visitMethodInsn(...) {
    if (opcode == Opcodes.INVOKESTATIC && owner.equals("org/sec/asm/core/ASM")
            && isInlineName(name)
            && isInlineDescriptor(descriptor)) {
        return;
    }
    super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
}

核心是处理INVOKEDYNAMIC指令的问题，这里省略一大堆判断，判断的意义是找到我们的__asm__精确对应的INVOKEDYNAMIC指令以及参数，通过handle.getName拿到的是自动生成新方法的名称，然后实例化一个新的类，反射调用该类的某个方法，这里稍后分析。注意最后直接返回表示删除当前的INVOKEDYNAMIC指令，配合上文删除下一条指令，使得原有的__asm__块代码完全被删除


xxxxxxxxxx
@Override
public void visitInvokeDynamicInsn(...) {
    // 省略一大堆判断
    try {
        String methodName = handle.getName();
        Class<?> klass = generateInlineClass(
                new MethodRef(methodName, Constants.BLOCK_TYPE_DESC));
        ASMOpcodes block = new ASMCoreVisitor(this);

        Method method = klass.getDeclaredMethod(methodName,
                Constants.BLOCK_TYPE.parameterArray());
        method.setAccessible(true);
        method.invoke(null, block);

        visitor.rewrite = true;
    } catch (Exception ex) {
        throw new RuntimeException(ex);
    }
    return;
}

实例化新类的方法如下，通过之前保存的Map拿到对应的唯一类ClassWriter对象，其中包含了写入的__asm__生成的新方法，得到字节码后通过Unsafe的defineClass方法得到这个凭空捏造的类型。接下来我们会反射调用这个类包含的方法


xxxxxxxxxx
private Class<?> generateInlineClass(MethodRef info) {
    ClassWriter writer = methods.get(info);
    if (writer == null) {
        throw new RuntimeException("error");
    }
    byte[] classBytes = writer.toByteArray();
    return ClassDefiner.INSTANCE.defineClass(loader, classBytes, 0, classBytes.length);
}

public Class<?> defineClass(ClassLoader loader, byte[] classBytes, int off, int len) {
    try {
        Method m = UNSAFE.getClass().getDeclaredMethod("defineClass",
                String.class, byte[].class, int.class, int.class,
                ClassLoader.class, ProtectionDomain.class);
        return (Class<?>) m.invoke(UNSAFE, null, classBytes, off, len, loader, null);
    } catch (Exception ex) {
        throw new RuntimeException(ex);
    }
}

反射调用这个凭空捏造类中包含的__asm__块自动生成的新方法


xxxxxxxxxx
ASMOpcodes block = new ASMCoreVisitor(this);

Method method = klass.getDeclaredMethod(methodName,
        Constants.BLOCK_TYPE.parameterArray());
method.setAccessible(true);
method.invoke(null, block);

visitor.rewrite = true;

ASMCoreVisitor中没有什么魔法，是简单的使用MethodVisitor构造字节码的方法，选取了其中两处很简单的指令代码，这也是ASM生成字节码的代码。


xxxxxxxxxx
@Override
public ASMOpcodes SALOAD() {
    return visitInsn(SALOAD);
}

@Override
public ASMOpcodes ISTORE(int idx) {
    return visitVarInsn(ISTORE, idx);
}

注意上文传入的初始化参数是this所以这里的visitXxx等方法通过这个MethodVisitor (this)构造出我们希望的ASM代码，反射调用后this对象中保存了调用这一套visit方法后的内容（可以理解为希望的指令已经写入了this中，等待ClassWriter.toBytesArray即可导出包含了希望指令的字节码）

可能会有一个疑问，对这里的this进行一系列的调用之后，是否会影响到后续正常的代码？不会，因为当前的操作全都在visitInvokeDynamicInsn方法中，且一系列调用后返回的是null以删除原来的__asm__块，后续的指令调用后续指令对应的visit方法，与这里将没有关联

处理畸形字节码

这样的反射调用，执行了通过UNSAFE定义的新类的新方法中的ASM代码，使this对象包含了希望的字节码，但这里引出了下一个问题，MethodVisitor是畸形的，并不能直接转字节码使用，因为只有指令没有符号表，常量池的内容有问题。可以测试直接导出字节码，分析后发现确实进入了指令，但是由于符号表和其它一些原因，这个字节码是无法正常运行的

通过上文ASMVisitor中的魔法操作，现在rewriter中的字节码指令是我们真正需要的，但这个字节码整体来看是畸形的，还需要后续处理，代码如下


xxxxxxxxxx
ClassWriter rewriter = new ClassWriter(reader, 0);
ASMVisitor inline = new ASMVisitor(rewriter, loader, methods);
reader.accept(inline, 0);
if (inline.rewrite) {
    // 一个新的writer
    ClassWriter writer = new LoaderWriter(ClassWriter.COMPUTE_FRAMES, loader);
    // 复制符号表（也许没有意义）
    ASMUtil.copySymbolTable(rewriter, writer);
    // 读取原有的畸形字节码生成正确的字节码
    new ClassReader(rewriter.toByteArray()).accept(writer, 0);
    // 得到合法的字节码
    return writer.toByteArray();
}

public static void copySymbolTable(ClassWriter source, ClassWriter target) {
    try {
        Field table = source.getClass().getDeclaredField("symbolTable");
        table.setAccessible(true);
        Object symbolTable = table.get(source);
        table.set(target, symbolTable);
    } catch (Exception ex) {
        ex.printStackTrace();
    }
}

老外代码使用一些魔法来复制symbolTable对象，我分析老外的魔法其实本质是一个反射操作，于是写了以上的代码。但是我对这个操作持怀疑态度，因为新的ClassWriter会根据指令自动生成symbolTable对象，复制未必有意义。删除这一处但复制后测试代码没有问题，所以这可能是多余的操作

况且除了符号表，还有一些结构可能存在问题，例如局部变量数组的大小，是否会溢出等问题

通过一个新的ClassWriter读取畸形字节码重新生成后，会得到正确的符号表，以及正确的字节码，拿到正确的字节码之后，一切就简单了，写回到对应的文件中，提示用户再次执行即可

Patch

在JDK 8中拿到的SystemClassLoader是URLClassLoader其中包含了我们编译生成的target/classes以及测试的test-classes目录，我收集到这些路径并调用Runner的方法进行Patch


xxxxxxxxxx
ClassLoader cl = ClassLoader.getSystemClassLoader();
// important: only jdk8 allow
URL[] urls = ((URLClassLoader) cl).getURLs();
List<URL> targetUrls = new ArrayList<>();
for (URL u : urls) {
    if (!u.toString().endsWith(".jar") &&
            u.toString().toLowerCase().contains("classes")) {
        // this must be class path
        targetUrls.add(u);
    }
}
if (targetUrls.size() == 0) {
    return;
}
try {
    for (URL targetUrl : targetUrls) {
        String temp;
        if(System.getProperty("os.name").toLowerCase().contains("windows")){
            temp = targetUrl.toURI().getPath().substring(1);
        }else{
            temp = targetUrl.toURI().getPath();
        }
        Runner.run(new String[]{temp});
    }
} catch (Exception e) {
    e.printStackTrace();
}

遍历目录找所有classes路径下但.class文件，进行Patch


xxxxxxxxxx
    public static void run(String[] args) throws Exception {
        // ...
        Files.walkFileTree(dir, new SimpleFileVisitor<Path>() {
            @Override
            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
                FileVisitResult result = super.visitFile(file, attrs);
                if (file.getFileName().toString().endsWith(".class")) {
                    byte[] bytes = Files.readAllBytes(file);
                    byte[] rewrite = Patch.patchBytes(loader, bytes);
                    if (rewrite != bytes) {
                        Files.write(file, rewrite, StandardOpenOption.TRUNCATE_EXISTING);
                    }
                }
                return result;
            }
        });
        loader.close();
    }

核心Patch的过程如下，首先收集得到一个Map对象，保存了所有的__asm__表达式生成的新方法以及对应的ClassWriter对象，接下来找到每一个方法中的__asm__块，删除已有的指令，反射将希望的ASM代码加入到MethodVisitor中，得到一个畸形的ClassWriter对象，畸形的字节码需要一个新的ClassWriter重新生成符号表，得到真正的字节码，然后写回到原来的文件


xxxxxxxxxx
    public static byte[] patchBytes(ClassLoader loader, byte[] bytes) {
        ClassReader reader = new ClassReader(bytes);
        Map<MethodRef, ClassWriter> methods = new HashMap<>();
        reader.accept(new CollectVisitor(methods), 0);
        if (methods.isEmpty()) {
            return bytes;
        }
        ClassWriter rewriter = new ClassWriter(reader, 0);
        ASMVisitor inline = new ASMVisitor(rewriter, loader, methods);
        reader.accept(inline, 0);
        if (inline.rewrite) {
            ClassWriter writer = new LoaderWriter(ClassWriter.COMPUTE_FRAMES, loader);
            ASMUtil.copySymbolTable(rewriter, writer);
            new ClassReader(rewriter.toByteArray()).accept(writer, 0);
            return writer.toByteArray();
        }
        return bytes;
    }

关于

代码地址在：github

其中有些地方可能我的理解还不到位，如果有错误之处，还请大佬教导，我进行修正