1
2
3
4
5
6 package mockit.internal.capturing;
7
8 import static mockit.internal.capturing.CapturedType.isNotToBeCaptured;
9
10 import edu.umd.cs.findbugs.annotations.NonNull;
11 import edu.umd.cs.findbugs.annotations.Nullable;
12
13 import java.lang.instrument.ClassFileTransformer;
14 import java.security.ProtectionDomain;
15 import java.util.Collections;
16 import java.util.HashMap;
17 import java.util.Map;
18
19 import mockit.asm.classes.ClassReader;
20 import mockit.asm.classes.ClassVisitor;
21 import mockit.asm.metadata.ClassMetadataReader;
22 import mockit.asm.types.JavaType;
23 import mockit.internal.ClassFile;
24 import mockit.internal.ClassIdentification;
25 import mockit.internal.startup.Startup;
26 import mockit.internal.state.TestRun;
27
28 public final class CaptureTransformer<M> implements ClassFileTransformer {
29 @NonNull
30 private final CapturedType capturedType;
31 @NonNull
32 private final String capturedTypeDesc;
33 @NonNull
34 private final CaptureOfImplementations<M> captureOfImplementations;
35 @NonNull
36 private final Map<ClassIdentification, byte[]> transformedClasses;
37 @NonNull
38 private final Map<String, Boolean> superTypesSearched;
39 @Nullable
40 private final M typeMetadata;
41 private boolean inactive;
42
43 CaptureTransformer(@NonNull CapturedType capturedType,
44 @NonNull CaptureOfImplementations<M> captureOfImplementations, boolean registerTransformedClasses,
45 @Nullable M typeMetadata) {
46 this.capturedType = capturedType;
47 capturedTypeDesc = JavaType.getInternalName(capturedType.baseType);
48 this.captureOfImplementations = captureOfImplementations;
49 transformedClasses = registerTransformedClasses ? new HashMap<>(2)
50 : Collections.<ClassIdentification, byte[]> emptyMap();
51 superTypesSearched = new HashMap<>();
52 this.typeMetadata = typeMetadata;
53 }
54
55 public void deactivate() {
56 inactive = true;
57
58 if (!transformedClasses.isEmpty()) {
59 for (Map.Entry<ClassIdentification, byte[]> classNameAndOriginalBytecode : transformedClasses.entrySet()) {
60 ClassIdentification classId = classNameAndOriginalBytecode.getKey();
61 byte[] originalBytecode = classNameAndOriginalBytecode.getValue();
62
63 Startup.redefineMethods(classId, originalBytecode);
64 }
65
66 transformedClasses.clear();
67 }
68 }
69
70 @Nullable
71 @Override
72 public byte[] transform(@Nullable ClassLoader loader, @NonNull String classDesc,
73 @Nullable Class<?> classBeingRedefined, @Nullable ProtectionDomain protectionDomain,
74 @NonNull byte[] classfileBuffer) {
75 if (classBeingRedefined != null || inactive || isNotToBeCaptured(protectionDomain, classDesc)) {
76 return null;
77 }
78
79 if (isClassToBeCaptured(loader, classfileBuffer)) {
80 String className = classDesc.replace('/', '.');
81 ClassReader cr = new ClassReader(classfileBuffer);
82 return modifyAndRegisterClass(loader, className, cr);
83 }
84
85 return null;
86 }
87
88 private boolean isClassToBeCaptured(@Nullable ClassLoader loader, @NonNull byte[] classfileBuffer) {
89 ClassMetadataReader cmr = new ClassMetadataReader(classfileBuffer);
90 String superName = cmr.getSuperClass();
91
92 if (capturedTypeDesc.equals(superName)) {
93 return true;
94 }
95
96 String[] interfaces = cmr.getInterfaces();
97
98 if (interfaces != null && isClassWhichImplementsACapturingInterface(interfaces)) {
99 return true;
100 }
101
102 return superName != null && searchSuperTypes(loader, superName, interfaces);
103 }
104
105 private boolean isClassWhichImplementsACapturingInterface(@NonNull String[] interfaces) {
106 for (String implementedInterface : interfaces) {
107 if (capturedTypeDesc.equals(implementedInterface)) {
108 return true;
109 }
110 }
111
112 return false;
113 }
114
115 private boolean searchSuperTypes(@Nullable ClassLoader loader, @NonNull String superName,
116 @Nullable String[] interfaces) {
117 if (!"java/lang/Object".equals(superName) && !superName.startsWith("mockit/")
118 && searchSuperType(loader, superName)) {
119 return true;
120 }
121
122 if (interfaces != null && interfaces.length > 0) {
123 for (String itf : interfaces) {
124 if (!itf.startsWith("java/") && !itf.startsWith("javax/") && !itf.startsWith("jakarta/")
125 && searchSuperType(loader, itf)) {
126 return true;
127 }
128 }
129 }
130
131 return false;
132 }
133
134 private boolean searchSuperType(@Nullable ClassLoader loader, @NonNull String superName) {
135 return superTypesSearched.computeIfAbsent(superName, name -> {
136 byte[] classfileBytes = ClassFile.getClassFile(loader, name);
137 return isClassToBeCaptured(loader, classfileBytes);
138 });
139 }
140
141 @NonNull
142 private byte[] modifyAndRegisterClass(@Nullable ClassLoader loader, @NonNull String className,
143 @NonNull ClassReader cr) {
144 ClassVisitor modifier = captureOfImplementations.createModifier(loader, cr, capturedType.baseType,
145 typeMetadata);
146 cr.accept(modifier);
147
148 ClassIdentification classId = new ClassIdentification(loader, className);
149 byte[] originalBytecode = cr.getBytecode();
150
151 if (transformedClasses == Collections.<ClassIdentification, byte[]> emptyMap()) {
152 TestRun.mockFixture().addTransformedClass(classId, originalBytecode);
153 } else {
154 transformedClasses.put(classId, originalBytecode);
155 }
156
157 TestRun.mockFixture().registerMockedClass(capturedType.baseType);
158 return modifier.toByteArray();
159 }
160
161 @Nullable
162 public <C extends CaptureOfImplementations<?>> C getCaptureOfImplementationsIfApplicable(@NonNull Class<?> aType) {
163 if (typeMetadata != null && capturedType.baseType.isAssignableFrom(aType)) {
164
165 return (C) captureOfImplementations;
166 }
167
168 return null;
169 }
170
171 public boolean areCapturedClasses(@NonNull Class<?> mockedClass1, @NonNull Class<?> mockedClass2) {
172 Class<?> baseType = capturedType.baseType;
173 return baseType.isAssignableFrom(mockedClass1) && baseType.isAssignableFrom(mockedClass2);
174 }
175 }