Manipulating Java Class Files with ASM 4 - Part Two: Tree API

What is ASM tree API: ASM Tree API is the part of ASM that lets you create/modify the class in memory. The class is viewed as a tree of information. Like the whole class is an instance of ClassNode, which contain a list of FieldNode objects, a list of MethodNode objects etc. This article assumes that the reader has already read the first part here.

A simple class through tree API: Let's use tree API to create our first class. Again I am going to jump right into a code example, because there is nothing better than a code example. The generated class has a main method that prints "Hello World!".

TreeAPIDemo.java

package com.geekyarticles.asm;

import java.io.DataOutputStream;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;

import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.FieldInsnNode;
import org.objectweb.asm.tree.InsnNode;
import org.objectweb.asm.tree.LdcInsnNode;
import org.objectweb.asm.tree.MethodInsnNode;
import org.objectweb.asm.tree.MethodNode;


public class TreeAPIDemo {
    public static void main(String [] args) throws Exception{
        
        ClassNode classNode=new ClassNode(4);//4 is just the API version number
        
        //These properties of the classNode must be set
        classNode.version=Opcodes.V1_6;//The generated class will only run on JRE 1.6 or above
        classNode.access=Opcodes.ACC_PUBLIC;
        classNode.signature="Lcom/geekyarticles/asm/Generated;";
        classNode.name="com/geekyarticles/asm/Generated";
        classNode.superName="java/lang/Object";
        
        //Create a method
        MethodNode mainMethod=new MethodNode(4,Opcodes.ACC_PUBLIC|Opcodes.ACC_STATIC,"main", "([Ljava/lang/String;)V",null, null);
        
        mainMethod.instructions.add(new FieldInsnNode(Opcodes.GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;"));
        mainMethod.instructions.add(new LdcInsnNode("Hello World!"));
        mainMethod.instructions.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V"));
        mainMethod.instructions.add(new InsnNode(Opcodes.RETURN));

        //Add the method to the classNode
        classNode.methods.add(mainMethod);
        
        //Write the class
        ClassWriter cw=new ClassWriter(ClassWriter.COMPUTE_MAXS|ClassWriter.COMPUTE_FRAMES);
        classNode.accept(cw);
        
        //Dump the class in a file
        File outDir=new File("out/com/geekyarticles/asm");
        outDir.mkdirs();
        DataOutputStream dout=new DataOutputStream(new FileOutputStream(new File(outDir,"Generated.class")));
        dout.write(cw.toByteArray());
        dout.flush();
        dout.close();
        
    }
}

As you can see, the code is very simple. A primary advantage over BCEL is that unlike BCEL, ASM does not require you to add every constant explicitly to the constant pool. Instead, ASM takes care of the constant pool itself.

Reading a class file: A ClassNode is a ClassVisitor. So, reading a class for use in tree API is as simple as creating a ClassReader object and using it to read a class file, while passing the ClassNode object in its accept method as a parameter. Once this is done, the ClassNode passed is fully initalized with all the information present in the class. In the following example, we will print all the methods in the class.

TreeAPIClassReaderDemo.java

package com.geekyarticles.asm;

import java.io.FileInputStream;
import java.io.InputStream;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.MethodNode;

public class TreeAPIClassReaderDemo {

    public static void main(String[] args) throws Exception{
        InputStream in=new FileInputStream("out/com/geekyarticles/asm/Generated.class");
        
        ClassReader cr=new ClassReader(in);
        ClassNode classNode=new ClassNode();
        
        //ClassNode is a ClassVisitor
        cr.accept(classNode, 0);
        
        //Let's move through all the methods
        
        for(MethodNode methodNodeclassNode.methods){
            System.out.println(methodNode.name+"  "+methodNode.desc);
        }

    }

}

Modifying a class file: Modifying a class file is a combination of the above two procedures. We first read the class in the usual way, make necessary changes to the data, and then write it back to a file. The following program implements an automatic injection of some logging code. Currently our Logger class only prints to the standard output. Every method annotated with @Loggable will be logged when they begin and when the return. In this we do not log the throw-exception. However that can also be implemented in the same manner by checking opcode ATHROW.

LoggingInsertion.java LoggingTest.java Logger.java Loggable.java

package com.geekyarticles.asm;

import java.io.DataOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.InputStream;
import java.util.Iterator;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.AnnotationNode;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.InsnList;
import org.objectweb.asm.tree.LdcInsnNode;
import org.objectweb.asm.tree.MethodInsnNode;
import org.objectweb.asm.tree.MethodNode;

public class LoggingInsertion {

    public static void main(String[] args) throws Exception{
        InputStream in=LoggingInsertion.class.getResourceAsStream("/com/geekyarticles/asm/LoggingTest.class");

        ClassReader cr=new ClassReader(in);
        ClassNode classNode=new ClassNode();
        cr.accept(classNode, 0);

        //Let's move through all the methods

        for(MethodNode methodNodeclassNode.methods){
            System.out.println(methodNode.name+"  "+methodNode.desc);
            boolean hasAnnotation=false;
            if(methodNode.visibleAnnotations!=null){
                for(AnnotationNode annotationNodemethodNode.visibleAnnotations){
                    if(annotationNode.desc.equals("Lcom/geekyarticles/asm/Loggable;")){
                        hasAnnotation=true;
                        break;
                    }
                }
            }
            if(hasAnnotation){
                //Lets insert the begin logger
                InsnList beginList=new InsnList();
                beginList.add(new LdcInsnNode(methodNode.name));
                beginList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "com/geekyarticles/asm/Logger", "logMethodStart", "(Ljava/lang/String;)V"));
                
                Iterator<AbstractInsnNode> insnNodes=methodNode.instructions.iterator();
                while(insnNodes.hasNext()){
                    System.out.println(insnNodes.next().getOpcode());
                }
                
                methodNode.instructions.insert(beginList);
                System.out.println(methodNode.instructions);
                
                //A method can have multiple places for return
                //All of them must be handled.
                insnNodes=methodNode.instructions.iterator();
                while(insnNodes.hasNext()){
                    AbstractInsnNode insn=insnNodes.next();
                    System.out.println(insn.getOpcode());
                    
                    if(insn.getOpcode()==Opcodes.IRETURN
                            ||insn.getOpcode()==Opcodes.RETURN
                            ||insn.getOpcode()==Opcodes.ARETURN
                            ||insn.getOpcode()==Opcodes.LRETURN
                            ||insn.getOpcode()==Opcodes.DRETURN){
                        InsnList endList=new InsnList();
                        endList.add(new LdcInsnNode(methodNode.name));
                        endList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "com/geekyarticles/asm/Logger", "logMethodReturn", "(Ljava/lang/String;)V"));
                        methodNode.instructions.insertBefore(insn, endList);
                    }
                    
                }
                
            }
        }

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


        File outDir=new File("out/com/geekyarticles/asm");
        outDir.mkdirs();
        DataOutputStream dout=new DataOutputStream(new FileOutputStream(new File(outDir,"LoggingTest.class")));
        dout.write(cw.toByteArray());
        dout.flush();
        dout.close();

    }

}

package com.geekyarticles.asm;

public class LoggingTest {
    public static void run1(){
        System.out.println("run 1");
    }
    @Loggable
    public static void run2(){
        System.out.println("run 2");
    }
    
    @Loggable
    public static void main(String [] args){
        run1();
        run2();
    }
}

package com.geekyarticles.asm;

public class Logger {
    public static void logMethodStart(String methodName){
        System.out.println("Starting method: "+methodName);
    }
    
    public static void logMethodReturn(String methodName){
        System.out.println("Ending method: "+methodName);
    }
}

package com.geekyarticles.asm;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface Loggable {

}

If you run this program, the generated file will have a dependency on the class Logger. Manually copy the Logger class to the correct package in the out directory. If you run the generated class (which is a modified version of LoggingTest class), the following would be the output.

bash-4.1$ java  com.geekyarticles.asm.LoggingTest
Starting method: main
run 1
Starting method: run2
run 2
Ending method: run2
Ending method: main

Note that unlike normal Lists, an InsnList object can be modified while we iterate over it. Any changes are immidiately reflected. So, if some instructions are inserted after the current position, that will also be iterated over.