Merge pull request #530 from Nizernizer/fix/property-disabled-features

fix: add property "disabled_features".
HXSecurity · Jun 7, 2023 · faaf558 · faaf558
2 parents 06b1088 + 1713cf4
commit faaf558
Show file tree

Hide file tree

Showing 6 changed files with 52 additions and 26 deletions.
diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/Agent.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/Agent.java
@@ -39,6 +39,7 @@ private static String[] parseAgentArgs(String[] args) throws ParseException {
         attachOptions.addOption(build("log_path", "log_path", "optional: DongTai agent log print path."));
         attachOptions.addOption(build("log_disable_collector", "log_disable_collector", "optional: DongTai agent disable log collector."));
         attachOptions.addOption(build("disabled_plugins", "disabled_plugins", "optional: DongTai agent disable plugins."));
+        attachOptions.addOption(build("disabled_features", "disabled_features", "optional: DongTai agent disable features."));
 
         CommandLineParser parser = new DefaultParser();
         HelpFormatter formatter = new HelpFormatter();

diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/IastProperties.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/IastProperties.java
@@ -34,6 +34,7 @@ public class IastProperties {
         put("log_disable_collector", PropertyConstant.PROPERTY_LOG_DISABLE_COLLECTOR);
         put("uuid_path", PropertyConstant.PROPERTY_UUID_PATH);
         put("disabled_plugins", PropertyConstant.PROPERTY_DISABLED_PLUGINS);
+        put("disabled_features", PropertyConstant.PROPERTY_DISABLED_FEATURES);
     }};
 
     private static IastProperties instance;

diff --git a/dongtai-common/src/main/java/io/dongtai/iast/common/constants/PropertyConstant.java b/dongtai-common/src/main/java/io/dongtai/iast/common/constants/PropertyConstant.java
@@ -25,12 +25,12 @@ public class PropertyConstant {
     public static final String PROPERTY_JAR_API_URL = "iast.jar.api.url";
     public static final String PROPERTY_LOG_ADDRESS = "dongtai.log.address";
     public static final String PROPERTY_LOG_PORT = "dongtai.log.port";
-    public static final String PROPERTY_FALLBACK_VERSION = "dongtai.fallback.version";
     public static final String PROPERTY_DUMP_CLASS_PATH = "iast.dump.class.path";
     public static final String PROPERTY_DUMP_CLASS_ENABLE = "iast.dump.class.enable";
     public static final String PROPERTY_SERVICE_HEARTBEAT_INTERVAL = "iast.service.heartbeat.interval";
     public static final String PROPERTY_RESPONSE_LENGTH = "dongtai.response.length";
     public static final String PROPERTY_POLICY_PATH = "dongtai.policy.path";
     public static final String PROPERTY_UUID_PATH = "dongtai.uuid.path";
     public static final String PROPERTY_DISABLED_PLUGINS = "dongtai.disabled.plugins";
+    public static final String PROPERTY_DISABLED_FEATURES = "dongtai.disabled_features";
 }
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/PluginRegister.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/PluginRegister.java
@@ -12,6 +12,7 @@
 import io.dongtai.iast.core.bytecode.enhance.plugin.spring.DispatchApiCollector;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.Policy;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.PolicyManager;
+import io.dongtai.iast.core.utils.PropertyUtils;
 import org.objectweb.asm.ClassVisitor;
 
 import java.util.*;
