fecth response for servlet PrintWriter

Author: lostsnow

dongtai-common/src/main/java/io/dongtai/iast/common/scope/Scope.java

@@ -4,7 +4,7 @@ public enum Scope {
     HTTP_REQUEST(1),
     HTTP_ENTRY(2),
     SERVLET_INPUT_STREAM_READ(3),
-    SERVLET_OUTPUT_STREAM_WRITE(4),
+    SERVLET_OUTPUT_WRITE(4),
     ;
 
     private final int id;

dongtai-common/src/main/java/io/dongtai/iast/common/scope/ScopeTracker.java

@@ -14,7 +14,7 @@ public GeneralScope getScope(Scope scope) {
                 return this.get().getHttpEntryScope();
             case SERVLET_INPUT_STREAM_READ:
                 return this.get().getServletInputStreamReadScope();
-            case SERVLET_OUTPUT_STREAM_WRITE:
+            case SERVLET_OUTPUT_WRITE:
                 return this.get().getServletOutputStreamWriteScope();
             default:
                 return null;

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

@@ -120,6 +120,18 @@ static Method getAsmMethod(final Class<?> clazz,
             int.class
     );
 
+    Method SPY$onPrintWriterWrite = InnerHelper.getAsmMethod(
+            SpyDispatcher.class,
+            "onPrintWriterWrite",
+            String.class,
+            Object.class,
+            int.class,
+            String.class,
+            char[].class,
+            int.class,
+            int.class
+    );
+
     Method SPY$enterSource = InnerHelper.getAsmMethod(
             SpyDispatcher.class,
             "enterSource"

dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/PluginRegister.java

@@ -50,7 +50,9 @@ public ClassVisitor initial(ClassVisitor classVisitor, ClassContext context, Pol
             if (pluginVisitor != classVisitor) {
                 classVisitor = pluginVisitor;
                 // TODO: need transform multiple times?
-                break;
+                if (!context.getClassName().equals(DispatchJ2ee.APACHE_COYOTE_WRITER)) {
+                    break;
+                }
             }
         }
         return classVisitor;

dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/j2ee/dispatch/DispatchJ2ee.java

@@ -22,7 +22,7 @@ public class DispatchJ2ee implements DispatchPlugin {
     public static final String JAKARTA_SERVLET_INPUT_STREAM = " jakarta.servlet.ServletInputStream".substring(1);
     public static final String JAVAX_SERVLET_OUTPUT_STREAM = " javax.servlet.ServletOutputStream".substring(1);
     public static final String JAKARTA_SERVLET_OUTPUT_STREAM = " jakarta.servlet.ServletOutputStream".substring(1);
-
+    public static final String APACHE_COYOTE_WRITER = " org.apache.catalina.connector.CoyoteWriter".substring(1);
 
     @Override
     public ClassVisitor dispatch(ClassVisitor classVisitor, ClassContext context, Policy policy) {
@@ -41,6 +41,8 @@ public ClassVisitor dispatch(ClassVisitor classVisitor, ClassContext context, Po
             classVisitor = new ServletOutputStreamAdapter(classVisitor, context);
         } else if (ancestors.contains(JAKARTA_SERVLET_OUTPUT_STREAM)) {
             classVisitor = new ServletOutputStreamAdapter(classVisitor, context);
+        } else if (APACHE_COYOTE_WRITER.equals(className)) {
+            classVisitor = new PrintWriterAdapter(classVisitor, context);
         }
         return classVisitor;
     }

dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/j2ee/dispatch/PrintWriterAdapter.java

@@ -0,0 +1,41 @@
+package io.dongtai.iast.core.bytecode.enhance.plugin.framework.j2ee.dispatch;
+
+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.*;
+
+import java.util.*;
+
+public class PrintWriterAdapter extends AbstractClassVisitor {
+    private static final Set<String> WRITE_DESC = new HashSet<>(Arrays.asList("(I)V", "([CII)V", "(Ljava/lang/String;II)V"));
+
+    PrintWriterAdapter(ClassVisitor classVisitor, ClassContext context) {
+        super(classVisitor, context);
+    }
+
+    @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);
+        Type[] typeOfArgs = Type.getArgumentTypes(desc);
+        String signCode = AsmUtils.buildSignature(context.getClassName(), name, desc);
+        if (isWrite(name, desc)) {
+            DongTaiLog.debug("Adding HTTP response PrintWriter by {}", signCode);
+            mv = new PrintWriterWriteAdviceAdapter(mv, access, name, desc, signCode, context);
+            setTransformed();
+        }
+        if (hasTransformed()) {
+            DongTaiLog.trace("rewrite method {}.{} for listener[match={}]", context.getClassName(), name, context.getMatchedClassName());
+        }
+
+        return mv;
+    }
+
+    private boolean isWrite(String name, String desc) {
+        if ("write".equals(name) && WRITE_DESC.contains(desc)) {
+            return true;
+        }
+        return false;
+    }
+}

dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/j2ee/dispatch/PrintWriterWriteAdviceAdapter.java

@@ -0,0 +1,68 @@
+package io.dongtai.iast.core.bytecode.enhance.plugin.framework.j2ee.dispatch;
+
+import io.dongtai.iast.common.scope.Scope;
+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;
+
+public class PrintWriterWriteAdviceAdapter extends AbstractAdviceAdapter {
+    public PrintWriterWriteAdviceAdapter(MethodVisitor mv, int access, String name, String desc, String signature,
+                                         ClassContext context) {
+        super(mv, access, name, desc, context, "j2ee", signature);
+    }
+
+    @Override
+    protected void before() {
+        mark(tryLabel);
+        enterScope(Scope.SERVLET_OUTPUT_WRITE);
+
+        Label elseLabel = new Label();
+
+        isFirstLevelScope(Scope.SERVLET_OUTPUT_WRITE);
+        mv.visitJumpInsn(EQ, elseLabel);
+
+        onPrintWriterWrite();
+
+        mark(elseLabel);
+    }
+
+    @Override
+    protected void after(final int opcode) {
+        leaveScope(Scope.SERVLET_OUTPUT_WRITE);
+    }
+
+    private void onPrintWriterWrite() {
+        if ("(I)V".equals(this.desc)) {
+            invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+            push(desc);
+            loadThis();
+            loadArg(0);
+            pushNull();
+            pushNull();
+            push(-1);
+            push(-1);
+            invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$onPrintWriterWrite);
+        } else if ("([CII)V".equals(this.desc)) {
+            invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+            push(desc);
+            loadThis();
+            push(-1);
+            pushNull();
+            loadArg(0);
+            loadArg(1);
+            loadArg(2);
+            invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$onPrintWriterWrite);
+        } else if ("(Ljava/lang/String;II)V".equals(this.desc)) {
+            invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+            push(desc);
+            loadThis();
+            push(-1);
+            loadArg(0);
+            pushNull();
+            loadArg(1);
+            loadArg(2);
+            invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$onPrintWriterWrite);
+        }
+    }
+}

dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/j2ee/dispatch/ServletOutputStreamWriteAdviceAdapter.java

@@ -15,11 +15,11 @@ public ServletOutputStreamWriteAdviceAdapter(MethodVisitor mv, int access, Strin
     @Override
     protected void before() {
         mark(tryLabel);
-        enterScope(Scope.SERVLET_OUTPUT_STREAM_WRITE);
+        enterScope(Scope.SERVLET_OUTPUT_WRITE);
 
         Label elseLabel = new Label();
 
-        isFirstLevelScope(Scope.SERVLET_OUTPUT_STREAM_WRITE);
+        isFirstLevelScope(Scope.SERVLET_OUTPUT_WRITE);
         mv.visitJumpInsn(EQ, elseLabel);
 
         onServletOutputStreamWrite();
@@ -29,7 +29,7 @@ protected void before() {
 
     @Override
     protected void after(final int opcode) {
-        leaveScope(Scope.SERVLET_OUTPUT_STREAM_WRITE);
+        leaveScope(Scope.SERVLET_OUTPUT_WRITE);
     }
 
     private void onServletOutputStreamWrite() {

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

@@ -245,6 +245,26 @@ public void onServletOutputStreamWrite(String desc, Object stream, int b, byte[]
         }
     }
 
+    @Override
+    public void onPrintWriterWrite(String desc, Object writer, int b, String s, char[] cs, int offset, int len) {
+        try {
+            ScopeManager.SCOPE_TRACKER.getPolicyScope().enterAgent();
+            if (!EngineManager.isEngineRunning()) {
+                return;
+            }
+
+            if (!ScopeManager.SCOPE_TRACKER.getScope(Scope.HTTP_ENTRY).in()) {
+                return;
+            }
+
+            HttpImpl.onPrintWriterWrite(desc, writer, b, s, cs, offset, len);
+        } catch (Throwable e) {
+            DongTaiLog.warn(ErrorCode.SPY_COLLECT_HTTP_FAILED, "response body", e);
+        } finally {
+            ScopeManager.SCOPE_TRACKER.getPolicyScope().leaveAgent();
+        }
+    }
+
     /**
      * mark for enter Source Entry Point
      *

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

@@ -227,29 +227,83 @@ public static void onServletOutputStreamWrite(String desc, Object stream, int b,
             maxLength = 50000;
         }
 
-        if ("(I)V".equals(desc)) {
-            if (b == -1) {
-                return;
-            }
-            ByteArrayOutputStream buff = EngineManager.BODY_BUFFER.getResponse();
-            if (buff.size() < maxLength) {
-                buff.write(b);
-            }
-        } else if ("([B)V".equals(desc)) {
-            if (bs == null) {
-                return;
+        try {
+            if ("(I)V".equals(desc)) {
+                if (b == -1) {
+                    return;
+                }
+                ByteArrayOutputStream buff = EngineManager.BODY_BUFFER.getResponse();
+                if (buff.size() < maxLength) {
+                    buff.write(b);
+                }
+            } else if ("([B)V".equals(desc)) {
+                if (bs == null) {
+                    return;
+                }
+                onServletOutputStreamWrite("([BII)V", stream, b, bs, 0, bs.length);
+            } else if ("([BII)V".equals(desc)) {
+                if (bs == null || offset < 0 || len < 0) {
+                    return;
+                }
+
+                ByteArrayOutputStream buff = EngineManager.BODY_BUFFER.getResponse();
+                int size = buff.size();
+                if (size < maxLength) {
+                    buff.write(bs, offset, Math.min(len, maxLength - size));
+                }
             }
-            onServletOutputStreamWrite("([BII)V", stream, b, bs, 0, bs.length);
-        } else if ("([BII)V".equals(desc)) {
-            if (bs == null || offset < 0 || len < 0) {
+        } catch (Throwable ignore) {
+        }
+    }
+
+    public static void onPrintWriterWrite(String desc, Object writer, int b, String s, char[] cs, int offset, int len) {
+        try {
+            boolean getBody = ((Config<Boolean>) ConfigBuilder.getInstance().getConfig(ConfigKey.REPORT_RESPONSE_BODY)).get();
+            if (!getBody) {
                 return;
             }
+        } catch (Throwable ignore) {
+            return;
+        }
 
-            ByteArrayOutputStream buff = EngineManager.BODY_BUFFER.getResponse();
-            int size = buff.size();
-            if (size < maxLength) {
-                buff.write(bs, offset, Math.min(len, maxLength - size));
+        Integer maxLength = PropertyUtils.getInstance().getResponseLength();
+        if (maxLength == 0) {
+            return;
+        } else if (maxLength < 0 || maxLength > 50000) {
+            maxLength = 50000;
+        }
+
+        try {
+            if ("(I)V".equals(desc)) {
+                if (b == -1) {
+                    return;
+                }
+                ByteArrayOutputStream buff = EngineManager.BODY_BUFFER.getResponse();
+                if (buff.size() < maxLength) {
+                    buff.write(b);
+                }
+            } else if ("([CII)V".equals(desc)) {
+                if (cs == null || offset < 0 || len < 0) {
+                    return;
+                }
+
+                ByteArrayOutputStream buff = EngineManager.BODY_BUFFER.getResponse();
+                int size = buff.size();
+                if (size < maxLength) {
+                    buff.write((new String(cs, offset, Math.min(len, maxLength - size))).getBytes());
+                }
+            } else if ("(Ljava/lang/String;II)V".equals(desc)) {
+                if (StringUtils.isEmpty(s) || offset < 0 || len < 0) {
+                    return;
+                }
+
+                ByteArrayOutputStream buff = EngineManager.BODY_BUFFER.getResponse();
+                int size = buff.size();
+                if (size < maxLength) {
+                    buff.write((new String(s.toCharArray(), offset, Math.min(len, maxLength - size))).getBytes());
+                }
             }
+        } catch (Throwable ignore) {
         }
     }