Implement HttpServerResponseCustomizer support for Undertow (#8265)

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<TypeDescription> 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);
+    }
+  }
+}

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<HttpServerExchange> {
+  INSTANCE;
+
+  UndertowHttpResponseMutator() {}
+
+  @Override
+  public void appendHeader(HttpServerExchange exchange, String name, String value) {
+    exchange.getResponseHeaders().add(HttpString.tryFromString(name), value);
+  }
+}

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<TypeInstrumentation> typeInstrumentations() {
-    return asList(new HandlerInstrumentation(), new HttpServerExchangeInstrumentation());
+    return asList(
+        new HandlerInstrumentation(),
+        new HttpServerExchangeInstrumentation(),
+        new HttpTransferEncodingInstrumentation());
   }
 }

instrumentation/undertow-1.4/javaagent/src/test/groovy/UndertowServerDispatchTest.groovy

@@ -34,6 +34,11 @@ class UndertowServerDispatchTest extends HttpServerTest<Undertow> implements Age
     return false
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    true
+  }
+
   @Override
   Undertow startServer(int port) {
     Undertow server = Undertow.builder()

instrumentation/undertow-1.4/javaagent/src/test/groovy/UndertowServerTest.groovy

@@ -113,6 +113,11 @@ class UndertowServerTest extends HttpServerTest<Undertow> implements AgentTestTr
     attributes
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    true
+  }
+
   def "test send response"() {
     setup:
     def uri = address.resolve("sendResponse")