From 2d7d3731091c10d60e27f0f85db61ddc3ecad325 Mon Sep 17 00:00:00 2001 From: lostsnow Date: Tue, 17 May 2022 21:16:10 +0800 Subject: [PATCH 1/2] fixes kafka service detection --- .../plugin/service/kafka/DispatchKafka.java | 3 ++ .../kafka/KafkaAbstractConfigAdapter.java | 27 +++++++++++ ...KafkaAbstractConfigInitAdviceAdapter.java} | 13 +++--- .../service/kafka/KafkaConsumerAdapter.java | 10 +---- .../kafka/KafkaConsumerInitAdviceAdapter.java | 36 --------------- .../service/kafka/KafkaProducerAdapter.java | 6 +-- .../kafka/KafkaProducerAdviceAdapter.java | 45 ------------------- 7 files changed, 38 insertions(+), 102 deletions(-) create mode 100644 dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaAbstractConfigAdapter.java rename dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/{KafkaConsumerAdviceAdapter.java => KafkaAbstractConfigInitAdviceAdapter.java} (73%) delete mode 100644 dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaConsumerInitAdviceAdapter.java delete mode 100644 dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaProducerAdviceAdapter.java diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/DispatchKafka.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/DispatchKafka.java index 749baba8b..22ad6413b 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/DispatchKafka.java +++ b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/DispatchKafka.java @@ -7,6 +7,7 @@ public class DispatchKafka implements DispatchPlugin { private final String classOfKafkaProducer = " org.apache.kafka.clients.producer.KafkaProducer".substring(1); private final String classOfKafkaConsumer = " org.apache.kafka.clients.consumer.KafkaConsumer".substring(1); + private final String classOfAbstractConfig = " org.apache.kafka.common.config.AbstractConfig".substring(1); private final String classOfSpringKafkaMessageListenerContainer = " org.springframework.kafka.listener.KafkaMessageListenerContainer$ListenerConsumer".substring(1); @Override @@ -17,6 +18,8 @@ public ClassVisitor dispatch(ClassVisitor classVisitor, IastContext context) { classVisitor = new KafkaProducerAdapter(classVisitor, context); } else if (classOfKafkaConsumer.equals(className)) { classVisitor = new KafkaConsumerAdapter(classVisitor, context); + } else if (classOfAbstractConfig.equals(className)) { + classVisitor = new KafkaAbstractConfigAdapter(classVisitor, context); } else if (classOfSpringKafkaMessageListenerContainer.equals(className)) { classVisitor = new SpringKafkaMessageListenerContainerAdapter(classVisitor, context); } diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaAbstractConfigAdapter.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaAbstractConfigAdapter.java new file mode 100644 index 000000000..abcdbb7f7 --- /dev/null +++ b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaAbstractConfigAdapter.java @@ -0,0 +1,27 @@ +package io.dongtai.iast.core.bytecode.enhance.plugin.service.kafka; + +import io.dongtai.iast.core.bytecode.enhance.IastContext; +import io.dongtai.iast.core.bytecode.enhance.plugin.AbstractClassVisitor; +import io.dongtai.log.DongTaiLog; +import org.objectweb.asm.*; + +public class KafkaAbstractConfigAdapter extends AbstractClassVisitor { + private String classDesc; + + public KafkaAbstractConfigAdapter(ClassVisitor classVisitor, IastContext context) { + super(classVisitor, context); + } + + @Override + public MethodVisitor visitMethod(final int access, final String name, final String desc, final String signature, final String[] exceptions) { + MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions); + int argCount = Type.getArgumentTypes(desc).length; + + if ("".equals(name) && argCount >= 3) { + DongTaiLog.debug("Adding kafka tracking for type {}.{}", context.getClassName(), name); + mv = new KafkaAbstractConfigInitAdviceAdapter(mv, access, name, desc); + setTransformed(); + } + return mv; + } +} diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaConsumerAdviceAdapter.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaAbstractConfigInitAdviceAdapter.java similarity index 73% rename from dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaConsumerAdviceAdapter.java rename to dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaAbstractConfigInitAdviceAdapter.java index 5ffa1091a..b1e82020f 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaConsumerAdviceAdapter.java +++ b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaAbstractConfigInitAdviceAdapter.java @@ -8,23 +8,22 @@ import org.objectweb.asm.Type; import org.objectweb.asm.commons.AdviceAdapter; -import java.util.ArrayList; +import java.util.List; -public class KafkaConsumerAdviceAdapter extends AdviceAdapter implements AsmTypes, AsmMethods { +public class KafkaAbstractConfigInitAdviceAdapter extends AdviceAdapter implements AsmTypes, AsmMethods { private int localServers; private int localServersString; - protected KafkaConsumerAdviceAdapter(MethodVisitor mv, int access, String name, String desc) { + protected KafkaAbstractConfigInitAdviceAdapter(MethodVisitor mv, int access, String name, String desc) { super(AsmUtils.api, mv, access, name, desc); } @Override protected void onMethodExit(int opcode) { if (opcode != ATHROW) { - localServers = newLocal(Type.getType(ArrayList.class)); - loadArg(0); + localServers = newLocal(Type.getType(List.class)); + loadThis(); push("bootstrap.servers"); - mv.visitMethodInsn(INVOKEINTERFACE, "java/util/Map", "get", "(Ljava/lang/Object;)Ljava/lang/Object;", true); - mv.visitTypeInsn(CHECKCAST, "java/util/ArrayList"); + mv.visitMethodInsn(INVOKEVIRTUAL, " org/apache/kafka/common/config/AbstractConfig".substring(1), "getList", "(Ljava/lang/String;)Ljava/util/List;", false); storeLocal(localServers); localServersString = newLocal(Type.getType(String.class)); diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaConsumerAdapter.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaConsumerAdapter.java index 06e6df88b..dc05fd3fa 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaConsumerAdapter.java +++ b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaConsumerAdapter.java @@ -17,15 +17,7 @@ public MethodVisitor visitMethod(final int access, final String name, final Stri MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions); int argCount = Type.getArgumentTypes(desc).length; - if ("".equals(name)) { - DongTaiLog.debug("Adding kafka tracking for type {}.{}", context.getClassName(), name); - if ("(Ljava/util/Map;Lorg/apache/kafka/common/serialization/Deserializer;Lorg/apache/kafka/common/serialization/Deserializer;)V".equals(desc)) { - mv = new KafkaConsumerAdviceAdapter(mv, access, name, desc); - } else if ("(Lorg/apache/kafka/clients/consumer/ConsumerConfig;Lorg/apache/kafka/common/serialization/Deserializer;Lorg/apache/kafka/common/serialization/Deserializer;)V".equals(desc)) { - mv = new KafkaConsumerInitAdviceAdapter(mv, access, name, desc); - } - setTransformed(); - } else if ("poll".equals(name) && argCount == 2) { + if ("poll".equals(name) && argCount == 2) { DongTaiLog.debug("Adding kafka tracking for type {}.{}", context.getClassName(), name); mv = new KafkaConsumerPollAdviceAdapter(mv, access, name, desc); diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaConsumerInitAdviceAdapter.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaConsumerInitAdviceAdapter.java deleted file mode 100644 index 58889f9b7..000000000 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaConsumerInitAdviceAdapter.java +++ /dev/null @@ -1,36 +0,0 @@ -package io.dongtai.iast.core.bytecode.enhance.plugin.service.kafka; - -import io.dongtai.iast.core.bytecode.enhance.asm.AsmMethods; -import io.dongtai.iast.core.bytecode.enhance.asm.AsmTypes; -import io.dongtai.iast.core.handler.hookpoint.service.ServiceType; -import io.dongtai.iast.core.utils.AsmUtils; -import org.objectweb.asm.MethodVisitor; -import org.objectweb.asm.Type; -import org.objectweb.asm.commons.AdviceAdapter; - -public class KafkaConsumerInitAdviceAdapter extends AdviceAdapter implements AsmTypes, AsmMethods { - private int localServersString; - protected KafkaConsumerInitAdviceAdapter(MethodVisitor mv, int access, String name, String desc) { - super(AsmUtils.api, mv, access, name, desc); - } - - @Override - protected void onMethodExit(int opcode) { - if (opcode != ATHROW) { - localServersString = newLocal(Type.getType(String.class)); - loadArg(0); - push("bootstrap.servers"); - mv.visitMethodInsn(INVOKEINTERFACE, " org/apache/kafka/clients/consumer/ConsumerConfig".substring(1), - "getString", "(Ljava/lang/String;)Ljava/lang/String;", false); - storeLocal(localServersString); - - invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher); - push(ServiceType.KAFKA.getCategory()); - push(ServiceType.KAFKA.getType()); - loadLocal(localServersString); - push(""); - push("KafkaUrlHandler"); - invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$reportService); - } - } -} diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaProducerAdapter.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaProducerAdapter.java index 35b2df28b..89640a348 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaProducerAdapter.java +++ b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaProducerAdapter.java @@ -15,11 +15,7 @@ public MethodVisitor visitMethod(final int access, final String name, final Stri MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions); int argCount = Type.getArgumentTypes(desc).length; - if ("".equals(name) && argCount == 7) { - DongTaiLog.debug("Adding kafka tracking for type {}.{}", context.getClassName(), name); - mv = new KafkaProducerAdviceAdapter(mv, access, name, desc); - setTransformed(); - } else if ("send".equals(name) && argCount == 2) { + if ("send".equals(name) && argCount == 2) { DongTaiLog.debug("Adding kafka tracking for type {}.{}", context.getClassName(), name); mv = new KafkaProducerSendAdviceAdapter(mv, access, name, desc); setTransformed(); diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaProducerAdviceAdapter.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaProducerAdviceAdapter.java deleted file mode 100644 index 0e7581d7b..000000000 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/service/kafka/KafkaProducerAdviceAdapter.java +++ /dev/null @@ -1,45 +0,0 @@ -package io.dongtai.iast.core.bytecode.enhance.plugin.service.kafka; - -import io.dongtai.iast.core.bytecode.enhance.asm.AsmMethods; -import io.dongtai.iast.core.bytecode.enhance.asm.AsmTypes; -import io.dongtai.iast.core.handler.hookpoint.service.ServiceType; -import io.dongtai.iast.core.utils.AsmUtils; -import org.objectweb.asm.MethodVisitor; -import org.objectweb.asm.Type; -import org.objectweb.asm.commons.AdviceAdapter; - -import java.util.ArrayList; - -public class KafkaProducerAdviceAdapter extends AdviceAdapter implements AsmTypes, AsmMethods { - private int localServers; - private int localServersString; - protected KafkaProducerAdviceAdapter(MethodVisitor mv, int access, String name, String desc) { - super(AsmUtils.api, mv, access, name, desc); - } - - @Override - protected void onMethodExit(int opcode) { - if (opcode != ATHROW) { - localServers = newLocal(Type.getType(ArrayList.class)); - loadArg(0); - push("bootstrap.servers"); - mv.visitMethodInsn(INVOKEINTERFACE, "java/util/Map", "get", "(Ljava/lang/Object;)Ljava/lang/Object;", true); - mv.visitTypeInsn(CHECKCAST, "java/util/ArrayList"); - storeLocal(localServers); - - localServersString = newLocal(Type.getType(String.class)); - push(","); - loadLocal(localServers); - mv.visitMethodInsn(INVOKESTATIC, "java/lang/String", "join", "(Ljava/lang/CharSequence;Ljava/lang/Iterable;)Ljava/lang/String;", false); - storeLocal(localServersString); - - invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher); - push(ServiceType.KAFKA.getCategory()); - push(ServiceType.KAFKA.getType()); - loadLocal(localServersString); - push(""); - push("KafkaUrlHandler"); - invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$reportService); - } - } -} From 411021f54f27fcb41061f497d901b32361fee2e0 Mon Sep 17 00:00:00 2001 From: lostsnow Date: Wed, 18 May 2022 12:11:02 +0800 Subject: [PATCH 2/2] fixes taint pool null exception --- .../handler/hookpoint/SpyDispatcherImpl.java | 2 +- .../controller/impl/PropagatorImpl.java | 12 +- .../framework/dubbo/DubboHandler.java | 64 +++---- .../hookpoint/framework/grpc/GrpcHandler.java | 161 +++++++++--------- .../hookpoint/service/kafka/KafkaHandler.java | 2 +- .../core/utils/threadlocal/IastTaintPool.java | 13 +- 6 files changed, 127 insertions(+), 127 deletions(-) diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/SpyDispatcherImpl.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/SpyDispatcherImpl.java index c824f59d8..a25011013 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/SpyDispatcherImpl.java +++ b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/SpyDispatcherImpl.java @@ -483,7 +483,7 @@ public boolean collectMethodPool(Object instance, Object[] argumentArray, Object HttpImpl.solveHttp(event); } else if (HookType.RPC.equals(hookType)) { solveRPC(framework, event); - } else if (HookType.PROPAGATOR.equals(hookType) && !EngineManager.TAINT_POOL.get().isEmpty()) { + } else if (HookType.PROPAGATOR.equals(hookType) && !EngineManager.TAINT_POOL.isEmpty()) { PropagatorImpl.solvePropagator(event, INVOKE_ID_SEQUENCER); } else if (HookType.SOURCE.equals(hookType)) { SourceImpl.solveSource(event, INVOKE_ID_SEQUENCER); diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/PropagatorImpl.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/PropagatorImpl.java index b45e49cc1..c3a1d5c37 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/PropagatorImpl.java +++ b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/PropagatorImpl.java @@ -1,19 +1,13 @@ package io.dongtai.iast.core.handler.hookpoint.controller.impl; import io.dongtai.iast.core.EngineManager; -import io.dongtai.iast.core.handler.hookpoint.models.IastHookRuleModel; -import io.dongtai.iast.core.handler.hookpoint.models.IastPropagatorModel; -import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent; +import io.dongtai.iast.core.handler.hookpoint.models.*; import io.dongtai.iast.core.handler.hookpoint.vulscan.dynamic.TrackUtils; import io.dongtai.iast.core.utils.StackUtils; import io.dongtai.iast.core.utils.TaintPoolUtils; -import io.dongtai.log.DongTaiLog; import java.lang.reflect.Method; -import java.util.ArrayList; -import java.util.HashSet; -import java.util.List; -import java.util.Set; +import java.util.*; import java.util.concurrent.atomic.AtomicInteger; import static io.dongtai.iast.core.utils.HashCode.isNotEmpty; @@ -35,7 +29,7 @@ public class PropagatorImpl { private static final String SPRING_OBJECT = " org.springframework.".substring(1); public static void solvePropagator(MethodEvent event, AtomicInteger invokeIdSequencer) { - if (!EngineManager.TAINT_POOL.get().isEmpty()) { + if (!EngineManager.TAINT_POOL.isEmpty()) { IastPropagatorModel propagator = IastHookRuleModel.getPropagatorByMethodSignature(event.signature); if (propagator != null) { auxiliaryPropagator(propagator, invokeIdSequencer, event); diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/framework/dubbo/DubboHandler.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/framework/dubbo/DubboHandler.java index 16eb43f8e..33b9eb109 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/framework/dubbo/DubboHandler.java +++ b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/framework/dubbo/DubboHandler.java @@ -192,7 +192,7 @@ private static boolean isNotEmpty(Object obj) { } public static void solveClientExit(Object invocation, Object rpcResult) { - if (EngineManager.TAINT_POOL.get().isEmpty()) { + if (EngineManager.TAINT_POOL.isEmpty()) { return; } @@ -259,36 +259,38 @@ public static void solveClientExit(Object invocation, Object rpcResult) { public static void solveServiceExit(Object invocation, Object rpcResult) { try { - if (null != EngineManager.TAINT_POOL.get() && !EngineManager.TAINT_POOL.get().isEmpty()) { - MethodEvent event = new MethodEvent( - 0, - 0, - "*.dubbo.monitor.support.MonitorFilter", - "*.dubbo.monitor.support.MonitorFilter", - "invoke", - "com.alibaba.dubbo.monitor.support.MonitorFilter#invoke", - "com.alibaba.dubbo.monitor.support.MonitorFilter#invoke", - null, - new Object[]{rpcResult}, - null, - "DUBBO", - false, - null - ); - Set modelItems = SourceImpl.parseCustomModel(rpcResult); - boolean isHitTaints = false; - for (Object item : modelItems) { - isHitTaints = isHitTaints || TaintPoolUtils.poolContains(item, event, false); - } - if (isHitTaints) { - int invokeId = SpyDispatcherImpl.INVOKE_ID_SEQUENCER.getAndIncrement(); - event.setInvokeId(invokeId); - event.setPlugin("DUBBO"); - event.setServiceName(""); - event.setProjectPropagatorClose(true); - event.setCallStack(StackUtils.getLatestStack(5)); - EngineManager.TRACK_MAP.addTrackMethod(invokeId, event); - } + if (EngineManager.TAINT_POOL.isEmpty()) { + return; + } + + MethodEvent event = new MethodEvent( + 0, + 0, + "*.dubbo.monitor.support.MonitorFilter", + "*.dubbo.monitor.support.MonitorFilter", + "invoke", + "com.alibaba.dubbo.monitor.support.MonitorFilter#invoke", + "com.alibaba.dubbo.monitor.support.MonitorFilter#invoke", + null, + new Object[]{rpcResult}, + null, + "DUBBO", + false, + null + ); + Set modelItems = SourceImpl.parseCustomModel(rpcResult); + boolean isHitTaints = false; + for (Object item : modelItems) { + isHitTaints = isHitTaints || TaintPoolUtils.poolContains(item, event, false); + } + if (isHitTaints) { + int invokeId = SpyDispatcherImpl.INVOKE_ID_SEQUENCER.getAndIncrement(); + event.setInvokeId(invokeId); + event.setPlugin("DUBBO"); + event.setServiceName(""); + event.setProjectPropagatorClose(true); + event.setCallStack(StackUtils.getLatestStack(5)); + EngineManager.TRACK_MAP.addTrackMethod(invokeId, event); } }catch (Exception e){ DongTaiLog.debug(e); diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/framework/grpc/GrpcHandler.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/framework/grpc/GrpcHandler.java index e8eb0546f..dad633cf0 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/framework/grpc/GrpcHandler.java +++ b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/framework/grpc/GrpcHandler.java @@ -155,90 +155,93 @@ public static void closeGrpcCall() { public static void blockingUnaryCall(Object req, Object res) { // todo: 判断 req 的相关自定义对象是否与污点有关 - if (!EngineManager.TAINT_POOL.get().isEmpty()) { - MethodEvent event = new MethodEvent( - 0, - 0, - "io.grpc.stub.ClientCalls", - "io.grpc.stub.ClientCalls", - "blockingUnaryCall", - "io.grpc.stub.ClientCalls.blockingUnaryCall(io.grpc.Channel, io.grpc.MethodDescriptor, io.grpc.CallOptions, ReqT)", - "io.grpc.stub.ClientCalls.blockingUnaryCall(io.grpc.Channel, io.grpc.MethodDescriptor, io.grpc.CallOptions, ReqT)", - null, - new Object[]{req}, - res, - "GRPC", - false, - null - ); - Set modelItems = SourceImpl.parseCustomModel(req); - boolean isHitTaints = false; - for (Object item : modelItems) { - isHitTaints = isHitTaints || TaintPoolUtils.poolContains(item, event, false); - } - if (isHitTaints) { - int invokeId = SpyDispatcherImpl.INVOKE_ID_SEQUENCER.getAndIncrement(); - event.setInvokeId(invokeId); - event.setPlugin("GRPC"); - // todo: 获取 service name - event.setServiceName(""); - // todo: 获取 traceId - event.setTraceId(sharedTraceId.get()); - event.setCallStack(StackUtils.getLatestStack(5)); - EngineManager.TRACK_MAP.addTrackMethod(invokeId, event); - Set resModelItems = SourceImpl.parseCustomModel(res); - sharedRespData.remove(); - Set taintPool = EngineManager.TAINT_POOL.get(); - Set resModelSet = new HashSet(); - for (Object obj : resModelItems) { - // fixme: 暂时只跟踪字符串相关内容 - if (obj instanceof String) { - resModelSet.add(obj); - addCustomResp.set(true); - taintPool.add(obj); - int identityHashCode = System.identityHashCode(obj); - event.addTargetHash(identityHashCode); - event.addTargetHashForRpc(obj.hashCode()); - EngineManager.TAINT_HASH_CODES.get().add(identityHashCode); - } + if (EngineManager.TAINT_POOL.isEmpty()) { + return; + } + + MethodEvent event = new MethodEvent( + 0, + 0, + "io.grpc.stub.ClientCalls", + "io.grpc.stub.ClientCalls", + "blockingUnaryCall", + "io.grpc.stub.ClientCalls.blockingUnaryCall(io.grpc.Channel, io.grpc.MethodDescriptor, io.grpc.CallOptions, ReqT)", + "io.grpc.stub.ClientCalls.blockingUnaryCall(io.grpc.Channel, io.grpc.MethodDescriptor, io.grpc.CallOptions, ReqT)", + null, + new Object[]{req}, + res, + "GRPC", + false, + null + ); + Set modelItems = SourceImpl.parseCustomModel(req); + boolean isHitTaints = false; + for (Object item : modelItems) { + isHitTaints = isHitTaints || TaintPoolUtils.poolContains(item, event, false); + } + if (isHitTaints) { + int invokeId = SpyDispatcherImpl.INVOKE_ID_SEQUENCER.getAndIncrement(); + event.setInvokeId(invokeId); + event.setPlugin("GRPC"); + // todo: 获取 service name + event.setServiceName(""); + // todo: 获取 traceId + event.setTraceId(sharedTraceId.get()); + event.setCallStack(StackUtils.getLatestStack(5)); + EngineManager.TRACK_MAP.addTrackMethod(invokeId, event); + Set resModelItems = SourceImpl.parseCustomModel(res); + sharedRespData.remove(); + Set taintPool = EngineManager.TAINT_POOL.get(); + Set resModelSet = new HashSet(); + for (Object obj : resModelItems) { + // fixme: 暂时只跟踪字符串相关内容 + if (obj instanceof String) { + resModelSet.add(obj); + addCustomResp.set(true); + taintPool.add(obj); + int identityHashCode = System.identityHashCode(obj); + event.addTargetHash(identityHashCode); + event.addTargetHashForRpc(obj.hashCode()); + EngineManager.TAINT_HASH_CODES.get().add(identityHashCode); } - sharedRespData.set(resModelSet); } + sharedRespData.set(resModelSet); } - } public static void sendMessage(Object message) { - if (!EngineManager.TAINT_POOL.get().isEmpty()) { - MethodEvent event = new MethodEvent( - 0, - 0, - "io.grpc.internal.ServerCallImpl", - "io.grpc.internal.ServerCallImpl", - "sendMessage", - "io.grpc.internal.ServerCallImpl.sendMessage(RespT)", - "io.grpc.internal.ServerCallImpl.sendMessage(RespT)", - null, - new Object[]{message}, - null, - "GRPC", - false, - null - ); - Set modelItems = SourceImpl.parseCustomModel(message); - boolean isHitTaints = false; - for (Object item : modelItems) { - isHitTaints = isHitTaints || TaintPoolUtils.poolContains(item, event, false); - } - if (isHitTaints) { - int invokeId = SpyDispatcherImpl.INVOKE_ID_SEQUENCER.getAndIncrement(); - event.setInvokeId(invokeId); - event.setPlugin("GRPC"); - event.setServiceName(""); - event.setProjectPropagatorClose(true); - event.setCallStack(StackUtils.getLatestStack(5)); - EngineManager.TRACK_MAP.addTrackMethod(invokeId, event); - } + if (EngineManager.TAINT_POOL.isEmpty()) { + return; + } + + MethodEvent event = new MethodEvent( + 0, + 0, + "io.grpc.internal.ServerCallImpl", + "io.grpc.internal.ServerCallImpl", + "sendMessage", + "io.grpc.internal.ServerCallImpl.sendMessage(RespT)", + "io.grpc.internal.ServerCallImpl.sendMessage(RespT)", + null, + new Object[]{message}, + null, + "GRPC", + false, + null + ); + Set modelItems = SourceImpl.parseCustomModel(message); + boolean isHitTaints = false; + for (Object item : modelItems) { + isHitTaints = isHitTaints || TaintPoolUtils.poolContains(item, event, false); + } + if (isHitTaints) { + int invokeId = SpyDispatcherImpl.INVOKE_ID_SEQUENCER.getAndIncrement(); + event.setInvokeId(invokeId); + event.setPlugin("GRPC"); + event.setServiceName(""); + event.setProjectPropagatorClose(true); + event.setCallStack(StackUtils.getLatestStack(5)); + EngineManager.TRACK_MAP.addTrackMethod(invokeId, event); } } @@ -246,7 +249,7 @@ public static void toStringUtf8(Object value) { Boolean added = addCustomResp.get(); if (added != null && added) { if (sharedRespData.get().contains(value)) { - EngineManager.TAINT_POOL.get().add(value); + EngineManager.TAINT_POOL.addToPool(value); } } } diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/service/kafka/KafkaHandler.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/service/kafka/KafkaHandler.java index f0cf46974..f98089582 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/service/kafka/KafkaHandler.java +++ b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/service/kafka/KafkaHandler.java @@ -29,7 +29,7 @@ public static void beforeSend(Object record) { } public static void trackSend(Object record) { - if (EngineManager.TAINT_POOL.get().isEmpty()) { + if (EngineManager.TAINT_POOL.isEmpty()) { return; } diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/utils/threadlocal/IastTaintPool.java b/dongtai-core/src/main/java/io/dongtai/iast/core/utils/threadlocal/IastTaintPool.java index 7ed121ee3..4f8242fc4 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/utils/threadlocal/IastTaintPool.java +++ b/dongtai-core/src/main/java/io/dongtai/iast/core/utils/threadlocal/IastTaintPool.java @@ -1,13 +1,11 @@ package io.dongtai.iast.core.utils.threadlocal; import io.dongtai.iast.core.EngineManager; -import io.dongtai.iast.core.utils.PropertyUtils; import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent; +import io.dongtai.iast.core.utils.PropertyUtils; import io.dongtai.log.DongTaiLog; -import java.util.HashSet; -import java.util.Map; -import java.util.Set; +import java.util.*; /** * @author dongzhiyong@huoxian.cn @@ -90,14 +88,17 @@ public void addTaintToPool(Object obj, MethodEvent event, boolean isSource) { } public void addToPool(Object obj) { + if (this.get() == null) { + return; + } this.get().add(obj); } public boolean isEmpty() { - return this.get().isEmpty(); + return this.get() == null || this.get().isEmpty(); } public boolean isNotEmpty() { - return !this.get().isEmpty(); + return this.get() != null && !this.get().isEmpty(); } }