Enhance ( HXSecurity#240 ): optimizate agent attach performance

exexute · Feb 15, 2022 · 1582ccb · 1582ccb
1 parent 1a72702
commit 1582ccb
Show file tree

Hide file tree

Showing 15 changed files with 249 additions and 234 deletions.
diff --git a/dongtai-agent/pom.xml b/dongtai-agent/pom.xml
@@ -124,6 +124,16 @@
                                     </excludes>
                                 </filter>
                             </filters>
+                            <relocations>
+                                <relocation>
+                                    <pattern>org.</pattern>
+                                    <shadedPattern>${shade-prefix}.org.</shadedPattern>
+                                </relocation>
+                                <relocation>
+                                    <pattern>oshi.</pattern>
+                                    <shadedPattern>${shade-prefix}.oshi.</shadedPattern>
+                                </relocation>
+                            </relocations>
                         </configuration>
                     </execution>
                 </executions>

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/IastClassFileTransformer.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/IastClassFileTransformer.java
diff --git a/...ecode/enhance/IastClassAncestorQuery.java → ...re/bytecode/enhance/IastClassDiagram.java b/...ecode/enhance/IastClassAncestorQuery.java → ...re/bytecode/enhance/IastClassDiagram.java
@@ -2,6 +2,7 @@
 
 import io.dongtai.iast.core.service.ErrorLogReport;
 import io.dongtai.iast.core.utils.ThrowableUtils;
+
 import java.io.InputStream;
 import java.util.Arrays;
 import java.util.Collections;
@@ -13,6 +14,7 @@
 import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+
 import org.json.JSONObject;
 import org.objectweb.asm.ClassReader;
 
@@ -21,9 +23,9 @@
  *
  * @author dongzhiyong@huoxian.cn
  */
-public class IastClassAncestorQuery {
+public class IastClassDiagram {
 
-    private final Map<String, Set<String>> classAncestorMap;
+    private final Map<String, Set<String>> diagrams;
     private static final Map<String, List<String>> DEFAULT_INTERFACE_LIST_MAP;
     private static final String BASE_CLASS = "java/lang/Object";
 
@@ -32,30 +34,40 @@ public synchronized void setLoader(ClassLoader loader) {
     }
 
     private ClassLoader loader;
-    private static IastClassAncestorQuery instance;
+    private static IastClassDiagram instance;
 
-    public static IastClassAncestorQuery getInstance() {
+    public static IastClassDiagram getInstance() {
         if (instance == null) {
-            instance = new IastClassAncestorQuery();
+            instance = new IastClassDiagram();
         }
         return instance;
     }
 
-    private IastClassAncestorQuery() {
-        this.classAncestorMap = new ConcurrentHashMap<String, Set<String>>();
+    public Set<String> getDiagram(String className) {
+        return diagrams.get(className);
+    }
+
+    public void setDiagram(String className, Set<String> diagram) {
+        diagrams.put(className, diagram);
+    }
+
+    private IastClassDiagram() {
+        this.diagrams = new ConcurrentHashMap<String, Set<String>>();
     }
 
     public synchronized void saveAncestors(String className, String superName, String[] interfaces) {
-        Set<String> ancestorSet = this.classAncestorMap.get(className);
+        Set<String> ancestorSet = this.diagrams.get(className);
         ancestorSet = ancestorSet == null ? new HashSet<String>() : ancestorSet;
 
+        ancestorSet.add(className);
         if (!BASE_CLASS.equals(superName)) {
-            ancestorSet.add(superName);
+            ancestorSet.add(superName.replace("/", "."));
+        }
+        for (String interfaceClazzName : interfaces) {
+            ancestorSet.add(interfaceClazzName.replace("/", "."));
         }
 
-        Collections.addAll(ancestorSet, interfaces);
-
-        this.classAncestorMap.put(className, ancestorSet);
+        this.diagrams.put(className, ancestorSet);
     }
 
     /**
@@ -67,7 +79,7 @@ public synchronized void saveAncestors(String className, String superName, Strin
      * @return 当前类的类族
      */
     public synchronized Set<String> getAncestors(String className, String superClassName, String[] interfaces) {
-        Set<String> ancestors = this.classAncestorMap.get(className);
+        Set<String> ancestors = this.diagrams.get(className);
 
         if (!isNullOrEmpty(superClassName) && !BASE_CLASS.equals(superClassName)) {
             addClassToAncestor(superClassName, ancestors);
@@ -87,7 +99,7 @@ public synchronized Set<String> getAncestors(String className, String superClass
     }
 
     private void addClassToAncestor(String className, Set<String> ancestors) {
-        Set<String> set = this.classAncestorMap.get(className);
+        Set<String> set = this.diagrams.get(className);
         if (null != set) {
             for (String subClassName : set) {
                 if (!ancestors.contains(subClassName)) {
@@ -105,12 +117,12 @@ private void addClassToAncestor(String className, Set<String> ancestors) {
                     if (tempDefaultMap != null) {
                         ancestors.addAll(tempDefaultMap);
                     }
-                    Set<String> tempClassMap = classAncestorMap.get(tempClass);
+                    Set<String> tempClassMap = diagrams.get(tempClass);
                     if (tempClassMap != null) {
                         ancestors.addAll(tempClassMap);
                     }
                 }
-                this.classAncestorMap.put(className, tempClassFamily);
+                this.diagrams.put(className, tempClassFamily);
             }
         }
         List<String> list = DEFAULT_INTERFACE_LIST_MAP.get(className);
@@ -177,11 +189,11 @@ public static boolean isNullOrEmpty(String className) {
      * todo 利用类名查找实现的接口列表、继承的父类
      */
     public static Set<String> getFamilyFromClass(String className) {
-        return instance == null ? null : instance.classAncestorMap.get(className);
+        return instance == null ? null : instance.diagrams.get(className);
     }
 
     static {
-        DEFAULT_INTERFACE_LIST_MAP = new HashMap();
+        DEFAULT_INTERFACE_LIST_MAP = new HashMap<String, List<String>>();
         DEFAULT_INTERFACE_LIST_MAP.put(" org/apache/jasper/runtime/HttpJspBase".substring(1),
                 Collections.singletonList(" javax/servlet/jsp/JspPage".substring(1)));
         DEFAULT_INTERFACE_LIST_MAP.put(" javax/servlet/http/HttpServletResponse".substring(1),

diff --git a/...i-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/IastClassHookPointMatcher.java b/...i-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/IastClassHookPointMatcher.java
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/IastContext.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/IastContext.java
@@ -45,23 +45,15 @@ private IastContext(String className, Set<String> ancestors, String[] interfaces
 
 
     public void setClassName(String className) {
-        this.className = className.replace('/', '.');
+        this.className = className;
     }
 
     public void setAncestor(Set<String> ancestors) {
-        Set<String> copyedAncestors = new HashSet<String>();
-        for (String className : ancestors) {
-            copyedAncestors.add(className.replace('/', '.'));
-        }
-        this.ancestors = copyedAncestors;
+        this.ancestors = ancestors;
     }
 
     public void setInterface(String[] interfaces) {
-        String[] copyedInterfaces = new String[interfaces.length];
-        for (int index = 0; index < interfaces.length; index++) {
-            copyedInterfaces[index] = interfaces[index].replace('/', '.');
-        }
-        this.interfaces = copyedInterfaces;
+        this.interfaces = interfaces;
     }
 
     public void setFlags(int flags) {

diff --git a/...va/io/dongtai/iast/core/bytecode/enhance/plugin/framework/j2ee/dispatch/DispatchJ2ee.java b/...va/io/dongtai/iast/core/bytecode/enhance/plugin/framework/j2ee/dispatch/DispatchJ2ee.java
@@ -24,11 +24,11 @@ public class DispatchJ2ee implements DispatchPlugin {
     @Override
     public ClassVisitor dispatch(ClassVisitor classVisitor, IastContext context) {
         String className = context.getClassName();
-        Set<String> ancestors = context.getAncestors();
+        Set<String> diagram = context.getAncestors();
 
         if (Modifier.isInterface(context.getFlags())) {
             DongTaiLog.trace("Ignoring interface " + className);
-        } else if (isServletDispatch(className, ancestors) || isJakartaServlet(className)) {
+        } else if (isServletDispatch(className, diagram) || isJakartaServlet(className)) {
             classVisitor = new ServletDispatcherAdapter(classVisitor, context);
         }
         return classVisitor;
@@ -39,10 +39,10 @@ public String isMatch() {
         return null;
     }
 
-    private boolean isServletDispatch(String className, Set<String> ancestors) {
+    private boolean isServletDispatch(String className, Set<String> diagram) {
         boolean isServlet = FACES_SERVLET.equals(className);
         isServlet = (isServlet || HTTP_SERVLET.equals(className));
-        return (isServlet || ancestors.contains(FILTER) || ancestors.contains(FILTER_CHAIN));
+        return (isServlet || diagram.contains(FILTER) || diagram.contains(FILTER_CHAIN));
     }
 
     private boolean isJakartaServlet(String className) {

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/graphy/GraphBuilder.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/graphy/GraphBuilder.java
@@ -2,7 +2,7 @@
 
 import io.dongtai.iast.core.EngineManager;
 import io.dongtai.iast.core.utils.PropertyUtils;
-import io.dongtai.iast.core.bytecode.enhance.IastClassAncestorQuery;
+import io.dongtai.iast.core.bytecode.enhance.IastClassDiagram;
 import io.dongtai.iast.core.handler.hookpoint.controller.impl.HttpImpl;
 import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
 import io.dongtai.iast.core.handler.hookpoint.vulscan.ReportConstant;
@@ -47,7 +47,7 @@ public static List<GraphNode> build() {
                             event.getCallerClass(),
                             event.getCallerMethod(),
                             event.getCallerLine(),
-                            event.object != null ? IastClassAncestorQuery
+                            event.object != null ? IastClassDiagram
                                     .getFamilyFromClass(event.object.getClass().getName().replace("\\.", "/")) : null,
                             event.getMatchClassName(),
                             event.getOriginClassName(),

diff --git a/...i-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/IastHookRuleModel.java b/...i-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/IastHookRuleModel.java
@@ -59,6 +59,14 @@ public static void buildModel() {
         ModelBuilder.buildRemote();
     }
 
+    public static IastHookRuleModel getInstance() {
+        return instance;
+    }
+
+    public boolean isHookClass(String className) {
+        return hookClassnames.contains(className) || hookSuperClassnames.contains(className);
+    }
+
     /**
      * 创建IASTHook规则单例对象
      *

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/init/impl/ConfigEngine.java b/dongtai-core/src/main/java/io/dongtai/iast/core/init/impl/ConfigEngine.java
@@ -16,13 +16,13 @@ public class ConfigEngine implements IEngine {
 
     @Override
     public void init(PropertyUtils cfg, Instrumentation inst) {
+        DongTaiLog.info("Initialize the core configuration of the engine");
+        IastHookRuleModel.buildModel();
+        DongTaiLog.info("The engine's core configuration is initialized successfully.");
     }
 
     @Override
     public void start() {
-        DongTaiLog.info("Initialize the core configuration of the engine");
-        IastHookRuleModel.buildModel();
-        DongTaiLog.info("The engine's core configuration is initialized successfully.");
     }
 
     @Override

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/service/ErrorLogReport.java b/dongtai-core/src/main/java/io/dongtai/iast/core/service/ErrorLogReport.java
@@ -3,6 +3,7 @@
 import io.dongtai.iast.core.EngineManager;
 import io.dongtai.iast.core.handler.hookpoint.vulscan.ReportConstant;
 import io.dongtai.iast.core.utils.Constants;
+
 import java.io.PrintWriter;
 import java.io.StringWriter;
 
@@ -23,6 +24,7 @@ public static void sendErrorLog(String errorLog) {
     }
 
     public static void sendErrorLog(Throwable t) {
+        t.printStackTrace();
         if (ENABLE_UPLOAD) {
             StringWriter sw = new StringWriter();
             PrintWriter pw = new PrintWriter(sw);

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/utils/ConfigUtils.java b/dongtai-core/src/main/java/io/dongtai/iast/core/utils/ConfigUtils.java
@@ -5,6 +5,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.HashSet;
+import java.util.Set;
 
 import org.apache.commons.io.IOUtils;
 import org.apache.commons.io.LineIterator;
@@ -25,7 +26,7 @@ public static InputStream getResourceAsStreamFromFilename(String filename) {
         return ConfigUtils.class.getClassLoader().getResourceAsStream(filename);
     }
 
-    public static HashSet[] loadConfigFromFile(String filename) {
+    public static Set<String>[] loadConfigFromFile(String filename) {
         HashSet<String> container = new HashSet<String>();
         HashSet<String> startWith = new HashSet<String>();
         HashSet<String> endWith = new HashSet<String>();
@@ -69,8 +70,8 @@ public static String[] loadExtConfigFromFile(String filename) {
         return extStringArray;
     }
 
-    public static HashSet loadConfigFromFileByLine(String filename) {
-        HashSet<String> container = new HashSet<String>();
+    public static Set<String> loadConfigFromFileByLine(String filename) {
+        Set<String> container = new HashSet<String>();
         InputStream fis = null;
         try {
             fis = getResourceAsStreamFromFilename(filename);