Fix sql interceptor

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

@@ -14,6 +14,9 @@ public class DispatchDubbo implements DispatchPlugin {
     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);
 
+    //com.caucho.hessian.client.HessianProxy.sendRequest(java.lang.String,java.lang.Object[])
+    public static final String DUBBO_PROXY_HESSIAN = " com.caucho.hessian.client.HessianProxy".substring(1);
+
     @Override
     public ClassVisitor dispatch(ClassVisitor classVisitor, ClassContext context, Policy policy) {
         String className = context.getClassName();
@@ -33,7 +36,10 @@ public ClassVisitor dispatch(ClassVisitor classVisitor, ClassContext context, Po
         } else if (APACHE_DUBBO_PROXY_HANDLER.equals(className)) {
             classVisitor = new DubboProxyHandlerAdapter(classVisitor, context, " org.apache".substring(1));
         }
-
+        if (DUBBO_PROXY_HESSIAN.equals(className)){
+            System.out.println("dispatch" + className);
+            classVisitor = new DubboHessianAdapter(classVisitor, context);
+        }
         return classVisitor;
     }
 

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

@@ -0,0 +1,34 @@
+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;
+
+/**
+ * @author mazepeng
+ * @date 2023/11/2 17:40
+ */
+public class DubboHessianAdapter extends AbstractClassVisitor {
+
+    private static final String HESSIAN_ADDREQUESTHEADERS = " com.caucho.hessian.client.HessianProxy.addRequestHeaders(com.caucho.hessian.client.HessianConnection)".substring(1);
+
+    public DubboHessianAdapter(ClassVisitor classVisitor, ClassContext context) {
+        super(classVisitor, context);
+    }
+
+
+    @Override
+    public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+        MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
+        String signCode = AsmUtils.buildSignature(context.getClassName(), name, descriptor);
+        if (HESSIAN_ADDREQUESTHEADERS.equals(signCode)) {
+            DongTaiLog.debug("Adding dubbo provider source tracking by {}", signCode);
+            System.out.println("hessian增强完成");
+            mv = new DubboHessianAddRequestHeadersAdapter(mv, access, name, descriptor,this.context,"hessian",signCode);
+            setTransformed();
+        }
+        return mv;    }
+}

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

@@ -0,0 +1,121 @@
+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.Label;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
+import org.objectweb.asm.commons.Method;
+
+/**
+ * @author mazepeng
+ * @date 2023/11/2 17:40
+ */
+public class DubboHessianAddRequestHeadersAdapter  extends AbstractAdviceAdapter {
+
+
+    private Label exHandler;
+
+    private final Type urlType;
+    private static final Method URL_TO_STRING_METHOD = Method.getMethod("java.lang.String toString()");
+
+
+    public DubboHessianAddRequestHeadersAdapter(MethodVisitor mv, int access, String name, String desc, ClassContext context, String type, String signCode) {
+        super(mv, access, name, desc, context, type, signCode);
+        this.urlType = Type.getObjectType("java/net/URL");
+
+
+    }
+    @Override
+    public void visitMaxs(int maxStack, int maxLocals) {
+        visitLabel(this.catchLabel);
+        visitLabel(this.exHandler);
+        leaveMethod(ATHROW);
+        throwException();
+        visitTryCatchBlock(this.tryLabel, this.catchLabel, this.exHandler, ASM_TYPE_THROWABLE.getInternalName());
+        super.visitMaxsNew(maxStack, maxLocals);
+    }
+
+    @Override
+    protected void onMethodEnter() {
+        this.tryLabel = new Label();
+        visitLabel(this.tryLabel);
+        enterMethod();
+        this.catchLabel = new Label();
+        this.exHandler = new Label();
+    }
+
+    @Override
+    protected void onMethodExit(int opcode) {
+        leaveMethod(opcode);
+    }
+    private void leaveMethod(int opcode) {
+        leaveScope();
+    }
+    private void leaveScope() {
+        invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+        push(false);
+        invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leavePropagator);
+    }
+
+
+    private void enterMethod() {
+        skipCollect();
+        enterScope();
+        Label endLabel = new Label();
+        traceMethod();
+        mark(endLabel);
+    }
+
+    @Override
+    protected void before() {
+    }
+
+    @Override
+    protected void after(int opcode) {
+    }
+
+    private void skipCollect() {
+        invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+        loadArg(0);
+        invokeInterface(ASM_TYPE_SPY_DISPATCHER,SPY$isSkipCollectDubbo);
+        pop();
+    }
+
+    private void enterScope() {
+        invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+        push(false);
+        invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterPropagator);
+    }
+
+    private void isFirstScope() {
+        invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+        invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$isFirstLevelPropagator);
+    }
+
+    private void traceMethod() {
+        invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+
+        //加载this 1
+        loadThis();
+        //2
+        dup();
+        visitFieldInsn(Opcodes.GETFIELD,"com/caucho/hessian/client/HessianProxy","_url","Ljava/net/URL;");
+        invokeVirtual(this.urlType,URL_TO_STRING_METHOD);
+        //参数 3
+        loadArg(0);
+        //arguments 4
+        pushNull();
+        //headers 5
+        pushNull();
+        //6
+        push(this.classContext.getClassName());
+        //7
+        push(this.name);
+        // 8
+        push(this.signature);
+        //获取静态getDispatcher
+        invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$traceDubboInvoke);
+    }
+}

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

