Skip to content

Commit 960d8f2

Browse files
lostsnowlostsnow
committed
add dubbo consumer invoke track
1 parent a2f01fd commit 960d8f2

File tree

15 files changed

+393
-179
lines changed

15 files changed

+393
-179
lines changed

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

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,7 @@
44

55
import java.lang.dongtai.SpyDispatcher;
66
import java.lang.dongtai.SpyDispatcherHandler;
7-
import java.util.Collection;
8-
import java.util.Enumeration;
7+
import java.util.*;
98

109
/**
1110
* 常用的ASM method 集合 省得我到处声明
@@ -209,6 +208,19 @@ static Method getAsmMethod(final Class<?> clazz,
209208
String.class
210209
);
211210

211+
Method SPY$traceDubboInvoke = InnerHelper.getAsmMethod(
212+
SpyDispatcher.class,
213+
"traceDubboInvoke",
214+
Object.class,
215+
String.class,
216+
Object.class,
217+
Object[].class,
218+
Map.class,
219+
String.class,
220+
String.class,
221+
String.class
222+
);
223+
212224
Method SPY$reportService = InnerHelper.getAsmMethod(
213225
SpyDispatcher.class,
214226
"reportService",

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import io.dongtai.iast.core.bytecode.enhance.ClassContext;
44
import io.dongtai.iast.core.bytecode.enhance.plugin.authentication.shiro.DispatchShiro;
55
import io.dongtai.iast.core.bytecode.enhance.plugin.core.DispatchClassPlugin;
6+
import io.dongtai.iast.core.bytecode.enhance.plugin.framework.dubbo.DispatchDubbo;
67
import io.dongtai.iast.core.bytecode.enhance.plugin.framework.feign.DispatchFeign;
78
import io.dongtai.iast.core.bytecode.enhance.plugin.framework.j2ee.dispatch.DispatchJ2ee;
89
import io.dongtai.iast.core.bytecode.enhance.plugin.hardcoded.DispatchHardcodedPlugin;
@@ -34,6 +35,7 @@ public PluginRegister() {
3435
this.plugins.add(new DispatchJdbc());
3536
this.plugins.add(new DispatchShiro());
3637
this.plugins.add(new DispatchFeign());
38+
this.plugins.add(new DispatchDubbo());
3739

3840
this.plugins.add(new DispatchClassPlugin());
3941
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/DispatchDubbo.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
package io.dongtai.iast.core.bytecode.enhance.plugin.framework.dubbo;
2+
3+
import io.dongtai.iast.core.bytecode.enhance.ClassContext;
4+
import io.dongtai.iast.core.bytecode.enhance.plugin.DispatchPlugin;
5+
import io.dongtai.iast.core.handler.hookpoint.models.policy.Policy;
6+
import org.objectweb.asm.ClassVisitor;
7+
8+
public class DispatchDubbo implements DispatchPlugin {
9+
public static final String LEGACY_DUBBO_SYNC_HANDLER = " com.alibaba.dubbo.rpc.listener.ListenerInvokerWrapper".substring(1);
10+
11+
@Override
12+
public ClassVisitor dispatch(ClassVisitor classVisitor, ClassContext context, Policy policy) {
13+
String className = context.getClassName();
14+
15+
if (LEGACY_DUBBO_SYNC_HANDLER.equals(className)) {
16+
classVisitor = new LegacyDubboSyncHandlerAdapter(classVisitor, context);
17+
}
18+
19+
return classVisitor;
20+
}
21+
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/LegacyDubboSyncHandlerAdapter.java‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package io.dongtai.iast.core.bytecode.enhance.plugin.framework.dubbo;
2+
3+
import io.dongtai.iast.core.bytecode.enhance.ClassContext;
4+
import io.dongtai.iast.core.bytecode.enhance.plugin.AbstractClassVisitor;
5+
import io.dongtai.iast.core.utils.AsmUtils;
6+
import io.dongtai.log.DongTaiLog;
7+
import org.objectweb.asm.ClassVisitor;
8+
import org.objectweb.asm.MethodVisitor;
9+
10+
public class LegacyDubboSyncHandlerAdapter extends AbstractClassVisitor {
11+
private static final String LEGACY_DUBBO_SYNC_HANDLER_INVOKE = " com.alibaba.dubbo.rpc.listener.ListenerInvokerWrapper.invoke(com.alibaba.dubbo.rpc.Invocation)".substring(1);
12+
13+
public LegacyDubboSyncHandlerAdapter(ClassVisitor classVisitor, ClassContext context) {
14+
super(classVisitor, context);
15+
}
16+
17+
@Override
18+
public MethodVisitor visitMethod(final int access, final String name, final String desc, final String signature, final String[] exceptions) {
19+
MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
20+
String signCode = AsmUtils.buildSignature(context.getClassName(), name, desc);
21+
22+
if (LEGACY_DUBBO_SYNC_HANDLER_INVOKE.equals(signCode)) {
23+
DongTaiLog.debug("Adding dubbo tracking by {}", signCode);
24+
mv = new LegacyDubboSyncHandlerInvokeAdviceAdapter(mv, access, name, desc, signCode, this.context);
25+
setTransformed();
26+
}
27+
return mv;
28+
}
29+
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/LegacyDubboSyncHandlerInvokeAdviceAdapter.java‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
package io.dongtai.iast.core.bytecode.enhance.plugin.framework.dubbo;
2+
3+
import io.dongtai.iast.core.bytecode.enhance.ClassContext;
4+
import io.dongtai.iast.core.bytecode.enhance.plugin.AbstractAdviceAdapter;
5+
import org.objectweb.asm.*;
6+
import org.objectweb.asm.commons.Method;
7+
8+
public class LegacyDubboSyncHandlerInvokeAdviceAdapter extends AbstractAdviceAdapter {
9+
private static final Method GET_URL_METHOD = Method.getMethod(" com.alibaba.dubbo.common.URL getUrl()".substring(1));
10+
private static final Method GET_ARGUMENTS_METHOD = Method.getMethod("java.lang.Object[] getArguments()");
11+
private static final Method GET_GETATTACHMENTS_METHOD = Method.getMethod("java.util.Map getAttachments()");
12+
private static final Method URL_TO_STRING_METHOD = Method.getMethod("java.lang.String toString()");
13+
14+
private Label exHandler;
15+
private final Type handlerType;
16+
private final Type urlType;
17+
private final Type invocationType;
18+
19+
protected LegacyDubboSyncHandlerInvokeAdviceAdapter(MethodVisitor mv, int access, String name, String desc, String signature, ClassContext context) {
20+
super(mv, access, name, desc, context, "dubbo", signature);
21+
this.handlerType = Type.getObjectType(" com/alibaba/dubbo/rpc/listener/ListenerInvokerWrapper".substring(1));
22+
this.invocationType = Type.getObjectType(" com/alibaba/dubbo/rpc/Invocation".substring(1));
23+
this.urlType = Type.getObjectType(" com/alibaba/dubbo/common/URL".substring(1));
24+
}
25+
26+
@Override
27+
protected void onMethodEnter() {
28+
this.tryLabel = new Label();
29+
visitLabel(this.tryLabel);
30+
enterMethod();
31+
this.catchLabel = new Label();
32+
this.exHandler = new Label();
33+
}
34+
35+
@Override
36+
protected void onMethodExit(int opcode) {
37+
leaveMethod(opcode);
38+
}
39+
40+
@Override
41+
protected void before() {
42+
}
43+
44+
@Override
45+
protected void after(int opcode) {
46+
}
47+
48+
@Override
49+
public void visitMaxs(int maxStack, int maxLocals) {
50+
visitLabel(this.catchLabel);
51+
visitLabel(this.exHandler);
52+
leaveMethod(ATHROW);
53+
throwException();
54+
visitTryCatchBlock(this.tryLabel, this.catchLabel, this.exHandler, ASM_TYPE_THROWABLE.getInternalName());
55+
super.visitMaxsNew(maxStack, maxLocals);
56+
}
57+
58+
private void enterMethod() {
59+
enterScope();
60+
61+
Label elseLabel = new Label();
62+
Label endLabel = new Label();
63+
64+
isFirstScope();
65+
mv.visitJumpInsn(Opcodes.IFEQ, elseLabel);
66+
67+
traceMethod();
68+
69+
mark(elseLabel);
70+
mark(endLabel);
71+
}
72+
73+
private void leaveMethod(int opcode) {
74+
leaveScope();
75+
}
76+
77+
private void enterScope() {
78+
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
79+
push(false);
80+
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterPropagator);
81+
}
82+
83+
private void leaveScope() {
84+
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
85+
push(false);
86+
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leavePropagator);
87+
}
88+
89+
private void isFirstScope() {
90+
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
91+
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$isFirstLevelPropagator);
92+
}
93+
94+
private void traceMethod() {
95+
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
96+
loadThisOrPushNullIfIsStatic();
97+
dup();
98+
invokeVirtual(this.handlerType, GET_URL_METHOD);
99+
invokeVirtual(this.urlType, URL_TO_STRING_METHOD);
100+
loadArg(0);
101+
loadArg(0);
102+
invokeInterface(this.invocationType, GET_ARGUMENTS_METHOD);
103+
loadArg(0);
104+
invokeInterface(this.invocationType, GET_GETATTACHMENTS_METHOD);
105+
push(this.classContext.getClassName());
106+
push(this.name);
107+
push(this.signature);
108+
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$traceDubboInvoke);
109+
pop();
110+
}
111+
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/feign/DispatchFeign.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import org.objectweb.asm.ClassVisitor;
77

88
public class DispatchFeign implements DispatchPlugin {
9-
private static final String FEIGN_SYNC_HANDLER = "feign.SynchronousMethodHandler";
9+
public static final String FEIGN_SYNC_HANDLER = "feign.SynchronousMethodHandler";
1010

1111
@Override
1212
public ClassVisitor dispatch(ClassVisitor classVisitor, ClassContext context, Policy policy) {

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import io.dongtai.iast.core.handler.hookpoint.graphy.GraphBuilder;
1212
import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
1313
import io.dongtai.iast.core.handler.hookpoint.models.policy.*;
14+
import io.dongtai.iast.core.handler.hookpoint.service.trace.DubboService;
1415
import io.dongtai.iast.core.handler.hookpoint.service.trace.FeignService;
1516
import io.dongtai.iast.core.utils.StringUtils;
1617
import io.dongtai.log.DongTaiLog;
@@ -535,6 +536,28 @@ public boolean traceFeignInvoke(Object instance, Object[] parameters,
535536
return false;
536537
}
537538

539+
@Override
540+
public boolean traceDubboInvoke(Object instance, String url, Object invocation, Object[] arguments,
541+
Map<String, String> headers, String className, String methodName,
542+
String signature) {
543+
try {
544+
ScopeManager.SCOPE_TRACKER.getPolicyScope().enterAgent();
545+
if (!isCollectAllowed(false)) {
546+
return false;
547+
}
548+
549+
MethodEvent event = new MethodEvent(className, className, methodName,
550+
signature, instance, arguments, null);
551+
552+
DubboService.solveSyncInvoke(event, invocation, url, headers, INVOKE_ID_SEQUENCER);
553+
} catch (Throwable e) {
554+
DongTaiLog.error(ErrorCode.SPY_TRACE_DUBBO_INVOKE_FAILED, e);
555+
} finally {
556+
ScopeManager.SCOPE_TRACKER.getPolicyScope().leaveAgent();
557+
}
558+
return false;
559+
}
560+
538561
@SuppressWarnings("unchecked")
539562
private boolean isCollectAllowed(boolean isEnterEntry) {
540563
if (!EngineManager.isEngineRunning()) {

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

Lines changed: 5 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,12 @@
44
import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
55
import io.dongtai.iast.core.handler.hookpoint.models.policy.SourceNode;
66
import io.dongtai.iast.core.handler.hookpoint.models.policy.TaintPosition;
7-
import io.dongtai.iast.core.handler.hookpoint.models.taint.range.*;
7+
import io.dongtai.iast.core.handler.hookpoint.models.taint.range.TaintRangesBuilder;
88
import io.dongtai.iast.core.utils.StackUtils;
99
import io.dongtai.iast.core.utils.TaintPoolUtils;
1010

11-
import java.lang.reflect.Array;
12-
import java.util.*;
11+
import java.util.ArrayList;
12+
import java.util.Set;
1313
import java.util.concurrent.atomic.AtomicInteger;
1414

1515
/**
@@ -81,99 +81,12 @@ private static boolean trackTarget(MethodEvent event, SourceNode sourceNode) {
8181
return false;
8282
}
8383

84-
trackObject(event, sourceNode, event.returnInstance, 0);
84+
TaintPoolUtils.trackObject(event, sourceNode, event.returnInstance, 0);
8585
// @TODO: hook json serializer for custom model
8686
handlerCustomModel(event, sourceNode);
8787
return true;
8888
}
8989

90-
private static void trackObject(MethodEvent event, SourceNode sourceNode, Object obj, int depth) {
91-
if (depth >= 10 || !TaintPoolUtils.isNotEmpty(obj) || !TaintPoolUtils.isAllowTaintType(obj)) {
92-
return;
93-
}
94-
95-
int hash = System.identityHashCode(obj);
96-
if (EngineManager.TAINT_HASH_CODES.contains(hash)) {
97-
return;
98-
}
99-
100-
Class<?> cls = obj.getClass();
101-
if (cls.isArray() && !cls.getComponentType().isPrimitive()) {
102-
trackArray(event, sourceNode, obj, depth);
103-
} else if (obj instanceof Iterator && !(obj instanceof Enumeration)) {
104-
trackIterator(event, sourceNode, (Iterator<?>) obj, depth);
105-
} else if (obj instanceof Map) {
106-
trackMap(event, sourceNode, (Map<?, ?>) obj, depth);
107-
} else if (obj instanceof Map.Entry) {
108-
trackMapEntry(event, sourceNode, (Map.Entry<?, ?>) obj, depth);
109-
} else if (obj instanceof Collection) {
110-
if (obj instanceof List) {
111-
trackList(event, sourceNode, (List<?>) obj, depth);
112-
} else {
113-
trackIterator(event, sourceNode, ((Collection<?>) obj).iterator(), depth);
114-
}
115-
} else if ("java.util.Optional".equals(obj.getClass().getName())) {
116-
trackOptional(event, sourceNode, obj, depth);
117-
} else {
118-
int len = TaintRangesBuilder.getLength(obj);
119-
if (len == 0) {
120-
return;
121-
}
122-
123-
TaintRanges tr = new TaintRanges(new TaintRange(0, len));
124-
if (sourceNode.hasTags()) {
125-
String[] tags = sourceNode.getTags();
126-
for (String tag : tags) {
127-
tr.add(new TaintRange(tag, 0, len));
128-
}
129-
}
130-
event.targetRanges.add(new MethodEvent.MethodEventTargetRange(hash, tr));
131-
132-
EngineManager.TAINT_HASH_CODES.add(hash);
133-
event.addTargetHash(hash);
134-
EngineManager.TAINT_RANGES_POOL.add(hash, tr);
135-
}
136-
}
137-
138-
private static void trackArray(MethodEvent event, SourceNode sourceNode, Object arr, int depth) {
139-
int length = Array.getLength(arr);
140-
for (int i = 0; i < length; i++) {
141-
trackObject(event, sourceNode, Array.get(arr, i), depth + 1);
142-
}
143-
}
144-
145-
private static void trackIterator(MethodEvent event, SourceNode sourceNode, Iterator<?> it, int depth) {
146-
while (it.hasNext()) {
147-
trackObject(event, sourceNode, it.next(), depth + 1);
148-
}
149-
}
150-
151-
private static void trackMap(MethodEvent event, SourceNode sourceNode, Map<?, ?> map, int depth) {
152-
for (Object key : map.keySet()) {
153-
trackObject(event, sourceNode, key, depth + 1);
154-
trackObject(event, sourceNode, map.get(key), depth + 1);
155-
}
156-
}
157-
158-
private static void trackMapEntry(MethodEvent event, SourceNode sourceNode, Map.Entry<?, ?> entry, int depth) {
159-
trackObject(event, sourceNode, entry.getKey(), depth + 1);
160-
trackObject(event, sourceNode, entry.getValue(), depth + 1);
161-
}
162-
163-
private static void trackList(MethodEvent event, SourceNode sourceNode, List<?> list, int depth) {
164-
for (Object obj : list) {
165-
trackObject(event, sourceNode, obj, depth + 1);
166-
}
167-
}
168-
169-
private static void trackOptional(MethodEvent event, SourceNode sourceNode, Object obj, int depth) {
170-
try {
171-
Object v = ((Optional<?>) obj).orElse(null);
172-
trackObject(event, sourceNode, v, depth + 1);
173-
} catch (Throwable ignore) {
174-
}
175-
}
176-
17790
/**
17891
* todo: 处理过程和结果需要细化
17992
*
@@ -183,7 +96,7 @@ public static void handlerCustomModel(MethodEvent event, SourceNode sourceNode)
18396
if (!"getSession".equals(event.getMethodName())) {
18497
Set<Object> modelValues = TaintPoolUtils.parseCustomModel(event.returnInstance);
18598
for (Object modelValue : modelValues) {
186-
trackObject(event, sourceNode, modelValue, 0);
99+
TaintPoolUtils.trackObject(event, sourceNode, modelValue, 0);
187100
}
188101
}
189102
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/policy/PolicyManager.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package io.dongtai.iast.core.handler.hookpoint.models.policy;
22

3+
import io.dongtai.iast.core.bytecode.enhance.plugin.framework.dubbo.DispatchDubbo;
4+
import io.dongtai.iast.core.bytecode.enhance.plugin.framework.feign.DispatchFeign;
35
import io.dongtai.iast.core.bytecode.enhance.plugin.framework.j2ee.dispatch.DispatchJ2ee;
46
import io.dongtai.iast.core.utils.StringUtils;
57
import io.dongtai.log.DongTaiLog;
@@ -27,7 +29,8 @@ public class PolicyManager {
2729
" org.springframework.web.servlet.FrameworkServlet".substring(1),
2830
" javax.servlet.http.Cookie".substring(1),
2931
" org/springframework/web/servlet/mvc/annotation/AnnotationMethodHandlerAdapter$ServletHandlerMethodInvoker".substring(1),
30-
" feign.SynchronousMethodHandler.invoke".substring(1)
32+
DispatchFeign.FEIGN_SYNC_HANDLER,
33+
DispatchDubbo.LEGACY_DUBBO_SYNC_HANDLER
3134
));
3235
private static final Set<String> HOOK_CLASS_SUFFIX_NAMES = new HashSet<String>(Collections.singletonList(
3336
".dubbo.monitor.support.MonitorFilter"

0 commit comments

Comments
 (0)