1/*
2 * Copyright (C) 2017 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
17package com.android.dx.mockito.inline;
18
19import org.mockito.exceptions.base.MockitoException;
20
21import java.lang.reflect.InvocationTargetException;
22import java.lang.reflect.Method;
23import java.lang.reflect.Modifier;
24import java.security.ProtectionDomain;
25import java.util.Arrays;
26import java.util.Collections;
27import java.util.HashSet;
28import java.util.Map;
29import java.util.Random;
30import java.util.Set;
31
32/**
33 * Adds entry hooks (that eventually call into
34 * {@link MockMethodAdvice#handle(Object, Method, Object[])} to all non-abstract methods of the
35 * supplied classes.
36 *
37 * <p></p>Transforming a class to adding entry hooks follow the following simple steps:
38 * <ol>
39 * <li>{@link #mockClass(MockFeatures)}</li>
40 * <li>{@link JvmtiAgent#requestTransformClasses(Class[])}</li>
41 * <li>{@link JvmtiAgent#nativeRetransformClasses(Class[])}</li>
42 * <li>agent.cc::Transform</li>
43 * <li>{@link JvmtiAgent#runTransformers(ClassLoader, String, Class, ProtectionDomain, byte[])}</li>
44 * <li>{@link #transform(Class, byte[])}</li>
45 * <li>{@link #nativeRedefine(String, byte[])}</li>
46 * </ol>
47 *
48 */
49class ClassTransformer {
50    // Some classes are so deeply optimized inside the runtime that they cannot be transformed
51    private static final Set<Class<? extends java.io.Serializable>> EXCLUDES = new HashSet<>(
52            Arrays.asList(Class.class,
53                    Boolean.class,
54                    Byte.class,
55                    Short.class,
56                    Character.class,
57                    Integer.class,
58                    Long.class,
59                    Float.class,
60                    Double.class,
61                    String.class));
62    private final static Random random = new Random();
63
64    /** Jvmti agent responsible for triggering transformation s*/
65    private final JvmtiAgent agent;
66
67    /** Types that have already be transformed */
68    private final Set<Class<?>> mockedTypes;
69
70    /**
71     * A unique identifier that is baked into the transformed classes. The entry hooks will then
72     * pass this identifier to
73     * {@code com.android.dx.mockito.inline.MockMethodDispatcher#get(String, Object)} to
74     * find the advice responsible for handling the method call interceptions.
75     */
76    private final String identifier;
77
78    /**
79     * We can only have a single transformation going on at a time, hence synchronize the
80     * transformation process via this lock.
81     *
82     * @see #mockClass(MockFeatures)
83     */
84    private final static Object lock = new Object();
85
86    /**
87     * Create a new generator.
88     *
89     * Creating more than one generator might cause transformations to overwrite each other.
90     *
91     * @param agent agent used to trigger transformations
92     * @param dispatcherClass {@code com.android.dx.mockito.inline.MockMethodDispatcher}
93     *                        that will dispatch method calls that might need to get intercepted.
94     * @param mocks list of mocked objects. As all objects of a class use the same transformed
95     *              bytecode the {@link MockMethodAdvice} needs to check this list if a object is
96     *              mocked or not.
97     */
98    ClassTransformer(JvmtiAgent agent, Class dispatcherClass,
99                     Map<Object, InvocationHandlerAdapter> mocks) {
100        this.agent = agent;
101        mockedTypes = Collections.synchronizedSet(new HashSet<Class<?>>());
102        identifier = Long.toString(random.nextLong());
103        MockMethodAdvice advice = new MockMethodAdvice(mocks);
104
105        try {
106            dispatcherClass.getMethod("set", String.class, Object.class).invoke(null, identifier,
107                    advice);
108        } catch (InvocationTargetException | IllegalAccessException | NoSuchMethodException e) {
109            throw new IllegalStateException(e);
110        }
111
112        agent.addTransformer(this);
113    }
114
115    /**
116     * Trigger the process to add entry hooks to a class (and all its parents).
117     *
118     * @param features specification what to mock
119     */
120    <T> void mockClass(MockFeatures<T> features) {
121        boolean subclassingRequired = !features.interfaces.isEmpty()
122                || Modifier.isAbstract(features.mockedType.getModifiers());
123
124        if (subclassingRequired
125                && !features.mockedType.isArray()
126                && !features.mockedType.isPrimitive()
127                && Modifier.isFinal(features.mockedType.getModifiers())) {
128            throw new MockitoException("Unsupported settings with this type '"
129                    + features.mockedType.getName() + "'");
130        }
131
132        synchronized (lock) {
133            Set<Class<?>> types = new HashSet<>();
134            Class<?> type = features.mockedType;
135
136            do {
137                boolean wasAdded = mockedTypes.add(type);
138
139                if (wasAdded) {
140                    if (!EXCLUDES.contains(type)) {
141                        types.add(type);
142                    }
143
144                    type = type.getSuperclass();
145                } else {
146                    break;
147                }
148            } while (type != null && !type.isInterface());
149
150            if (!types.isEmpty()) {
151                try {
152                    agent.requestTransformClasses(types.toArray(new Class<?>[types.size()]));
153                } catch (UnmodifiableClassException exception) {
154                    for (Class<?> failed : types) {
155                        mockedTypes.remove(failed);
156                    }
157
158                    throw new MockitoException("Could not modify all classes " + types, exception);
159                }
160            }
161        }
162    }
163
164    /**
165     * Add entry hooks to all methods of a class.
166     *
167     * <p>Called by the agent after triggering the transformation via
168     * {@link #mockClass(MockFeatures)}.
169     *
170     * @param classBeingRedefined class the hooks should be added to
171     * @param classfileBuffer original byte code of the class
172     *
173     * @return transformed class
174     */
175    byte[] transform(Class<?> classBeingRedefined, byte[] classfileBuffer) throws
176            IllegalClassFormatException {
177        if (classBeingRedefined == null
178                || !mockedTypes.contains(classBeingRedefined)) {
179            return null;
180        } else {
181            try {
182                return nativeRedefine(identifier, classfileBuffer);
183            } catch (Throwable throwable) {
184                throw new IllegalClassFormatException();
185            }
186        }
187    }
188
189    /**
190     * Check if the class should be transformed.
191     *
192     * @param classBeingRedefined The class that might need to transformed
193     *
194     * @return {@code true} iff the class needs to be transformed
195     */
196    boolean shouldTransform(Class<?> classBeingRedefined) {
197        return classBeingRedefined != null && mockedTypes.contains(classBeingRedefined);
198    }
199
200    private native byte[] nativeRedefine(String identifier, byte[] original);
201}
202