Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fixes policy node map key #407

Merged
merged 1 commit into from
Nov 8, 2022
Merged

fixes policy node map key #407

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ...ore/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/AbstractAdviceAdapter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
31 changes: 7 additions & 24 deletions .../src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/DispatchClassPlugin.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -97,29 +98,11 @@ private MethodVisitor lazyAop(MethodVisitor mv, int access, String name, String
MethodContext methodContext) {
Set<PolicyNode> matchedNodes = new HashSet<PolicyNode>();

List<SourceNode> sourceNodes = this.policy.getSources();
if (sourceNodes != null && sourceNodes.size() != 0) {
for (SourceNode sourceNode : sourceNodes) {
if (sourceNode.getMethodMatcher().match(methodContext)) {
matchedNodes.add(sourceNode);
}
}
}

List<PropagatorNode> propagatorNodes = this.policy.getPropagators();
if (sourceNodes != null && sourceNodes.size() != 0) {
for (PropagatorNode propagatorNode : propagatorNodes) {
if (propagatorNode.getMethodMatcher().match(methodContext)) {
matchedNodes.add(propagatorNode);
}
}
}

List<SinkNode> sinkNodes = this.policy.getSinks();
if (sourceNodes != null && sourceNodes.size() != 0) {
for (SinkNode sinkNode : sinkNodes) {
if (sinkNode.getMethodMatcher().match(methodContext)) {
matchedNodes.add(sinkNode);
Map<String, PolicyNode> policyNodesMap = this.policy.getPolicyNodesMap();
if (policyNodesMap != null && policyNodesMap.size() != 0) {
for (Map.Entry<String, PolicyNode> entry : policyNodesMap.entrySet()) {
if (entry.getValue().getMethodMatcher().match(methodContext)) {
matchedNodes.add(entry.getValue());
}
}
}
Expand Down
8 changes: 4 additions & 4 deletions dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/SpyDispatcherImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -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) {
Expand All @@ -499,6 +499,6 @@ private PolicyNode getPolicyNode(String methodMatcher) {
return null;
}

return policy.getPolicyNode(methodMatcher);
return policy.getPolicyNode(policyKey);
}
}
22 changes: 15 additions & 7 deletions dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/policy/Policy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ public List<SourceNode> getSources() {

public void addSource(SourceNode source) {
this.sources.add(source);
addHooks(source);
addPolicyNode(source);
}

public List<PropagatorNode> getPropagators() {
Expand All @@ -25,7 +25,7 @@ public List<PropagatorNode> getPropagators() {

public void addPropagator(PropagatorNode propagator) {
this.propagators.add(propagator);
addHooks(propagator);
addPolicyNode(propagator);
}

public List<SinkNode> getSinks() {
Expand All @@ -34,22 +34,22 @@ public List<SinkNode> 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<String, PolicyNode> 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());
}
}
Expand Down Expand Up @@ -78,4 +78,12 @@ public String getMatchedClass(String className, Set<String> ancestors) {
public boolean isMatchClass(String className) {
return this.classHooks.contains(className) || this.ancestorClassHooks.contains(className);
}

public Set<String> getClassHooks() {
return this.classHooks;
}

public Set<String> getAncestorClassHooks() {
return this.ancestorClassHooks;
}
}
0 ...policy-node-count-src0-p2-sink2-cls2.json → ...ode-count-src0-p2-sink2-policy4-cls2.json
View file
Open in desktop
File renamed without changes.
0 ...policy-node-count-src1-p3-sink1-cls5.json → ...ode-count-src1-p3-sink1-policy5-cls4.json
View file
Open in desktop
File renamed without changes.
10 changes: 7 additions & 3 deletions ...src/test/java/io/dongtai/iast/core/handler/hookpoint/models/policy/PolicyBuilderTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,8 +63,8 @@ public void run() throws PolicyException {
@Test
public void testBuild() throws PolicyException {
Map<String, List<Integer>> tests = new HashMap<String, List<Integer>>() {{
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<String, List<Integer>> entry : tests.entrySet()) {
JSONArray policyConfig = PolicyBuilder.fetchFromFile(POLICY_DIR + entry.getKey());
Expand All @@ -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<String> classes = policy.getClassHooks();
classes.addAll(policy.getAncestorClassHooks());
Assert.assertEquals("build hook class count" + entry.getKey(), entry.getValue().get(4).intValue(),
classes.size());
}

PolicyException exception;
Expand Down