1 /*
2  * Copyright (C) 2018 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.dx.mockito.inline;
18 
19 import android.os.Build;
20 
21 import org.mockito.Mockito;
22 import org.mockito.creation.instance.Instantiator;
23 import org.mockito.exceptions.base.MockitoException;
24 import org.mockito.invocation.MockHandler;
25 import org.mockito.mock.MockCreationSettings;
26 import org.mockito.plugins.InstantiatorProvider2;
27 import org.mockito.plugins.MockMaker;
28 
29 import java.io.IOException;
30 import java.lang.reflect.InvocationTargetException;
31 import java.lang.reflect.Method;
32 import java.lang.reflect.Modifier;
33 import java.util.HashMap;
34 import java.util.Map;
35 import java.util.Set;
36 import java.util.function.BiConsumer;
37 
38 /**
39  * Creates mock markers and adds stubbing hooks to static method
40  *
41  * <p>This is done by transforming the byte code of the classes to add method entry hooks.
42  */
43 public final class InlineStaticMockMaker implements MockMaker {
44     /**
45      * {@link StaticJvmtiAgent} set up during one time init
46      */
47     private static final StaticJvmtiAgent AGENT;
48 
49     /**
50      * Error  during one time init or {@code null} if init was successful
51      */
52     private static final Throwable INITIALIZATION_ERROR;
53     public static ThreadLocal<Class> mockingInProgressClass = new ThreadLocal<>();
54     public static ThreadLocal<BiConsumer<Class<?>, Method>> onMethodCallDuringStubbing
55             = new ThreadLocal<>();
56     public static ThreadLocal<BiConsumer<Class<?>, Method>> onMethodCallDuringVerification
57             = new ThreadLocal<>();
58 
59     /*
60      * One time setup to allow the system to mocking via this mock maker.
61      */
62     static {
63         StaticJvmtiAgent agent;
64         Throwable initializationError = null;
65 
66         try {
67             try {
68                 agent = new StaticJvmtiAgent();
69             } catch (IOException ioe) {
70                 throw new IllegalStateException("Mockito could not self-attach a jvmti agent to " +
71                         "the current VM. This feature is required for inline mocking.\nThis error" +
72                         " occured due to an I/O error during the creation of this agent: " + ioe
73                         + "\n\nPotentially, the current VM does not support the jvmti API " +
74                         "correctly", ioe);
75             }
76 
77             // Blacklisted APIs were introduced in Android P:
78             //
79             // https://android-developers.googleblog.com/2018/02/
80             // improving-stability-by-reducing-usage.html
81             //
82             // This feature prevents access to blacklisted fields and calling of blacklisted APIs
83             // if the calling class is not trusted.
84             Method allowHiddenApiReflectionFrom;
85             try {
86                 Class vmDebug = Class.forName("dalvik.system.VMDebug");
87                 allowHiddenApiReflectionFrom = vmDebug.getDeclaredMethod(
88                         "allowHiddenApiReflectionFrom", Class.class);
89             } catch (ClassNotFoundException | NoSuchMethodException e) {
90                 throw new IllegalStateException("Cannot find "
91                         + "VMDebug#allowHiddenApiReflectionFrom.");
92             }
93 
94             // The StaticMockMethodAdvice is used by methods of spies to call the real methods. As
95             // the real methods might be blacklisted, this class needs to be marked as trusted.
96             try {
allowHiddenApiReflectionFrom.invoke(null, StaticMockMethodAdvice.class)97                 allowHiddenApiReflectionFrom.invoke(null, StaticMockMethodAdvice.class);
98             } catch (InvocationTargetException e) {
99                 throw e.getCause();
100             }
101         } catch (Throwable throwable) {
102             agent = null;
103             initializationError = throwable;
104         }
105 
106         AGENT = agent;
107         INITIALIZATION_ERROR = initializationError;
108     }
109 
110     /**
111      * All currently active mock markers. We modify the class's byte code. Some objects of the class
112      * are modified, some are not. This list helps the {@link MockMethodAdvice} help figure out if a
113      * object's method calls should be intercepted.
114      */
115     private final Map<Object, InvocationHandlerAdapter> markerToHandler = new MarkerToHandlerMap();
116     private final Map<Class, Object> classToMarker = new HashMap<>();
117 
118     /**
119      * Class doing the actual byte code transformation.
120      */
121     private final StaticClassTransformer classTransformer;
122 
123     /**
124      * Create a new mock maker.
125      */
InlineStaticMockMaker()126     public InlineStaticMockMaker() {
127         if (INITIALIZATION_ERROR != null) {
128             throw new RuntimeException("Could not initialize inline mock maker.\n" + "\n" +
129                     "Release: Android " + Build.VERSION.RELEASE + " " + Build.VERSION.INCREMENTAL
130                     + "Device: " + Build.BRAND + " " + Build.MODEL, INITIALIZATION_ERROR);
131         }
132 
133         classTransformer = new StaticClassTransformer(AGENT, InlineDexmakerMockMaker
134                 .DISPATCHER_CLASS, markerToHandler, classToMarker);
135     }
136 
137     @Override
createMock(MockCreationSettings<T> settings, MockHandler handler)138     public <T> T createMock(MockCreationSettings<T> settings, MockHandler handler) {
139         Class<T> typeToMock = settings.getTypeToMock();
140         if (!typeToMock.equals(mockingInProgressClass.get()) || Modifier.isAbstract(typeToMock
141                 .getModifiers())) {
142             return null;
143         }
144 
145         Set<Class<?>> interfacesSet = settings.getExtraInterfaces();
146         InvocationHandlerAdapter handlerAdapter = new InvocationHandlerAdapter(handler);
147 
148         classTransformer.mockClass(MockFeatures.withMockFeatures(typeToMock, interfacesSet));
149 
150         Instantiator instantiator = Mockito.framework().getPlugins().getDefaultPlugin
151                 (InstantiatorProvider2.class).getInstantiator(settings);
152 
153         T mock;
154         try {
155             mock = instantiator.newInstance(typeToMock);
156         } catch (org.mockito.creation.instance.InstantiationException e) {
157             throw new MockitoException("Unable to create mock instance of type '" + typeToMock
158                     .getSimpleName() + "'", e);
159         }
160 
161         if (classToMarker.containsKey(typeToMock)) {
162             throw new MockitoException(typeToMock + " is already mocked");
163         }
164         classToMarker.put(typeToMock, mock);
165 
166         markerToHandler.put(mock, handlerAdapter);
167         return mock;
168     }
169 
170     @Override
resetMock(Object mock, MockHandler newHandler, MockCreationSettings settings)171     public void resetMock(Object mock, MockHandler newHandler, MockCreationSettings settings) {
172         InvocationHandlerAdapter adapter = getInvocationHandlerAdapter(mock);
173         if (adapter != null) {
174             if (mockingInProgressClass.get() == mock.getClass()) {
175                 markerToHandler.remove(mock);
176                 classToMarker.remove(mock.getClass());
177             } else {
178                 adapter.setHandler(newHandler);
179             }
180         }
181     }
182 
183     @Override
isTypeMockable(final Class<?> type)184     public TypeMockability isTypeMockable(final Class<?> type) {
185         if (mockingInProgressClass.get() == type) {
186             return new TypeMockability() {
187                 @Override
188                 public boolean mockable() {
189                     return !Modifier.isAbstract(type.getModifiers()) && !type.isPrimitive() && type
190                             != String.class;
191                 }
192 
193                 @Override
194                 public String nonMockableReason() {
195                     if (Modifier.isAbstract(type.getModifiers())) {
196                         return "abstract type";
197                     }
198 
199                     if (type.isPrimitive()) {
200                         return "primitive type";
201                     }
202 
203                     if (type == String.class) {
204                         return "string";
205                     }
206 
207                     return "not handled type";
208                 }
209             };
210         } else {
211             return null;
212         }
213     }
214 
215     @Override
216     public MockHandler getHandler(Object mock) {
217         InvocationHandlerAdapter adapter = getInvocationHandlerAdapter(mock);
218         return adapter != null ? adapter.getHandler() : null;
219     }
220 
221     /**
222      * Get the {@link InvocationHandlerAdapter} registered for a marker.
223      *
224      * @param marker marker of the class that might have mocking set up
225      * @return adapter for this class, or {@code null} if not mocked
226      */
227     private InvocationHandlerAdapter getInvocationHandlerAdapter(Object marker) {
228         if (marker == null) {
229             return null;
230         }
231 
232         return markerToHandler.get(marker);
233     }
234 }
235