sf.codegen.EntityEnhancerJavassist Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of sorm Show documentation
java jpa tool for spring
The newest version!
package sf.codegen;

import javassist.ClassPath;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.JavassistUtils;
import javassist.LoaderClassPath;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.smallframework.spring.path.ClassPathScanningUtils;
import sf.tools.IOUtils;
import sf.tools.SystemUtils;

import java.io.File;
import java.io.InputStream;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;

/**
 * Entity静态增强任务类,按URL来访问资源
 * 作用 这个类中提供了{@link #enhance(String...)}
 * 方法，可以对当前classpath下的Entity类进行字节码增强。
 * @author
 * @Date
 */
public class EntityEnhancerJavassist implements IEntityEnhancer {
    private static final Logger logger = LoggerFactory.getLogger(EntityEnhancerJavassist.class);

    private List roots;

    @Override
    public IEntityEnhancer addRoot(URL url) {
        if (url != null) {
            if (roots == null) {
                roots = new ArrayList<>();
            }
            roots.add(url);
        }
        return this;
    }

    /**
     * 在当前的classpath目录下扫描Entity类(.clsss文件)，使用字节码增强修改这些class文件。
     * @param pkgNames
     */
    @Override
    public void enhance(final String... pkgNames) {
        ClassPool pool = new ClassPool(true);//为了被gc回收,不可使用ClassPool.getDefault()
        //bugfix: javassit高版本类路径查找有问题.
        ClassLoader cl = Thread.currentThread().getContextClassLoader();
        LoaderClassPath lcp = new LoaderClassPath(cl);
        pool.appendClassPath(lcp);
        int n = 0;
        if (roots == null || roots.size() == 0) {
            Set resources = ClassPathScanningUtils.getPathScanning().getClassResourcesURL(pkgNames);
            for (URL url : resources) {
                if (url.getProtocol().contains(SystemUtils.URL_PROTOCOL_FILE)) {//是否是文件url
                    try {
                        if (enhanceJavassist(url, pool)) {
                            n++;
                        }
                    } catch (Exception e) {
                        logger.error("Enhance error: {}", url, e);
                    }
                }
            }
        } else {
            for (URL root : roots) {
                //设置类加载路径
                ClassPath cp = null;
                try {
                    cp = pool.appendClassPath(root.getFile());
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
                Set resources = ClassPathScanningUtils.getPathScanning().getClassResourcesURL(root, pkgNames);
                for (URL url : resources) {
                    if (url.getProtocol().contains(SystemUtils.URL_PROTOCOL_FILE)) {//是否是文件url
                        try {
                            if (enhanceJavassist(url, pool)) {
                                n++;
                            }
                        } catch (Exception e) {
                            logger.error("Enhance error: {}", url, e);
                        }
                    }
                }
                //回收路径资源.
                if (cp != null) {
                    pool.removeClassPath(cp);
                }
            }
        }
        pool.removeClassPath(lcp);
        JavassistUtils.removeAllClass(pool);//清理多余的类
        System.out.println(n + " classes enhanced.");
    }

    private boolean enhanceJavassist(URL resource, ClassPool pool) {
        EnhanceTaskJavassist enhancer = new EnhanceTaskJavassist();
        InputStream inputStream = null;
        try {
            inputStream = resource.openStream();
            CtClass cc = pool.makeClass(inputStream);
            byte[] result = enhancer.doEnhance(pool, cc);
            if (result != null) {
                if (result.length == 0) {
                    System.out.println(cc.getName() + " is already enhanced.");
                } else {
                    IOUtils.saveAsFile(new File(resource.toURI()), result);
                    System.out.println("enhanced class:" + cc.getName());// 增强完成
                    return true;
                }
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            IOUtils.closeQuietly(inputStream);
        }
        return false;
    }

}