diff --git a/dongtai-common/src/main/java/io/dongtai/iast/common/scope/PolicyScope.java b/dongtai-common/src/main/java/io/dongtai/iast/common/scope/PolicyScope.java index a49787c0c..b35eb864c 100644 --- a/dongtai-common/src/main/java/io/dongtai/iast/common/scope/PolicyScope.java +++ b/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 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 getSinkQueue() { + return sinkQueue; + } + + public void addSinkType(String sinkType) { + this.sinkQueue.push(sinkType); + } } diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/SpyDispatcherImpl.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/SpyDispatcherImpl.java index 8462f4c3b..59a383362 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/SpyDispatcherImpl.java +++ b/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; } diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/SinkImpl.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/SinkImpl.java index 143675eed..281bff3d5 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/SinkImpl.java +++ b/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); diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/dynamic/DynamicPropagatorScanner.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/dynamic/DynamicPropagatorScanner.java index 3104929a2..53d61d12a 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/dynamic/DynamicPropagatorScanner.java +++ b/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()); } } diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CookieFlagsMissingVulScan.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CookieFlagsMissingVulScan.java index 2d0f8fabd..18fb35ec1 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CookieFlagsMissingVulScan.java +++ b/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: {}, {}", diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CryptoBacCiphersVulScan.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CryptoBacCiphersVulScan.java index 1cdb5c9e7..6474008d7 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CryptoBacCiphersVulScan.java +++ b/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: {}, {}", diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CryptoBadMacVulScan.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CryptoBadMacVulScan.java index 97219c3e1..4bea90a98 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CryptoBadMacVulScan.java +++ b/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: {}, {}", diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CryptoWeakRandomnessVulScan.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CryptoWeakRandomnessVulScan.java index 67cb30ec0..0d05e5bf5 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/vulscan/normal/CryptoWeakRandomnessVulScan.java +++ b/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()); } }