1
2
3
4
5
6 package mockit.internal.capturing;
7
8 import edu.umd.cs.findbugs.annotations.NonNull;
9 import edu.umd.cs.findbugs.annotations.Nullable;
10
11 import mockit.asm.classes.ClassReader;
12 import mockit.internal.BaseClassModifier;
13 import mockit.internal.ClassFile;
14 import mockit.internal.startup.Startup;
15 import mockit.internal.state.TestRun;
16
17 public abstract class CaptureOfImplementations<M> {
18 protected CaptureOfImplementations() {
19 }
20
21 protected final void makeSureAllSubtypesAreModified(@NonNull Class<?> baseType, boolean registerCapturedClasses,
22 @Nullable M typeMetadata) {
23 CapturedType captureMetadata = new CapturedType(baseType);
24 redefineClassesAlreadyLoaded(captureMetadata, baseType, typeMetadata);
25 createCaptureTransformer(captureMetadata, registerCapturedClasses, typeMetadata);
26 }
27
28 private void redefineClassesAlreadyLoaded(@NonNull CapturedType captureMetadata, @NonNull Class<?> baseType,
29 @Nullable M typeMetadata) {
30 Class<?>[] classesLoaded = Startup.instrumentation().getAllLoadedClasses();
31
32 for (Class<?> aClass : classesLoaded) {
33 if (captureMetadata.isToBeCaptured(aClass)) {
34 redefineClass(aClass, baseType, typeMetadata);
35 }
36 }
37 }
38
39 protected final void redefineClass(@NonNull Class<?> realClass, @NonNull Class<?> baseType,
40 @Nullable M typeMetadata) {
41 if (!TestRun.mockFixture().containsRedefinedClass(realClass)) {
42 ClassReader classReader;
43
44 try {
45 classReader = ClassFile.createReaderOrGetFromCache(realClass);
46 } catch (ClassFile.NotFoundException ignore) {
47 return;
48 }
49
50 TestRun.ensureThatClassIsInitialized(realClass);
51
52 BaseClassModifier modifier = createModifier(realClass.getClassLoader(), classReader, baseType,
53 typeMetadata);
54 classReader.accept(modifier);
55
56 if (modifier.wasModified()) {
57 byte[] modifiedClass = modifier.toByteArray();
58 redefineClass(realClass, modifiedClass);
59 }
60 }
61 }
62
63 @NonNull
64 protected abstract BaseClassModifier createModifier(@Nullable ClassLoader cl, @NonNull ClassReader cr,
65 @NonNull Class<?> baseType, @Nullable M typeMetadata);
66
67 protected abstract void redefineClass(@NonNull Class<?> realClass, @NonNull byte[] modifiedClass);
68
69 private void createCaptureTransformer(@NonNull CapturedType captureMetadata, boolean registerCapturedClasses,
70 @Nullable M typeMetadata) {
71 CaptureTransformer<M> transformer = new CaptureTransformer<>(captureMetadata, this, registerCapturedClasses,
72 typeMetadata);
73 Startup.instrumentation().addTransformer(transformer, true);
74 TestRun.mockFixture().addCaptureTransformer(transformer);
75 }
76 }