javassist使用样例

javassist是jboss的一个子项目，其主要的优点，在于简单，而且快速。直接使用java编码的形式，而不需要了解虚拟机指令，就能动态改变类的结构，或者动态生成类。
本次样例是 通过注解以及文件列表 方式获取需要增强方法，动态织入方法的开始与结束时间。

public class MonitorTransformer implements ClassFileTransformer {

    private static final Logger LOGGER = Logger.getLogger(MonitorTransformer.class.getCanonicalName());
    private static final Map<ClassLoader, ClassPool> POOL_MAP = Collections.synchronizedMap(new WeakHashMap<ClassLoader, ClassPool>());
    final static String prefix ="\n methodBeginTime = System.currentTimeMillis();\n";
    final static String postfix ="\nlong endTime = System.currentTimeMillis();\n";
    final static char point_regex = ';';
    final static List<String> methodList =new ArrayList<String>();
    private static final Map<String, SignetClass> TRACE_INTER_CLASS = new ConcurrentHashMap<String, SignetClass>();
    private static final Map<String, SignetMethod> TRACE_INTER_METHOD = new ConcurrentHashMap<String, SignetMethod>();
    private static final Map<String, SignetMethod> METHOD_CACHE = new ConcurrentHashMap<String, SignetMethod>();

    private static final Map<ClassLoader, ClassLoader> CLASS_LOADER_CACHE = Collections.synchronizedMap(new WeakHashMap<ClassLoader, ClassLoader>());
    private static final ClassPool defaultPool = new ClassPool(null);
    static {
        defaultPool.appendSystemPath();
    }

