Add body capturing for reactive spring webclient

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

@@ -46,9 +46,11 @@ public void recycle(ByteBuffer object) {
         });
 
     private enum CaptureState {
-        NOT_ELIGIBLE,
-        ELIGIBLE,
-        STARTED
+        NOT_ELIGIBLE, // initial state
+        ELIGIBLE, // eligible but before preconditions evaluation
+        PRECONDITIONS_PASSED, // post preconditions (passed), can start capture
+        PRECONDITIONS_FAILED, // post preconditions (failed), no body will be captured
+        STARTED // the body capturing has been started, a buffer was acquired
     }
 
     private volatile CaptureState state;
@@ -95,17 +97,42 @@ public boolean isEligibleForCapturing() {
     }
 
     @Override
-    public boolean startCapture(@Nullable String requestCharset, int numBytesToCapture) {
+    public boolean havePreconditionsBeenChecked() {
+        return state == CaptureState.PRECONDITIONS_PASSED
+               || state == CaptureState.PRECONDITIONS_FAILED
+               || state == CaptureState.STARTED;
+    }
+
+    @Override
+    public void markPreconditionsFailed() {
+        synchronized (this) {
+            if (state == CaptureState.ELIGIBLE) {
+                state = CaptureState.PRECONDITIONS_FAILED;
+            }
+        }
+    }
+
+    @Override
+    public void markPreconditionsPassed(@Nullable String requestCharset, int numBytesToCapture) {
         if (numBytesToCapture > WebConfiguration.MAX_BODY_CAPTURE_BYTES) {
             throw new IllegalArgumentException("Capturing " + numBytesToCapture + " bytes is not supported, maximum is " + WebConfiguration.MAX_BODY_CAPTURE_BYTES + " bytes");
         }
-        if (state == CaptureState.ELIGIBLE) {
+        synchronized (this) {
+            if (state == CaptureState.ELIGIBLE) {
+                if (requestCharset != null) {
+                    this.charset.append(requestCharset);
+                }
+                this.numBytesToCapture = numBytesToCapture;
+                state = CaptureState.PRECONDITIONS_PASSED;
+            }
+        }
+    }
+
+    @Override
+    public boolean startCapture() {
+        if (state == CaptureState.PRECONDITIONS_PASSED) {
             synchronized (this) {
-                if (state == CaptureState.ELIGIBLE) {
-                    if (requestCharset != null) {
-                        this.charset.append(requestCharset);
-                    }
-                    this.numBytesToCapture = numBytesToCapture;
+                if (state == CaptureState.PRECONDITIONS_PASSED) {
                     state = CaptureState.STARTED;
                     return true;
                 }

apm-agent-core/src/test/java/co/elastic/apm/agent/impl/context/BodyCaptureImplTest.java

@@ -32,7 +32,8 @@ public class BodyCaptureImplTest {
     public void testAppendTruncation() {
         BodyCaptureImpl capture = new BodyCaptureImpl();
         capture.markEligibleForCapturing();
-        capture.startCapture("foobar", 10);
+        capture.markPreconditionsPassed("foobar", 10);
+        capture.startCapture();
         assertThat(capture.isFull()).isFalse();
 
         capture.append("123Hello World!".getBytes(StandardCharsets.UTF_8), 3, 5);
@@ -55,21 +56,28 @@ public void testLifecycle() {
         BodyCaptureImpl capture = new BodyCaptureImpl();
 
         assertThat(capture.isEligibleForCapturing()).isFalse();
-        assertThat(capture.startCapture("foobar", 42))
+        assertThat(capture.havePreconditionsBeenChecked()).isFalse();
+        assertThat(capture.startCapture())
             .isFalse();
         assertThatThrownBy(() -> capture.append((byte) 42)).isInstanceOf(IllegalStateException.class);
 
         capture.markEligibleForCapturing();
         assertThat(capture.isEligibleForCapturing()).isTrue();
+        assertThat(capture.havePreconditionsBeenChecked()).isFalse();
         assertThatThrownBy(() -> capture.append((byte) 42)).isInstanceOf(IllegalStateException.class);
 
-        assertThat(capture.startCapture("foobar", 42))
-            .isTrue();
+        capture.markPreconditionsPassed("foobar", 42);
+        assertThat(capture.isEligibleForCapturing()).isTrue();
+        assertThat(capture.havePreconditionsBeenChecked()).isTrue();
+
+
+        assertThat(capture.startCapture()).isTrue();
         capture.append((byte) 42); //ensure no exception thrown
 
         // startCapture should return true only once
-        assertThat(capture.startCapture("foobar", 42))
-            .isFalse();
+        assertThat(capture.havePreconditionsBeenChecked()).isTrue();
+        assertThat(capture.startCapture()).isFalse();
+        assertThat(capture.havePreconditionsBeenChecked()).isTrue();
 
         capture.resetState();
         assertThat(capture.getCharset()).isNull();

apm-agent-core/src/test/java/co/elastic/apm/agent/report/serialize/DslJsonSerializerTest.java

@@ -475,7 +475,8 @@ private SpanImpl createSpanWithRequestBody(@Nullable byte[] bodyBytes, @Nullable
         SpanImpl span = new SpanImpl(tracer);
         BodyCaptureImpl bodyCapture = span.getContext().getHttp().getRequestBody();
         bodyCapture.markEligibleForCapturing();
-        bodyCapture.startCapture(charset, WebConfiguration.MAX_BODY_CAPTURE_BYTES);
+        bodyCapture.markPreconditionsPassed(charset, WebConfiguration.MAX_BODY_CAPTURE_BYTES);
+        bodyCapture.startCapture();
 
         if (bodyBytes != null) {
             bodyCapture.append(bodyBytes, 0, bodyBytes.length);

apm-agent-plugins/apm-apache-httpclient/apm-apache-httpclient-common/src/main/java/co/elastic/apm/agent/httpclient/common/RequestBodyCaptureRegistry.java

@@ -56,7 +56,7 @@ public static <REQUEST, HTTPENTITY> void potentiallyCaptureRequestBody(
         ApacheHttpClientEntityAccessor<REQUEST, HTTPENTITY> adapter,
         TextHeaderGetter<REQUEST> headerGetter
     ) {
-        if (HttpClientHelper.startRequestBodyCapture(abstractSpan, request, headerGetter)) {
+        if (HttpClientHelper.checkAndStartRequestBodyCapture(abstractSpan, request, headerGetter)) {
             Span<?> span = (Span<?>) abstractSpan;
             byte[] simpleBytes = adapter.getSimpleBodyBytes(request);
             if (simpleBytes != null) {

apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/HttpClientHelper.java

@@ -27,6 +27,7 @@
 import co.elastic.apm.agent.tracer.TraceState;
 import co.elastic.apm.agent.tracer.configuration.WebConfiguration;
 import co.elastic.apm.agent.tracer.dispatch.TextHeaderGetter;
+import co.elastic.apm.agent.tracer.metadata.BodyCapture;
 
 import javax.annotation.Nullable;
 import java.net.URI;
@@ -75,29 +76,43 @@ public static Span<?> startHttpClientSpan(TraceState<?> activeContext, String me
         return span;
     }
 
-    public static <R> boolean startRequestBodyCapture(@Nullable AbstractSpan<?> abstractSpan, R request, TextHeaderGetter<R> headerGetter) {
+    public static <R> void checkBodyCapturePreconditions(@Nullable AbstractSpan<?> abstractSpan, R request, TextHeaderGetter<R> headerGetter) {
         if (!(abstractSpan instanceof Span<?>)) {
-            return false;
+            return;
         }
         Span<?> span = (Span<?>) abstractSpan;
-        if (!span.getContext().getHttp().getRequestBody().isEligibleForCapturing()) {
-            return false;
+        BodyCapture bodyCapture = span.getContext().getHttp().getRequestBody();
+        if (!bodyCapture.isEligibleForCapturing()) {
+            return;
+        }
+        if (bodyCapture.havePreconditionsBeenChecked()) {
+            return;
         }
         WebConfiguration webConfig = GlobalTracer.get().getConfig(WebConfiguration.class);
         int byteCount = webConfig.getCaptureClientRequestBytes();
         if (byteCount == 0) {
-            return false;
+            bodyCapture.markPreconditionsFailed();
+            return;
         }
         List<WildcardMatcher> contentTypes = webConfig.getCaptureContentTypes();
         String contentTypeHeader = headerGetter.getFirstHeader("Content-Type", request);
         if (contentTypeHeader == null) {
             contentTypeHeader = "";
         }
         if (WildcardMatcher.anyMatch(contentTypes, contentTypeHeader) == null) {
+            bodyCapture.markPreconditionsFailed();
+            return;
+        }
+        bodyCapture.markPreconditionsPassed(extractCharsetFromContentType(contentTypeHeader), byteCount);
+    }
+
+    public static <R> boolean checkAndStartRequestBodyCapture(@Nullable AbstractSpan<?> abstractSpan, R request, TextHeaderGetter<R> headerGetter) {
+        if (!(abstractSpan instanceof Span<?>)) {
             return false;
         }
-        String charset = extractCharsetFromContentType(contentTypeHeader);
-        return span.getContext().getHttp().getRequestBody().startCapture(charset, byteCount);
+        checkBodyCapturePreconditions(abstractSpan, request, headerGetter);
+        Span<?> span = (Span<?>) abstractSpan;
+        return span.getContext().getHttp().getRequestBody().startCapture();
     }
 
     //Visible for testing

apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelper.java

@@ -4,7 +4,7 @@
 import co.elastic.apm.agent.tracer.SpanEndListener;
 import co.elastic.apm.agent.tracer.metadata.BodyCapture;
 
-class RequestBodyRecordingHelper implements SpanEndListener<Span<?>> {
+public class RequestBodyRecordingHelper implements SpanEndListener<Span<?>> {
 
     /**
      * We do not need to participate in span reference counting here.
@@ -21,17 +21,25 @@ public RequestBodyRecordingHelper(Span<?> clientSpan) {
         }
     }
 
-    void appendToBody(byte b) {
+
+    /**
+     * @param b the byte to append
+     * @return false, if the body buffer is full and future calls would be no-op. True otherwise.
+     */
+    public boolean appendToBody(byte b) {
         if (clientSpan != null) {
             BodyCapture requestBody = clientSpan.getContext().getHttp().getRequestBody();
             requestBody.append(b);
             if (requestBody.isFull()) {
                 releaseSpan();
+            } else {
+                return true;
             }
         }
+        return false;
     }
 
-    void appendToBody(byte[] b, int off, int len) {
+    public void appendToBody(byte[] b, int off, int len) {
         if (clientSpan != null) {
             BodyCapture requestBody = clientSpan.getContext().getHttp().getRequestBody();
             requestBody.append(b, off, len);

apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelperTest.java

@@ -40,7 +40,8 @@ public void ensureNoModificationAfterSpanEnd() {
         SpanImpl span = rootTx.createSpan();
         BodyCaptureImpl spanBody = span.getContext().getHttp().getRequestBody();
         spanBody.markEligibleForCapturing();
-        spanBody.startCapture(null, 100);
+        spanBody.markPreconditionsPassed(null, 100);
+        spanBody.startCapture();
 
         RequestBodyRecordingHelper helper = new RequestBodyRecordingHelper(span);
         helper.appendToBody(new byte[]{1, 2, 3, 4}, 1, 2);
@@ -66,7 +67,8 @@ public void ensureLimitRespected() {
         SpanImpl span = rootTx.createSpan();
         BodyCaptureImpl spanBody = span.getContext().getHttp().getRequestBody();
         spanBody.markEligibleForCapturing();
-        spanBody.startCapture(null, 3);
+        spanBody.markPreconditionsPassed(null, 3);
+        spanBody.startCapture();
 
         RequestBodyRecordingHelper helper = new RequestBodyRecordingHelper(span);
         helper.appendToBody((byte) 1);

apm-agent-plugins/apm-spring-webflux/apm-spring-webclient-plugin/src/main/java/co/elastic/apm/agent/springwebclient/BodyCaptureRegistry.java

@@ -0,0 +1,30 @@
+package co.elastic.apm.agent.springwebclient;
+
+import co.elastic.apm.agent.httpclient.RequestBodyRecordingHelper;
+import co.elastic.apm.agent.sdk.weakconcurrent.WeakConcurrent;
+import co.elastic.apm.agent.sdk.weakconcurrent.WeakMap;
+import co.elastic.apm.agent.tracer.AbstractSpan;
+import co.elastic.apm.agent.tracer.Span;
+import org.springframework.http.client.reactive.ClientHttpRequest;
+
+import javax.annotation.Nullable;
+
+public class BodyCaptureRegistry {
+
+    private static final WeakMap<ClientHttpRequest, RequestBodyRecordingHelper> PENDING_RECORDINGS = WeakConcurrent.buildMap();
+
+    public static void maybeCaptureBodyFor(AbstractSpan<?> abstractSpan, ClientHttpRequest request) {
+        if (!(abstractSpan instanceof Span<?>)) {
+            return;
+        }
+        Span<?> span = (Span<?>) abstractSpan;
+        if (span.getContext().getHttp().getRequestBody().startCapture()) {
+            PENDING_RECORDINGS.put(request, new RequestBodyRecordingHelper(span));
+        }
+    }
+
+    @Nullable
+    public static RequestBodyRecordingHelper activateRecording(ClientHttpRequest request) {
+        return PENDING_RECORDINGS.remove(request);
+    }
+}

apm-agent-plugins/apm-spring-webflux/apm-spring-webclient-plugin/src/main/java/co/elastic/apm/agent/springwebclient/ClientHttpConnectorInstrumentation.java

@@ -0,0 +1,102 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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.
+ */
+package co.elastic.apm.agent.springwebclient;
+
+import co.elastic.apm.agent.sdk.ElasticApmInstrumentation;
+import co.elastic.apm.agent.tracer.GlobalTracer;
+import co.elastic.apm.agent.tracer.TraceState;
+import co.elastic.apm.agent.tracer.Tracer;
+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 org.springframework.http.HttpMethod;
+import org.springframework.http.client.reactive.ClientHttpRequest;
+import reactor.core.publisher.Mono;
+
+import javax.annotation.Nullable;
+import java.net.URI;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.function.Function;
+
+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.nameStartsWith;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+/**
+ * Instruments all {@link org.springframework.http.client.reactive.ClientHttpConnector} types, to preserve the span context
+ * within the callback passed to {@link org.springframework.http.client.reactive.ClientHttpConnector#connect(HttpMethod, URI, Function)}.
+ * If the span is ready for it, this will cause the request body to be captured.
+ */
+public class ClientHttpConnectorInstrumentation extends ElasticApmInstrumentation {
+
+    private static final Tracer tracer = GlobalTracer.get();
+
+    @Override
+    public ElementMatcher<? super NamedElement> getTypeMatcherPreFilter() {
+        return nameStartsWith("org.springframework.http.")
+            .and(nameContains("Connector"));
+    }
+
+    @Override
+    public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+        return hasSuperType(named("org.springframework.http.client.reactive.ClientHttpConnector"))
+            .and(not(isInterface()));
+    }
+
+    @Override
+    public ElementMatcher<? super MethodDescription> getMethodMatcher() {
+        return named("connect")
+            .and(takesArgument(2, named("java.util.function.Function")));
+    }
+
+    @Override
+    public Collection<String> getInstrumentationGroupNames() {
+        return Arrays.asList("http-client", "spring-webclient");
+    }
+
+    public static class AdviceClass {
+
+        @Nullable
+        @Advice.OnMethodEnter(suppress = Throwable.class, inline = false)
+        @Advice.AssignReturned.ToArguments(@Advice.AssignReturned.ToArguments.ToArgument(2))
+        public static Function<? super ClientHttpRequest, Mono<Void>> onBefore(@Advice.Argument(2) Function<? super ClientHttpRequest, Mono<Void>> requestCallback) {
+            TraceState<?> context = tracer.currentContext();
+            if (context.isEmpty()) {
+                return requestCallback;
+            }
+            return new Function<ClientHttpRequest, Mono<Void>>() {
+                @Override
+                public Mono<Void> apply(ClientHttpRequest clientHttpRequest) {
+                    // Note that even though ClientHttpRequest exposes headers via the interface, those are empty
+                    // therefore we check the span capturing pre-conditions in the WebClientExchangeFunctionInstrumentation instead
+                    BodyCaptureRegistry.maybeCaptureBodyFor(context.getSpan(), clientHttpRequest);
+                    return requestCallback.apply(clientHttpRequest);
+                }
+            };
+        }
+
+    }
+}