implement method record visitor

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 {
     }
 
 }
-

agent/src/main/java/com/alibaba/testable/TestableFileTransformer.java

@@ -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;
-    }
-
-}

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;
+    }
+}

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<Class> classesVector = (Vector<Class>)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;
+    }
+
+}

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<MethodInfo> methods = new ArrayList<MethodInfo>();
+    private boolean needTransform;
+
+    private static final String ENABLE_TESTABLE_INJECT = "Lcom/alibaba/testable/annotation/EnableTestableInject;";
+
+    public List<MethodInfo> 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);
+    }
+}

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<MethodInfo> methods;
+
+    public TestableVisitor(ClassWriter cw, List<MethodInfo> methods) {
         super(Opcodes.ASM8, cw);
+        this.methods = methods;
     }
 
     @Override

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 {

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 {