fix: dubbo blacklist skip collection.

Author: ‘niuerzhuang’

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

@@ -278,6 +278,12 @@ static Method getAsmMethod(final Class<?> clazz,
             String.class
     );
 
+    Method SPY$isSkipCollect = InnerHelper.getAsmMethod(
+            SpyDispatcher.class,
+            "isSkipCollect",
+            Object.class
+    );
+
     Method SPY$reportService = InnerHelper.getAsmMethod(
             SpyDispatcher.class,
             "reportService",

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

@@ -61,6 +61,7 @@ public void visitMaxs(int maxStack, int maxLocals) {
     }
 
     private void enterMethod() {
+        skipCollect();
         enterScope();
 
         Label elseLabel = new Label();
@@ -113,4 +114,11 @@ private void traceMethod() {
         invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$traceDubboInvoke);
         pop();
     }
+
+    private void skipCollect() {
+        invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+        loadArg(0);
+        invokeInterface(ASM_TYPE_SPY_DISPATCHER,SPY$isSkipCollect);
+        pop();
+    }
 }

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

@@ -7,6 +7,7 @@
 import io.dongtai.iast.common.scope.ScopeManager;
 import io.dongtai.iast.core.EngineManager;
 import io.dongtai.iast.core.bytecode.enhance.plugin.spring.SpringApplicationImpl;
+import io.dongtai.iast.core.handler.skip.BlackUrlSkipHandler;
 import io.dongtai.iast.core.handler.hookpoint.controller.HookType;
 import io.dongtai.iast.core.handler.hookpoint.controller.impl.*;
 import io.dongtai.iast.core.handler.hookpoint.graphy.GraphBuilder;
@@ -19,6 +20,8 @@
 import io.dongtai.log.ErrorCode;
 
 import java.lang.dongtai.SpyDispatcher;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
 import java.net.InetSocketAddress;
 import java.util.*;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -697,6 +700,21 @@ public boolean traceDubboInvoke(Object instance, String url, Object invocation,
         return false;
     }
 
+    @Override
+    public boolean isSkipCollect(Object invocation) {
+        if (BlackUrlSkipHandler.isBlackUrl()){
+            Method setAttachmentMethod = null;
+            try {
+                setAttachmentMethod = invocation.getClass().getMethod("setAttachment", String.class, String.class);
+                setAttachmentMethod.setAccessible(true);
+                setAttachmentMethod.invoke(invocation, BlackUrlSkipHandler.getHeaderKey(), BlackUrlSkipHandler.isBlackUrl().toString());
+            } catch (NoSuchMethodException | InvocationTargetException | IllegalAccessException e) {
+                DongTaiLog.error(ErrorCode.get("SPY_SKIP_COLLECT_DUBBO_FAILED"), e);
+            }
+        }
+        return false;
+    }
+
     private boolean isCollectAllowed(boolean isEnterEntry) {
         if (!EngineManager.isEngineRunning()) {
             return false;

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

@@ -4,6 +4,7 @@
 import io.dongtai.iast.common.config.ConfigBuilder;
 import io.dongtai.iast.common.config.ConfigKey;
 import io.dongtai.iast.core.EngineManager;
+import io.dongtai.iast.core.handler.skip.BlackUrlSkipHandler;
 import io.dongtai.iast.core.handler.context.ContextManager;
 import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.SourceNode;
@@ -54,6 +55,9 @@ public static void collectDubboRequestSource(Object handler, Object invocation,
         if (requestMeta == null) {
             return;
         }
+        if (null != headers.get(BlackUrlSkipHandler.getHeaderKey()) && headers.get(BlackUrlSkipHandler.getHeaderKey()).equals("true")){
+            return;
+        }
 
         String url = (String) requestMeta.get("requestURL") + "/" + methodName;
         String uri = (String) requestMeta.get("requestURI") + "/" + methodName;

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

@@ -3,6 +3,7 @@
 import io.dongtai.iast.common.config.*;
 import io.dongtai.iast.common.constants.AgentConstant;
 import io.dongtai.iast.core.EngineManager;
+import io.dongtai.iast.core.handler.skip.BlackUrlSkipHandler;
 import io.dongtai.iast.core.handler.hookpoint.IastClassLoader;
 import io.dongtai.iast.core.utils.*;
 import io.dongtai.iast.core.utils.matcher.ConfigMatcher;
@@ -74,6 +75,7 @@ public static void solveHttpRequest(Object obj, Object req, Object resp, Map<Str
             return;
         }
         if (ConfigMatcher.getInstance().getBlackUrl(requestMeta)) {
+            BlackUrlSkipHandler.setIsBlackUrl(true);
             return;
         }
 

dongtai-core/src/main/java/io/dongtai/iast/core/handler/skip/BlackUrlSkipHandler.java

@@ -0,0 +1,20 @@
+package io.dongtai.iast.core.handler.skip;
+
+import io.dongtai.iast.core.utils.threadlocal.BooleanThreadLocal;
+
+public class BlackUrlSkipHandler {
+
+    private static BooleanThreadLocal isBlackUrl = new BooleanThreadLocal(false);
+
+    public static Boolean isBlackUrl() {
+        return isBlackUrl.get();
+    }
+
+    public static void setIsBlackUrl(Boolean isBlackUrl) {
+        BlackUrlSkipHandler.isBlackUrl.set(isBlackUrl);
+    }
+
+    public static String getHeaderKey() {
+        return "dt-collect-skip";
+    }
+}

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

@@ -88,6 +88,7 @@ public enum ErrorCode {
     SPY_TRACE_DUBBO_CONSUMER_INVOKE_FAILED(20361, "hookpoint trace dubbo consumer invoke failed"),
     SPY_LEAVE_DUBBO_FAILED(20362, "hookpoint leave dubbo failed"),
     SPY_COLLECT_DUBBO_FAILED(20363, "hookpoint collect dubbo {} failed"),
+    SPY_SKIP_COLLECT_DUBBO_FAILED(20364, "hookpoint skip collect dubbo {} failed"),
 
     // report & replay
     REPORT_SEND_FAILED(20401, "send report to {} error, report: {}"),

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

@@ -261,4 +261,10 @@ public boolean traceDubboInvoke(Object instance, String url, Object invocation,
                                     String signature) {
         return false;
     }
+
+    @Override
+    public boolean isSkipCollect(Object invocation) {
+        return false;
+    }
+
 }

dongtai-spy/src/main/java/java/lang/dongtai/SpyDispatcher.java

@@ -170,4 +170,6 @@ boolean traceFeignInvoke(Object instance, Object[] parameters,
     boolean traceDubboInvoke(Object instance, String url, Object invocation, Object[] arguments,
                              Map<String, String> headers, String className, String methodName,
                              String signature);
+
+    boolean isSkipCollect(Object invocation);
 }