Capture text-based request bodies

apm-agent-core/src/main/java/co/elastic/apm/agent/impl/context/Request.java

@@ -42,7 +42,7 @@
 public class Request implements Recyclable {
 
 
-    private final ObjectPool<CharBuffer> charBufferPool = QueueBasedObjectPool.of(new MpmcAtomicArrayQueue<CharBuffer>(128), false,
+    private static final ObjectPool<CharBuffer> charBufferPool = QueueBasedObjectPool.of(new MpmcAtomicArrayQueue<CharBuffer>(128), false,
         new Allocator<CharBuffer>() {
             @Override
             public CharBuffer createInstance() {
@@ -71,6 +71,11 @@ public void recycle(CharBuffer object) {
      * A parsed key-value object of cookies
      */
     private final PotentiallyMultiValuedMap cookies = new PotentiallyMultiValuedMap();
+    /**
+     * Data should only contain the request body (not the query string). It can either be a dictionary (for standard HTTP requests) or a raw request body.
+     */
+    @Nullable
+    private String rawBody;
     /**
      * HTTP version.
      */
@@ -84,6 +89,7 @@ public void recycle(CharBuffer object) {
     private String method;
     @Nullable
     private CharBuffer bodyBuffer;
+    private boolean bodyBufferFinished = false;
 
     /**
      * Data should only contain the request body (not the query string). It can either be a dictionary (for standard HTTP requests) or a raw request body.
@@ -92,16 +98,34 @@ public void recycle(CharBuffer object) {
     public Object getBody() {
         if (!postParams.isEmpty()) {
             return postParams;
+        } else if (rawBody != null) {
+            return rawBody;
         } else {
             return bodyBuffer;
         }
     }
 
-    public void redactBody() {
+    @Nullable
+    public String getRawBody() {
+        return rawBody;
+    }
+
+    /**
+     * Sets the body as a raw string and removes any previously set {@link #postParams} or {@link #bodyBuffer}.
+     *
+     * @param rawBody the body as a raw string
+     */
+    public void setRawBody(String rawBody) {
         postParams.resetState();
         if (bodyBuffer != null) {
-            bodyBuffer.clear().append("[REDACTED]").flip();
+            charBufferPool.recycle(bodyBuffer);
+            bodyBuffer = null;
         }
+        this.rawBody = rawBody;
+    }
+
+    public void redactBody() {
+        setRawBody("[REDACTED]");
     }
 
     public Request addFormUrlEncodedParameter(String key, String value) {
@@ -132,6 +156,13 @@ public CharBuffer withBodyBuffer() {
         return this.bodyBuffer;
     }
 
+    public void endOfBufferInput() {
+        if (bodyBuffer != null && !bodyBufferFinished) {
+            bodyBufferFinished = true;
+            ((Buffer) bodyBuffer).flip();
+        }
+    }
+
     /**
      * Returns the associated pooled {@link CharBuffer} to record the request body.
      * <p>
@@ -142,6 +173,15 @@ public CharBuffer withBodyBuffer() {
      */
     @Nullable
     public CharBuffer getBodyBuffer() {
+        if (!bodyBufferFinished) {
+            return bodyBuffer;
+        } else {
+            return null;
+        }
+    }
+
+    @Nullable
+    public CharBuffer getBodyBufferForSerialization() {
         return bodyBuffer;
     }
 
@@ -231,6 +271,10 @@ public PotentiallyMultiValuedMap getCookies() {
         return cookies;
     }
 
+    void onTransactionEnd() {
+        endOfBufferInput();
+    }
+
     @Override
     public void resetState() {
         postParams.resetState();
@@ -240,10 +284,11 @@ public void resetState() {
         socket.resetState();
         url.resetState();
         cookies.resetState();
+        bodyBufferFinished = false;
         if (bodyBuffer != null) {
             charBufferPool.recycle(bodyBuffer);
+            bodyBuffer = null;
         }
-        bodyBuffer = null;
     }
 
     public void copyFrom(Request other) {
@@ -255,12 +300,12 @@ public void copyFrom(Request other) {
         this.url.copyFrom(other.url);
         this.cookies.copyFrom(other.cookies);
         if (other.bodyBuffer != null) {
-            final CharBuffer otherBuffer = other.getBodyBuffer();
+            final CharBuffer otherBuffer = other.bodyBuffer;
             final CharBuffer thisBuffer = this.withBodyBuffer();
             for (int i = 0; i < otherBuffer.length(); i++) {
                 thisBuffer.append(otherBuffer.charAt(i));
             }
-            thisBuffer.flip();
+            ((Buffer) thisBuffer).flip();
         }
     }
 

apm-agent-core/src/main/java/co/elastic/apm/agent/impl/context/TransactionContext.java

@@ -92,4 +92,8 @@ public void resetState() {
         request.resetState();
         user.resetState();
     }
+
+    public void onTransactionEnd() {
+        request.onTransactionEnd();
+    }
 }

apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/Db.java

@@ -38,7 +38,7 @@
  */
 public class Db implements Recyclable {
 
-    private final ObjectPool<CharBuffer> charBufferPool = QueueBasedObjectPool.of(new MpmcAtomicArrayQueue<CharBuffer>(128), false,
+    private static final ObjectPool<CharBuffer> charBufferPool = QueueBasedObjectPool.of(new MpmcAtomicArrayQueue<CharBuffer>(128), false,
         new Allocator<CharBuffer>() {
             @Override
             public CharBuffer createInstance() {

apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/Transaction.java

@@ -171,6 +171,7 @@ public void doEnd(long epochMicros) {
         if (type == null) {
             type = "custom";
         }
+        context.onTransactionEnd();
         this.tracer.endTransaction(this);
     }
 

apm-agent-core/src/main/java/co/elastic/apm/agent/matcher/WildcardMatcher.java

@@ -128,6 +128,18 @@ public static WildcardMatcher valueOf(final String wildcardString) {
         return new CompoundWildcardMatcher(wildcardString, matcher, matchers);
     }
 
+    /**
+     * Returns the first {@link WildcardMatcher} {@linkplain WildcardMatcher#matches(String) matching} the provided string.
+     *
+     * @param matchers the matchers which should be used to match the provided string
+     * @param s        the string to match against
+     * @return the first matching {@link WildcardMatcher}, or {@code null} if none match.
+     */
+    @Nullable
+    public static boolean isAnyMatch(List<WildcardMatcher> matchers, @Nullable String s) {
+        return anyMatch(matchers, s) != null;
+    }
+
     /**
      * Returns {@code true}, if any of the matchers match the provided string.
      *
@@ -136,17 +148,20 @@ public static WildcardMatcher valueOf(final String wildcardString) {
      * @return {@code true}, if any of the matchers match the provided string
      */
     @Nullable
-    public static WildcardMatcher anyMatch(List<WildcardMatcher> matchers, String s) {
+    public static WildcardMatcher anyMatch(List<WildcardMatcher> matchers, @Nullable String s) {
+        if (s == null) {
+            return null;
+        }
         return anyMatch(matchers, s, null);
     }
 
     /**
-     * Returns {@code true}, if any of the matchers match the provided partitioned string.
+     * Returns the first {@link WildcardMatcher} {@linkplain WildcardMatcher#matches(String) matching} the provided partitioned string.
      *
      * @param matchers   the matchers which should be used to match the provided string
      * @param firstPart  The first part of the string to match against.
      * @param secondPart The second part of the string to match against.
-     * @return {@code true}, if any of the matchers match the provided partitioned string
+     * @return the first matching {@link WildcardMatcher}, or {@code null} if none match.
      * @see #matches(String, String)
      */
     @Nullable

apm-agent-core/src/main/java/co/elastic/apm/agent/report/serialize/DslJsonSerializer.java

@@ -816,10 +816,15 @@ private void serializeRequest(final Request request) {
             // only one of those can be non-empty
             if (!request.getFormUrlEncodedParameters().isEmpty()) {
                 writeField("body", request.getFormUrlEncodedParameters());
-            } else if (request.getBodyBuffer() != null && request.getBodyBuffer().length() > 0) {
-                writeFieldName("body");
-                jw.writeString(request.getBodyBuffer());
-                jw.writeByte(COMMA);
+            } else if (request.getRawBody() != null) {
+                writeField("body", request.getRawBody());
+            } else {
+                final CharBuffer bodyBuffer = request.getBodyBufferForSerialization();
+                if (bodyBuffer != null && bodyBuffer.length() > 0) {
+                    writeFieldName("body");
+                    jw.writeString(bodyBuffer);
+                    jw.writeByte(COMMA);
+                }
             }
             if (request.getUrl().hasContent()) {
                 serializeUrl(request.getUrl());

apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/RequestStreamRecordingInstrumentation.java

@@ -0,0 +1,125 @@
+/*-
+ * #%L
+ * Elastic APM Java agent
+ * %%
+ * Copyright (C) 2018 - 2019 Elastic and contributors
+ * %%
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * #L%
+ */
+package co.elastic.apm.agent.servlet;
+
+import co.elastic.apm.agent.bci.ElasticApmInstrumentation;
+import co.elastic.apm.agent.bci.HelperClassManager;
+import co.elastic.apm.agent.bci.VisibleForAdvice;
+import co.elastic.apm.agent.impl.ElasticApmTracer;
+import co.elastic.apm.agent.impl.context.Request;
+import co.elastic.apm.agent.impl.transaction.Transaction;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.NamedElement;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+import javax.annotation.Nullable;
+import javax.servlet.ServletInputStream;
+import java.util.Arrays;
+import java.util.Collection;
+
+import static co.elastic.apm.agent.servlet.ServletInstrumentation.SERVLET_API;
+import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+import static net.bytebuddy.matcher.ElementMatchers.nameContains;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+
+public class RequestStreamRecordingInstrumentation extends ElasticApmInstrumentation {
+
+    @Nullable
+    @VisibleForAdvice
+    // referring to InputStreamWrapperFactory is legal because of type erasure
+    public static HelperClassManager<InputStreamWrapperFactory> wrapperHelperClassManager;
+
+    @Override
+    public void init(ElasticApmTracer tracer) {
+        wrapperHelperClassManager = HelperClassManager.ForSingleClassLoader.of(tracer,
+            "co.elastic.apm.agent.servlet.helper.InputStreamFactoryHelperImpl",
+            "co.elastic.apm.agent.servlet.helper.RecordingServletInputStreamWrapper");
+    }
+
+    @Override
+    public ElementMatcher<? super NamedElement> getTypeMatcherPreFilter() {
+        return nameContains("Request");
+    }
+
+    @Override
+    public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+        return hasSuperType(named("javax.servlet.ServletRequest")).and(not(isInterface()));
+    }
+
+    @Override
+    public ElementMatcher<? super MethodDescription> getMethodMatcher() {
+        return named("getInputStream").and(returns(named("javax.servlet.ServletInputStream")));
+    }
+
+    @Override
+    public Collection<String> getInstrumentationGroupNames() {
+        return Arrays.asList(SERVLET_API, "servlet-input-stream");
+    }
+
+    @Override
+    public Class<?> getAdviceClass() {
+        return GetInputStreamAdvice.class;
+    }
+
+    public interface InputStreamWrapperFactory {
+        ServletInputStream wrap(Request request, ServletInputStream servletInputStream);
+    }
+
+    public static class GetInputStreamAdvice {
+
+        @VisibleForAdvice
+        public static final ThreadLocal<Boolean> nestedThreadLocal = new ThreadLocal<Boolean>() {
+            @Override
+            protected Boolean initialValue() {
+                return Boolean.FALSE;
+            }
+        };
+
+        @Advice.OnMethodEnter(suppress = Throwable.class)
+        public static void onReadEnter(@Advice.This Object thiz,
+                                       @Advice.Local("transaction") Transaction transaction,
+                                       @Advice.Local("nested") boolean nested) {
+            nested = nestedThreadLocal.get();
+            nestedThreadLocal.set(Boolean.TRUE);
+        }
+
+        @Advice.OnMethodExit(suppress = Throwable.class)
+        public static void afterGetInputStream(@Advice.Return(readOnly = false) ServletInputStream inputStream,
+                                               @Advice.Local("nested") boolean nested) {
+            if (nested || tracer == null || wrapperHelperClassManager == null) {
+                return;
+            }
+            try {
+                final Transaction transaction = tracer.currentTransaction();
+                // only wrap if the body buffer has been initialized via ServletTransactionHelper.startCaptureBody
+                if (transaction != null && transaction.getContext().getRequest().getBodyBuffer() != null) {
+                    inputStream = wrapperHelperClassManager.getForClassLoaderOfClass(inputStream.getClass()).wrap(transaction.getContext().getRequest(), inputStream);
+                }
+            } finally {
+                nestedThreadLocal.set(Boolean.FALSE);
+            }
+        }
+    }
+}

apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/ServletApiAdvice.java

@@ -117,7 +117,7 @@ public static void onEnterServletService(@Advice.Argument(0) ServletRequest serv
 
             servletTransactionHelper.fillRequestContext(transaction, request.getProtocol(), request.getMethod(), request.isSecure(),
                 request.getScheme(), request.getServerName(), request.getServerPort(), request.getRequestURI(), request.getQueryString(),
-                request.getRemoteAddr());
+                request.getRemoteAddr(), request.getHeader("Content-Type"));
         }
     }