Skip to content

Commit e12f6cd

Browse files
author
‘niuerzhuang’
committed
fix: jdk9+ module inject.
1 parent 714de22 commit e12f6cd

File tree

11 files changed

+107
-23
lines changed

11 files changed

+107
-23
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: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -278,9 +278,15 @@ static Method getAsmMethod(final Class<?> clazz,
278278
String.class
279279
);
280280

281-
Method SPY$isSkipCollect = InnerHelper.getAsmMethod(
281+
Method SPY$isSkipCollectDubbo = InnerHelper.getAsmMethod(
282282
SpyDispatcher.class,
283-
"isSkipCollect",
283+
"isSkipCollectDubbo",
284+
Object.class
285+
);
286+
287+
Method SPY$isSkipCollectFeign = InnerHelper.getAsmMethod(
288+
SpyDispatcher.class,
289+
"isSkipCollectFeign",
284290
Object.class
285291
);
286292

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ private void traceMethod() {
118118
private void skipCollect() {
119119
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
120120
loadArg(0);
121-
invokeInterface(ASM_TYPE_SPY_DISPATCHER,SPY$isSkipCollect);
121+
invokeInterface(ASM_TYPE_SPY_DISPATCHER,SPY$isSkipCollectDubbo);
122122
pop();
123123
}
124124
}

‎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/skip/BlackUrlSkipHandler.java renamed to dongtai-core/src/main/java/io/dongtai/iast/core/handler/bypass/BlackUrlBypass.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1-
package io.dongtai.iast.core.handler.skip;
1+
package io.dongtai.iast.core.handler.bypass;
22

33
import io.dongtai.iast.core.utils.threadlocal.BooleanThreadLocal;
44

5-
public class BlackUrlSkipHandler {
5+
public class BlackUrlBypass {
66

77
private static BooleanThreadLocal isBlackUrl = new BooleanThreadLocal(false);
88

@@ -11,7 +11,7 @@ public static Boolean isBlackUrl() {
1111
}
1212

1313
public static void setIsBlackUrl(Boolean isBlackUrl) {
14-
BlackUrlSkipHandler.isBlackUrl.set(isBlackUrl);
14+
BlackUrlBypass.isBlackUrl.set(isBlackUrl);
1515
}
1616

1717
public static String getHeaderKey() {

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

Lines changed: 37 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77
import io.dongtai.iast.common.scope.ScopeManager;
88
import io.dongtai.iast.core.EngineManager;
99
import io.dongtai.iast.core.bytecode.enhance.plugin.spring.SpringApplicationImpl;
10-
import io.dongtai.iast.core.handler.skip.BlackUrlSkipHandler;
10+
import io.dongtai.iast.core.handler.bypass.BlackUrlBypass;
11+
import io.dongtai.iast.core.handler.context.ContextManager;
1112
import io.dongtai.iast.core.handler.hookpoint.controller.HookType;
1213
import io.dongtai.iast.core.handler.hookpoint.controller.impl.*;
1314
import io.dongtai.iast.core.handler.hookpoint.graphy.GraphBuilder;
@@ -16,10 +17,13 @@
1617
import io.dongtai.iast.core.handler.hookpoint.service.trace.DubboService;
1718
import io.dongtai.iast.core.handler.hookpoint.service.trace.FeignService;
1819
import io.dongtai.iast.core.utils.StringUtils;
20+
import io.dongtai.iast.core.utils.TaintPoolUtils;
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;
2327
import java.lang.reflect.InvocationTargetException;
2428
import java.lang.reflect.Method;
2529
import java.net.InetSocketAddress;
@@ -184,6 +188,14 @@ public void collectHttpRequest(Object obj, Object req, Object resp, StringBuffer
184188
put("headers", headers);
185189
put("replay-request", !StringUtils.isEmpty(headers.get("dongtai-replay-id")));
186190
}};
191+
if (ConfigMatcher.getInstance().getBlackUrl(requestMeta)) {
192+
BlackUrlBypass.setIsBlackUrl(true);
193+
return;
194+
}
195+
if (null != headers.get(BlackUrlBypass.getHeaderKey()) && headers.get(BlackUrlBypass.getHeaderKey()).equals("true")){
196+
BlackUrlBypass.setIsBlackUrl(true);
197+
return;
198+
}
187199
HttpImpl.solveHttpRequest(obj, req, resp, requestMeta);
188200
} catch (Throwable e) {
189201
DongTaiLog.warn(ErrorCode.get("SPY_COLLECT_HTTP_FAILED"), "request", e);
@@ -701,20 +713,40 @@ public boolean traceDubboInvoke(Object instance, String url, Object invocation,
701713
}
702714

703715
@Override
704-
public boolean isSkipCollect(Object invocation) {
705-
if (BlackUrlSkipHandler.isBlackUrl()){
716+
public boolean isSkipCollectDubbo(Object invocation) {
717+
if (BlackUrlBypass.isBlackUrl()){
706718
Method setAttachmentMethod = null;
707719
try {
708720
setAttachmentMethod = invocation.getClass().getMethod("setAttachment", String.class, String.class);
709721
setAttachmentMethod.setAccessible(true);
710-
setAttachmentMethod.invoke(invocation, BlackUrlSkipHandler.getHeaderKey(), BlackUrlSkipHandler.isBlackUrl().toString());
722+
setAttachmentMethod.invoke(invocation, BlackUrlBypass.getHeaderKey(), BlackUrlBypass.isBlackUrl().toString());
711723
} catch (NoSuchMethodException | InvocationTargetException | IllegalAccessException e) {
712-
DongTaiLog.error(ErrorCode.get("SPY_SKIP_COLLECT_DUBBO_FAILED"), e);
724+
DongTaiLog.error(ErrorCode.get("BYPASS_FAILED_DUBBO"), e);
713725
}
714726
}
715727
return false;
716728
}
717729

730+
@Override
731+
public boolean isSkipCollectFeign(Object instance) {
732+
Field metadataField = null;
733+
try {
734+
metadataField = instance.getClass().getDeclaredField("metadata");
735+
metadataField.setAccessible(true);
736+
Object metadata = metadataField.get(instance);
737+
Method templateMethod = metadata.getClass().getMethod("template");
738+
Object template = templateMethod.invoke(metadata);
739+
740+
Method addHeaderMethod = template.getClass().getDeclaredMethod("header", String.class, String[].class);
741+
addHeaderMethod.setAccessible(true);
742+
addHeaderMethod.invoke(template, BlackUrlBypass.getHeaderKey(), new String[]{});
743+
addHeaderMethod.invoke(template, BlackUrlBypass.getHeaderKey(), new String[]{BlackUrlBypass.isBlackUrl().toString()});
744+
} catch (NoSuchFieldException | NoSuchMethodException | InvocationTargetException | IllegalAccessException e) {
745+
DongTaiLog.error(ErrorCode.get("BYPASS_FAILED_FEIGN"), e);
746+
}
747+
return false;
748+
}
749+
718750
private boolean isCollectAllowed(boolean isEnterEntry) {
719751
if (!EngineManager.isEngineRunning()) {
720752
return false;

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +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.skip.BlackUrlSkipHandler;
7+
import io.dongtai.iast.core.handler.bypass.BlackUrlBypass;
88
import io.dongtai.iast.core.handler.context.ContextManager;
99
import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
1010
import io.dongtai.iast.core.handler.hookpoint.models.policy.SourceNode;
@@ -55,7 +55,8 @@ public static void collectDubboRequestSource(Object handler, Object invocation,
5555
if (requestMeta == null) {
5656
return;
5757
}
58-
if (null != headers.get(BlackUrlSkipHandler.getHeaderKey()) && headers.get(BlackUrlSkipHandler.getHeaderKey()).equals("true")){
58+
if (null != headers.get(BlackUrlBypass.getHeaderKey()) && headers.get(BlackUrlBypass.getHeaderKey()).equals("true")){
59+
BlackUrlBypass.setIsBlackUrl(true);
5960
return;
6061
}
6162

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

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
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.skip.BlackUrlSkipHandler;
76
import io.dongtai.iast.core.handler.hookpoint.IastClassLoader;
87
import io.dongtai.iast.core.utils.*;
98
import io.dongtai.iast.core.utils.matcher.ConfigMatcher;
@@ -74,10 +73,6 @@ public static void solveHttpRequest(Object obj, Object req, Object resp, Map<Str
7473
if (ConfigMatcher.getInstance().disableExtension((String) requestMeta.get("requestURI"))) {
7574
return;
7675
}
77-
if (ConfigMatcher.getInstance().getBlackUrl(requestMeta)) {
78-
BlackUrlSkipHandler.setIsBlackUrl(true);
79-
return;
80-
}
8176

8277
try {
8378
boolean enableVersionHeader = ConfigBuilder.getInstance().get(ConfigKey.ENABLE_VERSION_HEADER);

‎dongtai-log/src/main/java/io/dongtai/log/ErrorCode.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ public enum ErrorCode {
6565
TRANSFORM_ENGINE_DESTROY_REDEFINE_CLASSES_FAILED(20122, "transform engine failed to redefine classes when destroy"),
6666
ENGINE_PROPERTIES_INITIALIZE_FAILED(20131, "engine properties initialize failed"),
6767
CLASS_DIAGRAM_SCAN_JAR_ANCESTOR_FAILED(20141, "class diagram scan jar ancestor failed"),
68+
REDEFINE_MODULE_FAILED(20142, "There was a problem redefining the java.base module"),
6869

6970
// transform
7071
TRANSFORM_CLASS_FAILED(20201, "transform class {} failed"),
@@ -88,7 +89,8 @@ public enum ErrorCode {
8889
SPY_TRACE_DUBBO_CONSUMER_INVOKE_FAILED(20361, "hookpoint trace dubbo consumer invoke failed"),
8990
SPY_LEAVE_DUBBO_FAILED(20362, "hookpoint leave dubbo failed"),
9091
SPY_COLLECT_DUBBO_FAILED(20363, "hookpoint collect dubbo {} failed"),
91-
SPY_SKIP_COLLECT_DUBBO_FAILED(20364, "hookpoint skip collect dubbo {} failed"),
92+
BYPASS_FAILED_DUBBO(20364, "hookpoint skip collect dubbo {} failed"),
93+
BYPASS_FAILED_FEIGN(20365, "hookpoint skip collect feign {} failed"),
9294

9395
// report & replay
9496
REPORT_SEND_FAILED(20401, "send report to {} error, report: {}"),

‎dongtai-spy/src/main/java/java/lang/dongtai/NopSpy.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -263,7 +263,12 @@ public boolean traceDubboInvoke(Object instance, String url, Object invocation,
263263
}
264264

265265
@Override
266-
public boolean isSkipCollect(Object invocation) {
266+
public boolean isSkipCollectDubbo(Object invocation) {
267+
return false;
268+
}
269+
270+
@Override
271+
public boolean isSkipCollectFeign(Object instance) {
267272
return false;
268273
}
269274

0 commit comments

Comments
 (0)