fixes taint pool null exception

HXSecurity · May 18, 2022 · 7022fe1 · 7022fe1
1 parent 6a716a3
commit 7022fe1
Show file tree

Hide file tree

Showing 6 changed files with 128 additions and 128 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
@@ -483,11 +483,11 @@ 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);
-                        } else if (HookType.SINK.equals(hookType) && !EngineManager.TAINT_POOL.get().isEmpty()) {
+                        } else if (HookType.SINK.equals(hookType) && !EngineManager.TAINT_POOL.isEmpty()) {
                             SinkImpl.solveSink(event);
                         }
                     }

diff --git a/.../src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/PropagatorImpl.java b/.../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/...re/src/main/java/io/dongtai/iast/core/handler/hookpoint/framework/dubbo/DubboHandler.java b/...re/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<Object> 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<Object> 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/...core/src/main/java/io/dongtai/iast/core/handler/hookpoint/framework/grpc/GrpcHandler.java b/...core/src/main/java/io/dongtai/iast/core/handler/hookpoint/framework/grpc/GrpcHandler.java
@@ -155,98 +155,101 @@ 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<ReqT,RespT>, io.grpc.CallOptions, ReqT)",
-                    "io.grpc.stub.ClientCalls.blockingUnaryCall(io.grpc.Channel, io.grpc.MethodDescriptor<ReqT,RespT>, io.grpc.CallOptions, ReqT)",
-                    null,
-                    new Object[]{req},
-                    res,
-                    "GRPC",
-                    false,
-                    null
-            );
-            Set<Object> 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<Object> resModelItems = SourceImpl.parseCustomModel(res);
-                sharedRespData.remove();
-                Set<Object> taintPool = EngineManager.TAINT_POOL.get();
-                Set<Object> resModelSet = new HashSet<Object>();
-                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<ReqT,RespT>, io.grpc.CallOptions, ReqT)",
+                "io.grpc.stub.ClientCalls.blockingUnaryCall(io.grpc.Channel, io.grpc.MethodDescriptor<ReqT,RespT>, io.grpc.CallOptions, ReqT)",
+                null,
+                new Object[]{req},
+                res,
+                "GRPC",
+                false,
+                null
+        );
+        Set<Object> 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<Object> resModelItems = SourceImpl.parseCustomModel(res);
+            sharedRespData.remove();
+            Set<Object> taintPool = EngineManager.TAINT_POOL.get();
+            Set<Object> resModelSet = new HashSet<Object>();
+            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<Object> 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<Object> 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);
         }
     }
 
     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/...core/src/main/java/io/dongtai/iast/core/handler/hookpoint/service/kafka/KafkaHandler.java b/...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
@@ -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 [email protected]
@@ -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();
     }
 }
-Original file line number
+Diff line change
@@ Expand Up / @@ -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;
             }
@@ Expand Down @@