From f9498381d805cb956fb6ced9732637b4d823d51f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?=
 <manuel.alvarezalvarez@datadoghq.com>
Date: Fri, 31 May 2024 10:15:49 +0200
Subject: [PATCH] Add test for HandleVisitor in jetty 11

---
 .../test/groovy/HttpChannelAppSecTest.groovy  | 82 +++++++++++++++++++
 1 file changed, 82 insertions(+)
 create mode 100644 dd-java-agent/instrumentation/jetty-11/src/test/groovy/HttpChannelAppSecTest.groovy

diff --git a/dd-java-agent/instrumentation/jetty-11/src/test/groovy/HttpChannelAppSecTest.groovy b/dd-java-agent/instrumentation/jetty-11/src/test/groovy/HttpChannelAppSecTest.groovy
new file mode 100644
index 00000000000..37bc5c8c5bd
--- /dev/null
+++ b/dd-java-agent/instrumentation/jetty-11/src/test/groovy/HttpChannelAppSecTest.groovy
@@ -0,0 +1,82 @@
+import datadog.trace.agent.test.AgentTestRunner
+import net.bytebuddy.jar.asm.ClassReader
+import net.bytebuddy.jar.asm.ClassVisitor
+import net.bytebuddy.jar.asm.MethodVisitor
+import net.bytebuddy.utility.OpenedClassReader
+import org.eclipse.jetty.server.HttpChannel
+
+import java.lang.instrument.ClassFileTransformer
+import java.lang.instrument.IllegalClassFormatException
+import java.security.ProtectionDomain
+
+class HttpChannelAppSecTest extends AgentTestRunner {
+
+  void 'test blocking capabilities in HttpChannel'() {
+    given:
+    def target = HttpChannel
+    def name = target.name.replaceAll('\\.', '/')
+    def visitor = new BlockingClassVisitor()
+    def transformer = new ClassFileTransformer() {
+        @Override
+        byte[] transform(ClassLoader loader,
+          String className,
+          Class<?> classBeingRedefined,
+          ProtectionDomain protectionDomain,
+          byte[] classfileBuffer) throws IllegalClassFormatException {
+          if (name == className) {
+            final reader = new ClassReader(classfileBuffer)
+            reader.accept(visitor, 0)
+          }
+          return classfileBuffer
+        }
+      }
+
+    when:
+    INSTRUMENTATION.addTransformer(transformer, true)
+    INSTRUMENTATION.retransformClasses(target)
+
+    then:
+    visitor.blockApplied
+
+    cleanup:
+    INSTRUMENTATION.removeTransformer(transformer)
+  }
+
+  private static class BlockingClassVisitor extends ClassVisitor {
+
+    private BlockingMethodVisitor methodVisitor
+
+    protected BlockingClassVisitor() {
+      super(OpenedClassReader.ASM_API)
+    }
+
+    @Override
+    MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+      if ('handle' == name && '()Z' == descriptor) {
+        return methodVisitor = new BlockingMethodVisitor()
+      }
+      return super.visitMethod(access, name, descriptor, signature, exceptions)
+    }
+
+    boolean getBlockApplied() {
+      return methodVisitor?.blockApplied
+    }
+  }
+
+  private static class BlockingMethodVisitor extends MethodVisitor {
+
+    boolean blockApplied = false
+
+    protected BlockingMethodVisitor() {
+      super(OpenedClassReader.ASM_API)
+    }
+
+    @Override
+    void visitMethodInsn(int opcode, String owner, String name, String descriptor, boolean isInterface) {
+      if ('datadog/trace/bootstrap/instrumentation/api/AgentSpan' == owner && 'getRequestBlockingAction' == name) {
+        blockApplied = true
+      }
+      super.visitMethodInsn(opcode, owner, name, descriptor, isInterface)
+    }
+  }
+}