Skip to content

Commit a15a32a

Browse files
lostsnowlostsnow
authored
Merge pull request #540 from Nizernizer/fix/jdk-module
Fix/collection bypass and jdk9+ module inject
2 parents 09a90e9 + 2d4a13a commit a15a32a

File tree

14 files changed

+351
-5
lines changed

14 files changed

+351
-5
lines changed

‎dongtai-core/src/main/java/com/secnium/iast/core/AgentEngine.java‎

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,7 @@
1616

1717
import java.lang.dongtai.SpyDispatcherHandler;
1818
import java.lang.instrument.Instrumentation;
19-
import java.util.ArrayList;
20-
import java.util.ListIterator;
19+
import java.util.*;
2120

2221
/**
2322
* @author [email protected]
@@ -71,6 +70,7 @@ public static void install(String mode, String propertiesFilePath, Integer agent
7170
DongTaiLog.info("DongTai Engine is successfully installed to the JVM, and it takes {} s",
7271
stopWatch.getTime() / 1000);
7372
DongTaiLog.info("DongTai Agent Version: {}, DongTai Server: {}", AgentConstant.VERSION_VALUE, cfg.getBaseUrl());
73+
inject(inst);
7474
new ServiceDirReport().send();
7575
} catch (Throwable e) {
7676
DongTaiLog.error(ErrorCode.get("ENGINE_INSTALL_FAILED"), e);
@@ -134,4 +134,37 @@ private void destroy() {
134134
}
135135
}
136136

137+
138+
private static void redefineJavaBaseModule(Instrumentation instrumentation) {
139+
if (doesSupportModules()) {
140+
try {
141+
Instrumentation.class.getMethod("redefineModule", Class.forName("java.lang.Module"), Set.class, Map.class, Map.class, Set.class, Map.class).invoke(instrumentation, getModule(Object.class), Collections.emptySet(), Collections.emptyMap(), Collections.singletonMap("java.lang", Collections.singleton(getModule(EngineManager.class))), Collections.emptySet(), Collections.emptyMap());
142+
} catch (Exception e) {
143+
DongTaiLog.error(ErrorCode.REDEFINE_MODULE_FAILED,e);
144+
}
145+
}
146+
}
147+
148+
public static boolean doesSupportModules() {
149+
try {
150+
Class.forName("java.lang.Module");
151+
return true;
152+
} catch (ClassNotFoundException e) {
153+
return false;
154+
}
155+
}
156+
157+
private static Object getModule(Class<?> clazz) {
158+
try {
159+
return Class.class.getMethod("getModule", new Class[0]).invoke(clazz, new Object[0]);
160+
} catch (Exception e) {
161+
throw new IllegalStateException("There was a problem while getting the module of the class", e);
162+
}
163+
}
164+
public static void inject(Instrumentation inst) {
165+
if (doesSupportModules()) {
166+
redefineJavaBaseModule(inst);
167+
}
168+
}
169+
137170
}

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -256,6 +256,20 @@ static Method getAsmMethod(final Class<?> clazz,
256256
boolean.class
257257
);
258258

259+
Method SPY$skipCollect = InnerHelper.getAsmMethod(
260+
SpyDispatcher.class,
261+
"skipCollect",
262+
Object.class,
263+
Object[].class,
264+
Object.class,
265+
String.class,
266+
String.class,
267+
String.class,
268+
String.class,
269+
String.class,
270+
boolean.class
271+
);
272+
259273
Method SPY$traceFeignInvoke = InnerHelper.getAsmMethod(
260274
SpyDispatcher.class,
261275
"traceFeignInvoke",
@@ -279,6 +293,18 @@ static Method getAsmMethod(final Class<?> clazz,
279293
String.class
280294
);
281295

296+
Method SPY$isSkipCollectDubbo = InnerHelper.getAsmMethod(
297+
SpyDispatcher.class,
298+
"isSkipCollectDubbo",
299+
Object.class
300+
);
301+
302+
Method SPY$isSkipCollectFeign = InnerHelper.getAsmMethod(
303+
SpyDispatcher.class,
304+
"isSkipCollectFeign",
305+
Object.class
306+
);
307+
282308
Method SPY$reportService = InnerHelper.getAsmMethod(
283309
SpyDispatcher.class,
284310
"reportService",

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,32 @@ public void captureMethodState(
174174
pop();
175175
}
176176

177+
public void skipCollect(
178+
final int opcode,
179+
final PolicyNode policyNode,
180+
final boolean captureRet
181+
) {
182+
newLocal(ASM_TYPE_OBJECT);
183+
if (captureRet && !isThrow(opcode)) {
184+
loadReturn(opcode);
185+
} else {
186+
pushNull();
187+
}
188+
storeLocal(this.nextLocal - 1);
189+
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
190+
loadThisOrPushNullIfIsStatic();
191+
loadArgArray();
192+
loadLocal(this.nextLocal - 1);
193+
push(policyNode.toString());
194+
push(this.context.getClassName());
195+
push(this.context.getMatchedClassName());
196+
push(this.name);
197+
push(this.signature);
198+
push(Modifier.isStatic(this.access));
199+
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$skipCollect);
200+
pop();
201+
}
202+
177203
/**
178204
* 是否抛出异常返回(通过字节码判断)
179205
*

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,9 @@ public void onMethodEnter(MethodAdviceAdapter adapter, MethodVisitor mv, MethodC
1515
if (!(policyNode instanceof SinkNode)) {
1616
continue;
1717
}
18+
if ("ssrf".equals(((SinkNode) policyNode).getVulType())){
19+
adapter.skipCollect(-1, policyNode, false);
20+
}
1821

1922
enterScope(adapter, policyNode);
2023

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ public void visitMaxs(int maxStack, int maxLocals) {
6161
}
6262

6363
private void enterMethod() {
64+
skipCollect();
6465
enterScope();
6566

6667
Label elseLabel = new Label();
@@ -113,4 +114,11 @@ private void traceMethod() {
113114
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$traceDubboInvoke);
114115
pop();
115116
}
117+
118+
private void skipCollect() {
119+
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
120+
loadArg(0);
121+
invokeInterface(ASM_TYPE_SPY_DISPATCHER,SPY$isSkipCollectDubbo);
122+
pop();
123+
}
116124
}

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ public void visitMaxs(int maxStack, int maxLocals) {
4444
}
4545

4646
private void enterMethod() {
47+
skipCollect();
4748
enterScope();
4849

4950
Label elseLabel = new Label();
@@ -89,4 +90,11 @@ private void traceMethod() {
8990
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$traceFeignInvoke);
9091
pop();
9192
}
93+
94+
private void skipCollect() {
95+
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
96+
loadThisOrPushNullIfIsStatic();
97+
invokeInterface(ASM_TYPE_SPY_DISPATCHER,SPY$isSkipCollectFeign);
98+
pop();
99+
}
92100
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/handler/bypass/BlackUrlBypass.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
package io.dongtai.iast.core.handler.bypass;
2+
3+
import io.dongtai.iast.core.utils.threadlocal.BooleanThreadLocal;
4+
5+
public class BlackUrlBypass {
6+
7+
private static BooleanThreadLocal isBlackUrl = new BooleanThreadLocal(false);
8+
9+
public static Boolean isBlackUrl() {
10+
return isBlackUrl.get();
11+
}
12+
13+
public static void setIsBlackUrl(Boolean isBlackUrl) {
14+
BlackUrlBypass.isBlackUrl.set(isBlackUrl);
15+
}
16+
17+
public static String getHeaderKey() {
18+
return "dt-collect-skip";
19+
}
20+
}

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

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import io.dongtai.iast.common.scope.Scope;
77
import io.dongtai.iast.common.scope.ScopeManager;
88
import io.dongtai.iast.core.EngineManager;
9+
import io.dongtai.iast.core.handler.bypass.BlackUrlBypass;
910
import io.dongtai.iast.core.handler.hookpoint.api.DubboApiGatherThread;
1011
import io.dongtai.iast.core.handler.hookpoint.api.SpringGatherApiThread;
1112
import io.dongtai.iast.core.handler.hookpoint.controller.HookType;
@@ -15,11 +16,16 @@
1516
import io.dongtai.iast.core.handler.hookpoint.models.policy.*;
1617
import io.dongtai.iast.core.handler.hookpoint.service.trace.DubboService;
1718
import io.dongtai.iast.core.handler.hookpoint.service.trace.FeignService;
19+
import io.dongtai.iast.core.handler.hookpoint.service.trace.HttpService;
1820
import io.dongtai.iast.core.utils.StringUtils;
21+
import io.dongtai.iast.core.utils.matcher.ConfigMatcher;
1922
import io.dongtai.log.DongTaiLog;
2023
import io.dongtai.log.ErrorCode;
2124

2225
import java.lang.dongtai.SpyDispatcher;
26+
import java.lang.reflect.Field;
27+
import java.lang.reflect.InvocationTargetException;
28+
import java.lang.reflect.Method;
2329
import java.net.InetSocketAddress;
2430
import java.util.Collection;
2531
import java.util.Enumeration;
@@ -185,6 +191,14 @@ public void collectHttpRequest(Object obj, Object req, Object resp, StringBuffer
185191
put("headers", headers);
186192
put("replay-request", !StringUtils.isEmpty(headers.get("dongtai-replay-id")));
187193
}};
194+
if (ConfigMatcher.getInstance().getBlackUrl(requestMeta)) {
195+
BlackUrlBypass.setIsBlackUrl(true);
196+
return;
197+
}
198+
if (null != headers.get(BlackUrlBypass.getHeaderKey()) && headers.get(BlackUrlBypass.getHeaderKey()).equals("true")) {
199+
BlackUrlBypass.setIsBlackUrl(true);
200+
return;
201+
}
188202
HttpImpl.solveHttpRequest(obj, req, resp, requestMeta);
189203
} catch (Throwable e) {
190204
DongTaiLog.warn(ErrorCode.get("SPY_COLLECT_HTTP_FAILED"), "request", e);
@@ -714,6 +728,63 @@ public boolean traceDubboInvoke(Object instance, String url, Object invocation,
714728
return false;
715729
}
716730

731+
@Override
732+
public boolean isSkipCollectDubbo(Object invocation) {
733+
if (BlackUrlBypass.isBlackUrl()) {
734+
Method setAttachmentMethod = null;
735+
try {
736+
setAttachmentMethod = invocation.getClass().getMethod("setAttachment", String.class, String.class);
737+
setAttachmentMethod.setAccessible(true);
738+
setAttachmentMethod.invoke(invocation, BlackUrlBypass.getHeaderKey(), BlackUrlBypass.isBlackUrl().toString());
739+
} catch (NoSuchMethodException | InvocationTargetException | IllegalAccessException e) {
740+
DongTaiLog.error(ErrorCode.get("BYPASS_FAILED_DUBBO"), e);
741+
}
742+
}
743+
return false;
744+
}
745+
746+
@Override
747+
public boolean isSkipCollectFeign(Object instance) {
748+
if (BlackUrlBypass.isBlackUrl()) {
749+
Field metadataField = null;
750+
try {
751+
metadataField = instance.getClass().getDeclaredField("metadata");
752+
metadataField.setAccessible(true);
753+
Object metadata = metadataField.get(instance);
754+
Method templateMethod = metadata.getClass().getMethod("template");
755+
Object template = templateMethod.invoke(metadata);
756+
757+
Method addHeaderMethod = template.getClass().getDeclaredMethod("header", String.class, String[].class);
758+
addHeaderMethod.setAccessible(true);
759+
addHeaderMethod.invoke(template, BlackUrlBypass.getHeaderKey(), new String[]{});
760+
addHeaderMethod.invoke(template, BlackUrlBypass.getHeaderKey(), new String[]{BlackUrlBypass.isBlackUrl().toString()});
761+
} catch (NoSuchFieldException | NoSuchMethodException | InvocationTargetException | IllegalAccessException e) {
762+
DongTaiLog.error(ErrorCode.get("BYPASS_FAILED_FEIGN"), e);
763+
}
764+
}
765+
return false;
766+
}
767+
768+
@Override
769+
public boolean skipCollect(Object instance, Object[] parameters, Object retObject, String policyKey,
770+
String className, String matchedClassName, String methodName, String signature,
771+
boolean isStatic) {
772+
if (BlackUrlBypass.isBlackUrl()) {
773+
MethodEvent event = new MethodEvent(className, matchedClassName, methodName,
774+
signature, instance, parameters, retObject);
775+
PolicyNode policyNode = getPolicyNode(policyKey);
776+
if (policyNode == null) {
777+
return false;
778+
}
779+
HttpService httpService = new HttpService();
780+
if (httpService.match(event, policyNode)) {
781+
httpService.addBypass(event);
782+
return true;
783+
}
784+
}
785+
return false;
786+
}
787+
717788
private boolean isCollectAllowed(boolean isEnterEntry) {
718789
if (!EngineManager.isEngineRunning()) {
719790
return false;

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import io.dongtai.iast.common.config.ConfigBuilder;
55
import io.dongtai.iast.common.config.ConfigKey;
66
import io.dongtai.iast.core.EngineManager;
7+
import io.dongtai.iast.core.handler.bypass.BlackUrlBypass;
78
import io.dongtai.iast.core.handler.context.ContextManager;
89
import io.dongtai.iast.core.handler.hookpoint.IastClassLoader;
910
import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
@@ -94,6 +95,10 @@ public static void collectDubboRequestSource(Object handler, Object invocation,
9495
if (requestMeta == null) {
9596
return;
9697
}
98+
if (null != headers.get(BlackUrlBypass.getHeaderKey()) && headers.get(BlackUrlBypass.getHeaderKey()).equals("true")){
99+
BlackUrlBypass.setIsBlackUrl(true);
100+
return;
101+
}
97102

98103
String url = (String) requestMeta.get("requestURL") + "/" + methodName;
99104
String uri = (String) requestMeta.get("requestURI") + "/" + methodName;

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import io.dongtai.iast.common.config.*;
44
import io.dongtai.iast.common.constants.AgentConstant;
55
import io.dongtai.iast.core.EngineManager;
6+
import io.dongtai.iast.core.handler.bypass.BlackUrlBypass;
67
import io.dongtai.iast.core.handler.hookpoint.IastClassLoader;
78
import io.dongtai.iast.core.utils.*;
89
import io.dongtai.iast.core.utils.matcher.ConfigMatcher;
@@ -60,6 +61,7 @@ public static void solveHttpRequest(Object obj, Object req, Object resp, Map<Str
6061
String requestURL = ((StringBuffer) requestMeta.get("requestURL")).toString();
6162
Map<String, String> headers = (Map<String, String>) requestMeta.get("headers");
6263
if (requestDenyList.match(requestURL, headers)) {
64+
BlackUrlBypass.setIsBlackUrl(true);
6365
DongTaiLog.trace("HTTP Request {} deny to collect", requestURL);
6466
return;
6567
}
@@ -73,9 +75,6 @@ public static void solveHttpRequest(Object obj, Object req, Object resp, Map<Str
7375
if (ConfigMatcher.getInstance().disableExtension((String) requestMeta.get("requestURI"))) {
7476
return;
7577
}
76-
if (ConfigMatcher.getInstance().getBlackUrl(requestMeta)) {
77-
return;
78-
}
7978

8079
try {
8180
boolean enableVersionHeader = ConfigBuilder.getInstance().get(ConfigKey.ENABLE_VERSION_HEADER);

0 commit comments

Comments
 (0)