ASM method stack and local vars table introspection

gistfile1.txt

package com.offbynull.test.asmtest;

import java.io.DataOutputStream;
import java.io.FileOutputStream;
import java.io.InputStream;
import java.io.PrintWriter;
import java.util.HashMap;
import java.util.ListIterator;
import java.util.Map;
import java.util.Map.Entry;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.InsnList;
import org.objectweb.asm.tree.InsnNode;
import org.objectweb.asm.tree.LdcInsnNode;
import org.objectweb.asm.tree.MethodInsnNode;
import org.objectweb.asm.tree.MethodNode;
import org.objectweb.asm.tree.TypeInsnNode;
import org.objectweb.asm.tree.VarInsnNode;
import org.objectweb.asm.tree.analysis.Analyzer;
import org.objectweb.asm.tree.analysis.BasicInterpreter;
import org.objectweb.asm.tree.analysis.BasicValue;
import org.objectweb.asm.tree.analysis.Frame;
import org.objectweb.asm.util.TraceClassVisitor;

public class Main {

    public static void main(String[] args) throws Exception {
        InputStream in = Main.class.getResourceAsStream("/IntrospectionTest.class");

        ClassReader cr = new ClassReader(in);
        ClassNode classNode = new ClassNode();
        cr.accept(classNode, 0);
        cr.accept(new TraceClassVisitor(new PrintWriter(System.out)), 0);

        for (MethodNode methodNode : classNode.methods) {
            Analyzer<BasicValue> analyzer = new Analyzer<>(new BasicInterpreter());
            Frame<BasicValue>[] frames = analyzer.analyze(classNode.name, methodNode);

            ListIterator<AbstractInsnNode> it = methodNode.instructions.iterator();
            Map<MethodInsnNode, InsnList> invokeReplaceMap = new HashMap<>();
            while (it.hasNext()) {
                int nodeIdx = it.nextIndex();
                AbstractInsnNode node = it.next();
                if (node instanceof MethodInsnNode) {
                    MethodInsnNode methodInsnNode = (MethodInsnNode) node;
                    if (methodInsnNode.owner.equals("com/offbynull/test/asmtest/Introspecter")
                            && methodInsnNode.name.equals("stack")) {
                        InsnList replacementInstructions = stackToArray(methodNode, frames[nodeIdx]);
                        invokeReplaceMap.put(methodInsnNode, replacementInstructions);
                    }
                    if (methodInsnNode.owner.equals("com/offbynull/test/asmtest/Introspecter")
                            && methodInsnNode.name.equals("locals")) {
                        InsnList replacementInstructions = localsToArray(methodNode, frames[nodeIdx]);
                        invokeReplaceMap.put(methodInsnNode, replacementInstructions);
                    }
                }
            }
            
            for (Entry<MethodInsnNode, InsnList> entry : invokeReplaceMap.entrySet()) {
                methodNode.instructions.insert(entry.getKey(), entry.getValue());
                methodNode.instructions.remove(entry.getKey());
            }
        }

        //We are done now. so dump the class
        ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES);
        classNode.accept(cw);

