Merge pull request #529 from Nizernizer/fix/string-hash-enhance

fix: String hash enhance.
HXSecurity · Jun 7, 2023 · 06b1088 · 06b1088
2 parents 13729c7 + e6180bc
commit 06b1088
Show file tree

Hide file tree

Showing 12 changed files with 74 additions and 47 deletions.
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/EngineManager.java b/dongtai-core/src/main/java/io/dongtai/iast/core/EngineManager.java
@@ -136,16 +136,16 @@ public static void enterHttpEntry(Map<String, Object> requestMeta) {
         }
         REQUEST_CONTEXT.set(requestMeta);
         TRACK_MAP.set(new HashMap<Integer, MethodEvent>(1024));
-        TAINT_HASH_CODES.set(new HashSet<Integer>());
-        TAINT_RANGES_POOL.set(new HashMap<Integer, TaintRanges>());
+        TAINT_HASH_CODES.set(new HashSet<Long>());
+        TAINT_RANGES_POOL.set(new HashMap<Long, TaintRanges>());
         ScopeManager.SCOPE_TRACKER.getScope(Scope.HTTP_ENTRY).enter();
     }
 
     public static void enterDubboEntry(Map<String, Object> requestMeta) {
         REQUEST_CONTEXT.set(requestMeta);
         TRACK_MAP.set(new HashMap<Integer, MethodEvent>(1024));
-        TAINT_HASH_CODES.set(new HashSet<Integer>());
-        TAINT_RANGES_POOL.set(new HashMap<Integer, TaintRanges>());
+        TAINT_HASH_CODES.set(new HashSet<Long>());
+        TAINT_RANGES_POOL.set(new HashMap<Long, TaintRanges>());
         ScopeManager.SCOPE_TRACKER.getScope(Scope.DUBBO_ENTRY).enter();
     }
 }
