Search code examples
javainstrumentationjava-bytecode-asm

Avoiding variable slot collisions with ASM's LocalVariablesSorter


I'm having a hard time seeing how LocalVariablesSorter from ASM is able to keep variable slot collisions from happening. A variable might come from the original source, or I might create a variable with LocalVariablesSorter.newLocal(Type t). Later on, visitVarInsn aise going to come in for those slots. Both from the original code, and from my injected code. How could LocalVariablesSorter tell them apart. They will both have the same slot index, how does it move one to the proper slot? I don't see it happening in real life either.

Below is a program that shows off the problem. It instruments the Sample.sample method by injecting a local variable and uses it at the beginning and end of the method. In the middle the original source code provides its own variable. And as you can see, ASM gives them the same slot number, which is wrong. This is the whole point of what LocalVariablesSorter is supposed to stop, but I frankly don't see how it could possibly do so.

Here's the sample:

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.LocalVariablesSorter;
import org.objectweb.asm.util.CheckMethodAdapter;


public class LVSBug {

    public static void main(String[] args) throws IOException {

        try (InputStream is = new BufferedInputStream(LVSBug.class.getResourceAsStream("/Sample.class"))) {
            ClassReader cr = new ClassReader(is);
            ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_FRAMES|ClassWriter.COMPUTE_MAXS);
            ClassVisitor cv = new LVCInjector(cw);
            cr.accept(cv, ClassReader.EXPAND_FRAMES);

            try (OutputStream os = new BufferedOutputStream(new FileOutputStream(new File(System.getProperty("user.home"), "Sample.class")))) {
                os.write(cw.toByteArray());
            }

        }
    }
}

class LVCInjector extends ClassVisitor {

    public LVCInjector(ClassWriter cw) {
        super(Opcodes.ASM5, cw);
    }

    @Override
    public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
        MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
        if (name.equals("sample")) {
            CheckMethodAdapter checker = new CheckMethodAdapter(mv);
            return new LVMInjector(checker, access, desc);
        } else {
            return mv;
        }
    }
}

class LVMInjector extends LocalVariablesSorter {

    private int injectedReg;

    public LVMInjector(MethodVisitor mv, int access, String desc) {
        super(Opcodes.ASM5, access, desc, mv);
    }

    @Override
    public void visitCode() {
        super.visitCode();

        injectedReg = newLocal(Type.INT_TYPE);
        super.visitLdcInsn(Integer.valueOf(1));
        super.visitVarInsn(Opcodes.ISTORE, injectedReg);
    }

    @Override
    public void visitInsn(int opcode) {
        if (opcode == Opcodes.IRETURN) {
            super.visitInsn(Opcodes.POP);
            super.visitVarInsn(Opcodes.ILOAD, injectedReg);
        }

        super.visitInsn(opcode);
    }
}

class Sample {
    public static int sample(String s1) {
        Sample s = new Sample();
        return 0;
    }
} 

And here is the javap output of Sample.sample after instrumentation:

  public static int sample(java.lang.String);
    descriptor: (Ljava/lang/String;)I
    flags: ACC_PUBLIC, ACC_STATIC
    Code:
      stack=2, locals=3, args_size=1
         0: ldc           #22                 // int 1
         2: istore_2
         3: new           #1                  // class Sample
         6: dup
         7: invokespecial #16                 // Method "<init>":()V
        10: astore_2
        11: iconst_0
        12: pop
        13: iload_2
        14: ireturn
      LocalVariableTable:
        Start  Length  Slot  Name   Signature
            3      12     0    s1   Ljava/lang/String;
           11       4     2     s   LSample;
      LineNumberTable:
        line 85: 3
        line 86: 11

Notice that my injected variable gets slot 2, but the existing variable is given slot 2 as well, which is completely wrong.


Solution

  • OK, so I thought up a solution that is really simpler than what LocalvariablesSorter pretended to be. I believe that class is just conceptually broken. But here's how to easily add local variables to an existing method.

    1) Figure out what the slot number is of the last parameter. You can do this by looking at the method signature.

        int register = ((access & Opcodes.ACC_STATIC) != 0) ? 0 : 1;
        lastParmReg = register - 1;
        List<String> sigs = parseSignature(desc); // my own method (see below*)
        for (String sig : sigs) {
            lastParmReg = register;
            register += ("J".equals(sig) || "D".equals(sig)) ? 2 : 1;
        }
    
        int myNewInjectedReg = register;
        register += isMyNewRegADoubleOrLong() ? 2 : 1;
        ....
        ....
    

    2) You know how many local variables you are going to inject, so use the slots immediately above the lastParmReg for your locals.

    3) In your MethodVisitor, override the visitVarInsn, visitLocalVariable and visitLocalVariableAnnotation methods and do the following

    if the specified register slot is less than or equal to lastParmReq, just 
    call the super method passing the slot number as is.
    
    if it's larger than lastParmReg, then add the number of local variables you 
    are injecting to this value, and pass it along to super.
    

    Here's an example of visitVarInsn

    @Override
    public void visitVarInsn(int opcode, int var) {
        super.visitVarInsn(opcode, (var <= lastParmReg) ? var : var + numInjectedRegs);
    }
    

    Remember that your own injected locals will NOT be going through your MethodVisitor's visitXXX methods. You should only call super.XXXX on your local variables. The ones that come from the original source are the only ones that go thru your MethodVisitors methods.

    *Here's my parseSignature

    private static Pattern PARM_PATTERN = Pattern.compile("(\\[*(?:[ZCBSIJFD]|(?:L[^;]+;)))");
    
    private static List<String> parseSignature(String signature) {
        List<String> parms = new ArrayList<>();
    
        int openParenPos = signature.indexOf('(');
        int closeParenPos = signature.indexOf(')', openParenPos+1);
    
        String args = signature.substring(openParenPos + 1, closeParenPos);
        if (!args.isEmpty()) {
            Matcher m = PARM_PATTERN.matcher(args);
            while (m.find()) {
                parms.add(m.group(1));
            }
        }
    
        return parms;
    }