fix: custom model property.

HXSecurity · Jun 7, 2023 · 1713cf4 · 1713cf4
1 parent 3b805ed
commit 1713cf4
Showing 1 changed file with 21 additions and 17 deletions.
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) && !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);
-                            }
-                        }
-                        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++) {