Merge pull request #475 from lostsnow/fix/request-response-body-fetch

lostsnow · web-flow · commit 8ee423bce4f0 · 2023-02-24T17:14:30.000+08:00
Fix/request response body fetch
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
@@ -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;
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
@@ -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;
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
@@ -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"
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
@@ -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;
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/j2ee/dispatch/DispatchJ2ee.java b/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;
     }
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/j2ee/dispatch/PrintWriterAdapter.java b/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;
+    }
+}
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/j2ee/dispatch/PrintWriterWriteAdviceAdapter.java b/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);
+        }
+    }
+}
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/framework/j2ee/dispatch/ServletOutputStreamWriteAdviceAdapter.java b/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() {
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/SpyDispatcherImpl.java b/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
      *
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/HttpImpl.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/HttpImpl.java
@@ -95,6 +95,34 @@ public static void solveHttpRequest(Object obj, Object req, Object resp, Map<Str
         } catch (Throwable ignore) {
         }
 
+        try {
+            String contentType = ((Map<String, String>) requestMeta.get("headers")).get("Content-Type");
+            String method = (String) requestMeta.get("method");
+            if (("POST".equalsIgnoreCase(method) || "PUT".equalsIgnoreCase(method))
+                    && !isRawBody(contentType)) {
+                Method getParameterNamesMethod = ReflectUtils.getDeclaredMethodFromSuperClass(req.getClass(),
+                        "getParameterNames", null);
+                Method getParameterMethod = ReflectUtils.getDeclaredMethodFromSuperClass(req.getClass(),
+                        "getParameter", new Class[]{String.class});
+                Enumeration<?> parameterNames = (Enumeration<?>) getParameterNamesMethod.invoke(req);
+                StringBuilder postBody = new StringBuilder();
+                boolean first = true;
+                while (parameterNames.hasMoreElements()) {
+                    String key = (String) parameterNames.nextElement();
+                    if (first) {
+                        first = false;
+                        postBody.append(key).append("=").append((String) getParameterMethod.invoke(req, key));
+                    } else {
+                        postBody.append("&").append(key).append("=").append((String) getParameterMethod.invoke(req, key));
+                    }
+                }
+                if (postBody.length() > 0) {
+                    requestMeta.put("body", postBody.toString());
+                }
+            }
+        } catch (Throwable ignore) {
+        }
+
         EngineManager.enterHttpEntry(requestMeta);
         DongTaiLog.debug("HTTP Request:{} {} from: {}", requestMeta.get("method"), requestMeta.get("requestURI"),
                 obj.getClass().getName());
@@ -111,6 +139,9 @@ public static Map<String, String> parseRequestHeaders(Object req, Enumeration<?>
             try {
                 String key = (String) headerNames.nextElement();
                 String val = (String) getHeaderMethod.invoke(req, key);
+                if ("content-type".equalsIgnoreCase(key)) {
+                    key = "Content-Type";
+                }
                 headers.put(key, val);
             } catch (Throwable ignore) {
             }
@@ -119,6 +150,13 @@ public static Map<String, String> parseRequestHeaders(Object req, Enumeration<?>
     }
 
     public static void onServletInputStreamRead(int ret, String desc, Object stream, byte[] bs, int offset, int len) {
+        if (EngineManager.REQUEST_CONTEXT.get() != null
+                && EngineManager.REQUEST_CONTEXT.get().get("body") != null
+                && EngineManager.REQUEST_CONTEXT.get().get("body") != ""
+        ) {
+            return;
+        }
+
         if ("()I".equals(desc)) {
             if (ret == -1) {
                 return;
@@ -189,33 +227,92 @@ 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) {
         }
     }
 
     public static IastClassLoader getClassLoader() {
         return iastClassLoader;
     }
+
+    public static boolean isRawBody(String contentType) {
+        return contentType != null
+                && (contentType.contains("application/json") || contentType.contains("application/xml"));
+    }
 }
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/graphy/GraphBuilder.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/graphy/GraphBuilder.java
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/policy/PolicyManager.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/policy/PolicyManager.java
diff --git a/dongtai-core/src/main/resources/com.secnium.iast.resources/blackext.txt b/dongtai-core/src/main/resources/com.secnium.iast.resources/blackext.txt
diff --git a/dongtai-spy/src/main/java/java/lang/dongtai/NopSpy.java b/dongtai-spy/src/main/java/java/lang/dongtai/NopSpy.java
diff --git a/dongtai-spy/src/main/java/java/lang/dongtai/SpyDispatcher.java b/dongtai-spy/src/main/java/java/lang/dongtai/SpyDispatcher.java

Original file line number	Diff line number	Diff line change
`@@ -4,7 +4,7 @@ public enum Scope {`
`4`	`4`	`HTTP_REQUEST(1),`
`5`	`5`	`HTTP_ENTRY(2),`
`6`	`6`	`SERVLET_INPUT_STREAM_READ(3),`
`7`		`- SERVLET_OUTPUT_STREAM_WRITE(4),`
	`7`	`+ SERVLET_OUTPUT_WRITE(4),`
`8`	`8`	`;`
`9`	`9`
`10`	`10`	`private final int id;`
Original file line number	Diff line number	Diff line change
`@@ -50,7 +50,9 @@ public ClassVisitor initial(ClassVisitor classVisitor, ClassContext context, Pol`
`50`	`50`	`if (pluginVisitor != classVisitor) {`
`51`	`51`	`classVisitor = pluginVisitor;`
`52`	`52`	`// TODO: need transform multiple times?`
`53`		`- break;`
	`53`	`+ if (!context.getClassName().equals(DispatchJ2ee.APACHE_COYOTE_WRITER)) {`
	`54`	`+ break;`
	`55`	`+ }`
`54`	`56`	`}`
`55`	`57`	`}`
`56`	`58`	`return classVisitor;`