hypertrace · ryandens · Jun 25, 2021 · Jun 10, 2021 · Jun 21, 2021 · Jun 22, 2021
diff --git a/instrumentation/build.gradle.kts b/instrumentation/build.gradle.kts
@@ -40,6 +40,8 @@ dependencies{
     implementation(project(":instrumentation:apache-httpasyncclient-4.1"))
     implementation(project(":instrumentation:netty:netty-4.0"))
     implementation(project(":instrumentation:netty:netty-4.1"))
+    implementation(project(":instrumentation:undertow:undertow-1.4"))
+    implementation(project(":instrumentation:undertow:undertow-servlet-1.4"))
     implementation(project(":otel-extensions"))
 }
 

diff --git a/...ent/instrumentation/hypertrace/java/inputstream/InputStreamInstrumentationModuleTest.java b/...ent/instrumentation/hypertrace/java/inputstream/InputStreamInstrumentationModuleTest.java
@@ -20,13 +20,14 @@
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.sdk.trace.data.SpanData;
 import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 import org.ContextAccessor;
 import org.hypertrace.agent.core.instrumentation.SpanAndBuffer;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedBuffersFactory;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStream;
 import org.hypertrace.agent.testing.AbstractInstrumenterTest;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
@@ -94,7 +95,8 @@ public void readBytesOffset() {
   private void read(InputStream inputStream, Runnable read, String expected) {
     Span span = TEST_TRACER.spanBuilder("test-span").startSpan();
 
-    ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+    BoundedByteArrayOutputStream buffer =
+        BoundedBuffersFactory.createStream(StandardCharsets.ISO_8859_1);
     ContextAccessor.addToInputStreamContext(
         inputStream, new SpanAndBuffer(span, buffer, ATTRIBUTE_KEY, StandardCharsets.ISO_8859_1));
 

diff --git a/instrumentation/undertow/undertow-1.4/build.gradle.kts b/instrumentation/undertow/undertow-1.4/build.gradle.kts
@@ -0,0 +1,37 @@
+plugins {
+    `java-library`
+    id("net.bytebuddy.byte-buddy")
+    id("io.opentelemetry.instrumentation.auto-instrumentation")
+    muzzle
+}
+
+muzzle {
+    pass {
+        group = "io.undertow"
+        module = "undertow-core"
+        versions = "[1.4.0.Final,)"
+        assertInverse = true
+    }
+}
+
+afterEvaluate{
+    io.opentelemetry.instrumentation.gradle.bytebuddy.ByteBuddyPluginConfigurator(project,
+            sourceSets.main.get(),
+            "io.opentelemetry.javaagent.tooling.muzzle.collector.MuzzleCodeGenerationPlugin",
+            project(":javaagent-tooling").configurations["instrumentationMuzzle"] + configurations.runtimeClasspath
+    ).configure()
+}
+
+val versions: Map<String, String> by extra
+
+dependencies {
+    implementation("io.opentelemetry.javaagent.instrumentation:opentelemetry-javaagent-undertow-1.4:${versions["opentelemetry_java_agent"]}")
+    library("io.undertow:undertow-core:1.4.0.Final")
+    implementation(project(":instrumentation:undertow:undertow-common"))
+    testImplementation(testFixtures(project(":testing-common")))
+    testImplementation("javax.servlet:javax.servlet-api:3.1.0")
+    testImplementation("io.undertow:undertow-servlet:2.0.0.Final")
+    testRuntimeOnly(project(":instrumentation:servlet:servlet-3.0"))
+    testRuntimeOnly(project(":instrumentation:undertow:undertow-servlet-1.4"))
+}
+
diff --git a/...avaagent/instrumentation/hypertrace/undertow/v1_4/StreamSourceChannelInstrumentation.java b/...avaagent/instrumentation/hypertrace/undertow/v1_4/StreamSourceChannelInstrumentation.java
@@ -0,0 +1,101 @@
+/*
+ * Copyright The Hypertrace Authors
+ *
+ * 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.
+ */
+
+package io.opentelemetry.javaagent.instrumentation.hypertrace.undertow.v1_4;
+
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
+import io.opentelemetry.javaagent.instrumentation.hypertrace.undertow.v1_4.utils.Utils;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import io.opentelemetry.javaagent.tooling.bytebuddy.matcher.AgentElementMatchers;
+import io.undertow.server.HttpServerExchange;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.hypertrace.agent.core.instrumentation.HypertraceCallDepthThreadLocalMap;
+import org.hypertrace.agent.core.instrumentation.SpanAndBuffer;
+import org.xnio.channels.StreamSourceChannel;
+
+/** Instrumentation for {@link StreamSourceChannel} implementations */
+public final class StreamSourceChannelInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return AgentElementMatchers.safeHasSuperType(named("org.xnio.channels.StreamSourceChannel"));
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return Collections.singletonMap(
+        named("read")
+            .and(takesArguments(1))
+            .and(takesArgument(0, ByteBuffer.class))
+            .and(returns(int.class))
+            .and(isPublic()),
+        StreamSourceChannelInstrumentation.class.getName() + "$Read_advice");
+  }
+
+  /**
+   * Decorates the {@link StreamSourceChannel#read(ByteBuffer)} implementations with logic to
+   * capture data read into the request body {@link ByteBuffer} and report it. This instrumentation
+   * short-circuits if:
+   *
+   * <ul>
+   *   <li>We're in a nested {@link StreamSourceChannel#read(ByteBuffer)} call
+   *   <li>A {@link Throwable} was thrown in the context of the {@link
+   *       StreamSourceChannel#read(ByteBuffer)}, causing the method to exit
+   *   <li>The instrumented {@link StreamSourceChannel} was never put in the {@link
+   *       InstrumentationContext} by {@link
+   *       UndertowHttpServerExchangeInstrumentation.GetRequestChannel_advice#exit(HttpServerExchange,
+   *       StreamSourceChannel)}
+   * </ul>
+   */
+  static final class Read_advice {
+
+    @Advice.OnMethodEnter
+    public static void trackCallDepth() {
+      HypertraceCallDepthThreadLocalMap.incrementCallDepth(StreamSourceChannel.class);
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void exit(
+        @Advice.Return final int numBytesRead,
+        @Advice.This final StreamSourceChannel streamSourceChannel,
+        @Advice.Thrown final Throwable thrown,
+        @Advice.Argument(0) final ByteBuffer byteBuffer) {
+      if (HypertraceCallDepthThreadLocalMap.decrementCallDepth(StreamSourceChannel.class) > 0
+          || thrown != null) {
+        return;
+      }
+      final ContextStore<StreamSourceChannel, SpanAndBuffer> contextStore =
+          InstrumentationContext.get(StreamSourceChannel.class, SpanAndBuffer.class);
+      final SpanAndBuffer spanAndBuffer = contextStore.get(streamSourceChannel);
+      if (spanAndBuffer != null) {
+        Utils.handleRead(byteBuffer.asReadOnlyBuffer(), numBytesRead, spanAndBuffer);
+      }
+    }
+  }
+}
diff --git a/...t/instrumentation/hypertrace/undertow/v1_4/UndertowHttpServerExchangeInstrumentation.java b/...t/instrumentation/hypertrace/undertow/v1_4/UndertowHttpServerExchangeInstrumentation.java
@@ -0,0 +1,88 @@
+/*
+ * Copyright The Hypertrace Authors
+ *
+ * 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.
+ */
+
+package io.opentelemetry.javaagent.instrumentation.hypertrace.undertow.v1_4;
+
+import static net.bytebuddy.matcher.ElementMatchers.failSafe;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
+import io.opentelemetry.javaagent.instrumentation.hypertrace.undertow.common.RequestBodyCaptureMethod;
+import io.opentelemetry.javaagent.instrumentation.hypertrace.undertow.v1_4.utils.Utils;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import io.undertow.server.HttpServerExchange;
+import java.util.Collections;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.hypertrace.agent.core.instrumentation.SpanAndBuffer;
+import org.xnio.channels.StreamSourceChannel;
+
+/** Instrumentation for {@link HttpServerExchange} to capture request bodies */
+public final class UndertowHttpServerExchangeInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return failSafe(named("io.undertow.server.HttpServerExchange"));
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return Collections.singletonMap(
+        named("getRequestChannel")
+            .and(takesArguments(0))
+            .and(returns(named("org.xnio.channels.StreamSourceChannel")))
+            .and(isPublic()),
+        UndertowHttpServerExchangeInstrumentation.class.getName() + "$GetRequestChannel_advice");
+  }
+
+  /**
+   * Decorates {@link HttpServerExchange#getRequestChannel()} with instrumentation to store a {@link
+   * SpanAndBuffer} in the {@link InstrumentationContext}
+   */
+  static final class GetRequestChannel_advice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void exit(
+        @Advice.This final HttpServerExchange thizz,
+        @Advice.Return final StreamSourceChannel returnedChannel) {
+      final RequestBodyCaptureMethod requestBodyCaptureMethod =
+          InstrumentationContext.get(HttpServerExchange.class, RequestBodyCaptureMethod.class)
+              .get(thizz);
+      if (RequestBodyCaptureMethod.SERVLET.equals(requestBodyCaptureMethod)) {
+        // short circuit if we detect that we can capture the request body with servlet
+        // instrumentation
+        return;
+      }
+      final ContextStore<StreamSourceChannel, SpanAndBuffer> contextStore =
+          InstrumentationContext.get(StreamSourceChannel.class, SpanAndBuffer.class);
+      if (contextStore.get(returnedChannel) != null) {
+        // HttpServerExchange.getRequestChannel only creates a new channel the first time it is
+        // invoked.on subsequent invocations, we do not want to create a new buffer and put it in
+        // the context, as that would reset the state and could potential result in lost request
+        // bodies
+        return;
+      }
+      Utils.createAndStoreBufferForSpan(thizz, returnedChannel, contextStore);
+    }
+  }
+}
diff --git a/...try/javaagent/instrumentation/hypertrace/undertow/v1_4/UndertowInstrumentationModule.java b/...try/javaagent/instrumentation/hypertrace/undertow/v1_4/UndertowInstrumentationModule.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright The Hypertrace Authors
+ *
+ * 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.
+ */
+
+package io.opentelemetry.javaagent.instrumentation.hypertrace.undertow.v1_4;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.tooling.InstrumentationModule;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.Arrays;
+import java.util.List;
+
+@AutoService(InstrumentationModule.class)
+public final class UndertowInstrumentationModule extends InstrumentationModule {
+
+  public UndertowInstrumentationModule() {
+    super("undertow", "undertow-1.4-ht", "ht", "undertow-ht");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return Arrays.asList(
+        new UndertowHttpServerExchangeInstrumentation(), new StreamSourceChannelInstrumentation());
+  }
+}
diff --git a/...strumentation/hypertrace/undertow/v1_4/utils/BodyCapturingExchangeCompletionListener.java b/...strumentation/hypertrace/undertow/v1_4/utils/BodyCapturingExchangeCompletionListener.java
@@ -0,0 +1,50 @@
+/*
+ * Copyright The Hypertrace Authors
+ *
+ * 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.
+ */
+
+package io.opentelemetry.javaagent.instrumentation.hypertrace.undertow.v1_4.utils;
+
+import io.undertow.server.ExchangeCompletionListener;
+import io.undertow.server.HttpServerExchange;
+import java.io.UnsupportedEncodingException;
+import org.hypertrace.agent.core.instrumentation.HypertraceSemanticAttributes;
+import org.hypertrace.agent.core.instrumentation.SpanAndBuffer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public final class BodyCapturingExchangeCompletionListener implements ExchangeCompletionListener {
+
+  private static final Logger log =
+      LoggerFactory.getLogger(BodyCapturingExchangeCompletionListener.class);
+
+  private final SpanAndBuffer spanAndBuffer;
+
+  public BodyCapturingExchangeCompletionListener(final SpanAndBuffer spanAndBuffer) {
+    this.spanAndBuffer = spanAndBuffer;
+  }
+
+  @Override
+  public void exchangeEvent(HttpServerExchange exchange, NextListener nextListener) {
+    final String body;
+    try {
+      body = spanAndBuffer.byteArrayBuffer.toStringWithSuppliedCharset();
+    } catch (UnsupportedEncodingException e) {
+      log.error("illegal encoding", e);
+      return;
+    }
+    spanAndBuffer.span.setAttribute(HypertraceSemanticAttributes.HTTP_REQUEST_BODY, body);
+    nextListener.proceed();
+  }
+}