feature: Scope Handling of Sink Based on Vulnerability Types

dongtai-common/src/main/java/io/dongtai/iast/common/scope/PolicyScope.java

@@ -1,11 +1,15 @@
 package io.dongtai.iast.common.scope;
 
+import java.util.ArrayDeque;
+import java.util.Deque;
+
 public class PolicyScope {
     private int agentLevel;
     private int sourceLevel;
     private int propagatorLevel;
     private int propagatorSkipDepth;
     private int sinkLevel;
+    private final Deque<String> sinkQueue = new ArrayDeque<>();
     private int ignoreInternalLevel;
     /**
      * over max method pool size
@@ -66,10 +70,11 @@ public void enterSink() {
     public boolean isValidSink() {
         return this.agentLevel == 0
                 && this.ignoreInternalLevel == 0 && !this.overCapacity && this.sourceLevel == 0
-                && this.sinkLevel == 1;
+                && this.sinkLevel > 0;
     }
 
     public void leaveSink() {
+        this.sinkQueue.pop();
         this.sinkLevel = decrement(this.sinkLevel);
     }
 
@@ -95,4 +100,12 @@ private int decrement(int level) {
         }
         return 0;
     }
+
+    public Deque<String> getSinkQueue() {
+        return sinkQueue;
+    }
+
+    public void addSinkType(String sinkType) {
+        this.sinkQueue.push(sinkType);
+    }
 }

dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/SpyDispatcherImpl.java

@@ -541,7 +541,7 @@ public boolean isFirstLevelSink() {
                 return false;
             }
             return ScopeManager.SCOPE_TRACKER.inEnterEntry()
-                    && ScopeManager.SCOPE_TRACKER.getPolicyScope().isValidSink();
+                    && (ScopeManager.SCOPE_TRACKER.getPolicyScope().isValidSink());
         } catch (Throwable ignore) {
             return false;
         }

dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/SinkImpl.java

@@ -1,5 +1,6 @@
 package io.dongtai.iast.core.handler.hookpoint.controller.impl;
 
+import io.dongtai.iast.common.scope.ScopeManager;
 import io.dongtai.iast.core.EngineManager;
 import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.SinkNode;
@@ -24,6 +25,11 @@ public static void solveSink(MethodEvent event, SinkNode sinkNode) {
             return;
         }
 
+        if(ScopeManager.SCOPE_TRACKER.getPolicyScope().getSinkQueue().contains(sinkNode.getVulType())){
+            ScopeManager.SCOPE_TRACKER.getPolicyScope().addSinkType(sinkNode.getVulType());
+            return;
+        }
+
         String vulType = sinkNode.getVulType();
         if (VulnType.CRYPTO_WEAK_RANDOMNESS.equals(vulType)) {
             new CryptoWeakRandomnessVulScan().scan(event, sinkNode);

dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/dynamic/DynamicPropagatorScanner.java

@@ -1,5 +1,6 @@
 package io.dongtai.iast.core.handler.hookpoint.vulscan.dynamic;
 
+import io.dongtai.iast.common.scope.ScopeManager;
 import io.dongtai.iast.core.EngineManager;
 import io.dongtai.iast.core.handler.hookpoint.SpyDispatcherImpl;
 import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
@@ -113,6 +114,7 @@ public void scan(MethodEvent event, SinkNode sinkNode) {
             event.setStacks(stackTraceElements);
 
             EngineManager.TRACK_MAP.addTrackMethod(invokeId, event);
+            ScopeManager.SCOPE_TRACKER.getPolicyScope().addSinkType(sinkNode.getVulType());
         }
     }
 

dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CookieFlagsMissingVulScan.java

@@ -1,5 +1,6 @@
 package io.dongtai.iast.core.handler.hookpoint.vulscan.normal;
 
+import io.dongtai.iast.common.scope.ScopeManager;
 import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.SinkNode;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.TaintPosition;
@@ -34,6 +35,7 @@ public void scan(MethodEvent event, SinkNode sinkNode) {
                     continue;
                 }
                 sendReport(getLatestStack(), sinkNode.getVulType());
+                ScopeManager.SCOPE_TRACKER.getPolicyScope().addSinkType(sinkNode.getVulType());
                 break;
             } catch (Throwable e) {
                 DongTaiLog.trace("CookieFlagsMissingVulScan scan failed: {}, {}",

dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CryptoBacCiphersVulScan.java

@@ -1,5 +1,6 @@
 package io.dongtai.iast.core.handler.hookpoint.vulscan.normal;
 
+import io.dongtai.iast.common.scope.ScopeManager;
 import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.SinkNode;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.TaintPosition;
@@ -41,6 +42,7 @@ public void scan(MethodEvent event, SinkNode sinkNode) {
                     continue;
                 }
                 sendReport(getLatestStack(), sinkNode.getVulType());
+                ScopeManager.SCOPE_TRACKER.getPolicyScope().addSinkType(sinkNode.getVulType());
                 break;
             } catch (Throwable e) {
                 DongTaiLog.trace("CryptoBacCiphersVulScan scan failed: {}, {}",

dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CryptoBadMacVulScan.java

@@ -1,5 +1,6 @@
 package io.dongtai.iast.core.handler.hookpoint.vulscan.normal;
 
+import io.dongtai.iast.common.scope.ScopeManager;
 import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.SinkNode;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.TaintPosition;
@@ -48,6 +49,7 @@ public void scan(MethodEvent event, SinkNode sinkNode) {
                     }
                 }
                 sendReport(latestStack, sinkNode.getVulType());
+                ScopeManager.SCOPE_TRACKER.getPolicyScope().addSinkType(sinkNode.getVulType());
                 break;
             } catch (Throwable e) {
                 DongTaiLog.trace("CryptoBadMacVulScan scan failed: {}, {}",

dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CryptoWeakRandomnessVulScan.java

@@ -1,5 +1,6 @@
 package io.dongtai.iast.core.handler.hookpoint.vulscan.normal;
 
+import io.dongtai.iast.common.scope.ScopeManager;
 import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.SinkNode;
 
@@ -22,5 +23,6 @@ public void scan(MethodEvent event, SinkNode sinkNode) {
             return;
         }
         sendReport(getLatestStack(), sinkNode.getVulType());
+        ScopeManager.SCOPE_TRACKER.getPolicyScope().addSinkType(sinkNode.getVulType());
     }
 }