Merge pull request #577 from java-sec/fix-class-diagram-deadlock

feat: 重构IastClassDiagram

Author: CC11001100

dongtai-core/pom.xml

@@ -258,6 +258,11 @@
             <scope>test</scope>
         </dependency>
         -->
+        <dependency>
+            <groupId>org.openjdk.jmh</groupId>
+            <artifactId>jmh-generator-annprocess</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
 </project>

dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/IastClassFileTransformer.java

@@ -16,7 +16,9 @@
 import io.dongtai.log.DongTaiLog;
 import io.dongtai.log.ErrorCode;
 import org.apache.commons.lang3.time.StopWatch;
-import org.objectweb.asm.*;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.ClassWriter;
 
 import java.io.File;
 import java.lang.dongtai.SpyDispatcherHandler;
@@ -27,9 +29,12 @@
 import java.security.CodeSource;
 import java.security.ProtectionDomain;
 import java.util.*;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentSkipListSet;
 
 import static org.apache.commons.io.FileUtils.writeByteArrayToFile;
-import static org.objectweb.asm.ClassWriter.*;
+import static org.objectweb.asm.ClassWriter.COMPUTE_FRAMES;
+import static org.objectweb.asm.ClassWriter.COMPUTE_MAXS;
 
 /**
  * @author [email protected]
@@ -119,6 +124,8 @@ public byte[] transform(final ClassLoader loader,
                             final Class<?> classBeingRedefined,
                             final ProtectionDomain protectionDomain,
                             final byte[] srcByteCodeArray) {
+
+        // TODO 那下面这段逻辑是否就可以去掉了？
         String threadName = Thread.currentThread().getName();
         if (threadName.startsWith("DongTai-IAST-Core")) {
             return null;
@@ -172,12 +179,10 @@ public byte[] transform(final ClassLoader loader,
             }
             final String className = classContext.getClassName();
 
-            Set<String> ancestors = classDiagram.getDiagram(className);
+            // 设置类的祖先类
+            Set<String> ancestors = classDiagram.getClassAncestorSet(className);
             if (ancestors == null) {
-                classDiagram.setLoader(loader);
-                classDiagram.saveAncestors(className, classContext.getSuperClassName(), classContext.getInterfaces());
-                ancestors = classDiagram.getAncestors(className, classContext.getSuperClassName(),
-                        classContext.getInterfaces());
+                ancestors = classDiagram.updateAncestorsByClassContext(loader, classContext);
             }
             classContext.setAncestors(ancestors);
 
@@ -288,9 +293,9 @@ public Class<?>[] findForRetransform() {
                     continue;
                 }
                 String className = clazz.getName();
-                Set<String> diagram = classDiagram.getDiagram(className);
+                Set<String> diagram = classDiagram.getClassAncestorSet(className);
                 if (diagram == null) {
-                    diagram = new HashSet<String>();
+                    diagram = ConcurrentHashMap.newKeySet();
                     Queue<Class<?>> classQueue = new LinkedList<Class<?>>();
 
                     classQueue.add(clazz);
@@ -305,7 +310,7 @@ public Class<?>[] findForRetransform() {
                         Class<?>[] interfaces = currentClazz.getInterfaces();
                         Collections.addAll(classQueue, interfaces);
                     }
-                    classDiagram.setDiagram(className, diagram);
+                    classDiagram.setClassAncestorSet(className, diagram);
                 }
                 for (String clazzName : diagram) {
                     if (this.policyManager.isHookClass(clazzName) ||