Merge pull request #563 from Nizernizer/fix/deadlock

fix: agent deadlock.

Author: lostsnow

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

@@ -62,6 +62,13 @@ public static IastClassFileTransformer getInstance(Instrumentation inst, PolicyM
         return INSTANCE;
     }
 
+    public static IastClassFileTransformer getInstance() {
+        if (null != INSTANCE) {
+            return INSTANCE;
+        }
+        return null;
+    }
+
     IastClassFileTransformer(Instrumentation inst, PolicyManager policyManager) {
         this.inst = inst;
         this.isDumpClass = EngineManager.getInstance().isEnableDumpClass();
@@ -112,29 +119,39 @@ public byte[] transform(final ClassLoader loader,
                             final Class<?> classBeingRedefined,
                             final ProtectionDomain protectionDomain,
                             final byte[] srcByteCodeArray) {
+        String threadName = Thread.currentThread().getName();
+        if (threadName.startsWith("DongTai-IAST-Core")
+                || threadName.startsWith("DongTai-IAST-AgentStateMonitor")
+                || threadName.startsWith("DongTai-IAST-ConfigMonitor")
+                || threadName.startsWith("DongTai-IAST-FallbackConfigMonitor")
+                || threadName.startsWith("DongTai-IAST-HearBeatMonitor")
+                || threadName.startsWith("DongTai-IAST-PerformanceMonitor")) {
+            return null;
+        }
+
+        if (internalClassName == null
+                || internalClassName.startsWith("io/dongtai/")
+                || internalClassName.startsWith("com/secnium/iast/")
+                || internalClassName.startsWith("java/lang/iast/")
+                || internalClassName.startsWith("cn/huoxian/iast/")
+                || internalClassName.startsWith("META-INF/")
+                || "module-info".equals(internalClassName)) {
+            return null;
+        }
+
+        if (null != loader && loader.toString().toLowerCase().contains("rasp")) {
+            return null;
+        }
+
         try {
             ScopeManager.SCOPE_TRACKER.getPolicyScope().enterAgent();
 
-            if (internalClassName == null
-                    || internalClassName.startsWith("io/dongtai/")
-                    || internalClassName.startsWith("com/secnium/iast/")
-                    || internalClassName.startsWith("java/lang/iast/")
-                    || internalClassName.startsWith("cn/huoxian/iast/")
-                    || internalClassName.startsWith("META-INF/")
-                    || "module-info".equals(internalClassName)) {
-                return null;
-            }
-
             if (" com/alibaba/fastjson/JSON".substring(1).equals(internalClassName)) {
                 FastjsonCheck.setJsonClassLoader(loader);
             } else if (" com/alibaba/fastjson/parser/ParserConfig".substring(1).equals(internalClassName)) {
                 FastjsonCheck.setParseConfigClassLoader(loader);
             }
 
-            if (null != loader && loader.toString().toLowerCase().contains("rasp")) {
-                return null;
-            }
-
             if (loader != null && protectionDomain != null) {
                 final CodeSource codeSource = protectionDomain.getCodeSource();
                 if (codeSource == null) {
@@ -156,7 +173,6 @@ public byte[] transform(final ClassLoader loader,
 
             ClassContext classContext = new ClassContext(cr, loader);
             if (Modifier.isInterface(classContext.getModifier())) {
-                sourceCodeBak = null;
                 return null;
             }
             final String className = classContext.getClassName();
@@ -186,11 +202,9 @@ public byte[] transform(final ClassLoader loader,
                     return dumpClassIfNecessary(cr.getClassName(), cw.toByteArray(), srcByteCodeArray);
                 }
             }
-            sourceCodeBak = null;
         } catch (Throwable throwable) {
             DongTaiLog.warn(ErrorCode.get("TRANSFORM_CLASS_FAILED"), internalClassName, throwable);
         } finally {
-            classDiagram.setLoader(null);
             ScopeManager.SCOPE_TRACKER.getPolicyScope().leaveAgent();
         }
 
@@ -347,5 +361,9 @@ public void reTransform() {
     public static HashMap<Object, byte[]> getTransformMap() {
         return transformMap;
     }
+
+    public IastClassDiagram getClassDiagram() {
+        return classDiagram;
+    }
 }
 

dongtai-core/src/main/java/io/dongtai/iast/core/init/impl/TransformEngine.java

@@ -73,6 +73,9 @@ public void destroy() {
                 DongTaiLog.error(ErrorCode.get("TRANSFORM_ENGINE_DESTROY_REDEFINE_CLASSES_FAILED"), e);
             }
         }
+        if (IastClassFileTransformer.getInstance() != null) {
+            IastClassFileTransformer.getInstance().getClassDiagram().setLoader(null);
+        }
         inst = null;
         classFileTransformer = null;
     }