Merge pull request #483 from HXSecurity/beta

1.9.0-beta1
HXSecurity · Mar 10, 2023 · 47eda6d · 47eda6d
2 parents e8aedc7 + 3e0753a
commit 47eda6d
Show file tree

Hide file tree

Showing 37 changed files with 1,257 additions and 203 deletions.
diff --git a/.github/workflows/code-check.yml b/.github/workflows/code-check.yml
@@ -5,9 +5,13 @@ name: Build Agent and Upload To OSS
 
 on:
   push:
-    branches: [ main ]
+    branches:
+      - main
+      - beta
   pull_request:
-    branches: [ main ]
+    branches:
+      - main
+      - beta
     paths-ignore:
       - '.github/**'
       - 'changes/**'

diff --git a/.github/workflows/codeql-analysis.yml b/.github/workflows/codeql-analysis.yml
@@ -13,7 +13,9 @@ name: "CodeQL"
 
 on:
   push:
-    branches: [ main ]
+    branches:
+      - main
+      - beta
     paths-ignore:
       - '.github/**'
       - 'changes/**'
@@ -24,7 +26,9 @@ on:
       - 'LICENSE'
       - '.gitignore'
   pull_request:
-    branches: [ main ]
+    branches:
+      - main
+      - beta
     paths-ignore:
       - '.github/**'
       - 'changes/**'

