• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /*
2  * Copyright (c) 2016 Mockito contributors
3  * This program is made available under the terms of the MIT License.
4  */
5 package org.mockito.internal.creation.bytebuddy;
6 
7 import net.bytebuddy.ByteBuddy;
8 import net.bytebuddy.ClassFileVersion;
9 import net.bytebuddy.asm.Advice;
10 import net.bytebuddy.asm.AsmVisitorWrapper;
11 import net.bytebuddy.description.field.FieldDescription;
12 import net.bytebuddy.description.field.FieldList;
13 import net.bytebuddy.description.method.MethodDescription;
14 import net.bytebuddy.description.method.MethodList;
15 import net.bytebuddy.description.method.ParameterDescription;
16 import net.bytebuddy.description.type.TypeDescription;
17 import net.bytebuddy.dynamic.ClassFileLocator;
18 import net.bytebuddy.dynamic.scaffold.MethodGraph;
19 import net.bytebuddy.dynamic.scaffold.TypeValidation;
20 import net.bytebuddy.implementation.Implementation;
21 import net.bytebuddy.jar.asm.ClassVisitor;
22 import net.bytebuddy.jar.asm.MethodVisitor;
23 import net.bytebuddy.jar.asm.Opcodes;
24 import net.bytebuddy.matcher.ElementMatchers;
25 import net.bytebuddy.pool.TypePool;
26 import net.bytebuddy.utility.OpenedClassReader;
27 import net.bytebuddy.utility.RandomString;
28 import org.mockito.exceptions.base.MockitoException;
29 import org.mockito.internal.util.concurrent.WeakConcurrentMap;
30 import org.mockito.internal.util.concurrent.WeakConcurrentSet;
31 import org.mockito.mock.SerializableMode;
32 
33 import java.lang.instrument.ClassFileTransformer;
34 import java.lang.instrument.Instrumentation;
35 import java.lang.reflect.Modifier;
36 import java.security.ProtectionDomain;
37 import java.util.Arrays;
38 import java.util.HashSet;
39 import java.util.Set;
40 
41 import static net.bytebuddy.implementation.MethodDelegation.withDefaultConfiguration;
42 import static net.bytebuddy.implementation.bind.annotation.TargetMethodAnnotationDrivenBinder.ParameterBinder.ForFixedValue.OfConstant.of;
43 import static net.bytebuddy.matcher.ElementMatchers.*;
44 import static org.mockito.internal.util.StringUtil.join;
45 
46 public class InlineBytecodeGenerator implements BytecodeGenerator, ClassFileTransformer {
47 
48     private static final String PRELOAD = "org.mockito.inline.preload";
49 
50     @SuppressWarnings("unchecked")
51     static final Set<Class<?>> EXCLUDES = new HashSet<Class<?>>(Arrays.asList(Class.class,
52             Boolean.class,
53             Byte.class,
54             Short.class,
55             Character.class,
56             Integer.class,
57             Long.class,
58             Float.class,
59             Double.class,
60             String.class));
61 
62     private final Instrumentation instrumentation;
63 
64     private final ByteBuddy byteBuddy;
65 
66     private final WeakConcurrentSet<Class<?>> mocked;
67 
68     private final BytecodeGenerator subclassEngine;
69 
70     private final AsmVisitorWrapper mockTransformer;
71 
72     private volatile Throwable lastException;
73 
InlineBytecodeGenerator(Instrumentation instrumentation, WeakConcurrentMap<Object, MockMethodInterceptor> mocks)74     public InlineBytecodeGenerator(Instrumentation instrumentation, WeakConcurrentMap<Object, MockMethodInterceptor> mocks) {
75         preload();
76         this.instrumentation = instrumentation;
77         byteBuddy = new ByteBuddy()
78             .with(TypeValidation.DISABLED)
79             .with(Implementation.Context.Disabled.Factory.INSTANCE)
80             .with(MethodGraph.Compiler.ForDeclaredMethods.INSTANCE);
81         mocked = new WeakConcurrentSet<Class<?>>(WeakConcurrentSet.Cleaner.INLINE);
82         String identifier = RandomString.make();
83         subclassEngine = new TypeCachingBytecodeGenerator(new SubclassBytecodeGenerator(withDefaultConfiguration()
84             .withBinders(of(MockMethodAdvice.Identifier.class, identifier))
85             .to(MockMethodAdvice.ForReadObject.class), isAbstract().or(isNative()).or(isToString())), false);
86         mockTransformer = new AsmVisitorWrapper.ForDeclaredMethods()
87             .method(isVirtual()
88                     .and(not(isBridge().or(isHashCode()).or(isEquals()).or(isDefaultFinalizer())))
89                     .and(not(isDeclaredBy(nameStartsWith("java.")).<MethodDescription>and(isPackagePrivate()))),
90                 Advice.withCustomMapping()
91                     .bind(MockMethodAdvice.Identifier.class, identifier)
92                     .to(MockMethodAdvice.class))
93             .method(isHashCode(),
94                 Advice.withCustomMapping()
95                     .bind(MockMethodAdvice.Identifier.class, identifier)
96                     .to(MockMethodAdvice.ForHashCode.class))
97             .method(isEquals(),
98                 Advice.withCustomMapping()
99                     .bind(MockMethodAdvice.Identifier.class, identifier)
100                     .to(MockMethodAdvice.ForEquals.class));
101         MockMethodDispatcher.set(identifier, new MockMethodAdvice(mocks, identifier));
102         instrumentation.addTransformer(this, true);
103     }
104 
105     /**
106      * Mockito allows to mock about any type, including such types that we are relying on ourselves. This can cause a circularity:
107      * 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
108      * up this instance, we need to create key instances that rely on weak reference properties. Loading the later classes will happen before
109      * the key instances are completed what will cause Mockito to check if those key instances are themselves mocks what causes a loop which
110      * results in a circularity error. This is not normally a problem as we explicitly check if the instance that we investigate is one of
111      * our instance of which we hold a reference by reference equality what does not cause any code execution. But it seems like the load
112      * 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
113      * making sure that crucuial JVM types are loaded before we create the first inline mock. Unfortunately, these types dependant on a JVM's
114      * implementation and we can only maintain types that we know of from well-known JVM implementations such as HotSpot and extend this list
115      * once we learn of further problematic types for future Java versions. To allow users to whitelist their own types, we do not also offer
116      * a property that allows running problematic tests before a new Mockito version can be released and that allows us to ask users to
117      * easily validate that whitelisting actually solves a problem as circularities could also be caused by other problems.
118      */
preload()119     private static void preload() {
120         String preloads = System.getProperty(PRELOAD);
121         if (preloads == null) {
122             preloads = "java.lang.WeakPairMap,java.lang.WeakPairMap$Pair,java.lang.WeakPairMap$Pair$Weak";
123         }
124         for (String preload : preloads.split(",")) {
125             try {
126                 Class.forName(preload, false, null);
127             } catch (ClassNotFoundException ignored) {
128             }
129         }
130     }
131 
132     @Override
mockClass(MockFeatures<T> features)133     public <T> Class<? extends T> mockClass(MockFeatures<T> features) {
134         boolean subclassingRequired = !features.interfaces.isEmpty()
135             || features.serializableMode != SerializableMode.NONE
136             || Modifier.isAbstract(features.mockedType.getModifiers());
137 
138         checkSupportedCombination(subclassingRequired, features);
139 
140         synchronized (this) {
141             triggerRetransformation(features);
142         }
143 
144         return subclassingRequired ?
145                 subclassEngine.mockClass(features) :
146                 features.mockedType;
147     }
148 
triggerRetransformation(MockFeatures<T> features)149     private <T> void triggerRetransformation(MockFeatures<T> features) {
150         Set<Class<?>> types = new HashSet<Class<?>>();
151         Class<?> type = features.mockedType;
152         do {
153             if (mocked.add(type)) {
154                 types.add(type);
155                 addInterfaces(types, type.getInterfaces());
156             }
157             type = type.getSuperclass();
158         } while (type != null);
159         if (!types.isEmpty()) {
160             try {
161                 instrumentation.retransformClasses(types.toArray(new Class<?>[types.size()]));
162                 Throwable throwable = lastException;
163                 if (throwable != null) {
164                     throw new IllegalStateException(join("Byte Buddy could not instrument all classes within the mock's type hierarchy",
165                         "",
166                         "This problem should never occur for javac-compiled classes. This problem has been observed for classes that are:",
167                         " - Compiled by older versions of scalac",
168                         " - Classes that are part of the Android distribution"), throwable);
169                 }
170             } catch (Exception exception) {
171                 for (Class<?> failed : types) {
172                     mocked.remove(failed);
173                 }
174                 throw new MockitoException("Could not modify all classes " + types, exception);
175             } finally {
176                 lastException = null;
177             }
178         }
179     }
180 
checkSupportedCombination(boolean subclassingRequired, MockFeatures<T> features)181     private <T> void checkSupportedCombination(boolean subclassingRequired, MockFeatures<T> features) {
182         if (subclassingRequired
183                 && !features.mockedType.isArray()
184                 && !features.mockedType.isPrimitive()
185                 && Modifier.isFinal(features.mockedType.getModifiers())) {
186             throw new MockitoException("Unsupported settings with this type '" + features.mockedType.getName() + "'");
187         }
188     }
189 
addInterfaces(Set<Class<?>> types, Class<?>[] interfaces)190     private void addInterfaces(Set<Class<?>> types, Class<?>[] interfaces) {
191         for (Class<?> type : interfaces) {
192             if (mocked.add(type)) {
193                 types.add(type);
194                 addInterfaces(types, type.getInterfaces());
195             }
196         }
197     }
198 
199     @Override
transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer)200     public byte[] transform(ClassLoader loader,
201                             String className,
202                             Class<?> classBeingRedefined,
203                             ProtectionDomain protectionDomain,
204                             byte[] classfileBuffer) {
205         if (classBeingRedefined == null
206             || !mocked.contains(classBeingRedefined)
207             || EXCLUDES.contains(classBeingRedefined)) {
208             return null;
209         } else {
210             try {
211                 return byteBuddy.redefine(classBeingRedefined, ClassFileLocator.Simple.of(classBeingRedefined.getName(), classfileBuffer))
212                     // Note: The VM erases parameter meta data from the provided class file (bug). We just add this information manually.
213                     .visit(new ParameterWritingVisitorWrapper(classBeingRedefined))
214                     .visit(mockTransformer)
215                     .make()
216                     .getBytes();
217             } catch (Throwable throwable) {
218                 lastException = throwable;
219                 return null;
220             }
221         }
222     }
223 
224     private static class ParameterWritingVisitorWrapper extends AsmVisitorWrapper.AbstractBase {
225 
226         private final Class<?> type;
227 
ParameterWritingVisitorWrapper(Class<?> type)228         private ParameterWritingVisitorWrapper(Class<?> type) {
229             this.type = type;
230         }
231 
232         @Override
wrap(TypeDescription instrumentedType, ClassVisitor classVisitor, Implementation.Context implementationContext, TypePool typePool, FieldList<FieldDescription.InDefinedShape> fields, MethodList<?> methods, int writerFlags, int readerFlags)233         public ClassVisitor wrap(TypeDescription instrumentedType,
234                                  ClassVisitor classVisitor,
235                                  Implementation.Context implementationContext,
236                                  TypePool typePool,
237                                  FieldList<FieldDescription.InDefinedShape> fields,
238                                  MethodList<?> methods,
239                                  int writerFlags,
240                                  int readerFlags) {
241             return implementationContext.getClassFileVersion().isAtLeast(ClassFileVersion.JAVA_V8)
242                     ? new ParameterAddingClassVisitor(classVisitor, new TypeDescription.ForLoadedType(type))
243                     : classVisitor;
244         }
245 
246         private static class ParameterAddingClassVisitor extends ClassVisitor {
247 
248             private final TypeDescription typeDescription;
249 
ParameterAddingClassVisitor(ClassVisitor cv, TypeDescription typeDescription)250             private ParameterAddingClassVisitor(ClassVisitor cv, TypeDescription typeDescription) {
251                 super(OpenedClassReader.ASM_API, cv);
252                 this.typeDescription = typeDescription;
253             }
254 
255             @Override
visitMethod(int access, String name, String desc, String signature, String[] exceptions)256             public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
257                 MethodVisitor methodVisitor = super.visitMethod(access, name, desc, signature, exceptions);
258                 MethodList<?> methodList = typeDescription.getDeclaredMethods().filter((name.equals(MethodDescription.CONSTRUCTOR_INTERNAL_NAME)
259                         ? isConstructor()
260                         : ElementMatchers.<MethodDescription>named(name)).and(hasDescriptor(desc)));
261                 if (methodList.size() == 1 && methodList.getOnly().getParameters().hasExplicitMetaData()) {
262                     for (ParameterDescription parameterDescription : methodList.getOnly().getParameters()) {
263                         methodVisitor.visitParameter(parameterDescription.getName(), parameterDescription.getModifiers());
264                     }
265                     return new MethodParameterStrippingMethodVisitor(methodVisitor);
266                 } else {
267                     return methodVisitor;
268                 }
269             }
270         }
271 
272         private static class MethodParameterStrippingMethodVisitor extends MethodVisitor {
273 
MethodParameterStrippingMethodVisitor(MethodVisitor mv)274             public MethodParameterStrippingMethodVisitor(MethodVisitor mv) {
275                 super(Opcodes.ASM5, mv);
276             }
277 
278             @Override
visitParameter(String name, int access)279             public void visitParameter(String name, int access) {
280                 // suppress to avoid additional writing of the parameter if retained.
281             }
282         }
283     }
284 }
285