    public MonitorTransformer(){
        MonitorConfig config;
        try {
            //读取配置文件
            config = new MonitorConfigImpl();
            String methodStr = config.getStringValue("methodList", null);
            Iterable<String> it = Splitter.on(point_regex).split(methodStr);
            //将读取的配置文件加入要检测的方法列表
            if(null!=it){
                Iterator<String> itor = it.iterator();
                while (itor.hasNext()) {
                    methodList.add(itor.next());
                }
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }


    /**
     * 获取监控方法列表
     * @param ctClass
     */
    protected void getMonitorMethodByAnnotation(CtClass ctClass){
        try {
            SignetClass signetClass = getSignetClass(ctClass);
            if (!ctClass.isInterface() && signetClass != null) {
                CtMethod[] methods = ctClass.getDeclaredMethods();
                CtMethod[] extMethods = ctClass.getMethods();
                HashSet<CtMethod> total = new HashSet<CtMethod>();
                Collections.addAll(total, methods);
                Collections.addAll(total, extMethods);
                for (CtMethod m : total) {
                    SignetMethod signetMethod = getAutobahnenMethod(m, ctClass);
                    if (signetMethod != null) {
                        METHOD_CACHE.put(ctClass.getName() + "$" + m.getName() + "$" + m.getSignature(), signetMethod);
                        if(!methodList.contains(ctClass.getName() + "." + m.getName())){
                            methodList.add(ctClass.getName() + "." + m.getName());
                        }
                    }
                }
            }
        } catch (Throwable e) {
            e.printStackTrace();
        }
    }


    /**
     * 获取加@SignetClass 注解类
     * @param ctClass
     * @return
     */
    private  SignetClass getSignetClass(CtClass ctClass) {
        try {
            SignetClass monitorClass = (SignetClass) ctClass.getAnnotation(SignetClass.class);
            if (monitorClass == null && !ctClass.isInterface()) {
                CtClass[] ctClasses = ctClass.getInterfaces();
                for (CtClass ctClass1 : ctClasses) {
                    monitorClass = TRACE_INTER_CLASS.get(ctClass1.getName());
                    if (monitorClass == null) {
                        monitorClass = (SignetClass) ctClass1.getAnnotation(SignetClass.class);
                        if (monitorClass != null) {
                            TRACE_INTER_CLASS.put(ctClass1.getName(), monitorClass);
                        }
                    }
                    if (monitorClass != null) {
                        CtMethod[] methods = ctClass1.getDeclaredMethods();
                        for (CtMethod ctMethod : methods) {
                            SignetMethod autobahnen = (SignetMethod) ctMethod.getAnnotation(SignetMethod.class);
                            if (autobahnen != null) {
                                TRACE_INTER_METHOD.put(ctClass1.getName() + "$" + ctMethod.getName() + "$" + ctMethod.getSignature(), autobahnen);
                            }
                        }
                    }
                }
            }
            return monitorClass;
        } catch (Throwable e) {
            e.getStackTrace();
        }
        return null;
    }


    /**
     * 获取注解方法
     * @param m
     * @param c
     * @return
     */
    private SignetMethod getAutobahnenMethod(CtMethod m, CtClass c) {
        SignetMethod autobahnen = getDeclaredMonitorMethod(m);
        if (autobahnen == null) {
            try {
                CtClass[] ctClasses = c.getInterfaces();
                for (CtClass ctClass : ctClasses) {
                    String className = ctClass.getName();
                    if (TRACE_INTER_CLASS.containsKey(className)) {
                        autobahnen = TRACE_INTER_METHOD.get(className + "$" + m.getName() + "$" + m.getSignature());
                        if (autobahnen != null) {
                            return autobahnen;
                        }
                    }
                }
            } catch (Throwable e) {
                e.printStackTrace();
            }
        }
        return autobahnen;
    }

    /**
     * 获取加了@SignetMethod 注解方法
     * @param m
     * @return
     */
    private SignetMethod getDeclaredMonitorMethod(CtMethod m) {
        try {
            return (SignetMethod) m.getAnnotation(SignetMethod.class);
        } catch (Throwable e) {
            e.printStackTrace();
        }
        return null;
    }

    /* (non-Javadoc)
     * @see java.lang.instrument.ClassFileTransformer#transform(java.lang.ClassLoader, java.lang.String, java.lang.Class, java.security.ProtectionDomain, byte[])
     */
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                           ProtectionDomain protectionDomain,byte[] classfileBuffer)
            throws IllegalClassFormatException {
        if (classBeingRedefined != null) {
            LOGGER.info("redefined class [" + classBeingRedefined.getCanonicalName() + "] needless transform .");
            return null;
        }

        if (loader != null && !CLASS_LOADER_CACHE.containsKey(loader))
            CLASS_LOADER_CACHE.put(loader, null);

        //先判断下现在加载的class的包路径是不是需要增强的包，通过instrumentation进来的class路径用‘/’分割
        if(className.startsWith("com/xxx/change/web")){

            className = className.replace("/",".");
            CtClass ctclass = null;
            try {
                //用于取得字节码类，必须在当前的classpath中，使用全称 ,这部分是关于javassist的知识
               // ctclass = ClassPool.getDefault().get(className);
                ctclass =  getCtClass(classfileBuffer, loader);//defaultPool.get(className);

                //循环一下，看看哪些方法需要加时间监测
                getMonitorMethodByAnnotation(ctclass);
                for(String method :methodList){
                    if (method.startsWith(className)){
                        try{
                            ctclass.getDeclaredField("methodBeginTime");
                        }catch (NotFoundException e){
                            CtField ctField2 = new CtField(CtClass.longType, "methodBeginTime",ctclass);
                            ctField2.setModifiers(Modifier.PUBLIC);
                            ctclass.addField(ctField2);
                        }
                        //获取方法名
                        String methodName = method.substring(method.lastIndexOf('.')+1, method.length());
                        String outputStr ="\n logger.info(\"this method "+methodName+" cost:\" +(System.currentTimeMillis() - methodBeginTime) +\"ms.\");";
                        //得到这方法实例
                        CtMethod ctmethod = ctclass.getDeclaredMethod(methodName);
                        ctmethod.insertBefore(prefix);
                        ctmethod.insertAfter(outputStr, true);
                     //   CtMethod processed = CtNewMethod.copy(ctmethod, ctclass, null);
//                        //新定义一个方法叫做比如sayHello$impl
//                        String newMethodName = methodName +"$impl";
//                        //原来的方法改个名字
//                        ctmethod.setName(newMethodName);
//
//                        //创建新的方法，复制原来的方法 ，名字为原来的名字
//                        CtMethod newMethod = CtNewMethod.copy(ctmethod, methodName, ctclass, null);
//                        //构建新的方法体
//                        StringBuilder bodyStr =new StringBuilder();
//                        bodyStr.append("{");
//                        bodyStr.append(prefix);
//                        //调用原有代码，类似于method();($$)表示所有的参数
//                        bodyStr.append(newMethodName +"($$);\n");
//
//                        bodyStr.append(postfix);
//                        bodyStr.append(outputStr);
//
//                        bodyStr.append("}");
//                        //替换新方法
//                        newMethod.setBody(bodyStr.toString());
//                        //增加新方法
//                        ctclass.addMethod(newMethod);
                    }
                }
                return ctclass.toBytecode();
            } catch (IOException e) {
                //TODO Auto-generated catch block
                e.printStackTrace();
            } catch (CannotCompileException e) {
                //TODO Auto-generated catch block
                e.printStackTrace();
            } catch (NotFoundException e) {
                //TODO Auto-generated catch block
                e.printStackTrace();
            }
        }
        return null;
    }

    private CtClass getCtClass(byte[] classFileBuffer, ClassLoader classLoader) throws IOException {
        ClassPool classPool = getClassPool(classLoader);
        CtClass clazz = classPool.makeClass(new ByteArrayInputStream(classFileBuffer), false);
        clazz.defrost();
        return clazz;
    }

    public static ClassPool getClassPool(ClassLoader loader) {
        if (loader == null)
            return defaultPool;

        ClassPool pool = POOL_MAP.get(loader);
        if (pool == null) {
            pool = new ClassPool(true);
            pool.appendClassPath(new LoaderClassPath(loader));
            POOL_MAP.put(loader, pool);
        }
        return pool;
    }
}