        DataOutputStream dout = new DataOutputStream(new FileOutputStream("IntrospectionTest.class"));
        dout.write(cw.toByteArray());
        dout.flush();
        dout.close();
    }
    
    private static InsnList stackToArray(MethodNode methodNode, Frame<BasicValue> frameAtInstruction) {
        int stackSaveArrayLocalVarsIdx = methodNode.maxLocals + 1;
        int tempObjectLocalVarsIdx = methodNode.maxLocals + 2;
        
        InsnList insnList = new InsnList();
        
        // Create stack storage array and save it in local vars table
        insnList.add(new LdcInsnNode(frameAtInstruction.getStackSize()));
        insnList.add(new TypeInsnNode(Opcodes.ANEWARRAY, "java/lang/Object"));
        insnList.add(new VarInsnNode(Opcodes.ASTORE, stackSaveArrayLocalVarsIdx));
                        
        // Save the stack
        for (int i = frameAtInstruction.getStackSize() - 1; i >= 0; i--) {
            BasicValue basicValue = frameAtInstruction.getStack(i);
            Type type = basicValue.getType();
            
            
            // Convert the item to an object (if not already an object) and stores it in local vars table. Item removed from stack.
            switch (type.getSort()) {
                case Type.BOOLEAN:
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Boolean", "valueOf", "(Z)Ljava/lang/Boolean;"));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.BYTE:
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Byte", "valueOf", "(B)Ljava/lang/Byte;", false));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.SHORT:
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Short", "valueOf", "(S)Ljava/lang/Short;", false));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.CHAR:
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Character", "valueOf", "(C)Ljava/lang/Character;", false));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.INT:
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Integer", "valueOf", "(I)Ljava/lang/Integer;", false));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.FLOAT:
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Float", "valueOf", "(F)Ljava/lang/Float;", false));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.LONG:
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Long", "valueOf", "(J)Ljava/lang/Long;", false));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.DOUBLE:
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Double", "valueOf", "(D)Ljava/lang/Double;", false));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.ARRAY:
                case Type.OBJECT:
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.METHOD:
                case Type.VOID:
                default:
                    throw new IllegalArgumentException();
            }

            // Store item in to stack storage array
            insnList.add(new VarInsnNode(Opcodes.ALOAD, stackSaveArrayLocalVarsIdx));
            insnList.add(new LdcInsnNode(i));
            insnList.add(new VarInsnNode(Opcodes.ALOAD, tempObjectLocalVarsIdx));
            insnList.add(new InsnNode(Opcodes.AASTORE));
        }

        
        
        
        // Restore the stack
        for (int i = 0; i < frameAtInstruction.getStackSize(); i++) {
            BasicValue basicValue = frameAtInstruction.getStack(i);
            Type type = basicValue.getType();
            
            
            // Load item from stack storage array
            insnList.add(new VarInsnNode(Opcodes.ALOAD, stackSaveArrayLocalVarsIdx));
            insnList.add(new LdcInsnNode(i));
            insnList.add(new InsnNode(Opcodes.AALOAD));
            
            
            // Convert the item to an object (if not already an object) and stores it in local vars table. Item removed from stack.
            switch (type.getSort()) {
                case Type.BOOLEAN:
                    insnList.add(new TypeInsnNode(Opcodes.CHECKCAST, "java/lang/Boolean"));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/lang/Boolean", "booleanValue", "()Z", false));
                    break;
                case Type.BYTE:
                    insnList.add(new TypeInsnNode(Opcodes.CHECKCAST, "java/lang/Byte"));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/lang/Byte", "byteValue", "()B", false));
                    break;
                case Type.SHORT:
                    insnList.add(new TypeInsnNode(Opcodes.CHECKCAST, "java/lang/Short"));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/lang/Short", "shortValue", "()S", false));
                    break;
                case Type.CHAR:
                    insnList.add(new TypeInsnNode(Opcodes.CHECKCAST, "java/lang/Character"));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/lang/Character", "charValue", "()C", false));
                    break;
                case Type.INT:
                    insnList.add(new TypeInsnNode(Opcodes.CHECKCAST, "java/lang/Integer"));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/lang/Integer", "intValue", "()I", false));
                    break;
                case Type.FLOAT:
                    insnList.add(new TypeInsnNode(Opcodes.CHECKCAST, "java/lang/Float"));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/lang/Float", "floatValue", "()F", false));
                    break;
                case Type.LONG:
                    insnList.add(new TypeInsnNode(Opcodes.CHECKCAST, "java/lang/Long"));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/lang/Long", "longValue", "()J", false));
                    break;
                case Type.DOUBLE:
                    insnList.add(new TypeInsnNode(Opcodes.CHECKCAST, "java/lang/Double"));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/lang/Double", "doubleValue", "()D", false));
                    break;
                case Type.ARRAY:
                case Type.OBJECT:
                    break;
                case Type.METHOD:
                case Type.VOID:
                default:
                    throw new IllegalArgumentException();
            }
        }
        
        
        
        
        // Load up stack storage array as if it were returned by the invokestatic instruction
        insnList.add(new VarInsnNode(Opcodes.ALOAD, stackSaveArrayLocalVarsIdx));
        
        
        
        
        
        return insnList;
    }
    
    
    
    
    private static InsnList localsToArray(MethodNode methodNode, Frame<BasicValue> frameAtInstruction) {
        int localSaveArrayLocalVarsIdx = methodNode.maxLocals + 1;
        int tempObjectLocalVarsIdx = methodNode.maxLocals + 2;
        
        InsnList insnList = new InsnList();
        
        // Create local storage array and save it in local vars table
        insnList.add(new LdcInsnNode(frameAtInstruction.getLocals()));
        insnList.add(new TypeInsnNode(Opcodes.ANEWARRAY, "java/lang/Object"));
        insnList.add(new VarInsnNode(Opcodes.ASTORE, localSaveArrayLocalVarsIdx));
                        
        // Save the locals
        for (int i = 0; i < frameAtInstruction.getLocals(); i++) {
            BasicValue basicValue = frameAtInstruction.getLocal(i);
            Type type = basicValue.getType();
            
            if (type == null) {
                continue;
            }
            
            // Convert the item to an object (if not already an object) and stores it in local vars table.
            switch (type.getSort()) {
                case Type.BOOLEAN:
                    insnList.add(new VarInsnNode(Opcodes.ILOAD, i));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Boolean", "valueOf", "(Z)Ljava/lang/Boolean;"));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.BYTE:
                    insnList.add(new VarInsnNode(Opcodes.ILOAD, i));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Byte", "valueOf", "(B)Ljava/lang/Byte;", false));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.SHORT:
                    insnList.add(new VarInsnNode(Opcodes.ILOAD, i));
                    insnList.add(new InsnNode(Opcodes.DUP));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Short", "valueOf", "(S)Ljava/lang/Short;", false));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.CHAR:
                    insnList.add(new VarInsnNode(Opcodes.ILOAD, i));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Character", "valueOf", "(C)Ljava/lang/Character;", false));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.INT:
                    insnList.add(new VarInsnNode(Opcodes.ILOAD, i));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Integer", "valueOf", "(I)Ljava/lang/Integer;", false));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.FLOAT:
                    insnList.add(new VarInsnNode(Opcodes.FLOAD, i));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Float", "valueOf", "(F)Ljava/lang/Float;", false));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.LONG:
                    insnList.add(new VarInsnNode(Opcodes.LLOAD, i));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Double", "valueOf", "(D)Ljava/lang/Double;", false));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.DOUBLE:
                    insnList.add(new VarInsnNode(Opcodes.DLOAD, i));
                    insnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/Double", "valueOf", "(D)Ljava/lang/Double;", false));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.ARRAY:
                case Type.OBJECT:
                    insnList.add(new VarInsnNode(Opcodes.ALOAD, i));
                    insnList.add(new VarInsnNode(Opcodes.ASTORE, tempObjectLocalVarsIdx));
                    break;
                case Type.METHOD:
                case Type.VOID:
                default:
                    throw new IllegalArgumentException();
            }

            // Store item in to locals storage array
            insnList.add(new VarInsnNode(Opcodes.ALOAD, localSaveArrayLocalVarsIdx));
            insnList.add(new LdcInsnNode(i));
            insnList.add(new VarInsnNode(Opcodes.ALOAD, tempObjectLocalVarsIdx));
            insnList.add(new InsnNode(Opcodes.AASTORE));
        }
        
        
        
        
        // Load up locals storage array as if it were returned by the invokestatic instruction
        insnList.add(new VarInsnNode(Opcodes.ALOAD, localSaveArrayLocalVarsIdx));
        
        
        
        
        
        // Shove in to method nodes
        return insnList;
    }
}