@@ -145,9 +145,14 @@ public static void collectDubboRequestSource(Object handler, Object invocation,
             }
         }
 
+        if (handler.toString().startsWith("hessian")) {
+            Map<String, String> oldHeaders = (Map<String, String>) requestMeta.get("headers");
+            sHeaders.putAll(oldHeaders);
+        }
+
         if (!sHeaders.isEmpty()) {
             String traceIdKey = ContextManager.getHeaderKey();
-            if (headers.containsKey(traceIdKey)) {
+            if (sHeaders.containsKey(traceIdKey)) {
                 ContextManager.parseTraceId(sHeaders.get(traceIdKey));
             } else {
                 String newTraceId = ContextManager.currentTraceId();

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

@@ -31,7 +31,8 @@ public class PropagatorImpl {
             "java.net.URL.<init>(java.lang.String,java.lang.String,java.lang.String)", // indirect
             "java.net.URL.<init>(java.lang.String,java.lang.String,int,java.lang.String)", // indirect
             "java.net.URL.<init>(java.lang.String,java.lang.String,int,java.lang.String,java.net.URLStreamHandler)",
-            "com.fasterxml.jackson.databind.ObjectMapper.readValue(java.io.InputStream,com.fasterxml.jackson.databind.JavaType)"
+            "com.fasterxml.jackson.databind.ObjectMapper.readValue(java.io.InputStream,com.fasterxml.jackson.databind.JavaType)",
+            "com.baomidou.mybatisplus.core.override.MybatisMapperMethod.execute(org.apache.ibatis.session.SqlSession,java.lang.Object[])"
     ));
 
     public static void solvePropagator(MethodEvent event, PropagatorNode propagatorNode, AtomicInteger invokeIdSequencer) {

dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/service/trace/DubboService.java

@@ -15,7 +15,17 @@
 public class DubboService {
     public static void solveSyncInvoke(MethodEvent event, Object invocation, String url, Map<String, String> headers,
                                        AtomicInteger invokeIdSequencer) {
+
+
         try {
+            String traceId = ContextManager.nextTraceId();
+            //当类型为HessianURLConnection，只处理添加请求头即可
+            if (invocation.getClass().getSimpleName().equals("HessianURLConnection")) {
+                Method method = invocation.getClass().getMethod("addHeader", String.class, String.class);
+                method.setAccessible(true);
+                method.invoke(invocation, ContextManager.getHeaderKey(), traceId);
+                return;
+            }
             TaintPoolUtils.trackObject(event, null, event.parameterInstances, 0, false);
             boolean hasTaint = false;
             int sourceLen = 0;
@@ -34,10 +44,11 @@ public static void solveSyncInvoke(MethodEvent event, Object invocation, String
                 event.addParameterValue(1, headers, hasTaint);
             }
 
-            Method setAttachmentMethod = invocation.getClass().getMethod("setAttachment", String.class, String.class);
-            setAttachmentMethod.setAccessible(true);
-            String traceId = ContextManager.nextTraceId();
-            setAttachmentMethod.invoke(invocation, ContextManager.getHeaderKey(), traceId);
+                Method setAttachmentMethod = invocation.getClass().getMethod("setAttachment", String.class, String.class);
+                setAttachmentMethod.setAccessible(true);
+                setAttachmentMethod.invoke(invocation, ContextManager.getHeaderKey(), traceId);
+
+
 
             // add to method pool
             event.source = false;