/**
 * Simple BCEL bytecode example.
 *
 * Extend this example such that a program announces the following
 * instructions:  NEW, PUTFIELD, PUTSTATIC.
 * 
 * @author Godmar Back
 */
package intercept;

import java.util.Iterator;
import java.io.File;
import org.apache.bcel.Constants;
import org.apache.bcel.util.*;
import org.apache.bcel.classfile.*;
import org.apache.bcel.generic.*;

public class BCELInterceptor extends org.apache.bcel.util.ClassLoader {
    private InstructionFactory fac;
    private ClassGen cg;

    public BCELInterceptor(java.lang.ClassLoader parent) {
        super(parent, new String[] { "intercept.", "java.", "javax.", "sun." });
    }

    /* bcel.util.ClassLoader will invoke this method for every class
     * it loads, giving us a chance to modify it.
     */
    protected JavaClass modifyClass(JavaClass clazz) {
        // creates and then uses its own cloned ConstantPoolGen
        cg = new ClassGen(clazz);  

        for(Method m : clazz.getMethods())
            modifyMethod(m);

        clazz = cg.getJavaClass();
        /* output patched class for later inspection */
        try {
            clazz.dump(new File(clazz.getClassName() + "-patched.class"));
        } catch (Exception e) {
            System.out.println(e);
        }
        return clazz;
    }

    private void modifyMethod(Method oldmethod) {
        if (oldmethod.getCode() == null)        // native method
            return;

        MethodGen mg = new MethodGen(oldmethod, 
                                     cg.getClassName(), 
                                     cg.getConstantPool());
        InstructionList il = mg.getInstructionList();
        fac = new InstructionFactory(cg);
        InstructionFinder f = new InstructionFinder(il);

        // apply all patch actions for which the defined pattern matches. 
        for (PatchAction pa : actions) {
            Iterator it = f.search(pa.pattern(), pa.constraint());
            while (it.hasNext()) {
                InstructionHandle [] match = (InstructionHandle [])it.next();
                pa.doit(il, match, mg);
            }
        }

        mg.setMaxStack();
        cg.replaceMethod(oldmethod, mg.getMethod());
        il.dispose();
    }

    // an abstract class for simple instrumentation actions
    static abstract class PatchAction {
        // a helper
        protected static InstructionFinder.CodeConstraint alwaysTrue = 
            new InstructionFinder.CodeConstraint() {
                public boolean checkCode(InstructionHandle [] match) { 
                    return true; 
                }
            };

        protected String pattern;

        PatchAction(String pattern) {
            this.pattern = pattern;
        }

        String pattern() {
            return this.pattern;
        }

        InstructionFinder.CodeConstraint constraint() {
            return alwaysTrue;
        }

        abstract void doit(InstructionList il, InstructionHandle []match, MethodGen mg);
    }

    // short for if_acmpxx
    private PatchAction doStringCmp = new PatchAction("if_acmp") { 
        public void doit(InstructionList il, InstructionHandle []match, MethodGen mg) {
            if (match.length != 2)
                throw new Error("match length is not 2, it's " + match.length);

            /* insert 
                        DUP2    (x y -> x y x y)
                        INVOKE_STATIC intercept.cmp (x, y)
             */
            InstructionList intercept = new InstructionList();
            intercept.append(InstructionConstants.DUP2);
            intercept.append(fac.createInvoke("intercept.StringInterceptor", "cmp", 
                                Type.VOID, 
                                new Type [] { Type.OBJECT, Type.OBJECT }, 
                                Constants.INVOKESTATIC));
            System.out.println("before: " + il);
            il.insert(match[0], intercept);
            System.out.println("after: " + il);
            intercept.dispose();
        }
    };

    // matching all INVOKEVIRTUAL instructions
    // but also apply constraint, see below.
    private PatchAction wrongHashCode = new PatchAction("INVOKEVIRTUAL") {
        public void doit(InstructionList il, InstructionHandle []match, MethodGen mg) {
            if (match.length != 2)
                throw new Error("match length is not 2, it's " + match.length);

            // insert DUP2 before INVOKEVIRTUAL instruction to preserve arguments
            InstructionList intercept = new InstructionList(InstructionConstants.DUP2);
            il.insert(match[0], intercept);
            intercept.dispose();

            // after
            intercept = new InstructionList(
                fac.createInvoke("intercept.HashcodeInterceptor", "cmp", 
                                Type.BOOLEAN, 
                                new Type [] { Type.OBJECT, Type.OBJECT, Type.BOOLEAN }, 
                                Constants.INVOKESTATIC));
            il.append(match[0], intercept);
            intercept.dispose();
        }

        // specify additional constraint for match - we are only interested
        // in Object.equals() calls
        public InstructionFinder.CodeConstraint constraint() {
            return new InstructionFinder.CodeConstraint() {
                public boolean checkCode(InstructionHandle [] match) { 
                    INVOKEVIRTUAL iv = (INVOKEVIRTUAL)match[0].getInstruction();
                    String sig = iv.getSignature(cg.getConstantPool());
                    String name = iv.getName(cg.getConstantPool());
                    boolean rc = name.equals("equals") 
                            && sig.equals("(Ljava/lang/Object;)Z");
                    // System.out.println(sig + " " + name + " " + rc);
                    return rc;
                }
            };
        }
    };

    private PatchAction showException = new PatchAction("astore") {
        public void doit(InstructionList il, InstructionHandle []match, MethodGen mg) {
            boolean found = false;
            CodeExceptionGen ceg = null;
            for (CodeExceptionGen c : mg.getExceptionHandlers()) {
                if (match[0] == c.getHandlerPC()) {
                    ceg = c; found = true; break;
                }
            } 
            if (!found)
                return;

            InstructionList intercept = new InstructionList();
            intercept.append(InstructionConstants.DUP);
            intercept.append(fac.createInvoke("intercept.ExceptionInterceptor", "show", 
                                Type.VOID, 
                                new Type [] { Type.OBJECT }, 
                                Constants.INVOKESTATIC));
            System.out.println("before: " + il);
            InstructionHandle b = il.insert(match[0], intercept);
            ceg.setHandlerPC(b);
            System.out.println("after: " + il);
            intercept.dispose();
        }
    };

    PatchAction [] actions = new PatchAction[] { doStringCmp, 
                                                 wrongHashCode,
                                                 showException
                                                };
}