View Javadoc
1   /*
2    * MIT License
3    * Copyright (c) 2006-2025 JMockit developers
4    * See LICENSE file for full license text.
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             // noinspection unchecked
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 }