@@ -28,7 +29,7 @@ public class PluginRegister {
 
     public PluginRegister() {
         this.plugins = new ArrayList<>();
-        List<String> disabledPlugins = getdisabledPlugins();
+        List<String> disabledPlugins = PropertyUtils.getDisabledPlugins();
         List<DispatchPlugin> allPlugins = new ArrayList<>(Arrays.asList(
                 new DispatchApiCollector(),
                 new DispatchJ2ee(),
@@ -43,12 +44,6 @@ public PluginRegister() {
         this.plugins.add(new DispatchClassPlugin());
     }
 
-    private List<String> getdisabledPlugins() {
-        return Optional.ofNullable(System.getProperty("dongtai.disabled.plugins"))
-                .map(s -> Arrays.asList(s.split(",")))
-                .orElse(null);
-    }
-
     public ClassVisitor initial(ClassVisitor classVisitor, ClassContext context, PolicyManager policyManager) {
         Policy policy = policyManager.getPolicy();
         if (policy == null) {

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/utils/PropertyUtils.java b/dongtai-core/src/main/java/io/dongtai/iast/core/utils/PropertyUtils.java
@@ -6,7 +6,7 @@
 
 import java.io.File;
 import java.io.FileInputStream;
-import java.util.Properties;
+import java.util.*;
 
 /**
  * @author dongzhiyong@huoxian.cn
@@ -28,6 +28,8 @@ public class PropertyUtils {
     private String debugFlag;
     private Integer responseLength;
     private String policyPath;
+    private static List<String> disabledFeatureList;
+    private static Boolean isDisabledCustomModel;
 
     private final String propertiesFilePath;
 
@@ -198,4 +200,27 @@ public String getPolicyPath() {
         }
         return this.policyPath;
     }
+
+    public static List<String> getDisabledPlugins() {
+        return Optional.ofNullable(System.getProperty("dongtai.disabled.plugins"))
+                .map(s -> Arrays.asList(s.split(",")))
+                .orElse(null);
+    }
+
+    public static List<String> getDisabledFeatures() {
+        if (null == disabledFeatureList){
+            disabledFeatureList = Optional.ofNullable(System.getProperty("dongtai.disabled.features"))
+                    .map(s -> Arrays.asList(s.split(",")))
+                    .orElse(new ArrayList<>());
+        }
+        return disabledFeatureList;
+    }
+
+    public static Boolean isDisabledCustomModel() {
+        if (null == isDisabledCustomModel){
+            List<String> disabledFeatures = getDisabledFeatures();
+            isDisabledCustomModel = disabledFeatures.contains("custom-model-collection");
+        }
+        return isDisabledCustomModel;
+    }
 }
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/utils/TaintPoolUtils.java b/dongtai-core/src/main/java/io/dongtai/iast/core/utils/TaintPoolUtils.java
@@ -195,23 +195,7 @@ public static void trackObject(MethodEvent event, PolicyNode policyNode, Object
                 EngineManager.TAINT_HASH_CODES.add(hash);
                 event.addTargetHash(hash);
                 EngineManager.TAINT_RANGES_POOL.add(hash, tr);
-                if (isMicroservice && !(obj instanceof String)) {
-                    try {
-                        Field[] declaredFields = ReflectUtils.getDeclaredFieldsSecurity(cls);
-                        for (Field field : declaredFields) {
-                            if (!Modifier.isStatic(field.getModifiers())) {
-                                trackObject(event, policyNode, field.get(obj), depth + 1, isMicroservice);
-                            }
-                        }
-                        hash = System.identityHashCode(obj);
-                        if (EngineManager.TAINT_HASH_CODES.contains(hash)) {
-                            event.addSourceHash(hash);
-                        }
-                    } catch (Throwable e) {
-                        DongTaiLog.debug("solve model failed: {}, {}",
-                                e.getMessage(), e.getCause() != null ? e.getCause().getMessage() : "");
-                    }
-                }
+                TaintPoolUtils.customModel(isMicroservice,obj,cls,event,policyNode,depth);
             } else {
                 hash = getStringHash(obj);
                 if (EngineManager.TAINT_HASH_CODES.contains(hash)) {
@@ -221,6 +205,26 @@ public static void trackObject(MethodEvent event, PolicyNode policyNode, Object
         }
     }
 
+    private static void customModel(Boolean isMicroservice, Object obj, Class<?> cls, MethodEvent event,PolicyNode policyNode,int depth) {
+        if (isMicroservice && !(obj instanceof String) && !PropertyUtils.isDisabledCustomModel()) {
+            try {
+                Field[] declaredFields = ReflectUtils.getDeclaredFieldsSecurity(cls);
+                for (Field field : declaredFields) {
+                    if (!Modifier.isStatic(field.getModifiers())) {
+                        trackObject(event, policyNode, field.get(obj), depth + 1, isMicroservice);
+                    }
+                }
+                long hash = System.identityHashCode(obj);
+                if (EngineManager.TAINT_HASH_CODES.contains(hash)) {
+                    event.addSourceHash(hash);
+                }
+            } catch (Throwable e) {
+                DongTaiLog.debug("solve model failed: {}, {}",
+                        e.getMessage(), e.getCause() != null ? e.getCause().getMessage() : "");
+            }
+        }
+    }
+
     private static void trackArray(MethodEvent event, PolicyNode policyNode, Object arr, int depth, Boolean isMicroservice) {
         int length = Array.getLength(arr);
         for (int i = 0; i < length; i++) {