diff --git a/...-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/DubboImpl.java b/...-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/DubboImpl.java
@@ -101,7 +101,7 @@ public static void collectDubboRequestSource(Object handler, Object invocation,
 
         // for display taint range (full arguments value)
         String fv = event.parameterValues.get(0).getValue();
-        int hash = System.identityHashCode(fv);
+        long hash = TaintPoolUtils.toStringHash(fv.hashCode(),System.identityHashCode(fv));
         int len = TaintRangesBuilder.getLength(fv);
         TaintRanges tr = new TaintRanges(new TaintRange(0, len));
         event.targetRanges.add(0, new MethodEvent.MethodEventTargetRange(hash, tr));

diff --git a/.../src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/PropagatorImpl.java b/.../src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/PropagatorImpl.java
@@ -165,7 +165,7 @@ private static boolean setTarget(PropagatorNode propagatorNode, MethodEvent even
     }
 
     private static TaintRanges getTaintRanges(Object obj) {
-        int hash = System.identityHashCode(obj);
+        long hash = TaintPoolUtils.getStringHash(obj);
         TaintRanges tr = EngineManager.TAINT_RANGES_POOL.get(hash);
         if (tr == null) {
             tr = new TaintRanges();
@@ -209,7 +209,7 @@ private static void trackTaintRange(PropagatorNode propagatorNode, MethodEvent e
             }
         }
 
-        int tgtHash = 0;
+        long tgtHash = 0;
         Object tgt = null;
         Set<TaintPosition> targetLocs = propagatorNode.getTargets();
         // may have multiple targets?
@@ -218,17 +218,16 @@ private static void trackTaintRange(PropagatorNode propagatorNode, MethodEvent e
         }
         if (TaintPosition.hasObject(targetLocs)) {
             tgt = event.objectInstance;
-            tgtHash = System.identityHashCode(tgt);
+            tgtHash = TaintPoolUtils.getStringHash(tgt);
             oldTaintRanges = getTaintRanges(tgt);
         } else if (TaintPosition.hasReturn(targetLocs)) {
-            tgt = event.returnInstance;
-            tgtHash = System.identityHashCode(tgt);
+            tgtHash = TaintPoolUtils.getStringHash(tgt);
         } else if (TaintPosition.hasParameter(targetLocs)) {
             for (TaintPosition targetLoc : targetLocs) {
                 int parameterIndex = targetLoc.getParameterIndex();
                 if (event.parameterInstances.length > parameterIndex) {
                     tgt = event.parameterInstances[parameterIndex];
-                    tgtHash = System.identityHashCode(tgt);
+                    tgtHash = TaintPoolUtils.getStringHash(tgt);
                     oldTaintRanges = getTaintRanges(tgt);
                 }
             }

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/MethodEvent.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/MethodEvent.java
@@ -82,9 +82,9 @@ public class MethodEvent {
      */
     public String returnValue;
 
-    private final Set<Integer> sourceHashes = new HashSet<Integer>();
+    private final Set<Long> sourceHashes = new HashSet<Long>();
 
-    private final Set<Integer> targetHashes = new HashSet<Integer>();
+    private final Set<Long> targetHashes = new HashSet<Long>();
 
     public List<MethodEventTargetRange> targetRanges = new ArrayList<MethodEventTargetRange>();
 
@@ -118,10 +118,10 @@ public JSONObject toJson() {
     }
 
     public static class MethodEventSourceType {
-        private final Integer hash;
+        private final Long hash;
         private final String type;
 
-        public MethodEventSourceType(Integer hash, String type) {
+        public MethodEventSourceType(Long hash, String type) {
             this.hash = hash;
             this.type = type;
         }
@@ -135,10 +135,10 @@ public JSONObject toJson() {
     }
 
     public static class MethodEventTargetRange {
-        private final Integer hash;
+        private final Long hash;
         private final TaintRanges ranges;
 
-        public MethodEventTargetRange(Integer hash, TaintRanges ranges) {
+        public MethodEventTargetRange(Long hash, TaintRanges ranges) {
             this.hash = hash;
             this.ranges = ranges;
         }
@@ -234,19 +234,19 @@ private String formatValue(Object val, boolean hasTaint) {
                 + (hasTaint ? "*" : "") + String.valueOf(str.length());
     }
 
-    public Set<Integer> getSourceHashes() {
+    public Set<Long> getSourceHashes() {
         return sourceHashes;
     }
 
-    public void addSourceHash(int hashcode) {
+    public void addSourceHash(long hashcode) {
         this.sourceHashes.add(hashcode);
     }
 
-    public Set<Integer> getTargetHashes() {
+    public Set<Long> getTargetHashes() {
         return targetHashes;
     }
 
-    public void addTargetHash(int hashCode) {
+    public void addTargetHash(long hashCode) {
         this.targetHashes.add(hashCode);
     }
 

diff --git a/...java/io/dongtai/iast/core/handler/hookpoint/vulscan/dynamic/DynamicPropagatorScanner.java b/...java/io/dongtai/iast/core/handler/hookpoint/vulscan/dynamic/DynamicPropagatorScanner.java
@@ -122,7 +122,7 @@ private boolean sinkSourceHitTaintPool(MethodEvent event, SinkNode sinkNode) {
         if (VulnType.REFLECTED_XSS.equals(sinkNode.getVulType()) && !sourceInstances.isEmpty()) {
             boolean tagsHit = false;
             for (Object sourceInstance : sourceInstances) {
-                int hash = System.identityHashCode(sourceInstance);
+                long hash = TaintPoolUtils.getStringHash(sourceInstance);
                 TaintRanges tr = EngineManager.TAINT_RANGES_POOL.get(hash);
                 if (tr == null || tr.isEmpty()) {
                     continue;

diff --git a/.../main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/dynamic/PathTraversalCheck.java b/.../main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/dynamic/PathTraversalCheck.java
@@ -112,7 +112,7 @@ private boolean checkPath(String path, MethodEvent event) {
             return false;
         }
 
-        TaintRanges tr = EngineManager.TAINT_RANGES_POOL.get(System.identityHashCode(path));
+        TaintRanges tr = EngineManager.TAINT_RANGES_POOL.get(TaintPoolUtils.toStringHash(path.hashCode(),System.identityHashCode(path)));
         if (tr.isEmpty()) {
             return false;
         }

diff --git a/...src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/dynamic/SSRFSourceCheck.java b/...src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/dynamic/SSRFSourceCheck.java
@@ -361,7 +361,8 @@ private boolean addSourceType(MethodEvent event, Map<String, Object> sourceMap)
 
     private boolean checkTaintPool(MethodEvent event, String key, Object value) {
         if (!"".equals(value) && TaintPoolUtils.poolContains(value, event)) {
-            event.sourceTypes.add(new MethodEvent.MethodEventSourceType(System.identityHashCode(value), key));
+            long hash = TaintPoolUtils.getStringHash(value);
+            event.sourceTypes.add(new MethodEvent.MethodEventSourceType(hash, key));
             return true;
         }
         return false;

diff --git a/...java/io/dongtai/iast/core/handler/hookpoint/vulscan/dynamic/UnvalidatedRedirectCheck.java b/...java/io/dongtai/iast/core/handler/hookpoint/vulscan/dynamic/UnvalidatedRedirectCheck.java
@@ -125,8 +125,8 @@ private boolean checkValue(Object val, MethodEvent event) {
         if (!TaintPoolUtils.poolContains(val, event)) {
             return false;
         }
-
-        TaintRanges tr = EngineManager.TAINT_RANGES_POOL.get(System.identityHashCode(val));
+        long hash = TaintPoolUtils.getStringHash(val);
+        TaintRanges tr = EngineManager.TAINT_RANGES_POOL.get(hash);
         if (tr.isEmpty()) {
             return false;
         }

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/utils/HashCode.java b/dongtai-core/src/main/java/io/dongtai/iast/core/utils/HashCode.java
@@ -4,7 +4,7 @@
  * @author dongzhiyong@huoxian.cn
  */
 public class HashCode {
-    public static int calc(Object obj) {
+    public static long calc(Object obj) {
         if (obj instanceof String) {
             return ((String) obj).hashCode();
         } else {

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
@@ -29,11 +29,12 @@ public static boolean poolContains(Object obj, MethodEvent event) {
             return false;
         }
 
+        long hash = getStringHash(obj);
         boolean isContains;
         // check object hash exists
-        isContains = contains(obj);
+        isContains = contains(hash);
         if (isContains) {
-            event.addSourceHash(System.identityHashCode(obj));
+            event.addSourceHash(hash);
             return true;
         }
 
@@ -59,11 +60,11 @@ public static boolean poolContains(Object obj, MethodEvent event) {
     /**
      * 判断污点是否匹配
      *
-     * @param obj Object
+     * @param hash long
      * @return boolean
      */
-    private static boolean contains(Object obj) {
-        return EngineManager.TAINT_HASH_CODES.contains(System.identityHashCode(obj));
+    private static boolean contains(long hash) {
+        return EngineManager.TAINT_HASH_CODES.contains(hash);
     }
 
     /**
@@ -141,10 +142,17 @@ public static void trackObject(MethodEvent event, PolicyNode policyNode, Object
             return;
         }
 
-        int hash = 0;
+        long hash = 0;
+        long identityHash = 0;
         boolean isSourceNode = policyNode instanceof SourceNode;
         if (isSourceNode) {
-            hash = System.identityHashCode(obj);
+            if (obj instanceof String){
+                identityHash = System.identityHashCode(obj);
+                hash = toStringHash(obj.hashCode(),identityHash);
+            }else {
+                hash = System.identityHashCode(obj);
+                identityHash = hash;
+            }
             if (EngineManager.TAINT_HASH_CODES.contains(hash)) {
                 return;
             }
@@ -170,7 +178,7 @@ public static void trackObject(MethodEvent event, PolicyNode policyNode, Object
         } else {
             if (isSourceNode) {
                 int len = TaintRangesBuilder.getLength(obj);
-                if (hash == 0 || len == 0) {
+                if (identityHash == 0 || len == 0) {
                     return;
                 }
 
@@ -205,7 +213,7 @@ public static void trackObject(MethodEvent event, PolicyNode policyNode, Object
                     }
                 }
             } else {
-                hash = System.identityHashCode(obj);
+                hash = getStringHash(obj);
                 if (EngineManager.TAINT_HASH_CODES.contains(hash)) {
                     event.addSourceHash(hash);
                 }
@@ -251,4 +259,19 @@ private static void trackOptional(MethodEvent event, PolicyNode policyNode, Obje
         } catch (Throwable ignore) {
         }
     }
+
+    public static Long toStringHash(long objectHashCode,long identityHashCode) {
+        return (objectHashCode << 32) | (identityHashCode & 0xFFFFFFFFL);
+    }
+
+    public static Long getStringHash(Object obj) {
+        long hash;
+        if (obj instanceof String){
+            hash = TaintPoolUtils.toStringHash(obj.hashCode(),System.identityHashCode(obj));
+        }else {
+            hash = System.identityHashCode(obj);
+        }
+        return hash;
+    }
+
 }
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/utils/threadlocal/IastTaintHashCodes.java b/dongtai-core/src/main/java/io/dongtai/iast/core/utils/threadlocal/IastTaintHashCodes.java
@@ -11,24 +11,24 @@
 /**
  * @author dongzhiyong@huoxian.cn
  */
-public class IastTaintHashCodes extends ThreadLocal<HashSet<Integer>> {
+public class IastTaintHashCodes extends ThreadLocal<HashSet<Long>> {
     @Override
-    protected HashSet<Integer> initialValue() {
+    protected HashSet<Long> initialValue() {
         return null;
     }
 
     public boolean isEmpty() {
         return this.get() == null || this.get().isEmpty();
     }
 
-    public boolean contains(Integer hashCode) {
+    public boolean contains(Long hashCode) {
         if (this.get() == null) {
             return false;
         }
         return this.get().contains(hashCode);
     }
 
-    public void add(Integer hashCode) {
+    public void add(Long hashCode) {
         if (this.get() == null) {
             return;
         }
@@ -41,16 +41,20 @@ public void addObject(Object obj, MethodEvent event) {
         }
 
         try {
-            int subHashCode = 0;
+            long subHashCode = 0;
             if (obj instanceof String[]) {
                 String[] tempObjs = (String[]) obj;
                 for (String tempObj : tempObjs) {
-                    subHashCode = System.identityHashCode(tempObj);
+                    subHashCode = TaintPoolUtils.toStringHash(tempObj.hashCode(),System.identityHashCode(tempObj));
                     this.add(subHashCode);
                     event.addTargetHash(subHashCode);
                 }
             } else if (obj instanceof Map) {
-                int hashCode = System.identityHashCode(obj);
+                long hashCode = System.identityHashCode(obj);
+                this.add(hashCode);
+                event.addTargetHash(hashCode);
+            } else if (obj instanceof String){
+                long hashCode = TaintPoolUtils.toStringHash(obj.hashCode(),System.identityHashCode(obj));
                 this.add(hashCode);
                 event.addTargetHash(hashCode);
             } else if (obj.getClass().isArray() && !obj.getClass().getComponentType().isPrimitive()) {

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/utils/threadlocal/TaintRangesPool.java b/dongtai-core/src/main/java/io/dongtai/iast/core/utils/threadlocal/TaintRangesPool.java
@@ -4,17 +4,17 @@
 
 import java.util.Map;
 
-public class TaintRangesPool extends ThreadLocal<Map<Integer, TaintRanges>> {
+public class TaintRangesPool extends ThreadLocal<Map<Long, TaintRanges>> {
     @Override
-    protected Map<Integer, TaintRanges> initialValue() {
+    protected Map<Long, TaintRanges> initialValue() {
         return null;
     }
 
-    public void add(Integer hash, TaintRanges taintRanges) {
+    public void add(Long hash, TaintRanges taintRanges) {
         this.get().put(hash, taintRanges);
     }
 
-    public TaintRanges get(int hash) {
+    public TaintRanges get(long hash) {
         return this.get().get(hash);
     }
 }