mockit.internal.capturing.CaptureTransformer Maven / Gradle / Ivy

Go to download
/*
 * Copyright (c) 2006 JMockit developers
 * This file is subject to the terms of the MIT license (see LICENSE.txt).
 */
package mockit.internal.capturing;

import static mockit.internal.capturing.CapturedType.isNotToBeCaptured;

import edu.umd.cs.findbugs.annotations.NonNull;
import edu.umd.cs.findbugs.annotations.Nullable;

import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import mockit.asm.classes.ClassReader;
import mockit.asm.classes.ClassVisitor;
import mockit.asm.metadata.ClassMetadataReader;
import mockit.asm.types.JavaType;
import mockit.internal.ClassFile;
import mockit.internal.ClassIdentification;
import mockit.internal.startup.Startup;
import mockit.internal.state.TestRun;

public final class CaptureTransformer implements ClassFileTransformer {
    @NonNull
    private final CapturedType capturedType;
    @NonNull
    private final String capturedTypeDesc;
    @NonNull
    private final CaptureOfImplementations captureOfImplementations;
    @NonNull
    private final Map transformedClasses;
    @NonNull
    private final Map superTypesSearched;
    @Nullable
    private final M typeMetadata;
    private boolean inactive;

    CaptureTransformer(@NonNull CapturedType capturedType,
            @NonNull CaptureOfImplementations captureOfImplementations, boolean registerTransformedClasses,
            @Nullable M typeMetadata) {
        this.capturedType = capturedType;
        capturedTypeDesc = JavaType.getInternalName(capturedType.baseType);
        this.captureOfImplementations = captureOfImplementations;
        transformedClasses = registerTransformedClasses ? new HashMap<>(2)
                : Collections. emptyMap();
        superTypesSearched = new HashMap<>();
        this.typeMetadata = typeMetadata;
    }

    public void deactivate() {
        inactive = true;

        if (!transformedClasses.isEmpty()) {
            for (Map.Entry classNameAndOriginalBytecode : transformedClasses.entrySet()) {
                ClassIdentification classId = classNameAndOriginalBytecode.getKey();
                byte[] originalBytecode = classNameAndOriginalBytecode.getValue();

                Startup.redefineMethods(classId, originalBytecode);
            }

            transformedClasses.clear();
        }
    }

    @Nullable
    @Override
    public byte[] transform(@Nullable ClassLoader loader, @NonNull String classDesc,
            @Nullable Class classBeingRedefined, @Nullable ProtectionDomain protectionDomain,
            @NonNull byte[] classfileBuffer) {
        if (classBeingRedefined != null || inactive || isNotToBeCaptured(protectionDomain, classDesc)) {
            return null;
        }

        if (isClassToBeCaptured(loader, classfileBuffer)) {
            String className = classDesc.replace('/', '.');
            ClassReader cr = new ClassReader(classfileBuffer);
            return modifyAndRegisterClass(loader, className, cr);
        }

        return null;
    }

    private boolean isClassToBeCaptured(@Nullable ClassLoader loader, @NonNull byte[] classfileBuffer) {
        ClassMetadataReader cmr = new ClassMetadataReader(classfileBuffer);
        String superName = cmr.getSuperClass();

        if (capturedTypeDesc.equals(superName)) {
            return true;
        }

        String[] interfaces = cmr.getInterfaces();

        if (interfaces != null && isClassWhichImplementsACapturingInterface(interfaces)) {
            return true;
        }

        return superName != null && searchSuperTypes(loader, superName, interfaces);
    }

    private boolean isClassWhichImplementsACapturingInterface(@NonNull String[] interfaces) {
        for (String implementedInterface : interfaces) {
            if (capturedTypeDesc.equals(implementedInterface)) {
                return true;
            }
        }

        return false;
    }

    private boolean searchSuperTypes(@Nullable ClassLoader loader, @NonNull String superName,
            @Nullable String[] interfaces) {
        if (!"java/lang/Object".equals(superName) && !superName.startsWith("mockit/")
                && searchSuperType(loader, superName)) {
            return true;
        }

        if (interfaces != null && interfaces.length > 0) {
            for (String itf : interfaces) {
                if (!itf.startsWith("java/") && !itf.startsWith("javax/") && searchSuperType(loader, itf)) {
                    return true;
                }
            }
        }

        return false;
    }

    private boolean searchSuperType(@Nullable ClassLoader loader, @NonNull String superName) {
        Boolean extendsCapturedType = superTypesSearched.get(superName);

        if (extendsCapturedType == null) {
            byte[] classfileBytes = ClassFile.getClassFile(loader, superName);
            extendsCapturedType = isClassToBeCaptured(loader, classfileBytes);
            superTypesSearched.put(superName, extendsCapturedType);
        }

        return extendsCapturedType;
    }

    @NonNull
    private byte[] modifyAndRegisterClass(@Nullable ClassLoader loader, @NonNull String className,
            @NonNull ClassReader cr) {
        ClassVisitor modifier = captureOfImplementations.createModifier(loader, cr, capturedType.baseType,
                typeMetadata);
        cr.accept(modifier);

        ClassIdentification classId = new ClassIdentification(loader, className);
        byte[] originalBytecode = cr.getBytecode();

        if (transformedClasses == Collections. emptyMap()) {
            TestRun.mockFixture().addTransformedClass(classId, originalBytecode);
        } else {
            transformedClasses.put(classId, originalBytecode);
        }

        TestRun.mockFixture().registerMockedClass(capturedType.baseType);
        return modifier.toByteArray();
    }

    @Nullable
    public > C getCaptureOfImplementationsIfApplicable(@NonNull Class aType) {
        if (typeMetadata != null && capturedType.baseType.isAssignableFrom(aType)) {
            // noinspection unchecked
            return (C) captureOfImplementations;
        }

        return null;
    }

    public boolean areCapturedClasses(@NonNull Class mockedClass1, @NonNull Class mockedClass2) {
        Class baseType = capturedType.baseType;
        return baseType.isAssignableFrom(mockedClass1) && baseType.isAssignableFrom(mockedClass2);
    }
}