mockit.internal.expectations.transformation.ExpectationsTransformer 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.expectations.transformation;

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

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

import mockit.asm.classes.ClassReader;
import mockit.asm.classes.ClassVisitor;
import mockit.asm.classes.ClassWriter;
import mockit.asm.classes.WrappingClassVisitor;
import mockit.asm.methods.MethodVisitor;
import mockit.asm.methods.MethodWriter;
import mockit.internal.util.ClassNaming;
import mockit.internal.util.VisitInterruptedException;

public final class ExpectationsTransformer implements ClassFileTransformer {
    private static final String BASE_CLASSES = "mockit/Expectations mockit/Verifications mockit/VerificationsInOrder mockit/FullVerifications";

    @Nullable
    @Override
    public byte[] transform(@Nullable ClassLoader loader, @NonNull String className,
            @Nullable Class classBeingRedefined, @Nullable ProtectionDomain protectionDomain,
            @NonNull byte[] classfileBuffer) {
        if (classBeingRedefined == null && protectionDomain != null && className != null) {
            boolean anonymousClass = ClassNaming.isAnonymousClass(className);

            if (anonymousClass && !isJMockitClass(className) && !className.startsWith("org/junit/")) {
                ClassReader cr = new ClassReader(classfileBuffer);
                String superClassName = cr.getSuperName();

                if (!BASE_CLASSES.contains(superClassName)) {
                    return null;
                }

                return modifyInvocationsSubclass(cr, className);
            }
        }

        return null;
    }

    private static boolean isJMockitClass(@NonNull String classDesc) {
        return classDesc.startsWith("mockit/") && (classDesc.startsWith("mockit/internal/")
                || classDesc.startsWith("mockit/coverage/") || classDesc.startsWith("mockit/integration/"));
    }

    @Nullable
    private static byte[] modifyInvocationsSubclass(@NonNull ClassReader cr, @NonNull final String classDesc) {
        ClassWriter cw = new ClassWriter(cr);

        ClassVisitor modifier = new WrappingClassVisitor(cw) {
            @Override
            public MethodVisitor visitMethod(int access, @NonNull String name, @NonNull String desc,
                    @Nullable String signature, @Nullable String[] exceptions) {
                MethodWriter mw = cw.visitMethod(access, name, desc, signature, exceptions);

                if (!"".equals(name)) {
                    return mw;
                }

                return new InvocationBlockModifier(mw, classDesc);
            }
        };

        try {
            cr.accept(modifier);
            return modifier.toByteArray();
        } catch (VisitInterruptedException ignore) {
        } catch (Throwable e) {
            e.printStackTrace();
        }

        return null;
    }
}