src/main/java/org/mockito/internal/creation/bytebuddy/InlineBytecodeGenerator.java

/*
 * Copyright (c) 2016 Mockito contributors
 * This program is made available under the terms of the MIT License.
 */
package org.mockito.internal.creation.bytebuddy;

import net.bytebuddy.ByteBuddy;
import net.bytebuddy.ClassFileVersion;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.asm.AsmVisitorWrapper;
import net.bytebuddy.description.field.FieldDescription;
import net.bytebuddy.description.field.FieldList;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.method.MethodList;
import net.bytebuddy.description.method.ParameterDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.ClassFileLocator;
import net.bytebuddy.dynamic.scaffold.MethodGraph;
import net.bytebuddy.dynamic.scaffold.TypeValidation;
import net.bytebuddy.implementation.Implementation;
import net.bytebuddy.jar.asm.ClassVisitor;
import net.bytebuddy.jar.asm.MethodVisitor;
import net.bytebuddy.jar.asm.Opcodes;
import net.bytebuddy.matcher.ElementMatchers;
import net.bytebuddy.pool.TypePool;
import net.bytebuddy.utility.RandomString;
import org.mockito.exceptions.base.MockitoException;
import org.mockito.internal.util.concurrent.WeakConcurrentMap;
import org.mockito.internal.util.concurrent.WeakConcurrentSet;
import org.mockito.mock.SerializableMode;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Modifier;
import java.security.ProtectionDomain;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;

import static net.bytebuddy.implementation.MethodDelegation.withDefaultConfiguration;
import static net.bytebuddy.implementation.bind.annotation.TargetMethodAnnotationDrivenBinder.ParameterBinder.ForFixedValue.OfConstant.of;
import static net.bytebuddy.matcher.ElementMatchers.*;
import static org.mockito.internal.util.StringUtil.join;

public class InlineBytecodeGenerator implements BytecodeGenerator, ClassFileTransformer {

    private static final String PRELOAD = "org.mockito.inline.preload";

    @SuppressWarnings("unchecked")
    static final Set<Class<?>> EXCLUDES = new HashSet<Class<?>>(Arrays.asList(Class.class,
            Boolean.class,
            Byte.class,
            Short.class,
            Character.class,
            Integer.class,
            Long.class,
            Float.class,
            Double.class,
            String.class));

    private final Instrumentation instrumentation;

    private final ByteBuddy byteBuddy;

    private final WeakConcurrentSet<Class<?>> mocked;

    private final BytecodeGenerator subclassEngine;

    private final AsmVisitorWrapper mockTransformer;

    private volatile Throwable lastException;

    public InlineBytecodeGenerator(Instrumentation instrumentation, WeakConcurrentMap<Object, MockMethodInterceptor> mocks) {
        preload();
        this.instrumentation = instrumentation;
        byteBuddy = new ByteBuddy()
            .with(TypeValidation.DISABLED)
            .with(Implementation.Context.Disabled.Factory.INSTANCE)
            .with(MethodGraph.Compiler.ForDeclaredMethods.INSTANCE);
        mocked = new WeakConcurrentSet<Class<?>>(WeakConcurrentSet.Cleaner.INLINE);
        String identifier = RandomString.make();
        subclassEngine = new TypeCachingBytecodeGenerator(new SubclassBytecodeGenerator(withDefaultConfiguration()
            .withBinders(of(MockMethodAdvice.Identifier.class, identifier))
            .to(MockMethodAdvice.ForReadObject.class), isAbstract().or(isNative()).or(isToString())), false);
        mockTransformer = new AsmVisitorWrapper.ForDeclaredMethods()
            .method(isVirtual()
                    .and(not(isBridge().or(isHashCode()).or(isEquals()).or(isDefaultFinalizer())))
                    .and(not(isDeclaredBy(nameStartsWith("java.")).<MethodDescription>and(isPackagePrivate()))),
                Advice.withCustomMapping()
                    .bind(MockMethodAdvice.Identifier.class, identifier)
                    .to(MockMethodAdvice.class))
            .method(isHashCode(),
                Advice.withCustomMapping()
                    .bind(MockMethodAdvice.Identifier.class, identifier)
                    .to(MockMethodAdvice.ForHashCode.class))
            .method(isEquals(),
                Advice.withCustomMapping()
                    .bind(MockMethodAdvice.Identifier.class, identifier)
                    .to(MockMethodAdvice.ForEquals.class));
        MockMethodDispatcher.set(identifier, new MockMethodAdvice(mocks, identifier));
        instrumentation.addTransformer(this, true);
    }

    /**
     * Mockito allows to mock about any type, including such types that we are relying on ourselves. This can cause a circularity:
     * In order to check if an instance is a mock we need to look up if this instance is registered in the {@code mocked} set. But to look
     * up this instance, we need to create key instances that rely on weak reference properties. Loading the later classes will happen before
     * the key instances are completed what will cause Mockito to check if those key instances are themselves mocks what causes a loop which
     * results in a circularity error. This is not normally a problem as we explicitly check if the instance that we investigate is one of
     * our instance of which we hold a reference by reference equality what does not cause any code execution. But it seems like the load
     * order plays a role here with unloaded types being loaded before we even get to check the mock instance property. To avoid this, we are
     * making sure that crucuial JVM types are loaded before we create the first inline mock. Unfortunately, these types dependant on a JVM's
     * implementation and we can only maintain types that we know of from well-known JVM implementations such as HotSpot and extend this list
     * once we learn of further problematic types for future Java versions. To allow users to whitelist their own types, we do not also offer
     * a property that allows running problematic tests before a new Mockito version can be released and that allows us to ask users to
     * easily validate that whitelisting actually solves a problem as circularities could also be caused by other problems.
     */
    private static void preload() {
        String preloads = System.getProperty(PRELOAD);
        if (preloads == null) {
            preloads = "java.lang.WeakPairMap,java.lang.WeakPairMap$Pair,java.lang.WeakPairMap$Pair$Weak";
        }
        for (String preload : preloads.split(",")) {
            try {
                Class.forName(preload, false, null);
            } catch (ClassNotFoundException ignored) {
            }
        }
    }

