diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/DispatchDubbo.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/DispatchDubbo.java index 4bddf9c4c..edc289f7c 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/DispatchDubbo.java +++ b/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; } diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/DubboHessianAdapter.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/DubboHessianAdapter.java new file mode 100644 index 000000000..0027a88d7 --- /dev/null +++ b/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; } +} diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/DubboHessianAddRequestHeadersAdapter.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/dubbo/DubboHessianAddRequestHeadersAdapter.java new file mode 100644 index 000000000..43082a7bd --- /dev/null +++ b/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); + } +} diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/DubboImpl.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/DubboImpl.java index fe2cfe570..bca439d96 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/DubboImpl.java +++ b/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 oldHeaders = (Map) 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(); 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 a5c934112..f9c0ee283 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 @@ -31,7 +31,8 @@ public class PropagatorImpl { "java.net.URL.(java.lang.String,java.lang.String,java.lang.String)", // indirect "java.net.URL.(java.lang.String,java.lang.String,int,java.lang.String)", // indirect "java.net.URL.(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) { diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/service/trace/DubboService.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/service/trace/DubboService.java index 83e3c0657..e127198f5 100644 --- a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/service/trace/DubboService.java +++ b/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 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;