Skip to content

Commit ac33750

Browse files
author
‘niuerzhuang’
committed
feature: add validator handler.
1 parent c5817d5 commit ac33750

File tree

23 files changed

+325
-15
lines changed

23 files changed

+325
-15
lines changed

‎dongtai-common/src/main/java/io/dongtai/iast/common/constants/PropertyConstant.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,5 +32,5 @@ public class PropertyConstant {
3232
public static final String PROPERTY_POLICY_PATH = "dongtai.policy.path";
3333
public static final String PROPERTY_UUID_PATH = "dongtai.uuid.path";
3434
public static final String PROPERTY_DISABLED_PLUGINS = "dongtai.disabled.plugins";
35-
public static final String PROPERTY_DISABLED_FEATURES = "dongtai.disabled_features";
35+
public static final String PROPERTY_DISABLED_FEATURES = "dongtai.disabled.features";
3636
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/asm/AsmMethods.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,12 @@ static Method getAsmMethod(final Class<?> clazz,
219219
SpyDispatcher.class,
220220
"isFirstLevelSink"
221221
);
222+
223+
Method SPY$enterValidator = InnerHelper.getAsmMethod(
224+
SpyDispatcher.class,
225+
"enterValidator"
226+
);
227+
222228
Method SPY$enterIgnoreInternal = InnerHelper.getAsmMethod(
223229
SpyDispatcher.class,
224230
"enterIgnoreInternal"

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/DispatchClassPlugin.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ public class ClassVisit extends AbstractClassVisitor {
5454
new SourceAdapter(),
5555
new PropagatorAdapter(),
5656
new SinkAdapter(),
57+
new ValidatorAdapter(),
5758
};
5859
}
5960

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/adapter/ValidatorAdapter.java‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
package io.dongtai.iast.core.bytecode.enhance.plugin.core.adapter;
2+
3+
import io.dongtai.iast.core.bytecode.enhance.MethodContext;
4+
import io.dongtai.iast.core.handler.hookpoint.models.policy.PolicyNode;
5+
import io.dongtai.iast.core.handler.hookpoint.models.policy.ValidatorNode;
6+
import org.objectweb.asm.Label;
7+
import org.objectweb.asm.MethodVisitor;
8+
import org.objectweb.asm.Opcodes;
9+
10+
import java.util.Set;
11+
12+
public class ValidatorAdapter extends MethodAdapter {
13+
/**
14+
* @param adapter
15+
* @param mv
16+
* @param context
17+
* @param policyNodes
18+
*/
19+
@Override
20+
public void onMethodEnter(MethodAdviceAdapter adapter, MethodVisitor mv, MethodContext context, Set<PolicyNode> policyNodes) {
21+
}
22+
23+
/**
24+
* @param adapter
25+
* @param mv
26+
* @param opcode
27+
* @param context
28+
* @param policyNodes
29+
*/
30+
@Override
31+
public void onMethodExit(MethodAdviceAdapter adapter, MethodVisitor mv, int opcode, MethodContext context, Set<PolicyNode> policyNodes) {
32+
for (PolicyNode policyNode : policyNodes) {
33+
if (!(policyNode instanceof ValidatorNode)) {
34+
continue;
35+
}
36+
37+
Label elseLabel = new Label();
38+
Label endLabel = new Label();
39+
40+
isEnterScope(adapter);
41+
mv.visitJumpInsn(Opcodes.IFEQ, elseLabel);
42+
43+
adapter.trackMethod(opcode, policyNode, true);
44+
45+
adapter.mark(elseLabel);
46+
adapter.mark(endLabel);
47+
}
48+
}
49+
50+
private void isEnterScope(MethodAdviceAdapter adapter) {
51+
adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
52+
adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterValidator);
53+
}
54+
}

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

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -394,7 +394,8 @@ public void collectDubboResponse(Object result, byte status) {
394394
}
395395

396396
if (!ScopeManager.SCOPE_TRACKER.getScope(Scope.DUBBO_REQUEST).isFirst()
397-
|| !ScopeManager.SCOPE_TRACKER.getScope(Scope.DUBBO_ENTRY).in()) {
397+
|| !ScopeManager.SCOPE_TRACKER.getScope(Scope.DUBBO_ENTRY).in()
398+
|| ScopeManager.SCOPE_TRACKER.getScope(Scope.HTTP_REQUEST).in()) {
398399
return;
399400
}
400401

@@ -558,6 +559,17 @@ public void leaveSink() {
558559
}
559560
}
560561

