Skip to content

Commit 8219898

Browse files
lostsnowlostsnow
committed
add dubbo provider track
1 parent 960d8f2 commit 8219898

File tree

25 files changed

+622
-14
lines changed

25 files changed

+622
-14
lines changed

‎dongtai-common/src/main/java/io/dongtai/iast/common/scope/Scope.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,9 @@ public enum Scope {
55
HTTP_ENTRY(2),
66
SERVLET_INPUT_STREAM_READ(3),
77
SERVLET_OUTPUT_WRITE(4),
8+
DUBBO_REQUEST(5),
9+
DUBBO_ENTRY(6),
10+
DUBBO_SOURCE(7),
811
;
912

1013
private final int id;

‎dongtai-common/src/main/java/io/dongtai/iast/common/scope/ScopeAggregator.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,9 @@
33
public class ScopeAggregator {
44
private final GeneralScope httpRequestScope = new GeneralScope();
55
private final GeneralScope httpEntryScope = new GeneralScope();
6+
private final GeneralScope dubboRequestScope = new GeneralScope();
7+
private final GeneralScope dubboEntryScope = new GeneralScope();
8+
private final GeneralScope dubboSourceScope = new GeneralScope();
69
private final GeneralScope servletInputStreamReadScope = new GeneralScope();
710
private final GeneralScope servletOutputStreamWriteScope = new GeneralScope();
811
private final PolicyScope policyScope = new PolicyScope();
@@ -15,6 +18,18 @@ public GeneralScope getHttpEntryScope() {
1518
return httpEntryScope;
1619
}
1720

21+
public GeneralScope getDubboRequestScope() {
22+
return dubboRequestScope;
23+
}
24+
25+
public GeneralScope getDubboEntryScope() {
26+
return dubboEntryScope;
27+
}
28+
29+
public GeneralScope getDubboSourceScope() {
30+
return dubboSourceScope;
31+
}
32+
1833
public GeneralScope getServletInputStreamReadScope() {
1934
return servletInputStreamReadScope;
2035
}

‎dongtai-common/src/main/java/io/dongtai/iast/common/scope/ScopeTracker.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,19 @@ public GeneralScope getScope(Scope scope) {
1616
return this.get().getServletInputStreamReadScope();
1717
case SERVLET_OUTPUT_WRITE:
1818
return this.get().getServletOutputStreamWriteScope();
19+
case DUBBO_REQUEST:
20+
return this.get().getDubboRequestScope();
21+
case DUBBO_ENTRY:
22+
return this.get().getDubboEntryScope();
23+
case DUBBO_SOURCE:
24+
return this.get().getDubboSourceScope();
1925
default:
2026
return null;
2127
}
2228
}
2329

2430
public boolean inEnterEntry() {
25-
return this.get().getHttpEntryScope().in();
31+
return this.get().getHttpEntryScope().in() || this.get().getDubboRequestScope().in();
2632
}
2733

2834
public PolicyScope getPolicyScope() {

‎dongtai-core/src/main/java/io/dongtai/iast/core/EngineManager.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,4 +140,12 @@ public static void enterHttpEntry(Map<String, Object> requestMeta) {
140140
TAINT_RANGES_POOL.set(new HashMap<Integer, TaintRanges>());
141141
ScopeManager.SCOPE_TRACKER.getScope(Scope.HTTP_ENTRY).enter();
142142
}
143+
144+
public static void enterDubboEntry(Map<String, Object> requestMeta) {
145+
REQUEST_CONTEXT.set(requestMeta);
146+
TRACK_MAP.set(new HashMap<Integer, MethodEvent>(1024));
147+
TAINT_HASH_CODES.set(new HashSet<Integer>());
148+
TAINT_RANGES_POOL.set(new HashMap<Integer, TaintRanges>());
149+
ScopeManager.SCOPE_TRACKER.getScope(Scope.DUBBO_ENTRY).enter();
150+
}
143151
}

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

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

55
import java.lang.dongtai.SpyDispatcher;
66
import java.lang.dongtai.SpyDispatcherHandler;
7+
import java.net.InetSocketAddress;
78
import java.util.*;
89

910
/**
@@ -131,6 +132,49 @@ static Method getAsmMethod(final Class<?> clazz,
131132
int.class
132133
);
133134

135+
Method SPY$enterDubbo = InnerHelper.getAsmMethod(
136+
SpyDispatcher.class,
137+
"enterDubbo"
138+
);
139+
Method SPY$leaveDubbo = InnerHelper.getAsmMethod(
140+
SpyDispatcher.class,
141+
"leaveDubbo",
142+
Object.class,
143+
Object.class,
144+
Object.class,
145+
byte.class
146+
);
147+
Method SPY$isFirstLevelDubbo = InnerHelper.getAsmMethod(
148+
SpyDispatcher.class,
149+
"isFirstLevelDubbo"
150+
);
151+
Method SPY$collectDubboRequest = InnerHelper.getAsmMethod(
152+
SpyDispatcher.class,
153+
"collectDubboRequest",
154+
Object.class,
155+
Object.class,
156+
Object.class,
157+
String.class,
158+
InetSocketAddress.class,
159+
boolean.class,
160+
boolean.class,
161+
boolean.class,
162+
boolean.class
163+
);
164+
165+
Method SPY$collectDubboRequestSource = InnerHelper.getAsmMethod(
166+
SpyDispatcher.class,
167+
"collectDubboRequestSource",
168+
Object.class,
169+
Object.class,
170+
String.class,
171+
Object[].class,
172+
Map.class,
173+
String.class,
174+
String.class,
175+
String.class
176+
);
177+
134178
Method SPY$enterSource = InnerHelper.getAsmMethod(
135179
SpyDispatcher.class,
136180
"enterSource"

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,19 @@
77

88
public class DispatchDubbo implements DispatchPlugin {
99
public static final String LEGACY_DUBBO_SYNC_HANDLER = " com.alibaba.dubbo.rpc.listener.ListenerInvokerWrapper".substring(1);
10+
public static final String LEGACY_DUBBO_EXCHANGE_HANDLER = " com.alibaba.dubbo.remoting.exchange.support.header.HeaderExchangeHandler".substring(1);
11+
public static final String LEGACY_DUBBO_PROXY_HANDLER = " com.alibaba.dubbo.rpc.proxy.AbstractProxyInvoker".substring(1);
1012

1113
@Override
1214
public ClassVisitor dispatch(ClassVisitor classVisitor, ClassContext context, Policy policy) {
1315
String className = context.getClassName();
1416

1517
if (LEGACY_DUBBO_SYNC_HANDLER.equals(className)) {
1618
classVisitor = new LegacyDubboSyncHandlerAdapter(classVisitor, context);
19+
} else if (LEGACY_DUBBO_EXCHANGE_HANDLER.equals(className)) {
20+
classVisitor = new LegacyDubboExchangeHandlerAdapter(classVisitor, context);
21+
} else if (LEGACY_DUBBO_PROXY_HANDLER.equals(className)) {
22+
classVisitor = new LegacyDubboProxyHandlerAdapter(classVisitor, context);
1723
}
1824

1925
return classVisitor;

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

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
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 LegacyDubboExchangeHandleRequestAdviceAdapter 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 URL_TO_STRING_METHOD = Method.getMethod("java.lang.String toString()");
11+
private static final Method GET_REMOTE_ADDRESS_METHOD = Method.getMethod("java.net.InetSocketAddress getRemoteAddress()");
12+
private static final Method IS_TWO_WAY_METHOD = Method.getMethod("boolean isTwoWay()");
13+
private static final Method IS_EVENT_METHOD = Method.getMethod("boolean isEvent()");
14+
private static final Method IS_BROKEN_METHOD = Method.getMethod("boolean isBroken()");
15+
private static final Method IS_HEARTBEAT_METHOD = Method.getMethod("boolean isHeartbeat()");
16+
private static final Method GET_RESULT_METHOD = Method.getMethod("java.lang.Object getResult()");
17+
private static final Method GET_STATUS_METHOD = Method.getMethod("byte getStatus()");
18+
19+
private final Type endpointType;
20+
private final Type urlType;
21+
private final Type channelType;
22+
private final Type requestType;
23+
private final Type responseType;
24+
25+
protected LegacyDubboExchangeHandleRequestAdviceAdapter(MethodVisitor mv, int access, String name, String desc, String signature, ClassContext context) {
26+
super(mv, access, name, desc, context, "dubbo", signature);
27+
this.endpointType = Type.getObjectType(" com/alibaba/dubbo/remoting/Endpoint".substring(1));
28+
this.urlType = Type.getObjectType(" com/alibaba/dubbo/common/URL".substring(1));
29+
this.channelType = Type.getObjectType(" com/alibaba/dubbo/remoting/Channel".substring(1));
30+
this.requestType = Type.getObjectType(" com/alibaba/dubbo/remoting/exchange/Request".substring(1));
31+
this.responseType = Type.getObjectType(" com/alibaba/dubbo/remoting/exchange/Response".substring(1));
32+
}
33+
34+
@Override
35+
protected void before() {
36+
mark(tryLabel);
37+
Label elseLabel = new Label();
38+
39+
enterDubbo();
40+
isFirstLevelDubbo();
41+
mv.visitJumpInsn(EQ, elseLabel);
42+
43+
collectDubboRequest();
44+
45+
mark(elseLabel);
46+
}
47+
48+
@Override
49+
protected void after(int opcode) {
50+
// if (opcode != ATHROW) {
51+
// Label elseLabel = new Label();
52+
// isFirstLevelDubbo();
53+
// mv.visitJumpInsn(EQ, elseLabel);
54+
// collectDubboResponse(opcode);
55+
// mark(elseLabel);
56+
// }
57+
58+
leaveDubbo(opcode);
59+
}
60+
61+
private void enterDubbo() {
62+
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
63+
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterDubbo);
64+
}
65+
66+
private void leaveDubbo(int opcode) {
67+
int retLocal = newLocal(ASM_TYPE_OBJECT);
68+
if (!isThrow(opcode)) {
69+
loadReturn(opcode);
70+
} else {
71+
pushNull();
72+
}
73+
storeLocal(retLocal);
74+
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
75+
loadArg(1);
76+
loadLocal(retLocal);
77+
if (!isThrow(opcode)) {
78+
loadLocal(retLocal);
79+
invokeVirtual(this.responseType, GET_RESULT_METHOD);
80+
loadLocal(retLocal);
81+
invokeVirtual(this.responseType, GET_STATUS_METHOD);
82+
} else {
83+
pushNull();
84+
byte b = 0;
85+
push(b);
86+
}
87+
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leaveDubbo);
88+
}
89+
90+
private void isFirstLevelDubbo() {
91+
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
92+
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$isFirstLevelDubbo);
93+
}
94+
95+
private void collectDubboRequest() {
96+
Label tryL = new Label();
97+
Label catchL = new Label();
98+
Label exHandlerL = new Label();
99+
visitTryCatchBlock(tryL, catchL, exHandlerL, ASM_TYPE_THROWABLE.getInternalName());
100+
visitLabel(tryL);
101+
102+
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
103+
loadThis();
104+
loadArg(0);
105+
loadArg(1);
106+
loadArg(0);
107+
invokeInterface(this.endpointType, GET_URL_METHOD);
108+
invokeVirtual(this.urlType, URL_TO_STRING_METHOD);
109+
loadArg(0);
110+
invokeInterface(this.channelType, GET_REMOTE_ADDRESS_METHOD);
111+
loadArg(1);
112+
invokeVirtual(this.requestType, IS_TWO_WAY_METHOD);
113+
loadArg(1);
114+
invokeVirtual(this.requestType, IS_EVENT_METHOD);
115+
loadArg(1);
116+
invokeVirtual(this.requestType, IS_BROKEN_METHOD);
117+
loadArg(1);
118+
invokeVirtual(this.requestType, IS_HEARTBEAT_METHOD);
119+
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$collectDubboRequest);
120+
121+
visitLabel(catchL);
122+
Label endL = new Label();
123+
visitJumpInsn(GOTO, endL);
124+
visitLabel(exHandlerL);
125+
visitVarInsn(ASTORE, this.nextLocal);
126+
visitLabel(endL);
127+
}
128+
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/LegacyDubboExchangeHandlerAdapter.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 LegacyDubboExchangeHandlerAdapter extends AbstractClassVisitor {
11+
public static final String LEGACY_DUBBO_EXCHANGE_HANDLE_REQUEST = " com.alibaba.dubbo.remoting.exchange.support.header.HeaderExchangeHandler.handleRequest(com.alibaba.dubbo.remoting.exchange.ExchangeChannel,com.alibaba.dubbo.remoting.exchange.Request)".substring(1);
12+
13+
public LegacyDubboExchangeHandlerAdapter(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_EXCHANGE_HANDLE_REQUEST.equals(signCode)) {
23+
DongTaiLog.debug("Adding dubbo provider tracking by {}", signCode);
24+
mv = new LegacyDubboExchangeHandleRequestAdviceAdapter(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/LegacyDubboProxyHandlerAdapter.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 LegacyDubboProxyHandlerAdapter extends AbstractClassVisitor {
11+
public static final String LEGACY_DUBBO_PROXY_HANDLER_INVOKE = " com.alibaba.dubbo.rpc.proxy.AbstractProxyInvoker.invoke(com.alibaba.dubbo.rpc.Invocation)".substring(1);
12+
13+
public LegacyDubboProxyHandlerAdapter(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_PROXY_HANDLER_INVOKE.equals(signCode)) {
23+
DongTaiLog.debug("Adding dubbo provider source tracking by {}", signCode);
24+
mv = new LegacyDubboProxyHandlerInvokeAdviceAdapter(mv, access, name, desc, signCode, this.context);
25+
setTransformed();
26+
}
27+
return mv;
28+
}
29+
}

0 commit comments

Comments
 (0)