Instrument Jetty websocket pojo (#8562)

* Instrument Jetty websocket pojo

* spotbugs and muzzle

* fix muzzle and other tests

* capital S

* exclude jetty 12.1 because of api changes

* add isLatestTest for all latest tests

* Code review comments

Author: amarziali

dd-java-agent/instrumentation/jetty-11/build.gradle

@@ -39,6 +39,8 @@ dependencies {
   testFixturesCompileOnly "org.eclipse.jetty:jetty-server:11.0.0"
   testFixturesCompileOnly "org.eclipse.jetty:jetty-servlet:11.0.0"
   testFixturesCompileOnly "org.eclipse.jetty.websocket:websocket-jakarta-server:11.0.0"
+  testFixturesImplementation group: 'jakarta.websocket', name: 'jakarta.websocket-client-api', version: '2.0.0'
+
 
 
   testFixturesImplementation(project(':dd-java-agent:testing')) {
@@ -64,6 +66,8 @@ dependencies {
   testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-5')
   testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:javax-websocket-1.0')
   testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:jakarta-websocket-2.0')
+  testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-10")
+  testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-11")
 
   latestDepTestImplementation group: 'org.eclipse.jetty', name: 'jetty-server', version: '11.+',  {
     exclude group: 'org.slf4j', module: 'slf4j-api'

dd-java-agent/instrumentation/jetty-11/src/test/groovy/Jetty11Test.groovy

@@ -8,7 +8,7 @@ import org.eclipse.jetty.server.Server
 abstract class Jetty11Test extends HttpServerTest<Server> {
   @Override
   HttpServer server() {
-    new JettyServer(handler())
+    new JettyServer(handler(), useWebsocketPojoEndpoint())
   }
 
   protected Handler handler() {
@@ -89,10 +89,17 @@ abstract class Jetty11Test extends HttpServerTest<Server> {
   boolean testSessionId() {
     true
   }
+
+  protected boolean useWebsocketPojoEndpoint() {
+    false
+  }
 }
 
 class Jetty11V0ForkedTest extends Jetty11Test implements TestingGenericHttpNamingConventions.ServerV0 {
 }
 
 class Jetty11V1ForkedTest extends Jetty11Test implements TestingGenericHttpNamingConventions.ServerV1 {
+  protected boolean useWebsocketPojoEndpoint() {
+    true
+  }
 }

dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy

@@ -11,7 +11,11 @@ import jakarta.websocket.CloseReason
 import jakarta.websocket.Endpoint
 import jakarta.websocket.EndpointConfig
 import jakarta.websocket.MessageHandler
+import jakarta.websocket.OnClose
+import jakarta.websocket.OnMessage
+import jakarta.websocket.OnOpen
 import jakarta.websocket.Session
+import jakarta.websocket.server.ServerEndpoint
 import jakarta.websocket.server.ServerEndpointConfig
 import org.eclipse.jetty.server.Handler
 import org.eclipse.jetty.server.Server
@@ -29,13 +33,15 @@ class JettyServer implements WebsocketServer {
   def port = 0
   final server = new Server(0) // select random open port
 
-  JettyServer(Handler handler) {
+  JettyServer(Handler handler, usePojoWebsocketHandler = false) {
     server.handler = handler
     server.addBean(errorHandler)
+    def endpointClass = usePojoWebsocketHandler ? PojoEndpoint : WsEndpoint
+
     if (handler instanceof ServletContextHandler) {
       try {
         JakartaWebSocketServletContainerInitializer.configure(handler, (servletContext, container) -> {
-          container.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.class, "/websocket").build())
+          container.addEndpoint(ServerEndpointConfig.Builder.create(endpointClass, "/websocket").build())
         })
       } catch (Throwable ignored) {
       }
@@ -98,9 +104,9 @@ class JettyServer implements WebsocketServer {
   @Override
   void serverSendText(String[] messages) {
     if (messages.length == 1) {
-      WsEndpoint.activeSession.getBasicRemote().sendText(messages[0])
+      Lock.activeSession.getBasicRemote().sendText(messages[0])
     } else {
-      def remoteEndpoint = WsEndpoint.activeSession.getBasicRemote()
+      def remoteEndpoint = Lock.activeSession.getBasicRemote()
       for (int i = 0; i < messages.length; i++) {
         remoteEndpoint.sendText(messages[i], i == messages.length - 1)
       }
@@ -110,20 +116,20 @@ class JettyServer implements WebsocketServer {
   @Override
   void serverSendBinary(byte[][] binaries) {
     if (binaries.length == 1) {
-      WsEndpoint.activeSession.getBasicRemote().sendBinary(ByteBuffer.wrap(binaries[0]))
+      Lock.activeSession.getBasicRemote().sendBinary(ByteBuffer.wrap(binaries[0]))
     } else {
-      try (def stream = WsEndpoint.activeSession.getBasicRemote().getSendStream()) {
+      try (def stream = Lock.activeSession.getBasicRemote().getSendStream()) {
         binaries.each { stream.write(it) }
       }
     }
   }
 
   @Override
   synchronized void awaitConnected() {
-    synchronized (WsEndpoint) {
+    synchronized (Lock) {
       try {
-        while (WsEndpoint.activeSession == null) {
-          WsEndpoint.wait()
+        while (Lock.activeSession == null) {
+          Lock.wait()
         }
       } catch (InterruptedException ie) {
         Thread.currentThread().interrupt()
@@ -133,48 +139,79 @@ class JettyServer implements WebsocketServer {
 
   @Override
   void serverClose() {
-    WsEndpoint.activeSession?.close()
-    WsEndpoint.activeSession = null
+    Lock.activeSession?.close()
+    Lock.activeSession = null
   }
 
   @Override
   void setMaxPayloadSize(int size) {
-    WsEndpoint.activeSession?.setMaxTextMessageBufferSize(size)
-    WsEndpoint.activeSession?.setMaxBinaryMessageBufferSize(size)
+    Lock.activeSession?.setMaxTextMessageBufferSize(size)
+    Lock.activeSession?.setMaxBinaryMessageBufferSize(size)
   }
 
   @Override
   boolean canSplitLargeWebsocketPayloads() {
     false
   }
 
-  static class WsEndpoint extends Endpoint {
+  static class Lock {
     static Session activeSession
+  }
+
+  @ServerEndpoint("/websocket")
+  static class PojoEndpoint {
+
+    @OnOpen
+    void onOpen(Session session) {
+      Lock.activeSession = session
+      synchronized (Lock) {
+        Lock.notifyAll()
+      }
+    }
+
+    @OnMessage
+    void onText(String text, Session session, boolean last) {
+      runUnderTrace("onRead", {})
+
+    }
+
+    @OnMessage
+    void onBytes(byte[] binary) {
+      runUnderTrace("onRead", {})
+    }
+
+    @OnClose
+    void onClose() {
+      Lock.activeSession = null
+    }
+  }
+
+  static class WsEndpoint extends Endpoint {
 
     @Override
     void onOpen(Session session, EndpointConfig endpointConfig) {
       session.addMessageHandler(new MessageHandler.Partial<String>() {
           @Override
-          void onMessage(String s, boolean last) {
-            // jetty does not respect at all limiting the payload so we have to simulate it in few tests
+          void onMessage(String s, boolean b) {
             runUnderTrace("onRead", {})
           }
         })
-      session.addMessageHandler(new MessageHandler.Partial<byte[]>() {
+      session.addMessageHandler(new MessageHandler.Whole<ByteBuffer>() {
+
           @Override
-          void onMessage(byte[] b, boolean last) {
+          void onMessage(ByteBuffer buffer) {
             runUnderTrace("onRead", {})
           }
         })
-      activeSession = session
-      synchronized (WsEndpoint) {
-        WsEndpoint.notifyAll()
+      Lock.activeSession = session
+      synchronized (Lock) {
+        Lock.notifyAll()
       }
     }
 
     @Override
     void onClose(Session session, CloseReason closeReason) {
-      activeSession = null
+      Lock.activeSession = null
     }
   }
 }

dd-java-agent/instrumentation/jetty-12/build.gradle

@@ -38,12 +38,15 @@ addTestSuiteExtendingForDir('ee10LatestDepTest', 'latestDepTest', 'test/ee10')
     targetCompatibility = JavaVersion.VERSION_1_8
   }
 }
-
+[ee8LatestDepTest, ee9LatestDepTest, ee10LatestDepTest].each {
+  it.configure {
+    it.jvmArgs += ['-Dtest.dd.latestDepTest=true']
+  }
+}
 forbiddenApisMain_java17 {
   failOnMissingClasses = false
 }
-
-compileTestGroovy {
+tasks.withType(GroovyCompile).configureEach {
   javaLauncher = getJavaLauncherFor(17)
 }
 
@@ -70,6 +73,9 @@ dependencies {
   testImplementation testFixtures(project(':dd-java-agent:instrumentation:servlet:request-3'))
   testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:javax-websocket-1.0')
   testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:jakarta-websocket-2.0')
+  testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-10")
+  testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-11")
+  testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-12")
   testImplementation testFixtures(project(':dd-java-agent:appsec'))
   testRuntimeOnly project(':dd-java-agent:instrumentation:jetty-9')
   testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-5')

dd-java-agent/instrumentation/jetty-12/src/test/ee10/groovy/Jetty12Test.groovy

@@ -7,7 +7,7 @@ import org.eclipse.jetty.server.Server
 class Jetty12Test extends HttpServerTest<Server> implements TestingGenericHttpNamingConventions.ServerV0 {
   @Override
   HttpServer server() {
-    new JettyServer(JettyServer.servletHandler(TestServlet5))
+    new JettyServer(JettyServer.servletHandler(TestServlet5), useWebsocketPojoEndpoint())
   }
 
   @Override
@@ -49,4 +49,15 @@ class Jetty12Test extends HttpServerTest<Server> implements TestingGenericHttpNa
   boolean hasExtraErrorInformation() {
     true
   }
+
+  protected boolean useWebsocketPojoEndpoint() {
+    false
+  }
+}
+
+class Jetty12PojoWebsocketTest extends Jetty12Test {
+  protected boolean useWebsocketPojoEndpoint() {
+    // advices for pojo won't apply for latest alpha 12.1.+. It has to be adapted once jetty codebase will be stable
+    !isLatestDepTest
+  }
 }