    @Override
    public <T> Class<? extends T> mockClass(MockFeatures<T> features) {
        boolean subclassingRequired = !features.interfaces.isEmpty()
            || features.serializableMode != SerializableMode.NONE
            || Modifier.isAbstract(features.mockedType.getModifiers());

        checkSupportedCombination(subclassingRequired, features);

        synchronized (this) {
            triggerRetransformation(features);
        }

        return subclassingRequired ?
                subclassEngine.mockClass(features) :
                features.mockedType;
    }

    private <T> void triggerRetransformation(MockFeatures<T> features) {
        Set<Class<?>> types = new HashSet<Class<?>>();
        Class<?> type = features.mockedType;
        do {
            if (mocked.add(type)) {
                types.add(type);
                addInterfaces(types, type.getInterfaces());
            }
            type = type.getSuperclass();
        } while (type != null);
        if (!types.isEmpty()) {
            try {
                instrumentation.retransformClasses(types.toArray(new Class<?>[types.size()]));
                Throwable throwable = lastException;
                if (throwable != null) {
                    throw new IllegalStateException(join("Byte Buddy could not instrument all classes within the mock's type hierarchy",
                        "",
                        "This problem should never occur for javac-compiled classes. This problem has been observed for classes that are:",
                        " - Compiled by older versions of scalac",
                        " - Classes that are part of the Android distribution"), throwable);
                }
            } catch (Exception exception) {
                for (Class<?> failed : types) {
                    mocked.remove(failed);
                }
                throw new MockitoException("Could not modify all classes " + types, exception);
            } finally {
                lastException = null;
            }
        }
    }

    private <T> void checkSupportedCombination(boolean subclassingRequired, MockFeatures<T> features) {
        if (subclassingRequired
                && !features.mockedType.isArray()
                && !features.mockedType.isPrimitive()
                && Modifier.isFinal(features.mockedType.getModifiers())) {
            throw new MockitoException("Unsupported settings with this type '" + features.mockedType.getName() + "'");
        }
    }

    private void addInterfaces(Set<Class<?>> types, Class<?>[] interfaces) {
        for (Class<?> type : interfaces) {
            if (mocked.add(type)) {
                types.add(type);
                addInterfaces(types, type.getInterfaces());
            }
        }
    }

    @Override
    public byte[] transform(ClassLoader loader,
                            String className,
                            Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain,
                            byte[] classfileBuffer) {
        if (classBeingRedefined == null
            || !mocked.contains(classBeingRedefined)
            || EXCLUDES.contains(classBeingRedefined)) {
            return null;
        } else {
            try {
                return byteBuddy.redefine(classBeingRedefined, ClassFileLocator.Simple.of(classBeingRedefined.getName(), classfileBuffer))
                    // Note: The VM erases parameter meta data from the provided class file (bug). We just add this information manually.
                    .visit(new ParameterWritingVisitorWrapper(classBeingRedefined))
                    .visit(mockTransformer)
                    .make()
                    .getBytes();
            } catch (Throwable throwable) {
                lastException = throwable;
                return null;
            }
        }
    }

    private static class ParameterWritingVisitorWrapper extends AsmVisitorWrapper.AbstractBase {

        private final Class<?> type;

        private ParameterWritingVisitorWrapper(Class<?> type) {
            this.type = type;
        }

        @Override
        public ClassVisitor wrap(TypeDescription instrumentedType,
                                 ClassVisitor classVisitor,
                                 Implementation.Context implementationContext,
                                 TypePool typePool,
                                 FieldList<FieldDescription.InDefinedShape> fields,
                                 MethodList<?> methods,
                                 int writerFlags,
                                 int readerFlags) {
            return implementationContext.getClassFileVersion().isAtLeast(ClassFileVersion.JAVA_V8)
                    ? new ParameterAddingClassVisitor(classVisitor, new TypeDescription.ForLoadedType(type))
                    : classVisitor;
        }

        private static class ParameterAddingClassVisitor extends ClassVisitor {

            private final TypeDescription typeDescription;

            private ParameterAddingClassVisitor(ClassVisitor cv, TypeDescription typeDescription) {
                super(Opcodes.ASM6, cv);
                this.typeDescription = typeDescription;
            }

            @Override
            public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
                MethodVisitor methodVisitor = super.visitMethod(access, name, desc, signature, exceptions);
                MethodList<?> methodList = typeDescription.getDeclaredMethods().filter((name.equals(MethodDescription.CONSTRUCTOR_INTERNAL_NAME)
                        ? isConstructor()
                        : ElementMatchers.<MethodDescription>named(name)).and(hasDescriptor(desc)));
                if (methodList.size() == 1 && methodList.getOnly().getParameters().hasExplicitMetaData()) {
                    for (ParameterDescription parameterDescription : methodList.getOnly().getParameters()) {
                        methodVisitor.visitParameter(parameterDescription.getName(), parameterDescription.getModifiers());
                    }
                    return new MethodParameterStrippingMethodVisitor(methodVisitor);
                } else {
                    return methodVisitor;
                }
            }
        }

        private static class MethodParameterStrippingMethodVisitor extends MethodVisitor {

            public MethodParameterStrippingMethodVisitor(MethodVisitor mv) {
                super(Opcodes.ASM5, mv);
            }

            @Override
            public void visitParameter(String name, int access) {
                // suppress to avoid additional writing of the parameter if retained.
            }
        }
    }
}