562+
/**
563+
* mark for enter validator entry point
564+
*/
565+
@Override
566+
public boolean enterValidator() {
567+
if (!EngineManager.isEngineRunning()) {
568+
return false;
569+
}
570+
return !ScopeManager.SCOPE_TRACKER.inAgent() && ScopeManager.SCOPE_TRACKER.inEnterEntry();
571+
}
572+
561573
/**
562574
* Determines whether it is a layer 1 Sink entry
563575
*
@@ -674,6 +686,9 @@ public boolean collectMethod(Object instance, Object[] parameters, Object retObj
674686
} else if ((policyNode instanceof SinkNode)) {
675687
SinkImpl.solveSink(event, (SinkNode) policyNode);
676688
return true;
689+
} else if ((policyNode instanceof ValidatorNode)) {
690+
ValidatorImpl.solveValidator(event,(ValidatorNode)policyNode, INVOKE_ID_SEQUENCER);
691+
return true;
677692
}
678693

679694
return false;
@@ -731,7 +746,7 @@ public boolean traceDubboInvoke(Object instance, String url, Object invocation,
731746
@Override
732747
public boolean isSkipCollectDubbo(Object invocation) {
733748
if (BlackUrlBypass.isBlackUrl()) {
734-
Method setAttachmentMethod = null;
749+
Method setAttachmentMethod;
735750
try {
736751
setAttachmentMethod = invocation.getClass().getMethod("setAttachment", String.class, String.class);
737752
setAttachmentMethod.setAccessible(true);
@@ -746,7 +761,7 @@ public boolean isSkipCollectDubbo(Object invocation) {
746761
@Override
747762
public boolean isSkipCollectFeign(Object instance) {
748763
if (BlackUrlBypass.isBlackUrl()) {
749-
Field metadataField = null;
764+
Field metadataField;
750765
try {
751766
metadataField = instance.getClass().getDeclaredField("metadata");
752767
metadataField.setAccessible(true);

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import io.dongtai.iast.core.handler.context.ContextManager;
99
import io.dongtai.iast.core.handler.hookpoint.IastClassLoader;
1010
import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
11+
import io.dongtai.iast.core.handler.hookpoint.models.policy.PolicyNodeType;
1112
import io.dongtai.iast.core.handler.hookpoint.models.policy.SourceNode;
1213
import io.dongtai.iast.core.handler.hookpoint.models.policy.TaintPosition;
1314
import io.dongtai.iast.core.handler.hookpoint.models.taint.range.TaintRange;
@@ -178,6 +179,7 @@ public static void collectDubboRequestSource(Object handler, Object invocation,
178179

179180
int invokeId = invokeIdSequencer.getAndIncrement();
180181
event.setInvokeId(invokeId);
182+
event.setPolicyType(PolicyNodeType.SOURCE.getName());
181183

182184
event.source = true;
183185
event.setCallStacks(StackUtils.createCallStack(4));

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import io.dongtai.iast.core.EngineManager;
44
import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
5+
import io.dongtai.iast.core.handler.hookpoint.models.policy.PolicyNodeType;
56
import io.dongtai.iast.core.handler.hookpoint.models.policy.PropagatorNode;
67
import io.dongtai.iast.core.handler.hookpoint.models.policy.TaintPosition;
78
import io.dongtai.iast.core.handler.hookpoint.models.taint.range.*;
@@ -63,6 +64,7 @@ private static void addPropagator(PropagatorNode propagatorNode, MethodEvent eve
6364
event.setCallStacks(StackUtils.createCallStack(6));
6465
int invokeId = invokeIdSequencer.getAndIncrement();
6566
event.setInvokeId(invokeId);
67+
event.setPolicyType(PolicyNodeType.PROPAGATOR.getName());
6668
EngineManager.TRACK_MAP.get().put(invokeId, event);
6769
}
6870

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import io.dongtai.iast.core.EngineManager;
44
import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
5+
import io.dongtai.iast.core.handler.hookpoint.models.policy.PolicyNodeType;
56
import io.dongtai.iast.core.handler.hookpoint.models.policy.SourceNode;
67
import io.dongtai.iast.core.handler.hookpoint.models.policy.TaintPosition;
78
import io.dongtai.iast.core.handler.hookpoint.models.taint.range.TaintRangesBuilder;
@@ -37,6 +38,7 @@ public static void solveSource(MethodEvent event, SourceNode sourceNode, AtomicI
3738

3839
int invokeId = invokeIdSequencer.getAndIncrement();
3940
event.setInvokeId(invokeId);
41+
event.setPolicyType(PolicyNodeType.SOURCE.getName());
4042

4143
boolean valid = trackTarget(event, sourceNode);
4244
if (!valid) {

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

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
package io.dongtai.iast.core.handler.hookpoint.controller.impl;
2+
3+
import io.dongtai.iast.core.EngineManager;
4+
import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
5+
import io.dongtai.iast.core.handler.hookpoint.models.policy.PolicyNodeType;
6+
import io.dongtai.iast.core.handler.hookpoint.models.policy.TaintPosition;
7+
import io.dongtai.iast.core.handler.hookpoint.models.policy.ValidatorNode;
8+
import io.dongtai.iast.core.handler.hookpoint.models.taint.range.TaintRange;
9+
import io.dongtai.iast.core.handler.hookpoint.models.taint.range.TaintRanges;
10+
import io.dongtai.iast.core.handler.hookpoint.models.taint.range.TaintRangesBuilder;
11+
import io.dongtai.iast.core.utils.StackUtils;
12+
import io.dongtai.iast.core.utils.TaintPoolUtils;
13+
14+
import java.util.Set;
15+
import java.util.concurrent.atomic.AtomicInteger;
16+
17+
import static io.dongtai.iast.core.utils.TaintPoolUtils.getStringHash;
18+
19+
public class ValidatorImpl {
20+
21+
/**
22+
* 处理 Validator 点的事件
23+
*
24+
* @param event Validator 点事件
25+
*/
26+
public static void solveValidator(MethodEvent event, ValidatorNode validatorNode, AtomicInteger invokeIdSequencer) {
27+
if (EngineManager.TAINT_HASH_CODES.isEmpty()) {
28+
return;
29+
}
30+
Set<TaintPosition> sources = validatorNode.getSources();
31+
if (sources.isEmpty()) {
32+
return;
33+
}
34+
35+
for (TaintPosition position : sources) {
36+
Long hash = null;
37+
Integer len = null;
38+
if (position.isObject()) {
39+
if (TaintPoolUtils.isNotEmpty(event.objectInstance)
40+
&& TaintPoolUtils.isAllowTaintType(event.objectInstance)
41+
&& TaintPoolUtils.poolContains(event.objectInstance, event)) {
42+
hash = getStringHash(event.objectInstance);
43+
len = TaintRangesBuilder.getLength(event.objectInstance);
44+
}
45+
} else if (position.isParameter()) {
46+
int parameterIndex = position.getParameterIndex();
47+
if (parameterIndex >= event.parameterInstances.length) {
48+
continue;
49+
}
50+
Object parameter = event.parameterInstances[parameterIndex];
51+
if (TaintPoolUtils.isNotEmpty(parameter)
52+
&& TaintPoolUtils.isAllowTaintType(parameter)
53+
&& TaintPoolUtils.poolContains(parameter, event)) {
54+
hash = getStringHash(parameter);
55+
len = TaintRangesBuilder.getLength(parameter);
56+
}
57+
}
58+
59+
if (null != len && null != hash){
60+
TaintRanges tr = new TaintRanges(new TaintRange("validated", 0, len));
61+
if (validatorNode.hasTags()) {
62+
String[] tags = validatorNode.getTags();
63+
for (String tag : tags) {
64+
tr.add(new TaintRange(tag, 0, len));
65+
}
66+
}
67+
event.sourceRanges.add(new MethodEvent.MethodEventTargetRange(hash, tr));
68+
TaintRanges taintRanges = EngineManager.TAINT_RANGES_POOL.get().get(hash);
69+
if (null == taintRanges){
70+
EngineManager.TAINT_RANGES_POOL.add(hash, tr);
71+
}else {
72+
taintRanges.addAll(tr);
73+
}
74+
}
75+
}
76+
77+
event.source = false;
78+
event.setCallStacks(StackUtils.createCallStack(4));
79+
80+
int invokeId = invokeIdSequencer.getAndIncrement();
81+
event.setInvokeId(invokeId);
82+
event.setPolicyType(PolicyNodeType.VALIDATOR.getName());
83+
EngineManager.TRACK_MAP.addTrackMethod(invokeId, event);
84+
}
85+
86+
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/graphy/GraphBuilder.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,7 @@ public static JSONObject toJson(MethodEvent event) {
110110
List<String> targetPositions = new ArrayList<String>();
111111

112112
value.put("invokeId", event.getInvokeId());
113+
value.put("policyType", event.getPolicyType());
113114
value.put("source", event.isSource());
114115
value.put("originClassName", event.getOriginClassName());
115116
value.put("className", event.getMatchedClassName());
@@ -167,6 +168,14 @@ public static JSONObject toJson(MethodEvent event) {
167168
}
168169
}
169170

171+
if (event.sourceRanges.size() > 0) {
172+
JSONArray tr = new JSONArray();
173+
value.put("sourceRange", tr);
174+
for (MethodEvent.MethodEventTargetRange range : event.sourceRanges) {
175+
tr.add(range.toJson());
176+
}
177+
}
178+
170179
if (event.sourceTypes != null && event.sourceTypes.size() > 0) {
171180
JSONArray st = new JSONArray();
172181
value.put("sourceType", st);

0 commit comments

Comments
 (0)