From 592731a3e69a30cfff2bde78ab27862081347cc8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=87=91=E6=88=9F?= Date: Tue, 21 Jul 2020 00:08:05 +0800 Subject: [PATCH] implement method record visitor --- .../java/com/alibaba/testable/PreMain.java | 3 +- .../testable/TestableFileTransformer.java | 28 ------- .../alibaba/testable/model/MethodInfo.java | 38 +++++++++ .../transformer/TestableFileTransformer.java | 80 +++++++++++++++++++ .../testable/visitor/MethodRecordVisitor.java | 45 +++++++++++ .../TestableVisitor.java} | 12 ++- .../testable/annotation/EnableTestable.java | 2 +- .../annotation/EnableTestableInject.java | 2 +- 8 files changed, 176 insertions(+), 34 deletions(-) delete mode 100644 agent/src/main/java/com/alibaba/testable/TestableFileTransformer.java create mode 100644 agent/src/main/java/com/alibaba/testable/model/MethodInfo.java create mode 100644 agent/src/main/java/com/alibaba/testable/transformer/TestableFileTransformer.java create mode 100755 agent/src/main/java/com/alibaba/testable/visitor/MethodRecordVisitor.java rename agent/src/main/java/com/alibaba/testable/{ClassPrinter.java => visitor/TestableVisitor.java} (73%) diff --git a/agent/src/main/java/com/alibaba/testable/PreMain.java b/agent/src/main/java/com/alibaba/testable/PreMain.java index fa91dc5..29ef73b 100755 --- a/agent/src/main/java/com/alibaba/testable/PreMain.java +++ b/agent/src/main/java/com/alibaba/testable/PreMain.java @@ -1,5 +1,7 @@ package com.alibaba.testable; +import com.alibaba.testable.transformer.TestableFileTransformer; + import java.lang.instrument.Instrumentation; public class PreMain { @@ -9,4 +11,3 @@ public class PreMain { } } - diff --git a/agent/src/main/java/com/alibaba/testable/TestableFileTransformer.java b/agent/src/main/java/com/alibaba/testable/TestableFileTransformer.java deleted file mode 100644 index b2818e6..0000000 --- a/agent/src/main/java/com/alibaba/testable/TestableFileTransformer.java +++ /dev/null @@ -1,28 +0,0 @@ -package com.alibaba.testable; - -import org.objectweb.asm.ClassReader; -import org.objectweb.asm.ClassWriter; -import java.lang.instrument.ClassFileTransformer; -import java.lang.instrument.IllegalClassFormatException; -import java.security.ProtectionDomain; - -public class TestableFileTransformer implements ClassFileTransformer { - - private static String targetClass = "com/alibaba/testable/TransClass"; - - @Override - public byte[] transform(ClassLoader loader, String className, Class classBeingRedefined, - ProtectionDomain protectionDomain, byte[] classfileBuffer) - throws IllegalClassFormatException { - if (targetClass.equals(className)) { - // Print class structure with ASM - ClassReader reader = new ClassReader(classfileBuffer); - ClassWriter writer = new ClassWriter(reader, 0); - ClassPrinter visitor = new ClassPrinter(writer); - reader.accept(visitor, 0); - return writer.toByteArray(); - } - return null; - } - -} diff --git a/agent/src/main/java/com/alibaba/testable/model/MethodInfo.java b/agent/src/main/java/com/alibaba/testable/model/MethodInfo.java new file mode 100644 index 0000000..3da61a8 --- /dev/null +++ b/agent/src/main/java/com/alibaba/testable/model/MethodInfo.java @@ -0,0 +1,38 @@ +package com.alibaba.testable.model; + +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; + this.name = name; + this.desc = desc; + this.signature = signature; + this.exceptions = exceptions; + } + + public int getAccess() { + return access; + } + + public String getName() { + return name; + } + + public String getDesc() { + return desc; + } + + public String getSignature() { + return signature; + } + + public String[] getExceptions() { + return exceptions; + } +} diff --git a/agent/src/main/java/com/alibaba/testable/transformer/TestableFileTransformer.java b/agent/src/main/java/com/alibaba/testable/transformer/TestableFileTransformer.java new file mode 100644 index 0000000..f4d07cf --- /dev/null +++ b/agent/src/main/java/com/alibaba/testable/transformer/TestableFileTransformer.java @@ -0,0 +1,80 @@ +package com.alibaba.testable.transformer; + +import com.alibaba.testable.model.MethodInfo; +import com.alibaba.testable.visitor.MethodRecordVisitor; +import com.alibaba.testable.visitor.TestableVisitor; +import org.objectweb.asm.ClassReader; +import org.objectweb.asm.ClassWriter; + +import java.lang.annotation.Annotation; +import java.lang.instrument.ClassFileTransformer; +import java.lang.reflect.Field; +import java.security.ProtectionDomain; +import java.util.List; +import java.util.Vector; + +public class TestableFileTransformer implements ClassFileTransformer { + + private static final String ENABLE_TESTABLE = "com.alibaba.testable.annotation.EnableTestable"; + private static final String DOT = "."; + private static final String SLASH = "/"; + private static final String TEST_POSTFIX = "Test"; + + public byte[] transform(ClassLoader loader, String className, Class classBeingRedefined, + ProtectionDomain protectionDomain, byte[] classfileBuffer) { + if (null == loader || null == className) { + // Ignore system class + return null; + } + String dotClassName = className.replace(SLASH, DOT); + MethodRecordVisitor methodRecordVisitor = getMemberMethods(classfileBuffer, + isTestClassTestable(loader, dotClassName)); + if (!methodRecordVisitor.isNeedTransform()) { + // Neither EnableTestable on test class, nor EnableTestableInject on source class + return null; + } + + ClassReader reader = new ClassReader(classfileBuffer); + ClassWriter writer = new ClassWriter(reader, 0); + reader.accept(new TestableVisitor(writer, methodRecordVisitor.getMethods()), 0); + return writer.toByteArray(); + } + + private boolean isTestClassTestable(ClassLoader loader, String dotClassName) { + boolean needTransform = false; + try { + Field classesField = ClassLoader.class.getDeclaredField("classes"); + classesField.setAccessible(true); + Vector classesVector = (Vector)classesField.get(loader); + if (null != classesVector) { + for (Class c : classesVector) { + String testClassName = dotClassName + TEST_POSTFIX; + if (c.getName().endsWith(testClassName)) { + Class testClazz = Class.forName(testClassName); + for (Annotation a : testClazz.getAnnotations()) { + if (a.annotationType().getName().equals(ENABLE_TESTABLE)) { + needTransform = true; + } + } + } + } + } + } catch (NoSuchFieldException e) { + return false; + } catch (IllegalAccessException e) { + return false; + } catch (ClassNotFoundException e) { + return false; + } + return needTransform; + } + + private MethodRecordVisitor getMemberMethods(byte[] classfileBuffer, boolean needTransform) { + ClassReader reader = new ClassReader(classfileBuffer); + ClassWriter writer = new ClassWriter(reader, 0); + MethodRecordVisitor visitor = new MethodRecordVisitor(writer, needTransform); + reader.accept(visitor, 0); + return visitor; + } + +} diff --git a/agent/src/main/java/com/alibaba/testable/visitor/MethodRecordVisitor.java b/agent/src/main/java/com/alibaba/testable/visitor/MethodRecordVisitor.java new file mode 100755 index 0000000..c9c4acd --- /dev/null +++ b/agent/src/main/java/com/alibaba/testable/visitor/MethodRecordVisitor.java @@ -0,0 +1,45 @@ +package com.alibaba.testable.visitor; + +import com.alibaba.testable.model.MethodInfo; +import org.objectweb.asm.*; + +import java.util.ArrayList; +import java.util.List; + +public class MethodRecordVisitor extends ClassVisitor { + + /** + * Member methods + */ + private List methods = new ArrayList(); + private boolean needTransform; + + private static final String ENABLE_TESTABLE_INJECT = "Lcom/alibaba/testable/annotation/EnableTestableInject;"; + + public List getMethods() { + return methods; + } + + public boolean isNeedTransform() { + return needTransform; + } + + public MethodRecordVisitor(ClassWriter cw, boolean needTransform) { + super(Opcodes.ASM8, cw); + this.needTransform = needTransform; + } + + @Override + public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) { + if (descriptor.equals(ENABLE_TESTABLE_INJECT)) { + needTransform = true; + } + return super.visitAnnotation(descriptor, visible); + } + + @Override + public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) { + methods.add(new MethodInfo(access, name, desc, signature, exceptions)); + return super.visitMethod(access, name, desc, signature, exceptions); + } +} diff --git a/agent/src/main/java/com/alibaba/testable/ClassPrinter.java b/agent/src/main/java/com/alibaba/testable/visitor/TestableVisitor.java similarity index 73% rename from agent/src/main/java/com/alibaba/testable/ClassPrinter.java rename to agent/src/main/java/com/alibaba/testable/visitor/TestableVisitor.java index 7af3e3d..0768615 100755 --- a/agent/src/main/java/com/alibaba/testable/ClassPrinter.java +++ b/agent/src/main/java/com/alibaba/testable/visitor/TestableVisitor.java @@ -1,14 +1,20 @@ -package com.alibaba.testable; +package com.alibaba.testable.visitor; +import com.alibaba.testable.model.MethodInfo; import org.objectweb.asm.ClassVisitor; import org.objectweb.asm.ClassWriter; import org.objectweb.asm.MethodVisitor; import org.objectweb.asm.Opcodes; -public class ClassPrinter extends ClassVisitor { +import java.util.List; - public ClassPrinter(ClassWriter cw) { +public class TestableVisitor extends ClassVisitor { + + private List methods; + + public TestableVisitor(ClassWriter cw, List methods) { super(Opcodes.ASM8, cw); + this.methods = methods; } @Override diff --git a/core/src/main/java/com/alibaba/testable/annotation/EnableTestable.java b/core/src/main/java/com/alibaba/testable/annotation/EnableTestable.java index 4bb27d6..2417d56 100644 --- a/core/src/main/java/com/alibaba/testable/annotation/EnableTestable.java +++ b/core/src/main/java/com/alibaba/testable/annotation/EnableTestable.java @@ -7,7 +7,7 @@ import java.lang.annotation.*; * * @author flin */ -@Retention(RetentionPolicy.SOURCE) +@Retention(RetentionPolicy.RUNTIME) @Target(ElementType.TYPE) @Documented public @interface EnableTestable { diff --git a/core/src/main/java/com/alibaba/testable/annotation/EnableTestableInject.java b/core/src/main/java/com/alibaba/testable/annotation/EnableTestableInject.java index 24fe925..8b8eac5 100644 --- a/core/src/main/java/com/alibaba/testable/annotation/EnableTestableInject.java +++ b/core/src/main/java/com/alibaba/testable/annotation/EnableTestableInject.java @@ -7,7 +7,7 @@ import java.lang.annotation.*; * * @author flin */ -@Retention(RetentionPolicy.SOURCE) +@Retention(RetentionPolicy.RUNTIME) @Target(ElementType.TYPE) @Documented public @interface EnableTestableInject {