From e43de92d3eb28ce87e2b37690da6ca2c4c2199ad Mon Sep 17 00:00:00 2001 From: lostsnow Date: Mon, 7 Nov 2022 19:18:08 +0800 Subject: [PATCH] fixes policy node map key --- .../enhance/plugin/AbstractAdviceAdapter.java | 2 +- .../plugin/core/DispatchClassPlugin.java | 31 +++++-------------- .../handler/hookpoint/SpyDispatcherImpl.java | 8 ++--- .../hookpoint/models/policy/Policy.java | 22 ++++++++----- ...ode-count-src0-p2-sink2-policy4-cls2.json} | 0 ...ode-count-src1-p3-sink1-policy5-cls4.json} | 0 .../models/policy/PolicyBuilderTest.java | 10 ++++-- 7 files changed, 34 insertions(+), 39 deletions(-) rename dongtai-core/src/test/fixture/policy/{policy-node-count-src0-p2-sink2-cls2.json => policy-node-count-src0-p2-sink2-policy4-cls2.json} (100%) rename dongtai-core/src/test/fixture/policy/{policy-node-count-src1-p3-sink1-cls5.json => policy-node-count-src1-p3-sink1-policy5-cls4.json} (100%) diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/AbstractAdviceAdapter.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/AbstractAdviceAdapter.java index 334dc61dd..4caba41fc 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/AbstractAdviceAdapter.java +++ b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/AbstractAdviceAdapter.java @@ -136,7 +136,7 @@ public void trackMethod( loadThisOrPushNullIfIsStatic(); loadArgArray(); loadLocal(this.nextLocal - 1); - push(policyNode.getMethodMatcher().toString()); + push(policyNode.toString()); push(this.context.getClassName()); push(this.context.getMatchedClassName()); push(this.name); diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/DispatchClassPlugin.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/DispatchClassPlugin.java index 01d7ef914..8ba66807f 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/DispatchClassPlugin.java +++ b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/DispatchClassPlugin.java @@ -5,7 +5,8 @@ import io.dongtai.iast.core.bytecode.enhance.plugin.AbstractClassVisitor; import io.dongtai.iast.core.bytecode.enhance.plugin.DispatchPlugin; import io.dongtai.iast.core.bytecode.enhance.plugin.core.adapter.*; -import io.dongtai.iast.core.handler.hookpoint.models.policy.*; +import io.dongtai.iast.core.handler.hookpoint.models.policy.Policy; +import io.dongtai.iast.core.handler.hookpoint.models.policy.PolicyNode; import io.dongtai.iast.core.utils.AsmUtils; import io.dongtai.log.DongTaiLog; import org.objectweb.asm.ClassVisitor; @@ -97,29 +98,11 @@ private MethodVisitor lazyAop(MethodVisitor mv, int access, String name, String MethodContext methodContext) { Set matchedNodes = new HashSet(); - List sourceNodes = this.policy.getSources(); - if (sourceNodes != null && sourceNodes.size() != 0) { - for (SourceNode sourceNode : sourceNodes) { - if (sourceNode.getMethodMatcher().match(methodContext)) { - matchedNodes.add(sourceNode); - } - } - } - - List propagatorNodes = this.policy.getPropagators(); - if (sourceNodes != null && sourceNodes.size() != 0) { - for (PropagatorNode propagatorNode : propagatorNodes) { - if (propagatorNode.getMethodMatcher().match(methodContext)) { - matchedNodes.add(propagatorNode); - } - } - } - - List sinkNodes = this.policy.getSinks(); - if (sourceNodes != null && sourceNodes.size() != 0) { - for (SinkNode sinkNode : sinkNodes) { - if (sinkNode.getMethodMatcher().match(methodContext)) { - matchedNodes.add(sinkNode); + Map policyNodesMap = this.policy.getPolicyNodesMap(); + if (policyNodesMap != null && policyNodesMap.size() != 0) { + for (Map.Entry entry : policyNodesMap.entrySet()) { + if (entry.getValue().getMethodMatcher().match(methodContext)) { + matchedNodes.add(entry.getValue()); } } } 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 4c951a8ed..68210c83a 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 @@ -429,12 +429,12 @@ private void solveRPC(String framework, MethodEvent event) { } @Override - public boolean collectMethod(Object instance, Object[] parameters, Object retObject, String methodMatcher, + public boolean collectMethod(Object instance, Object[] parameters, Object retObject, String policyKey, String className, String matchedClassName, String methodName, String signature, boolean isStatic) { try { ScopeManager.SCOPE_TRACKER.getPolicyScope().enterAgent(); - PolicyNode policyNode = getPolicyNode(methodMatcher); + PolicyNode policyNode = getPolicyNode(policyKey); if (policyNode == null) { return false; } @@ -488,7 +488,7 @@ private boolean isCollectAllowed(String className, String methodName, String sig return true; } - private PolicyNode getPolicyNode(String methodMatcher) { + private PolicyNode getPolicyNode(String policyKey) { AgentEngine agentEngine = AgentEngine.getInstance(); PolicyManager policyManager = agentEngine.getPolicyManager(); if (policyManager == null) { @@ -499,6 +499,6 @@ private PolicyNode getPolicyNode(String methodMatcher) { return null; } - return policy.getPolicyNode(methodMatcher); + return policy.getPolicyNode(policyKey); } } diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/policy/Policy.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/policy/Policy.java index 5e36f2cc5..368d6b377 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/policy/Policy.java +++ b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/policy/Policy.java @@ -16,7 +16,7 @@ public List getSources() { public void addSource(SourceNode source) { this.sources.add(source); - addHooks(source); + addPolicyNode(source); } public List getPropagators() { @@ -25,7 +25,7 @@ public List getPropagators() { public void addPropagator(PropagatorNode propagator) { this.propagators.add(propagator); - addHooks(propagator); + addPolicyNode(propagator); } public List getSinks() { @@ -34,22 +34,22 @@ public List getSinks() { public void addSink(SinkNode sink) { this.sinks.add(sink); - addHooks(sink); + addPolicyNode(sink); } - public PolicyNode getPolicyNode(String methodMatcher) { - return this.policyNodesMap.get(methodMatcher); + public PolicyNode getPolicyNode(String policyKey) { + return this.policyNodesMap.get(policyKey); } public Map getPolicyNodesMap() { return this.policyNodesMap; } - public void addHooks(PolicyNode node) { + public void addPolicyNode(PolicyNode node) { SignatureMethodMatcher methodMatcher; if (node.getMethodMatcher() instanceof SignatureMethodMatcher) { methodMatcher = (SignatureMethodMatcher) node.getMethodMatcher(); - this.policyNodesMap.put(methodMatcher.toString(), node); + this.policyNodesMap.put(node.toString(), node); addHooks(methodMatcher.getSignature().getClassName(), node.getInheritable()); } } @@ -78,4 +78,12 @@ public String getMatchedClass(String className, Set ancestors) { public boolean isMatchClass(String className) { return this.classHooks.contains(className) || this.ancestorClassHooks.contains(className); } + + public Set getClassHooks() { + return this.classHooks; + } + + public Set getAncestorClassHooks() { + return this.ancestorClassHooks; + } } diff --git a/dongtai-core/src/test/fixture/policy/policy-node-count-src0-p2-sink2-cls2.json b/dongtai-core/src/test/fixture/policy/policy-node-count-src0-p2-sink2-policy4-cls2.json similarity index 100% rename from dongtai-core/src/test/fixture/policy/policy-node-count-src0-p2-sink2-cls2.json rename to dongtai-core/src/test/fixture/policy/policy-node-count-src0-p2-sink2-policy4-cls2.json diff --git a/dongtai-core/src/test/fixture/policy/policy-node-count-src1-p3-sink1-cls5.json b/dongtai-core/src/test/fixture/policy/policy-node-count-src1-p3-sink1-policy5-cls4.json similarity index 100% rename from dongtai-core/src/test/fixture/policy/policy-node-count-src1-p3-sink1-cls5.json rename to dongtai-core/src/test/fixture/policy/policy-node-count-src1-p3-sink1-policy5-cls4.json diff --git a/dongtai-core/src/test/java/io/dongtai/iast/core/handler/hookpoint/models/policy/PolicyBuilderTest.java b/dongtai-core/src/test/java/io/dongtai/iast/core/handler/hookpoint/models/policy/PolicyBuilderTest.java index e73c93644..59fa16cde 100644 --- a/dongtai-core/src/test/java/io/dongtai/iast/core/handler/hookpoint/models/policy/PolicyBuilderTest.java +++ b/dongtai-core/src/test/java/io/dongtai/iast/core/handler/hookpoint/models/policy/PolicyBuilderTest.java @@ -63,8 +63,8 @@ public void run() throws PolicyException { @Test public void testBuild() throws PolicyException { Map> tests = new HashMap>() {{ - put("policy-node-count-src1-p3-sink1-cls5.json", Arrays.asList(1, 3, 1, 5)); - put("policy-node-count-src0-p2-sink2-cls2.json", Arrays.asList(0, 2, 2, 2)); + put("policy-node-count-src0-p2-sink2-policy4-cls2.json", Arrays.asList(0, 2, 2, 4, 2)); + put("policy-node-count-src1-p3-sink1-policy5-cls4.json", Arrays.asList(1, 3, 1, 5, 4)); }}; for (Map.Entry> entry : tests.entrySet()) { JSONArray policyConfig = PolicyBuilder.fetchFromFile(POLICY_DIR + entry.getKey()); @@ -75,8 +75,12 @@ public void testBuild() throws PolicyException { policy.getPropagators().size()); Assert.assertEquals("build sink count " + entry.getKey(), entry.getValue().get(2).intValue(), policy.getSinks().size()); - Assert.assertEquals("build hook class count" + entry.getKey(), entry.getValue().get(3).intValue(), + Assert.assertEquals("build hook policy count" + entry.getKey(), entry.getValue().get(3).intValue(), policy.getPolicyNodesMap().size()); + Set classes = policy.getClassHooks(); + classes.addAll(policy.getAncestorClassHooks()); + Assert.assertEquals("build hook class count" + entry.getKey(), entry.getValue().get(4).intValue(), + classes.size()); } PolicyException exception;