From 6103a691ce360f428f65e5719872e48b45d1beba Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=87=91=E6=88=9F?= <jinji.lf@alibaba-inc.com>
Date: Tue, 21 Jul 2020 12:19:08 +0800
Subject: [PATCH] more effective check

---
 .../transformer/TestableFileTransformer.java  | 46 +++++++------------
 1 file changed, 17 insertions(+), 29 deletions(-)

diff --git a/agent/src/main/java/com/alibaba/testable/transformer/TestableFileTransformer.java b/agent/src/main/java/com/alibaba/testable/transformer/TestableFileTransformer.java
index f4d07cf..e8fe8d5 100644
--- a/agent/src/main/java/com/alibaba/testable/transformer/TestableFileTransformer.java
+++ b/agent/src/main/java/com/alibaba/testable/transformer/TestableFileTransformer.java
@@ -1,6 +1,5 @@
 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;
@@ -8,10 +7,9 @@ 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;
+import java.util.HashSet;
+import java.util.Set;
 
 public class TestableFileTransformer implements ClassFileTransformer {
 
@@ -20,6 +18,8 @@ public class TestableFileTransformer implements ClassFileTransformer {
     private static final String SLASH = "/";
     private static final String TEST_POSTFIX = "Test";
 
+    private static Set<String> loadedClassNames = new HashSet<String>();
+
     public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                             ProtectionDomain protectionDomain, byte[] classfileBuffer) {
         if (null == loader || null == className) {
@@ -27,8 +27,8 @@ public class TestableFileTransformer implements ClassFileTransformer {
             return null;
         }
         String dotClassName = className.replace(SLASH, DOT);
-        MethodRecordVisitor methodRecordVisitor = getMemberMethods(classfileBuffer,
-            isTestClassTestable(loader, dotClassName));
+        loadedClassNames.add(dotClassName);
+        MethodRecordVisitor methodRecordVisitor = getMemberMethods(classfileBuffer, checkTestClass(dotClassName));
         if (!methodRecordVisitor.isNeedTransform()) {
             // Neither EnableTestable on test class, nor EnableTestableInject on source class
             return null;
@@ -40,33 +40,21 @@ public class TestableFileTransformer implements ClassFileTransformer {
         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;
-                            }
-                        }
+    private boolean checkTestClass(String dotClassName) {
+        String testClassName = dotClassName + TEST_POSTFIX;
+        if (loadedClassNames.contains(testClassName)) {
+            try {
+                Class<?> testClazz = Class.forName(testClassName);
+                for (Annotation a : testClazz.getAnnotations()) {
+                    if (a.annotationType().getName().equals(ENABLE_TESTABLE)) {
+                        return true;
                     }
                 }
+            } catch (ClassNotFoundException e) {
+                return false;
             }
-        } catch (NoSuchFieldException e) {
-            return false;
-        } catch (IllegalAccessException e) {
-            return false;
-        } catch (ClassNotFoundException e) {
-            return false;
         }
-        return needTransform;
+        return false;
     }
 
     private MethodRecordVisitor getMemberMethods(byte[] classfileBuffer, boolean needTransform) {