mockit.internal.state.CachedClassfiles Maven / Gradle / Ivy

Go to download
/*
 * Copyright (c) 2006-2013 Rogério Liesenfeld
 * This file is subject to the terms of the MIT license (see LICENSE.txt).
 */
package mockit.internal.state;

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

import org.jetbrains.annotations.*;

/**
 * Holds a map of internal class names to the corresponding class files (bytecode arrays), for the classes
 * that have already been loaded during the test run.
 * These classfiles are not necessarily the same as those stored in the corresponding ".class" files
 * available from the runtime classpath.
 * If any third-party {@link java.lang.instrument.ClassFileTransformer}s are active, those original classfiles
 * may have been modified before being loaded by the JVM.
 * JMockit installs a {@code ClassFileTransformer} of its own which saves all potentially modified classfiles
 * here.
 * 
 * This bytecode cache allows classes to be mocked and un-mocked correctly, even in the presence of other
 * bytecode modification agents such as the AspectJ load-time weaver.
 */
public final class CachedClassfiles implements ClassFileTransformer
{
   public static final CachedClassfiles INSTANCE = new CachedClassfiles();

   private final Map> classLoadersAndClassfiles =
      new WeakHashMap>(2);

   @Nullable private ClassDefinition[] classesBeingMocked;
   public void setClassesBeingMocked(@Nullable ClassDefinition[] classDefs) { classesBeingMocked = classDefs; }

   private CachedClassfiles() {}

   @Nullable
   public byte[] transform(
      @Nullable ClassLoader loader, @NotNull String classDesc, @Nullable Class classBeingRedefined,
      @NotNull ProtectionDomain protectionDomain, @NotNull byte[] classfileBuffer)
   {
      if (classBeingRedefined == null) { // class definition
         if (!isExcluded(classDesc)) {
            addClassfile(loader, classDesc, classfileBuffer);
         }
      }
      else if (!isBeingMocked(classBeingRedefined) && !isExcluded(classDesc)) { // class redefinition
         addClassfileIfNotYetPresent(loader, classDesc, classfileBuffer);
      }

      return null;
   }

   private boolean isExcluded(@NotNull String classDesc)
   {
      return
         classDesc.startsWith("org/junit/") ||
         classDesc.startsWith("sun/") && !classDesc.startsWith("sun/proxy/$") ||
         classDesc.startsWith("com/intellij/") ||
         classDesc.startsWith("mockit/internal/") || classDesc.startsWith("mockit/integration/") ||
         classDesc.startsWith("mockit/coverage/") ||
         classDesc.startsWith("mockit/") && classDesc.indexOf('$') < 0;
   }

   private boolean isBeingMocked(@NotNull Class classBeingRedefined)
   {
      if (classesBeingMocked != null) {
         for (ClassDefinition classDef : classesBeingMocked) {
            if (classDef.getDefinitionClass() == classBeingRedefined) {
               return true;
            }
         }
      }

      return false;
   }

   private synchronized void addClassfile(
      @Nullable ClassLoader loader, @NotNull String classDesc, @NotNull byte[] classfile)
   {
      Map classfiles = getClassfiles(loader);
      classfiles.put(classDesc, classfile);
   }

   @NotNull private Map getClassfiles(@Nullable ClassLoader loader)
   {
      Map classfiles = classLoadersAndClassfiles.get(loader);

      if (classfiles == null) {
         classfiles = new HashMap(100);
         classLoadersAndClassfiles.put(loader, classfiles);
      }

      return classfiles;
   }

   private synchronized void addClassfileIfNotYetPresent(
      @Nullable ClassLoader loader, @NotNull String classDesc, @NotNull byte[] classfile)
   {
      Map classfiles = getClassfiles(loader);

      if (!classfiles.containsKey(classDesc)) {
         classfiles.put(classDesc, classfile);
      }
   }

   @Nullable private synchronized byte[] findClassfile(@NotNull Class aClass)
   {
      Map classfiles = getClassfiles(aClass.getClassLoader());
      return classfiles.get(aClass.getName().replace('.', '/'));
   }

   @Nullable private synchronized byte[] findClassfile(@Nullable ClassLoader loader, @NotNull String classDesc)
   {
      Map classfiles = getClassfiles(loader);
      return classfiles.get(classDesc);
   }

   @Nullable public static byte[] getClassfile(@NotNull Class aClass) { return INSTANCE.findClassfile(aClass); }

   @Nullable public static byte[] getClassfile(@Nullable ClassLoader loader, @NotNull String internalClassName)
   {
      return INSTANCE.findClassfile(loader, internalClassName);
   }

   public static void addClassfile(@NotNull Class aClass, @NotNull byte[] classfile)
   {
      INSTANCE.addClassfile(aClass.getClassLoader(), aClass.getName().replace('.', '/'), classfile);
   }
}