diff --git a/instrumentation/undertow-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/HttpTransferEncodingInstrumentation.java b/instrumentation/undertow-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/HttpTransferEncodingInstrumentation.java new file mode 100644 index 0000000000..4f6b6e46b8 --- /dev/null +++ b/instrumentation/undertow-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/HttpTransferEncodingInstrumentation.java @@ -0,0 +1,45 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.opentelemetry.javaagent.instrumentation.undertow; + +import static io.opentelemetry.javaagent.instrumentation.undertow.UndertowSingletons.helper; +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; + +import io.opentelemetry.context.Context; +import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder; +import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; +import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; +import io.undertow.server.HttpServerExchange; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +public class HttpTransferEncodingInstrumentation implements TypeInstrumentation { + @Override + public ElementMatcher typeMatcher() { + return named("io.undertow.server.protocol.http.HttpTransferEncoding"); + } + + @Override + public void transform(TypeTransformer transformer) { + transformer.applyAdviceToMethod( + named("createSinkConduit") + .and(takesArgument(0, named("io.undertow.server.HttpServerExchange"))), + this.getClass().getName() + "$ResponseAdvice"); + } + + @SuppressWarnings("unused") + public static class ResponseAdvice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter(@Advice.Argument(0) HttpServerExchange exchange) { + Context context = helper().getServerContext(exchange); + HttpServerResponseCustomizerHolder.getCustomizer() + .customize(context, exchange, UndertowHttpResponseMutator.INSTANCE); + } + } +} diff --git a/instrumentation/undertow-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/UndertowHttpResponseMutator.java b/instrumentation/undertow-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/UndertowHttpResponseMutator.java new file mode 100644 index 0000000000..ce12a57117 --- /dev/null +++ b/instrumentation/undertow-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/UndertowHttpResponseMutator.java @@ -0,0 +1,21 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.opentelemetry.javaagent.instrumentation.undertow; + +import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator; +import io.undertow.server.HttpServerExchange; +import io.undertow.util.HttpString; + +public enum UndertowHttpResponseMutator implements HttpServerResponseMutator { + INSTANCE; + + UndertowHttpResponseMutator() {} + + @Override + public void appendHeader(HttpServerExchange exchange, String name, String value) { + exchange.getResponseHeaders().add(HttpString.tryFromString(name), value); + } +} diff --git a/instrumentation/undertow-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/UndertowInstrumentationModule.java b/instrumentation/undertow-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/UndertowInstrumentationModule.java index 4810eca4dd..bbe18133e2 100644 --- a/instrumentation/undertow-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/UndertowInstrumentationModule.java +++ b/instrumentation/undertow-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/UndertowInstrumentationModule.java @@ -21,6 +21,9 @@ public class UndertowInstrumentationModule extends InstrumentationModule { @Override public List typeInstrumentations() { - return asList(new HandlerInstrumentation(), new HttpServerExchangeInstrumentation()); + return asList( + new HandlerInstrumentation(), + new HttpServerExchangeInstrumentation(), + new HttpTransferEncodingInstrumentation()); } } diff --git a/instrumentation/undertow-1.4/javaagent/src/test/groovy/UndertowServerDispatchTest.groovy b/instrumentation/undertow-1.4/javaagent/src/test/groovy/UndertowServerDispatchTest.groovy index 40d3a14a11..aea4d09b07 100644 --- a/instrumentation/undertow-1.4/javaagent/src/test/groovy/UndertowServerDispatchTest.groovy +++ b/instrumentation/undertow-1.4/javaagent/src/test/groovy/UndertowServerDispatchTest.groovy @@ -34,6 +34,11 @@ class UndertowServerDispatchTest extends HttpServerTest implements Age return false } + @Override + boolean hasResponseCustomizer(ServerEndpoint endpoint) { + true + } + @Override Undertow startServer(int port) { Undertow server = Undertow.builder() diff --git a/instrumentation/undertow-1.4/javaagent/src/test/groovy/UndertowServerTest.groovy b/instrumentation/undertow-1.4/javaagent/src/test/groovy/UndertowServerTest.groovy index 42cb713aaa..e1ff9aff60 100644 --- a/instrumentation/undertow-1.4/javaagent/src/test/groovy/UndertowServerTest.groovy +++ b/instrumentation/undertow-1.4/javaagent/src/test/groovy/UndertowServerTest.groovy @@ -113,6 +113,11 @@ class UndertowServerTest extends HttpServerTest implements AgentTestTr attributes } + @Override + boolean hasResponseCustomizer(ServerEndpoint endpoint) { + true + } + def "test send response"() { setup: def uri = address.resolve("sendResponse")