diff --git a/dongtai-common/src/main/java/io/dongtai/iast/common/constants/AgentConstant.java b/dongtai-common/src/main/java/io/dongtai/iast/common/constants/AgentConstant.java
@@ -1,7 +1,7 @@
 package io.dongtai.iast.common.constants;
 
 public class AgentConstant {
-    public static final String VERSION_VALUE = "v1.8.2";
+    public static final String VERSION_VALUE = "v1.9.0-beta1";
     public static final String LANGUAGE = "JAVA";
     public static final String THREAD_NAME_PREFIX = "DongTai-IAST-";
     public static final String THREAD_NAME_PREFIX_CORE = "DongTai-IAST-Core-";

diff --git a/dongtai-common/src/main/java/io/dongtai/iast/common/scope/Scope.java b/dongtai-common/src/main/java/io/dongtai/iast/common/scope/Scope.java
@@ -5,6 +5,9 @@ public enum Scope {
     HTTP_ENTRY(2),
     SERVLET_INPUT_STREAM_READ(3),
     SERVLET_OUTPUT_WRITE(4),
+    DUBBO_REQUEST(5),
+    DUBBO_ENTRY(6),
+    DUBBO_SOURCE(7),
     ;
 
     private final int id;

diff --git a/dongtai-common/src/main/java/io/dongtai/iast/common/scope/ScopeAggregator.java b/dongtai-common/src/main/java/io/dongtai/iast/common/scope/ScopeAggregator.java
@@ -3,6 +3,9 @@
 public class ScopeAggregator {
     private final GeneralScope httpRequestScope = new GeneralScope();
     private final GeneralScope httpEntryScope = new GeneralScope();
+    private final GeneralScope dubboRequestScope = new GeneralScope();
+    private final GeneralScope dubboEntryScope = new GeneralScope();
+    private final GeneralScope dubboSourceScope = new GeneralScope();
     private final GeneralScope servletInputStreamReadScope = new GeneralScope();
     private final GeneralScope servletOutputStreamWriteScope = new GeneralScope();
     private final PolicyScope policyScope = new PolicyScope();
@@ -15,6 +18,18 @@ public GeneralScope getHttpEntryScope() {
         return httpEntryScope;
     }
 
+    public GeneralScope getDubboRequestScope() {
+        return dubboRequestScope;
+    }
+
+    public GeneralScope getDubboEntryScope() {
+        return dubboEntryScope;
+    }
+
+    public GeneralScope getDubboSourceScope() {
+        return dubboSourceScope;
+    }
+
     public GeneralScope getServletInputStreamReadScope() {
         return servletInputStreamReadScope;
     }

diff --git a/dongtai-common/src/main/java/io/dongtai/iast/common/scope/ScopeTracker.java b/dongtai-common/src/main/java/io/dongtai/iast/common/scope/ScopeTracker.java
@@ -16,13 +16,19 @@ public GeneralScope getScope(Scope scope) {
                 return this.get().getServletInputStreamReadScope();
             case SERVLET_OUTPUT_WRITE:
                 return this.get().getServletOutputStreamWriteScope();
+            case DUBBO_REQUEST:
+                return this.get().getDubboRequestScope();
+            case DUBBO_ENTRY:
+                return this.get().getDubboEntryScope();
+            case DUBBO_SOURCE:
+                return this.get().getDubboSourceScope();
             default:
                 return null;
         }
     }
 
     public boolean inEnterEntry() {
-        return this.get().getHttpEntryScope().in();
+        return this.get().getHttpEntryScope().in() || this.get().getDubboRequestScope().in();
     }
 
     public PolicyScope getPolicyScope() {

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/EngineManager.java b/dongtai-core/src/main/java/io/dongtai/iast/core/EngineManager.java
@@ -140,4 +140,12 @@ public static void enterHttpEntry(Map<String, Object> requestMeta) {
         TAINT_RANGES_POOL.set(new HashMap<Integer, TaintRanges>());
         ScopeManager.SCOPE_TRACKER.getScope(Scope.HTTP_ENTRY).enter();
     }
+
+    public static void enterDubboEntry(Map<String, Object> requestMeta) {
+        REQUEST_CONTEXT.set(requestMeta);
+        TRACK_MAP.set(new HashMap<Integer, MethodEvent>(1024));
+        TAINT_HASH_CODES.set(new HashSet<Integer>());
+        TAINT_RANGES_POOL.set(new HashMap<Integer, TaintRanges>());
+        ScopeManager.SCOPE_TRACKER.getScope(Scope.DUBBO_ENTRY).enter();
+    }
 }
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/asm/AsmMethods.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/asm/AsmMethods.java
@@ -4,8 +4,8 @@
 
 import java.lang.dongtai.SpyDispatcher;
 import java.lang.dongtai.SpyDispatcherHandler;
-import java.util.Collection;
-import java.util.Enumeration;
+import java.net.InetSocketAddress;
+import java.util.*;
 
 /**
  * 常用的ASM method 集合 省得我到处声明
@@ -132,6 +132,54 @@ static Method getAsmMethod(final Class<?> clazz,
             int.class
     );
 
+    Method SPY$enterDubbo = InnerHelper.getAsmMethod(
+            SpyDispatcher.class,
+            "enterDubbo"
+    );
+    Method SPY$leaveDubbo = InnerHelper.getAsmMethod(
+            SpyDispatcher.class,
+            "leaveDubbo",
+            Object.class,
+            Object.class
+    );
+    Method SPY$isFirstLevelDubbo = InnerHelper.getAsmMethod(
+            SpyDispatcher.class,
+            "isFirstLevelDubbo"
+    );
+    Method SPY$collectDubboRequest = InnerHelper.getAsmMethod(
+            SpyDispatcher.class,
+            "collectDubboRequest",
+            Object.class,
+            Object.class,
+            Object.class,
+            String.class,
+            InetSocketAddress.class,
+            boolean.class,
+            boolean.class,
+            boolean.class,
+            boolean.class
+    );
+
+    Method SPY$collectDubboRequestSource = InnerHelper.getAsmMethod(
+            SpyDispatcher.class,
+            "collectDubboRequestSource",
+            Object.class,
+            Object.class,
+            String.class,
+            Object[].class,
+            Map.class,
+            String.class,
+            String.class,
+            String.class
+    );
+
+    Method SPY$collectDubboResponse = InnerHelper.getAsmMethod(
+            SpyDispatcher.class,
+            "collectDubboResponse",
+            Object.class,
+            byte.class
+    );
+
     Method SPY$enterSource = InnerHelper.getAsmMethod(
             SpyDispatcher.class,
             "enterSource"
@@ -209,6 +257,19 @@ static Method getAsmMethod(final Class<?> clazz,
             String.class
     );
 
+    Method SPY$traceDubboInvoke = InnerHelper.getAsmMethod(
+            SpyDispatcher.class,
+            "traceDubboInvoke",
+            Object.class,
+            String.class,
+            Object.class,
+            Object[].class,
+            Map.class,
+            String.class,
+            String.class,
+            String.class
+    );
+
     Method SPY$reportService = InnerHelper.getAsmMethod(
             SpyDispatcher.class,
             "reportService",

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/PluginRegister.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/PluginRegister.java
@@ -3,6 +3,7 @@
 import io.dongtai.iast.core.bytecode.enhance.ClassContext;
 import io.dongtai.iast.core.bytecode.enhance.plugin.authentication.shiro.DispatchShiro;
 import io.dongtai.iast.core.bytecode.enhance.plugin.core.DispatchClassPlugin;
+import io.dongtai.iast.core.bytecode.enhance.plugin.framework.dubbo.DispatchDubbo;
 import io.dongtai.iast.core.bytecode.enhance.plugin.framework.feign.DispatchFeign;
 import io.dongtai.iast.core.bytecode.enhance.plugin.framework.j2ee.dispatch.DispatchJ2ee;
 import io.dongtai.iast.core.bytecode.enhance.plugin.hardcoded.DispatchHardcodedPlugin;
@@ -34,6 +35,7 @@ public PluginRegister() {
         this.plugins.add(new DispatchJdbc());
         this.plugins.add(new DispatchShiro());
         this.plugins.add(new DispatchFeign());
+        this.plugins.add(new DispatchDubbo());
 
         this.plugins.add(new DispatchClassPlugin());
     }

diff --git a/...main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/DispatchDubbo.java b/...main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/DispatchDubbo.java
@@ -0,0 +1,39 @@
+package io.dongtai.iast.core.bytecode.enhance.plugin.framework.dubbo;
+
+import io.dongtai.iast.core.bytecode.enhance.ClassContext;
+import io.dongtai.iast.core.bytecode.enhance.plugin.DispatchPlugin;
+import io.dongtai.iast.core.handler.hookpoint.models.policy.Policy;
+import org.objectweb.asm.ClassVisitor;
+
+public class DispatchDubbo implements DispatchPlugin {
+    public static final String ALIBABA_DUBBO_SYNC_HANDLER = " com.alibaba.dubbo.rpc.listener.ListenerInvokerWrapper".substring(1);
+    public static final String APACHE_DUBBO_SYNC_HANDLER = " org.apache.dubbo.rpc.listener.ListenerInvokerWrapper".substring(1);
+    public static final String ALIBABA_DUBBO_EXCHANGE_HANDLER = " com.alibaba.dubbo.remoting.exchange.support.header.HeaderExchangeHandler".substring(1);
+    public static final String APACHE_DUBBO_EXCHANGE_HANDLER = " org.apache.dubbo.remoting.exchange.support.header.HeaderExchangeHandler".substring(1);
+    public static final String APACHE_DUBBO_EXCHANGE_CHANNEL = " org.apache.dubbo.remoting.exchange.support.header.HeaderExchangeChannel".substring(1);
+    public static final String ALIBABA_DUBBO_PROXY_HANDLER = " com.alibaba.dubbo.rpc.proxy.AbstractProxyInvoker".substring(1);
+    public static final String APACHE_DUBBO_PROXY_HANDLER = " org.apache.dubbo.rpc.proxy.AbstractProxyInvoker".substring(1);
+
+    @Override
+    public ClassVisitor dispatch(ClassVisitor classVisitor, ClassContext context, Policy policy) {
+        String className = context.getClassName();
+
+        if (ALIBABA_DUBBO_SYNC_HANDLER.equals(className)) {
+            classVisitor = new DubboSyncHandlerAdapter(classVisitor, context, " com.alibaba".substring(1));
+        } else if (APACHE_DUBBO_SYNC_HANDLER.equals(className)) {
+            classVisitor = new DubboSyncHandlerAdapter(classVisitor, context, " org.apache".substring(1));
+        } else if (ALIBABA_DUBBO_EXCHANGE_HANDLER.equals(className)) {
+            classVisitor = new DubboExchangeHandlerAdapter(classVisitor, context, " com.alibaba".substring(1));
+        } else if (APACHE_DUBBO_EXCHANGE_HANDLER.equals(className)) {
+            classVisitor = new DubboExchangeHandlerAdapter(classVisitor, context, " org.apache".substring(1));
+        } else if (APACHE_DUBBO_EXCHANGE_CHANNEL.equals(className)) {
+            classVisitor = new DubboExchangeChannelAdapter(classVisitor, context, " org.apache".substring(1));
+        } else if (ALIBABA_DUBBO_PROXY_HANDLER.equals(className)) {
+            classVisitor = new DubboProxyHandlerAdapter(classVisitor, context, " com.alibaba".substring(1));
+        } else if (APACHE_DUBBO_PROXY_HANDLER.equals(className)) {
+            classVisitor = new DubboProxyHandlerAdapter(classVisitor, context, " org.apache".substring(1));
+        }
+
+        return classVisitor;
+    }
+}
diff --git a/...ongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/DubboExchangeChannelAdapter.java b/...ongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/DubboExchangeChannelAdapter.java
@@ -0,0 +1,35 @@
+package io.dongtai.iast.core.bytecode.enhance.plugin.framework.dubbo;
+
+import io.dongtai.iast.core.bytecode.enhance.ClassContext;
+import io.dongtai.iast.core.bytecode.enhance.plugin.AbstractClassVisitor;
+import io.dongtai.iast.core.utils.AsmUtils;
+import io.dongtai.log.DongTaiLog;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.MethodVisitor;
+
+public class DubboExchangeChannelAdapter extends AbstractClassVisitor {
+    public static final String DUBBO_EXCHANGE_CHANNEL_SEND = "{package}.dubbo.remoting.exchange.support.header.HeaderExchangeChannel.send(java.lang.Object)";
+
+    private final String packageName;
+    private final String sendSign;
+
+    public DubboExchangeChannelAdapter(ClassVisitor classVisitor, ClassContext context, String packageName) {
+        super(classVisitor, context);
+        this.packageName = packageName;
+        this.sendSign = DUBBO_EXCHANGE_CHANNEL_SEND.replace("{package}", this.packageName);
+    }
+
+    @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);
+        String signCode = AsmUtils.buildSignature(context.getClassName(), name, desc);
+
+        if (this.sendSign.equals(signCode)) {
+            DongTaiLog.debug("Adding dubbo provider response tracking by {}", signCode);
+            mv = new DubboExchangeChannelSendAdviceAdapter(mv, access, name, desc, signCode,
+                    this.context, this.packageName);
+            setTransformed();
+        }
+        return mv;
+    }
+}
diff --git a/...t/core/bytecode/enhance/plugin/framework/dubbo/DubboExchangeChannelSendAdviceAdapter.java b/...t/core/bytecode/enhance/plugin/framework/dubbo/DubboExchangeChannelSendAdviceAdapter.java
@@ -0,0 +1,74 @@
+package io.dongtai.iast.core.bytecode.enhance.plugin.framework.dubbo;
+
+import io.dongtai.iast.core.bytecode.enhance.ClassContext;
+import io.dongtai.iast.core.bytecode.enhance.plugin.AbstractAdviceAdapter;
+import org.objectweb.asm.*;
+import org.objectweb.asm.commons.Method;
+
+public class DubboExchangeChannelSendAdviceAdapter extends AbstractAdviceAdapter {
+    private static final Method GET_RESULT_METHOD = Method.getMethod("java.lang.Object getResult()");
+    private static final Method GET_STATUS_METHOD = Method.getMethod("byte getStatus()");
+
+    private final String packageName;
+    private final Type objectType;
+    private final Type responseType;
+
+    protected DubboExchangeChannelSendAdviceAdapter(MethodVisitor mv, int access, String name, String desc,
+                                                    String signature, ClassContext context, String packageName) {
+        super(mv, access, name, desc, context, "dubbo", signature);
+        this.packageName = packageName;
+        String packageDesc = packageName.replace(".", "/");
+        this.responseType = Type.getObjectType(packageDesc + "/dubbo/remoting/exchange/Response");
+        this.objectType = Type.getObjectType("java/lang/Object");
+    }
+
+    @Override
+    protected void before() {
+        mark(tryLabel);
+        Label elseLabel = new Label();
+
+        isFirstLevelDubbo();
+        mv.visitJumpInsn(EQ, elseLabel);
+
+        collectDubboResponse();
+
+        mark(elseLabel);
+    }
+
+    @Override
+    protected void after(int opcode) {
+    }
+
+    private void isFirstLevelDubbo() {
+        invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+        invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$isFirstLevelDubbo);
+    }
+
+    private void collectDubboResponse() {
+        Label tryL = new Label();
+        Label catchL = new Label();
+        Label exHandlerL = new Label();
+        visitTryCatchBlock(tryL, catchL, exHandlerL, ASM_TYPE_THROWABLE.getInternalName());
+        visitLabel(tryL);
+
+
+        int respLocal = newLocal(this.responseType);
+        loadArg(0);
+        checkCast(this.responseType);
+        storeLocal(respLocal);
+
+        invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+        loadLocal(respLocal);
+        invokeVirtual(this.responseType, GET_RESULT_METHOD);
+        loadLocal(respLocal);
+        invokeVirtual(this.responseType, GET_STATUS_METHOD);
+        invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$collectDubboResponse);
+
+        visitLabel(catchL);
+        Label endL = new Label();
+        visitJumpInsn(GOTO, endL);
+        visitLabel(exHandlerL);
+        visitVarInsn(ASTORE, this.nextLocal);
+        visitLabel(endL);
+    }
+}