Skip to content

Commit 7a0200e

Browse files
lostsnowlostsnow
authored
Merge pull request #291 from lostsnow/fix/kafka-bytecode
fixes kafka bytecode, update service report
2 parents 1d33977 + 0cdbfad commit 7a0200e

18 files changed

+96
-108
lines changed

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

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
11
package io.dongtai.iast.core.bytecode.enhance.asm;
22

3-
import java.lang.dongtai.ServiceUrlHandler;
3+
import org.objectweb.asm.commons.Method;
4+
45
import java.lang.dongtai.SpyDispatcher;
56
import java.lang.dongtai.SpyDispatcherHandler;
67

7-
import org.objectweb.asm.commons.Method;
8-
98
/**
109
* 常用的ASM method 集合 省得我到处声明
1110
*
@@ -81,12 +80,6 @@ static Method getAsmMethod(final Class<?> clazz,
8180
"kafkaBeforeSend",
8281
Object.class
8382
);
84-
Method SPY$kafkaAfterSend = InnerHelper.getAsmMethod(
85-
SpyDispatcher.class,
86-
"kafkaAfterSend",
87-
Object.class,
88-
Object.class
89-
);
9083
Method SPY$kafkaAfterPoll = InnerHelper.getAsmMethod(
9184
SpyDispatcher.class,
9285
"kafkaAfterPoll",
@@ -187,6 +180,6 @@ static Method getAsmMethod(final Class<?> clazz,
187180
String.class,
188181
String.class,
189182
String.class,
190-
ServiceUrlHandler.class
183+
String.class
191184
);
192185
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/grpc/ServerStreamListenerImplClosedAdviceAdapter.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.grpc;
2+
3+
import io.dongtai.iast.core.bytecode.enhance.asm.AsmMethods;
4+
import io.dongtai.iast.core.bytecode.enhance.asm.AsmTypes;
5+
import io.dongtai.iast.core.utils.AsmUtils;
6+
import org.objectweb.asm.MethodVisitor;
7+
import org.objectweb.asm.commons.AdviceAdapter;
8+
9+
public class ServerStreamListenerImplClosedAdviceAdapter extends AdviceAdapter implements AsmTypes, AsmMethods {
10+
protected ServerStreamListenerImplClosedAdviceAdapter(MethodVisitor methodVisitor, int access, String name, String descriptor) {
11+
super(AsmUtils.api, methodVisitor, access, name, descriptor);
12+
}
13+
14+
@Override
15+
protected void onMethodExit(int opcode) {
16+
if (opcode != ATHROW) {
17+
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
18+
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$closeGrpcCall);
19+
}
20+
}
21+
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/jdbc/MysqlHostInfoAdviceAdapter.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ protected void onMethodExit(int opcode) {
2121
loadArg(1);
2222
loadArg(2);
2323
mv.visitMethodInsn(INVOKESTATIC, "java/lang/String", "valueOf", "(I)Ljava/lang/String;", false);
24-
mv.visitInsn(ACONST_NULL);
24+
push("");
2525
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$reportService);
2626
}
2727
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/jdbc/MysqlJdbcDriverParseUrlAdviceAdapter.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ protected void onMethodExit(int opcode) {
3131
loadLocal(nextLocal - 1);
3232
push("PORT");
3333
mv.visitMethodInsn(INVOKEVIRTUAL, "java/util/Properties", "getProperty", "(Ljava/lang/String;)Ljava/lang/String;", false);
34-
mv.visitInsn(ACONST_NULL);
34+
push("");
3535
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$reportService);
3636
mark(nonNullLabel);
3737
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/jdbc/PostgresqlDriverParseUrlAdviceAdapter.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ protected void onMethodExit(int opcode) {
3131
loadLocal(nextLocal - 1);
3232
push("PGPORT");
3333
mv.visitMethodInsn(INVOKEVIRTUAL, "java/util/Properties", "getProperty", "(Ljava/lang/String;)Ljava/lang/String;", false);
34-
mv.visitInsn(ACONST_NULL);
34+
push("");
3535
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$reportService);
3636
mark(nonNullLabel);
3737
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaConsumerAdapter.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,16 @@ public MethodVisitor visitMethod(final int access, final String name, final Stri
1717
MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
1818
int argCount = Type.getArgumentTypes(desc).length;
1919

20-
if ("<init>".equals(name) && "(Ljava/util/Map;Lorg/apache/kafka/common/serialization/Deserializer;Lorg/apache/kafka/common/serialization/Deserializer;)V".equals(desc)) {
20+
if ("<init>".equals(name)) {
2121
if (DongTaiLog.isDebugEnabled()) {
2222
DongTaiLog.debug("Adding kafka tracking for type {}.{}", context.getClassName(), name);
2323
}
2424

25-
mv = new KafkaConsumerAdviceAdapter(mv, access, name, desc);
25+
if ("(Ljava/util/Map;Lorg/apache/kafka/common/serialization/Deserializer;Lorg/apache/kafka/common/serialization/Deserializer;)V".equals(desc)) {
26+
mv = new KafkaConsumerAdviceAdapter(mv, access, name, desc);
27+
} else if ("(Lorg/apache/kafka/clients/consumer/ConsumerConfig;Lorg/apache/kafka/common/serialization/Deserializer;Lorg/apache/kafka/common/serialization/Deserializer;)V".equals(desc)) {
28+
mv = new KafkaConsumerInitAdviceAdapter(mv, access, name, desc);
29+
}
2630
setTransformed();
2731
} else if ("poll".equals(name) && argCount == 2) {
2832
if (DongTaiLog.isDebugEnabled()) {

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaConsumerAdviceAdapter.java‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
public class KafkaConsumerAdviceAdapter extends AdviceAdapter implements AsmTypes, AsmMethods {
1414
private int localServers;
1515
private int localServersString;
16-
private int localUrlHandler;
1716
protected KafkaConsumerAdviceAdapter(MethodVisitor mv, int access, String name, String desc) {
1817
super(AsmUtils.api, mv, access, name, desc);
1918
}
@@ -34,19 +33,12 @@ protected void onMethodExit(int opcode) {
3433
mv.visitMethodInsn(INVOKESTATIC, "java/lang/String", "join", "(Ljava/lang/CharSequence;Ljava/lang/Iterable;)Ljava/lang/String;", false);
3534
storeLocal(localServersString);
3635

37-
localUrlHandler = newLocal(ASM_TYPE_OBJECT);
38-
mv.visitTypeInsn(NEW, "io/dongtai/iast/core/handler/hookpoint/service/url/KafkaUrlHandler");
39-
dup();
40-
mv.visitMethodInsn(INVOKESPECIAL, "io/dongtai/iast/core/handler/hookpoint/service/url/KafkaUrlHandler",
41-
"<init>", "()V", false);
42-
storeLocal(localUrlHandler);
43-
4436
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
4537
push(ServiceType.KAFKA.getCategory());
4638
push(ServiceType.KAFKA.getType());
4739
loadLocal(localServersString);
4840
push("");
49-
loadLocal(localUrlHandler);
41+
push("KafkaUrlHandler");
5042
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$reportService);
5143
}
5244
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaConsumerInitAdviceAdapter.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
package io.dongtai.iast.core.bytecode.enhance.plugin.service.kafka;
2+
3+
import io.dongtai.iast.core.bytecode.enhance.asm.AsmMethods;
4+
import io.dongtai.iast.core.bytecode.enhance.asm.AsmTypes;
5+
import io.dongtai.iast.core.handler.hookpoint.service.ServiceType;
6+
import io.dongtai.iast.core.utils.AsmUtils;
7+
import org.objectweb.asm.MethodVisitor;
8+
import org.objectweb.asm.Type;
9+
import org.objectweb.asm.commons.AdviceAdapter;
10+
11+
public class KafkaConsumerInitAdviceAdapter extends AdviceAdapter implements AsmTypes, AsmMethods {
12+
private int localServersString;
13+
protected KafkaConsumerInitAdviceAdapter(MethodVisitor mv, int access, String name, String desc) {
14+
super(AsmUtils.api, mv, access, name, desc);
15+
}
16+
17+
@Override
18+
protected void onMethodExit(int opcode) {
19+
if (opcode != ATHROW) {
20+
localServersString = newLocal(Type.getType(String.class));
21+
loadArg(0);
22+
push("bootstrap.servers");
23+
mv.visitMethodInsn(INVOKEINTERFACE, " org/apache/kafka/clients/consumer/ConsumerConfig".substring(1),
24+
"getString", "(Ljava/lang/String;)Ljava/lang/String;", false);
25+
storeLocal(localServersString);
26+
27+
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
28+
push(ServiceType.KAFKA.getCategory());
29+
push(ServiceType.KAFKA.getType());
30+
loadLocal(localServersString);
31+
push("");
32+
push("KafkaUrlHandler");
33+
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$reportService);
34+
}
35+
}
36+
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaProducerAdviceAdapter.java‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
public class KafkaProducerAdviceAdapter extends AdviceAdapter implements AsmTypes, AsmMethods {
1414
private int localServers;
1515
private int localServersString;
16-
private int localUrlHandler;
1716
protected KafkaProducerAdviceAdapter(MethodVisitor mv, int access, String name, String desc) {
1817
super(AsmUtils.api, mv, access, name, desc);
1918
}
@@ -34,19 +33,12 @@ protected void onMethodExit(int opcode) {
3433
mv.visitMethodInsn(INVOKESTATIC, "java/lang/String", "join", "(Ljava/lang/CharSequence;Ljava/lang/Iterable;)Ljava/lang/String;", false);
3534
storeLocal(localServersString);
3635

37-
localUrlHandler = newLocal(ASM_TYPE_OBJECT);
38-
mv.visitTypeInsn(NEW, "io/dongtai/iast/core/handler/hookpoint/service/url/KafkaUrlHandler");
39-
dup();
40-
mv.visitMethodInsn(INVOKESPECIAL, "io/dongtai/iast/core/handler/hookpoint/service/url/KafkaUrlHandler",
41-
"<init>", "()V", false);
42-
storeLocal(localUrlHandler);
43-
4436
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
4537
push(ServiceType.KAFKA.getCategory());
4638
push(ServiceType.KAFKA.getType());
4739
loadLocal(localServersString);
4840
push("");
49-
loadLocal(localUrlHandler);
41+
push("KafkaUrlHandler");
5042
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$reportService);
5143
}
5244
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaProducerSendAdviceAdapter.java‎

Lines changed: 0 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -7,33 +7,14 @@
77
import org.objectweb.asm.commons.AdviceAdapter;
88

99
public class KafkaProducerSendAdviceAdapter extends AdviceAdapter implements AsmTypes, AsmMethods {
10-
private int localRecord;
1110
protected KafkaProducerSendAdviceAdapter(MethodVisitor mv, int access, String name, String desc) {
1211
super(AsmUtils.api, mv, access, name, desc);
1312
}
1413

1514
@Override
1615
protected void onMethodEnter() {
17-
loadArg(0);
18-
localRecord = newLocal(ASM_TYPE_OBJECT);
19-
dup();
20-
storeLocal(localRecord);
2116
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
2217
loadArg(0);
2318
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$kafkaBeforeSend);
24-
storeArg(0);
25-
}
26-
27-
@Override
28-
protected void onMethodExit(int opcode) {
29-
if (opcode != ATHROW) {
30-
newLocal(ASM_TYPE_OBJECT);
31-
dup();
32-
storeLocal(nextLocal - 1);
33-
invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
34-
loadLocal(localRecord);
35-
loadLocal(nextLocal - 1);
36-
invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$kafkaAfterSend);
37-
}
3819
}
3920
}

0 commit comments

Comments
 (0)