org.test4j.mock.faking.modifier.FakeTransformer Maven / Gradle / Ivy

Go to download
package org.test4j.mock.faking.modifier;

import g_asm.org.objectweb.asm.ClassReader;
import org.test4j.mock.faking.util.TypeUtility;
import org.test4j.mock.startup.Startup;

import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;
import java.util.*;

import static org.test4j.mock.faking.util.AsmConstant.acceptOptions;
import static org.test4j.mock.faking.util.ClassFile.writeBytes4Debug;

/**
 * 对需要进行mock增强的类进行重定义处理
 *
 * @author darui.wu
 */
public final class FakeTransformer implements ClassFileTransformer {

    public static final FakeTransformer INSTANCE = new FakeTransformer();
    /**
     * key: class desc
     * value: static method list (name+desc): isStatic
     */
    final static transient Map> fakedMethods = new HashMap<>(32);

    private Class beingCached;

    private FakeTransformer() {
    }

    public static Boolean findMethodInFaked(String classDesc, String methodDesc) {
        Map map = fakedMethods.get(classDesc);
        return map == null ? null : map.get(methodDesc);
    }

    @Override
    public byte[] transform(ClassLoader loader, String classDesc, Class aClass, ProtectionDomain domain, byte[] bytes) {
        byte[] faked = null;
        if (this.beenCaching(classDesc, aClass)) {
            faked = this.fake(classDesc, aClass, bytes);
            beingCached = null;
        }
        return faked;
    }

    private boolean beenCaching(String classDesc, Class beingRedefined) {
        if (classDesc == null) {
            return false;
        } else if (fakedMethods.containsKey(classDesc)) {
            return false;
        } else {
            return beingRedefined != null && beingRedefined == beingCached;
        }
    }

    /**
     * 更新变更逻辑
     *
     * @param classDesc
     * @param classToFake
     * @param bytes
     */
    private static byte[] fake(String classDesc, Class classToFake, byte[] bytes) {
        if (fakedMethods.containsKey(classDesc)) {
            return bytes;
        }
        synchronized (INSTANCE) {
            if (fakedMethods.containsKey(classDesc) || notMockType(classDesc)) {
                return bytes;
            }
            FakeClassModifier cv = new FakeClassModifier();
            new ClassReader(bytes).accept(cv, acceptOptions);

            byte[] faked = cv.toByteArray();
            writeBytes4Debug(classToFake.getName(), faked);
            INSTANCE.fakedMethods.put(classDesc, cv.getFakedMethods());
            return faked;
        }
    }

    /**
     * 不能被mock的类
     *
     * @param typeName
     * @return
     */
    public static boolean notMockType(String typeName) {
        String name = typeName.replace('/', '.');
        for (String prefix : Not_Mock_Packages) {
            if (name.startsWith(prefix)) {
                return true;
            }
        }
        return false;
    }

    private static final List Not_Mock_Packages = Arrays.asList(
        "java.lang.",
        "java.util.concurrent.",
        "java.util.jar.",
        "java.util.regex.",
        "java.util.zip.",
        "jdk.internal.",
        "sun.",
        "com.sun."
    );

    /**
     * 对declaredToFake及其子类进行mock增强
     */
    public static void applyFakes(Class declaredToFake) {
        if (declaredToFake == null) {
            return;
        }
        Set toBeFakes = TypeUtility.findAllClass(declaredToFake);
        for (Class classToModify : toBeFakes) {
            FakeTransformer.fakeClass(classToModify);
        }
    }

    private static void fakeClass(Class aClass) {
        String classDesc = TypeUtility.classPath(aClass);
        if (INSTANCE.fakedMethods.containsKey(classDesc)) {
            return;
        }
        INSTANCE.beingCached = aClass;
        Startup.reTransformClass(aClass);
    }
}