diff --git a/agent/src/main/java/com/alibaba/testable/agent/constant/ConstPool.java b/agent/src/main/java/com/alibaba/testable/agent/constant/ConstPool.java index b1d2755..f8cd50b 100644 --- a/agent/src/main/java/com/alibaba/testable/agent/constant/ConstPool.java +++ b/agent/src/main/java/com/alibaba/testable/agent/constant/ConstPool.java @@ -1,8 +1,5 @@ package com.alibaba.testable.agent.constant; -import java.util.ArrayList; -import java.util.List; - /** * @author flin */ @@ -13,9 +10,6 @@ public class ConstPool { public static final String TEST_POSTFIX = "Test"; - public static final List SYS_CLASSES = new ArrayList(); - static { - SYS_CLASSES.add("java/lang/StringBuilder"); - } - + public static final String ENABLE_TESTABLE = "com.alibaba.testable.core.annotation.EnableTestable"; + public static final String TESTABLE_INJECT = "com.alibaba.testable.core.annotation.TestableInject"; } diff --git a/agent/src/main/java/com/alibaba/testable/agent/handler/SourceClassHandler.java b/agent/src/main/java/com/alibaba/testable/agent/handler/SourceClassHandler.java index 102f13e..e76e1b3 100644 --- a/agent/src/main/java/com/alibaba/testable/agent/handler/SourceClassHandler.java +++ b/agent/src/main/java/com/alibaba/testable/agent/handler/SourceClassHandler.java @@ -1,12 +1,15 @@ package com.alibaba.testable.agent.handler; import com.alibaba.testable.agent.constant.ConstPool; +import com.alibaba.testable.agent.model.MethodInfo; import com.alibaba.testable.agent.util.ClassUtil; +import com.alibaba.testable.agent.util.CollectionUtil; import com.alibaba.testable.agent.util.StringUtil; import org.objectweb.asm.Opcodes; import org.objectweb.asm.Type; import org.objectweb.asm.tree.*; +import java.util.ArrayList; import java.util.HashSet; import java.util.List; import java.util.Set; @@ -24,33 +27,42 @@ public class SourceClassHandler extends ClassHandler { private static final String METHOD_DESC_PREFIX = "(Ljava/lang/Object;Ljava/lang/String;"; private static final String OBJECT_DESC = "Ljava/lang/Object;"; private static final String METHOD_DESC_POSTFIX = ")Ljava/lang/Object;"; + private List injectMethods; + + public SourceClassHandler(List injectMethods) { + this.injectMethods = injectMethods; + } @Override protected void transform(ClassNode cn) { - Set methodNames = new HashSet(); + List methods = new ArrayList(); for (MethodNode m : cn.methods) { if (!CONSTRUCTOR.equals(m.name)) { - methodNames.add(m.name); + methods.add(new MethodInfo(m.name, m.desc)); } } + Set memberInjectMethods = CollectionUtil.getCrossSet(methods, injectMethods); + Set newOperatorInjectMethods = CollectionUtil.getMinusSet(injectMethods, memberInjectMethods); for (MethodNode m : cn.methods) { - transformMethod(cn, m, methodNames); + transformMethod(cn, m, memberInjectMethods, MethodInfo.descSet(newOperatorInjectMethods)); } } - private void transformMethod(ClassNode cn, MethodNode mn, Set methodNames) { + private void transformMethod(ClassNode cn, MethodNode mn, Set memberInjectMethods, + Set newOperatorInjectDesc) { AbstractInsnNode[] instructions = mn.instructions.toArray(); int i = 0; do { if (instructions[i].getOpcode() == Opcodes.INVOKESPECIAL) { MethodInsnNode node = (MethodInsnNode)instructions[i]; - if (cn.name.equals(node.owner) && methodNames.contains(node.name)) { + if (cn.name.equals(node.owner) && memberInjectMethods.contains(new MethodInfo(node.name, node.desc))) { int rangeStart = getMemberMethodStart(instructions, i); if (rangeStart >= 0) { instructions = replaceMemberCallOps(mn, instructions, rangeStart, i); i = rangeStart; } - } else if (CONSTRUCTOR.equals(node.name) && !ConstPool.SYS_CLASSES.contains(node.owner)) { + } else if (CONSTRUCTOR.equals(node.name) && + newOperatorInjectDesc.contains(getConstructorInjectDesc(node))) { int rangeStart = getConstructorStart(instructions, node.owner, i); if (rangeStart >= 0) { instructions = replaceNewOps(mn, instructions, rangeStart, i); @@ -62,6 +74,11 @@ public class SourceClassHandler extends ClassHandler { } while (i < instructions.length); } + private String getConstructorInjectDesc(MethodInsnNode constructorNode) { + return constructorNode.desc.substring(0, constructorNode.desc.length() - 1) + + ClassUtil.toByteCodeClassName(constructorNode.owner); + } + private int getConstructorStart(AbstractInsnNode[] instructions, String target, int rangeEnd) { for (int i = rangeEnd - 1; i >= 0; i--) { if (instructions[i].getOpcode() == Opcodes.NEW && ((TypeInsnNode)instructions[i]).desc.equals(target)) { diff --git a/agent/src/main/java/com/alibaba/testable/agent/model/MethodInfo.java b/agent/src/main/java/com/alibaba/testable/agent/model/MethodInfo.java index d2ee60b..78716c4 100644 --- a/agent/src/main/java/com/alibaba/testable/agent/model/MethodInfo.java +++ b/agent/src/main/java/com/alibaba/testable/agent/model/MethodInfo.java @@ -1,26 +1,28 @@ package com.alibaba.testable.agent.model; +import java.util.Collection; +import java.util.HashSet; +import java.util.Set; + /** * @author flin */ public class MethodInfo { - private int access; private String name; private String desc; - private String signature; - private String[] exceptions; - public MethodInfo(int access, String name, String desc, String signature, String[] exceptions) { - this.access = access; + public MethodInfo(String name, String desc) { this.name = name; this.desc = desc; - this.signature = signature; - this.exceptions = exceptions; } - public int getAccess() { - return access; + public static Set descSet(Collection methodInfos) { + Set set = new HashSet(); + for (MethodInfo m : methodInfos) { + set.add(m.desc); + } + return set; } public String getName() { @@ -31,11 +33,21 @@ public class MethodInfo { return desc; } - public String getSignature() { - return signature; + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + MethodInfo that = (MethodInfo)o; + return name.equals(that.name) && desc.equals(that.desc); } - public String[] getExceptions() { - return exceptions; + @Override + public int hashCode() { + return 31 * name.hashCode() + desc.hashCode(); } + } diff --git a/agent/src/main/java/com/alibaba/testable/agent/transformer/TestableClassTransformer.java b/agent/src/main/java/com/alibaba/testable/agent/transformer/TestableClassTransformer.java index 131ae05..98e7ccc 100644 --- a/agent/src/main/java/com/alibaba/testable/agent/transformer/TestableClassTransformer.java +++ b/agent/src/main/java/com/alibaba/testable/agent/transformer/TestableClassTransformer.java @@ -3,6 +3,7 @@ package com.alibaba.testable.agent.transformer; import com.alibaba.testable.agent.constant.ConstPool; import com.alibaba.testable.agent.handler.SourceClassHandler; import com.alibaba.testable.agent.handler.TestClassHandler; +import com.alibaba.testable.agent.model.MethodInfo; import com.alibaba.testable.agent.util.ClassUtil; import java.io.IOException; @@ -18,9 +19,6 @@ import java.util.Set; */ public class TestableClassTransformer implements ClassFileTransformer { - private static final String ENABLE_TESTABLE = "com.alibaba.testable.core.annotation.EnableTestable"; - private static final String ENABLE_TESTABLE_INJECT = "com.alibaba.testable.core.annotation.EnableTestableInject"; - private static final Set loadedClassNames = new HashSet(); public byte[] transform(ClassLoader loader, String className, Class classBeingRedefined, @@ -33,10 +31,11 @@ public class TestableClassTransformer implements ClassFileTransformer { List annotations = ClassUtil.getAnnotations(className); List testAnnotations = ClassUtil.getAnnotations(className + ConstPool.TEST_POSTFIX); try { - if (annotations.contains(ENABLE_TESTABLE_INJECT) || testAnnotations.contains(ENABLE_TESTABLE)) { + if (testAnnotations.contains(ConstPool.ENABLE_TESTABLE)) { loadedClassNames.add(className); - return new SourceClassHandler().getBytes(className); - } else if (annotations.contains(ENABLE_TESTABLE)) { + List injectMethods = ClassUtil.getTestableInjectMethods(className + ConstPool.TEST_POSTFIX); + return new SourceClassHandler(injectMethods).getBytes(className); + } else if (annotations.contains(ConstPool.ENABLE_TESTABLE)) { loadedClassNames.add(className); return new TestClassHandler().getBytes(className); } diff --git a/agent/src/main/java/com/alibaba/testable/agent/util/ClassUtil.java b/agent/src/main/java/com/alibaba/testable/agent/util/ClassUtil.java index 386fdfb..54d1df0 100644 --- a/agent/src/main/java/com/alibaba/testable/agent/util/ClassUtil.java +++ b/agent/src/main/java/com/alibaba/testable/agent/util/ClassUtil.java @@ -1,9 +1,11 @@ package com.alibaba.testable.agent.util; import com.alibaba.testable.agent.constant.ConstPool; +import com.alibaba.testable.agent.model.MethodInfo; import org.objectweb.asm.ClassReader; import org.objectweb.asm.tree.AnnotationNode; import org.objectweb.asm.tree.ClassNode; +import org.objectweb.asm.tree.MethodNode; import java.io.IOException; import java.util.ArrayList; @@ -37,8 +39,7 @@ public class ClassUtil { ClassNode cn = new ClassNode(); new ClassReader(className).accept(cn, 0); for (AnnotationNode an : cn.visibleAnnotations) { - String annotationName = an.desc.replace(ConstPool.SLASH, ConstPool.DOT).substring(1, an.desc.length() - 1); - annotations.add(annotationName); + annotations.add(toDotSeparateFullClassName(an.desc)); } return annotations; } catch (Exception e) { @@ -46,6 +47,32 @@ public class ClassUtil { } } + public static List getTestableInjectMethods(String className) { + try { + List methodInfos = new ArrayList(); + ClassNode cn = new ClassNode(); + new ClassReader(className).accept(cn, 0); + for (MethodNode mn : cn.methods) { + checkMethodAnnotation(methodInfos, mn); + } + return methodInfos; + } catch (Exception e) { + return new ArrayList(); + } + } + + private static void checkMethodAnnotation(List methodInfos, MethodNode mn) { + if (mn.visibleAnnotations == null) { + return; + } + for (AnnotationNode an : mn.visibleAnnotations) { + if (toDotSeparateFullClassName(an.desc).equals(ConstPool.TESTABLE_INJECT)) { + methodInfos.add(new MethodInfo(mn.name, mn.desc)); + break; + } + } + } + public static List getParameterTypes(String desc) { List parameterTypes = new ArrayList(); boolean travelingClass = false; @@ -101,4 +128,9 @@ public class ClassUtil { public static String toByteCodeClassName(String className) { return TYPE_CLASS + className.replace(ConstPool.DOT, ConstPool.SLASH) + CLASS_END; } + + public static String toDotSeparateFullClassName(String className) { + return className.replace(ConstPool.SLASH, ConstPool.DOT).substring(1, className.length() - 1); + } + } diff --git a/agent/src/main/java/com/alibaba/testable/agent/util/CollectionUtil.java b/agent/src/main/java/com/alibaba/testable/agent/util/CollectionUtil.java index 2450776..7656e09 100644 --- a/agent/src/main/java/com/alibaba/testable/agent/util/CollectionUtil.java +++ b/agent/src/main/java/com/alibaba/testable/agent/util/CollectionUtil.java @@ -1,18 +1,20 @@ package com.alibaba.testable.agent.util; -import java.util.ArrayList; -import java.util.Collection; -import java.util.Collections; -import java.util.List; +import java.util.*; /** * @author flin */ public class CollectionUtil { - public static boolean containsAny(Collection hostContainer, Collection itemsToFind) { - for (Object o : hostContainer) { - for (Object i : itemsToFind) { + /** + * Check two collection has any equaled item + * @param collectionLeft the first collection + * @param collectionRight the second collection + */ + public static boolean containsAny(Collection collectionLeft, Collection collectionRight) { + for (Object o : collectionLeft) { + for (Object i : collectionRight) { if (o.equals(i)) { return true; } @@ -21,10 +23,44 @@ public class CollectionUtil { return false; } + /** + * Generate a list of item + * @param items elements to add + */ public static List listOf(T... items) { List list = new ArrayList(items.length); Collections.addAll(list, items); return list; } + /** + * Get cross set of two collections + * @param collectionLeft the first collection + * @param collectionRight the second collection + */ + public static Set getCrossSet(Collection collectionLeft, Collection collectionRight) { + Set crossSet = new HashSet(); + for (T i : collectionLeft) { + if (collectionRight.contains(i)) { + crossSet.add(i); + } + } + return crossSet; + } + + /** + * Get minus set of two collections + * @param collectionRaw original collection + * @param collectionMinus items to remove + */ + public static Set getMinusSet(Collection collectionRaw, Collection collectionMinus) { + Set crossSet = new HashSet(); + for (T i : collectionRaw) { + if (!collectionMinus.contains(i)) { + crossSet.add(i); + } + } + return crossSet; + } + } diff --git a/core/src/main/java/com/alibaba/testable/core/annotation/EnableTestableInject.java b/core/src/main/java/com/alibaba/testable/core/annotation/EnableTestableInject.java index ead675b..16b1b1d 100644 --- a/core/src/main/java/com/alibaba/testable/core/annotation/EnableTestableInject.java +++ b/core/src/main/java/com/alibaba/testable/core/annotation/EnableTestableInject.java @@ -12,4 +12,9 @@ import java.lang.annotation.*; @Documented public @interface EnableTestableInject { + /** + * Test class names + */ + String[] value(); + } diff --git a/core/src/main/java/com/alibaba/testable/core/annotation/TestableInject.java b/core/src/main/java/com/alibaba/testable/core/annotation/TestableInject.java index 74381d6..e047dec 100644 --- a/core/src/main/java/com/alibaba/testable/core/annotation/TestableInject.java +++ b/core/src/main/java/com/alibaba/testable/core/annotation/TestableInject.java @@ -7,7 +7,7 @@ import java.lang.annotation.*; * * @author flin */ -@Retention(RetentionPolicy.SOURCE) +@Retention(RetentionPolicy.RUNTIME) @Target(ElementType.METHOD) @Documented public @interface TestableInject {