Skip to content

Commit 184045b

Browse files
amarzialiamarziali
committed
Instrument Jetty websocket pojo
1 parent a464b9b commit 184045b

File tree

49 files changed

+3435
-139
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

49 files changed

+3435
-139
lines changed

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AdviceShader.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,10 @@ public byte[] shadeClass(byte[] bytecode) {
6565
public String uniqueHelper(String dottedName) {
6666
int packageEnd = dottedName.lastIndexOf('.');
6767
if (packageEnd > 0) {
68-
return dottedName.substring(0, packageEnd + 1) + "shaded" + dottedName.substring(packageEnd);
68+
return dottedName.substring(0, packageEnd + 1)
69+
+ "shaded"
70+
+ relocations.hashCode()
71+
+ dottedName.substring(packageEnd);
6972
}
7073
return dottedName;
7174
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,8 @@ dependencies {
6464
testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-5')
6565
testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:javax-websocket-1.0')
6666
testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:jakarta-websocket-2.0')
67+
testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-10")
68+
testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-11")
6769

6870
latestDepTestImplementation group: 'org.eclipse.jetty', name: 'jetty-server', version: '11.+', {
6971
exclude group: 'org.slf4j', module: 'slf4j-api'

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import org.eclipse.jetty.server.Server
88
abstract class Jetty11Test extends HttpServerTest<Server> {
99
@Override
1010
HttpServer server() {
11-
new JettyServer(handler())
11+
new JettyServer(handler(), useWebsocketPojoEndpoint())
1212
}
1313

1414
protected Handler handler() {
@@ -89,10 +89,17 @@ abstract class Jetty11Test extends HttpServerTest<Server> {
8989
boolean testSessionId() {
9090
true
9191
}
92+
93+
protected boolean useWebsocketPojoEndpoint() {
94+
false
95+
}
9296
}
9397

9498
class Jetty11V0ForkedTest extends Jetty11Test implements TestingGenericHttpNamingConventions.ServerV0 {
9599
}
96100

97101
class Jetty11V1ForkedTest extends Jetty11Test implements TestingGenericHttpNamingConventions.ServerV1 {
102+
protected boolean useWebsocketPojoEndpoint() {
103+
true
104+
}
98105
}

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

Lines changed: 59 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,11 @@ import jakarta.websocket.CloseReason
1111
import jakarta.websocket.Endpoint
1212
import jakarta.websocket.EndpointConfig
1313
import jakarta.websocket.MessageHandler
14+
import jakarta.websocket.OnClose
15+
import jakarta.websocket.OnMessage
16+
import jakarta.websocket.OnOpen
1417
import jakarta.websocket.Session
18+
import jakarta.websocket.server.ServerEndpoint
1519
import jakarta.websocket.server.ServerEndpointConfig
1620
import org.eclipse.jetty.server.Handler
1721
import org.eclipse.jetty.server.Server
@@ -29,13 +33,15 @@ class JettyServer implements WebsocketServer {
2933
def port = 0
3034
final server = new Server(0) // select random open port
3135

32-
JettyServer(Handler handler) {
36+
JettyServer(Handler handler, usePojoWebsocketHandler = false) {
3337
server.handler = handler
3438
server.addBean(errorHandler)
39+
def endpointClass = usePojoWebsocketHandler ? PojoEndpoint : WsEndpoint
40+
3541
if (handler instanceof ServletContextHandler) {
3642
try {
3743
JakartaWebSocketServletContainerInitializer.configure(handler, (servletContext, container) -> {
38-
container.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.class, "/websocket").build())
44+
container.addEndpoint(ServerEndpointConfig.Builder.create(endpointClass, "/websocket").build())
3945
})
4046
} catch (Throwable ignored) {
4147
}
@@ -98,9 +104,9 @@ class JettyServer implements WebsocketServer {
98104
@Override
99105
void serverSendText(String[] messages) {
100106
if (messages.length == 1) {
101-
WsEndpoint.activeSession.getBasicRemote().sendText(messages[0])
107+
Lock.activeSession.getBasicRemote().sendText(messages[0])
102108
} else {
103-
def remoteEndpoint = WsEndpoint.activeSession.getBasicRemote()
109+
def remoteEndpoint = Lock.activeSession.getBasicRemote()
104110
for (int i = 0; i < messages.length; i++) {
105111
remoteEndpoint.sendText(messages[i], i == messages.length - 1)
106112
}
@@ -110,20 +116,20 @@ class JettyServer implements WebsocketServer {
110116
@Override
111117
void serverSendBinary(byte[][] binaries) {
112118
if (binaries.length == 1) {
113-
WsEndpoint.activeSession.getBasicRemote().sendBinary(ByteBuffer.wrap(binaries[0]))
119+
Lock.activeSession.getBasicRemote().sendBinary(ByteBuffer.wrap(binaries[0]))
114120
} else {
115-
try (def stream = WsEndpoint.activeSession.getBasicRemote().getSendStream()) {
121+
try (def stream = Lock.activeSession.getBasicRemote().getSendStream()) {
116122
binaries.each { stream.write(it) }
117123
}
118124
}
119125
}
120126

121127
@Override
122128
synchronized void awaitConnected() {
123-
synchronized (WsEndpoint) {
129+
synchronized (Lock) {
124130
try {
125-
while (WsEndpoint.activeSession == null) {
126-
WsEndpoint.wait()
131+
while (Lock.activeSession == null) {
132+
Lock.wait()
127133
}
128134
} catch (InterruptedException ie) {
129135
Thread.currentThread().interrupt()
@@ -133,48 +139,79 @@ class JettyServer implements WebsocketServer {
133139

134140
@Override
135141
void serverClose() {
136-
WsEndpoint.activeSession?.close()
137-
WsEndpoint.activeSession = null
142+
Lock.activeSession?.close()
143+
Lock.activeSession = null
138144
}
139145

140146
@Override
141147
void setMaxPayloadSize(int size) {
142-
WsEndpoint.activeSession?.setMaxTextMessageBufferSize(size)
143-
WsEndpoint.activeSession?.setMaxBinaryMessageBufferSize(size)
148+
Lock.activeSession?.setMaxTextMessageBufferSize(size)
149+
Lock.activeSession?.setMaxBinaryMessageBufferSize(size)
144150
}
145151

146152
@Override
147153
boolean canSplitLargeWebsocketPayloads() {
148154
false
149155
}
150156

151-
static class WsEndpoint extends Endpoint {
157+
static class Lock {
152158
static Session activeSession
159+
}
160+
161+
@ServerEndpoint("/websocket")
162+
static class PojoEndpoint {
163+
164+
@OnOpen
165+
void onOpen(Session session) {
166+
Lock.activeSession = session
167+
synchronized (Lock) {
168+
Lock.notifyAll()
169+
}
170+
}
171+
172+
@OnMessage
173+
void onText(String text, Session session, boolean last) {
174+
runUnderTrace("onRead", {})
175+
176+
}
177+
178+
@OnMessage
179+
void onBytes(byte[] binary) {
180+
runUnderTrace("onRead", {})
181+
}
182+
183+
@OnClose
184+
void onClose() {
185+
Lock.activeSession = null
186+
}
187+
}
188+
189+
static class WsEndpoint extends Endpoint {
153190

154191
@Override
155192
void onOpen(Session session, EndpointConfig endpointConfig) {
156193
session.addMessageHandler(new MessageHandler.Partial<String>() {
157194
@Override
158-
void onMessage(String s, boolean last) {
159-
// jetty does not respect at all limiting the payload so we have to simulate it in few tests
195+
void onMessage(String s, boolean b) {
160196
runUnderTrace("onRead", {})
161197
}
162198
})
163-
session.addMessageHandler(new MessageHandler.Partial<byte[]>() {
199+
session.addMessageHandler(new MessageHandler.Whole<ByteBuffer>() {
200+
164201
@Override
165-
void onMessage(byte[] b, boolean last) {
202+
void onMessage(ByteBuffer buffer) {
166203
runUnderTrace("onRead", {})
167204
}
168205
})
169-
activeSession = session
170-
synchronized (WsEndpoint) {
171-
WsEndpoint.notifyAll()
206+
Lock.activeSession = session
207+
synchronized (Lock) {
208+
Lock.notifyAll()
172209
}
173210
}
174211

175212
@Override
176213
void onClose(Session session, CloseReason closeReason) {
177-
activeSession = null
214+
Lock.activeSession = null
178215
}
179216
}
180217
}

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,8 +42,7 @@ addTestSuiteExtendingForDir('ee10LatestDepTest', 'latestDepTest', 'test/ee10')
4242
forbiddenApisMain_java17 {
4343
failOnMissingClasses = false
4444
}
45-
46-
compileTestGroovy {
45+
tasks.withType(GroovyCompile).configureEach {
4746
javaLauncher = getJavaLauncherFor(17)
4847
}
4948

@@ -70,6 +69,9 @@ dependencies {
7069
testImplementation testFixtures(project(':dd-java-agent:instrumentation:servlet:request-3'))
7170
testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:javax-websocket-1.0')
7271
testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:jakarta-websocket-2.0')
72+
testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-10")
73+
testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-11")
74+
testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-12")
7375
testImplementation testFixtures(project(':dd-java-agent:appsec'))
7476
testRuntimeOnly project(':dd-java-agent:instrumentation:jetty-9')
7577
testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-5')

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import org.eclipse.jetty.server.Server
77
class Jetty12Test extends HttpServerTest<Server> implements TestingGenericHttpNamingConventions.ServerV0 {
88
@Override
99
HttpServer server() {
10-
new JettyServer(JettyServer.servletHandler(TestServlet5))
10+
new JettyServer(JettyServer.servletHandler(TestServlet5), useWebsocketPojoEndpoint())
1111
}
1212

1313
@Override
@@ -49,4 +49,13 @@ class Jetty12Test extends HttpServerTest<Server> implements TestingGenericHttpNa
4949
boolean hasExtraErrorInformation() {
5050
true
5151
}
52+
protected boolean useWebsocketPojoEndpoint() {
53+
false
54+
}
55+
}
56+
57+
class Jetty12PojoWebsocketTest extends Jetty12Test {
58+
protected boolean useWebsocketPojoEndpoint() {
59+
true
60+
}
5261
}

0 commit comments

Comments
 (0)