1
2
3
4
5
6 package mockit.coverage.modification;
7
8 import edu.umd.cs.findbugs.annotations.NonNull;
9 import edu.umd.cs.findbugs.annotations.Nullable;
10
11 import java.security.ProtectionDomain;
12 import java.util.ArrayList;
13 import java.util.HashSet;
14 import java.util.List;
15 import java.util.Set;
16
17 import mockit.asm.classes.ClassReader;
18
19 public final class ClassModification {
20 @NonNull
21 private final Set<String> modifiedClasses;
22 @NonNull
23 final List<ProtectionDomain> protectionDomainsWithUniqueLocations;
24 @NonNull
25 private final ClassSelection classSelection;
26
27 public ClassModification() {
28 modifiedClasses = new HashSet<>();
29 protectionDomainsWithUniqueLocations = new ArrayList<>();
30 classSelection = new ClassSelection();
31 }
32
33 public boolean shouldConsiderClassesNotLoaded() {
34 return !classSelection.loadedOnly;
35 }
36
37 boolean isToBeConsideredForCoverage(@NonNull String className, @NonNull ProtectionDomain protectionDomain) {
38 return !modifiedClasses.contains(className) && classSelection.isSelected(className, protectionDomain);
39 }
40
41 @Nullable
42 public byte[] modifyClass(@NonNull String className, @NonNull ProtectionDomain protectionDomain,
43 @NonNull byte[] originalClassfile) {
44 if (isToBeConsideredForCoverage(className, protectionDomain)) {
45 try {
46 byte[] modifiedClassfile = modifyClassForCoverage(className, originalClassfile);
47 registerModifiedClass(className, protectionDomain);
48 return modifiedClassfile;
49 } catch (VisitInterruptedException ignore) {
50
51 } catch (RuntimeException | AssertionError | ClassCircularityError e) {
52 e.printStackTrace();
53 }
54 }
55
56 return null;
57 }
58
59 @NonNull
60 private static byte[] modifyClassForCoverage(@NonNull String className, @NonNull byte[] classBytecode) {
61 byte[] modifiedBytecode = CoverageModifier.recoverModifiedByteCodeIfAvailable(className);
62
63 if (modifiedBytecode != null) {
64 return modifiedBytecode;
65 }
66
67 ClassReader cr = new ClassReader(classBytecode);
68 CoverageModifier modifier = new CoverageModifier(cr);
69 cr.accept(modifier);
70 return modifier.toByteArray();
71 }
72
73 private void registerModifiedClass(@NonNull String className, @NonNull ProtectionDomain pd) {
74 modifiedClasses.add(className);
75
76 if (pd.getClassLoader() != null && pd.getCodeSource() != null && pd.getCodeSource().getLocation() != null) {
77 addProtectionDomainIfHasUniqueNewPath(pd);
78 }
79 }
80
81 private void addProtectionDomainIfHasUniqueNewPath(@NonNull ProtectionDomain newPD) {
82 String newPath = newPD.getCodeSource().getLocation().getPath();
83
84 for (int i = protectionDomainsWithUniqueLocations.size() - 1; i >= 0; i--) {
85 ProtectionDomain previousPD = protectionDomainsWithUniqueLocations.get(i);
86 String previousPath = previousPD.getCodeSource().getLocation().getPath();
87
88 if (previousPath.startsWith(newPath)) {
89 return;
90 }
91 if (newPath.startsWith(previousPath)) {
92 protectionDomainsWithUniqueLocations.set(i, newPD);
93 return;
94 }
95 }
96
97 protectionDomainsWithUniqueLocations.add(newPD);
98 }
99 }