Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Selecting and modifying `if` statement with ASM

I want to update if statement in already existing class on particular line without changing the whole method.

Here is the target code (names of classes, methods and some code changed because they're not relevant):

public class Target extends Something {
   public Target(){
        super();
        //some code...
   }

    public Result targetMethod(Data firstPar, Data secondPar){
        if(someStatement()) {
            return Result.FAIL;
        } else {
            if(firstPar.check()){ //here is the line I want to change
                 firstPar.doSomething()
            }
            secondPar.doSomething();
            return Result.SUCCESS;
        }

    }
}

In this code I want to change if(firstPar.check()) to something like this:

if(firstPar.check() && !Utilities.someOtherCheck(firstPar, secondPar))

Here is the way I tried to solve this:

ClassNode node = new ClassNode();
ClassReader reader = new ClassReader(bytes); //bytes - original target class
reader.accept(node, 0);
ClassWriter cw = new ClassWriter(0);
node.accept(cw);
MethodVisitor visitor = cw.visitMethod(ACC_PUBLIC, "targetMethod", "<it's signature>", null, null);
visitor.visitCode();
//No idea what to do next

Things I don't understand:

  1. How do I properly choose a line in MethodVisitor?

  2. How do I change only a half of the if statement?

  3. How do I get bytecode of the new class (which will be generated after we change target class)?

It would be nice if you could provide some kind of example code.

Thanks!

like image 328
Daniil Dubrovsky Avatar asked Mar 13 '23 09:03

Daniil Dubrovsky


1 Answers

Question 1 is perhaps the most difficult. You'll need to find out the position to insert the instructions by recognizing some pattern. If you assume that firstPar.check() is called only once, then you can look for the following bytecode instructions for if(firstPar.check()):

ALOAD 1
INVOKEVIRTUAL Data.check ()Z
IFEQ L3

where L3 is the jump label if check return false.

For question 2, notice that bytecode instructions for if(firstPar.check() && !Utilities.someOtherCheck(firstPar, secondPar)) is:

ALOAD 1
INVOKEVIRTUAL Data.check ()Z
IFEQ L3
ALOAD 1
ALOAD 2
INVOKESTATIC Utilities.someOtherCheck (LData;LData;)Z
IFNE L3

So, you will need to insert 4 new instructions right after IFEQ L3.

You can do this by using the Tree API, where you create an adapter for targetMethod by subclassing ClassVisitor and MethodNode:

private static class ClassAdapter extends ClassVisitor {
    public ClassAdapter(ClassVisitor cv) {
        super(Opcodes.ASM5, cv);
    }

    @Override
    public MethodVisitor visitMethod(int access, String name, String desc,
            String signature, String[] exceptions) {
        MethodVisitor mv = cv.visitMethod(access, name, desc, signature, exceptions);
        if (name.equals("targetMethod"))
            return new MethodAdapter(access, name, desc, signature, exceptions, mv);
        else
            return mv;
    }
}

private static class MethodAdapter extends MethodNode {
    public MethodAdapter(int access, String name, String desc,
            String signature, String[] exceptions, MethodVisitor mv) {
        super(Opcodes.ASM5, access, name, desc, signature, exceptions);
        this.mv = mv;
    }
    // More to come ...
}

Inside MethodAdapter, you can override visitEnd to iterate all instructions inside the method, and try to detect the above 3 instructions and insert 4 new instructions after them:

@Override
public void visitEnd() {
    // Iterates all instructions in the method
    ListIterator<AbstractInsnNode> itr = instructions.iterator();
    while (itr.hasNext()) {
        // Checks whether the instruction is ALOAD 1
        AbstractInsnNode node = itr.next();
        if (node.getOpcode() != Opcodes.ALOAD
                || ((VarInsnNode) node).var != 1)
            continue;

        // Checks whether the next instruction is INVOKEVIRTUAL
        if (node.getNext() == null
                || node.getNext().getOpcode() != Opcodes.INVOKEVIRTUAL)
            continue;

        // Checks the invoked method name and signature
        MethodInsnNode next = (MethodInsnNode) node.getNext();
        if (!next.owner.equals("Data")
                || !next.name.equals("check")
                || !next.desc.equals("()Z"))
            continue;

        // Checks whether the next of the next instruction is IFEQ
        AbstractInsnNode next2 = next.getNext();
        if (next2 == null
                || next2.getOpcode() != Opcodes.IFEQ)
            continue;

        // Creates a list instructions to be inserted
        InsnList list = new InsnList();
        list.add(new VarInsnNode(Opcodes.ALOAD, 1));
        list.add(new VarInsnNode(Opcodes.ALOAD, 2));
        list.add(new MethodInsnNode(Opcodes.INVOKESTATIC, 
            "Utilities", "someOtherCheck", 
            "(LData;LData;)Z", false));
        list.add(new JumpInsnNode(Opcodes.IFNE, ((JumpInsnNode) next2).label));

        // Inserts the list, updates maxStack to at least 2, and we are done
        instructions.insert(next2, list);
        maxStack = Math.max(2, maxStack);
        break;
    }
    accept(mv);
}

To use the adapter, you can chain it with a ClassReader and a ClassWriter. Below, I also chain a TraceClassVisitor to print out a log file in the tmp directory:

ClassReader reader = new ClassReader("Target");
ClassWriter writer = new ClassWriter(reader, 0);
TraceClassVisitor printer = new TraceClassVisitor(writer, 
        new PrintWriter(System.getProperty("java.io.tmpdir") + File.separator + "Target.log"));
ClassAdapter adapter = new ClassAdapter(printer);
reader.accept(adapter, 0);
byte[] b = writer.toByteArray(); // The modified bytecode
like image 61
dejvuth Avatar answered Apr 26 '23 08:04

dejvuth