From 2139fd8b2deea8b001b52b6cf4b8c7eb117690c2 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Thu, 6 Feb 2025 11:04:51 +0100 Subject: [PATCH 01/15] Instrument Java Websocket API (JSR356) --- .../decorator/WebsocketDecorator.java | 168 +++++ .../websocket/HandlerContext.java | 127 ++++ .../websocket/HandlersExtractor.java | 91 +++ .../websocket/ResourceNameExtractor.java | 36 ++ .../bytebuddy/matcher/ignored_class_name.trie | 2 + .../groovy/ArmeriaJetty9ServerTest.groovy | 4 +- .../instrumentation/jetty-11/build.gradle | 10 + .../testFixtures/groovy/JettyServer.groovy | 108 +++- .../instrumentation/jetty-12/build.gradle | 21 + .../src/test/ee10/groovy/JettyServer.groovy | 108 +++- .../src/test/ee8/groovy/JettyServer.groovy | 115 +++- .../src/test/ee9/groovy/JettyServer.groovy | 113 +++- .../instrumentation/jetty-9/build.gradle | 9 + .../jetty10/JettyServerInstrumentation.java | 1 + .../WebSocketSessionInstrumentation.java | 88 +++ .../jetty10/OnCompletedAdvice.java | 29 + .../jetty9/Jetty9InactiveAppSecTest.groovy | 2 + .../instrumentation/jetty9/Jetty9Test.groovy | 12 +- .../JettyContinuationHandlerTest.groovy | 7 +- .../instrumentation/jetty9/JettyServer.groovy | 37 -- .../groovy/test/JettyServer.groovy | 144 +++++ .../jetty9 => test}/TestHandler.groovy | 74 ++- .../spring-webmvc-3.1/build.gradle | 5 +- .../boot/CustomBeanClassloaderTest.groovy | 2 +- .../test/boot/SpringBootBasedTest.groovy | 15 +- .../groovy/test/boot/SpringBootServer.groovy | 67 +- .../groovy/test/boot/WebsocketConfig.groovy | 21 + .../groovy/test/boot/WebsocketEndpoint.groovy | 42 ++ .../spring-webmvc-6.0/build.gradle | 4 + .../boot/SpringBootBasedTest.groovy | 111 +++- .../springweb6/boot/WebsocketConfig.groovy | 21 + .../springweb6/boot/WebsocketEndpoint.groovy | 42 ++ .../instrumentation/tomcat-5.5/build.gradle | 47 +- ...AppSecInactiveHttpServerForkedTest.groovy} | 4 +- .../TomcatNoPropagationForkedTest.groovy | 3 +- .../latestDepTest/groovy/TomcatServer.groovy | 93 ++- .../groovy/TomcatServletTest.groovy | 39 +- .../latestDepTest/groovy/WsEndpoint.groovy | 76 +++ .../WsHandshakeRequestInstrumentation.java | 52 ++ .../WsHttpUpgradeHandlerInstrumentation.java | 56 ++ .../src/test/groovy/TomcatServletTest.groovy | 2 +- .../tomcat9Test/groovy/TomcatServer.groovy | 161 +++++ .../groovy/TomcatWebsocketTest.groovy | 374 ++++++++++++ .../src/tomcat9Test/groovy/WsEndpoint.groovy | 79 +++ .../undertow/undertow-2.0/build.gradle | 3 + .../src/test/groovy/TestEndpoint.groovy | 33 + .../test/groovy/UndertowServletTest.groovy | 64 +- .../undertow/undertow-2.2/build.gradle | 7 + .../src/test/groovy/TestEndpoint.groovy | 35 ++ .../test/groovy/UndertowServletTest.groovy | 60 +- .../instrumentation/websocket/build.gradle | 4 + .../jakarta-websocket-2.0/build.gradle | 25 + .../jsr256/JakartaWebsocketModule.java | 25 + .../src/test/groovy/EndpointWrapper.groovy | 39 ++ .../src/test/groovy/Endpoints.groovy | 130 ++++ .../src/test/groovy/WebsocketTest.groovy | 577 ++++++++++++++++++ .../javax-websocket-1.0/build.gradle | 19 + .../AsyncRemoteEndpointInstrumentation.java | 245 ++++++++ .../BasicRemoteEndpointInstrumentation.java | 238 ++++++++ .../jsr256/EndpointInstrumentation.java | 109 ++++ .../jsr256/JavaxWebsocketModule.java | 64 ++ .../jsr256/MessageHandlerInstrumentation.java | 98 +++ .../jsr256/SessionInstrumentation.java | 203 ++++++ .../websocket/jsr256/TracingOutputStream.java | 68 +++ .../websocket/jsr256/TracingSendHandler.java | 33 + .../websocket/jsr256/TracingWriter.java | 58 ++ .../src/test/groovy/EndpointWrapper.groovy | 39 ++ .../src/test/groovy/Endpoints.groovy | 130 ++++ .../src/test/groovy/WebsocketTest.groovy | 577 ++++++++++++++++++ .../agent/test/asserts/LinksAssert.groovy | 50 ++ .../agent/test/asserts/SpanAssert.groovy | 21 +- .../agent/test/asserts/TagsAssert.groovy | 1 + .../agent/test/base/HttpServerTest.groovy | 238 +++++++- .../agent/test/base/WebsocketServer.groovy | 39 ++ .../datadog/trace/api/ConfigDefaults.java | 4 + .../java/datadog/trace/api/DDSpanTypes.java | 1 + .../main/java/datadog/trace/api/DDTags.java | 3 + .../config/TraceInstrumentationConfig.java | 7 + .../main/java/datadog/trace/core/DDSpan.java | 2 +- .../main/java/datadog/trace/api/Config.java | 25 +- .../datadog/trace/api/InstrumenterConfig.java | 12 + .../api/InstrumentationTags.java | 8 + .../api/InternalSpanTypes.java | 1 + .../instrumentation/api/SpanAttributes.java | 13 + .../instrumentation/api/SpanLink.java | 23 +- settings.gradle | 3 + 86 files changed, 5760 insertions(+), 196 deletions(-) create mode 100644 dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java create mode 100644 dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/HandlerContext.java create mode 100644 dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/HandlersExtractor.java create mode 100644 dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/ResourceNameExtractor.java create mode 100644 dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/WebSocketSessionInstrumentation.java create mode 100644 dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/OnCompletedAdvice.java delete mode 100644 dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/datadog/trace/instrumentation/jetty9/JettyServer.groovy create mode 100644 dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/test/JettyServer.groovy rename dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/{datadog/trace/instrumentation/jetty9 => test}/TestHandler.groovy (50%) create mode 100644 dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/WebsocketConfig.groovy create mode 100644 dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/WebsocketEndpoint.groovy create mode 100644 dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/WebsocketConfig.groovy create mode 100644 dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/WebsocketEndpoint.groovy rename dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/{TomcatAppSecInactiveHttpServerTest.groovy => TomcatAppSecInactiveHttpServerForkedTest.groovy} (90%) create mode 100644 dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/WsEndpoint.groovy create mode 100644 dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java create mode 100644 dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHttpUpgradeHandlerInstrumentation.java create mode 100644 dd-java-agent/instrumentation/tomcat-5.5/src/tomcat9Test/groovy/TomcatServer.groovy create mode 100644 dd-java-agent/instrumentation/tomcat-5.5/src/tomcat9Test/groovy/TomcatWebsocketTest.groovy create mode 100644 dd-java-agent/instrumentation/tomcat-5.5/src/tomcat9Test/groovy/WsEndpoint.groovy create mode 100644 dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/TestEndpoint.groovy create mode 100644 dd-java-agent/instrumentation/undertow/undertow-2.2/src/test/groovy/TestEndpoint.groovy create mode 100644 dd-java-agent/instrumentation/websocket/build.gradle create mode 100644 dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/build.gradle create mode 100644 dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/JakartaWebsocketModule.java create mode 100644 dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/EndpointWrapper.groovy create mode 100644 dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/Endpoints.groovy create mode 100644 dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/WebsocketTest.groovy create mode 100644 dd-java-agent/instrumentation/websocket/javax-websocket-1.0/build.gradle create mode 100644 dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/AsyncRemoteEndpointInstrumentation.java create mode 100644 dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/BasicRemoteEndpointInstrumentation.java create mode 100644 dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/EndpointInstrumentation.java create mode 100644 dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/JavaxWebsocketModule.java create mode 100644 dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/MessageHandlerInstrumentation.java create mode 100644 dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/SessionInstrumentation.java create mode 100644 dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/TracingOutputStream.java create mode 100644 dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/TracingSendHandler.java create mode 100644 dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/TracingWriter.java create mode 100644 dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/EndpointWrapper.groovy create mode 100644 dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/Endpoints.groovy create mode 100644 dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/WebsocketTest.groovy create mode 100644 dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy create mode 100644 dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/WebsocketServer.groovy diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java new file mode 100644 index 00000000000..63a2100bc62 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java @@ -0,0 +1,168 @@ +package datadog.trace.bootstrap.instrumentation.decorator; + +import static datadog.trace.api.DDTags.DECISION_MAKER_INHERITED; +import static datadog.trace.api.DDTags.DECISION_MAKER_RESOURCE; +import static datadog.trace.api.DDTags.DECISION_MAKER_SERVICE; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_CLOSE_CODE; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_CLOSE_REASON; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_FRAMES; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_LENGTH; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_RECEIVE_TIME; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_TYPE; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_SESSION_ID; +import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND; +import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_CONSUMER; +import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_PRODUCER; + +import datadog.trace.api.Config; +import datadog.trace.api.time.SystemTimeSource; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.SpanAttributes; +import datadog.trace.bootstrap.instrumentation.api.SpanLink; +import datadog.trace.bootstrap.instrumentation.api.Tags; +import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; +import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; +import javax.annotation.Nonnull; + +public class WebsocketDecorator extends BaseDecorator { + private static final CharSequence WEBSOCKET = UTF8BytesString.create("websocket"); + private static final String[] INSTRUMENTATION_NAMES = {WEBSOCKET.toString()}; + private static final CharSequence WEBSOCKET_RECEIVE = UTF8BytesString.create("websocket.receive"); + private static final CharSequence WEBSOCKET_SEND = UTF8BytesString.create("websocket.send"); + private static final CharSequence WEBSOCKET_CLOSE = UTF8BytesString.create("websocket.close"); + + private static final SpanAttributes SPAN_ATTRIBUTES_RECEIVE = + SpanAttributes.builder().put("dd.kind", "executed_from").build(); + private static final SpanAttributes SPAN_ATTRIBUTES_SEND = + SpanAttributes.builder().put("dd.kind", "resuming").build(); + + public static final WebsocketDecorator DECORATE = new WebsocketDecorator(); + + @Override + protected String[] instrumentationNames() { + return INSTRUMENTATION_NAMES; + } + + @Override + protected CharSequence spanType() { + return InternalSpanTypes.WEBSOCKET; + } + + @Override + protected CharSequence component() { + return WEBSOCKET; + } + + @Override + public AgentSpan afterStart(AgentSpan span) { + return super.afterStart(span).setMeasured(true); + } + + @Nonnull + public AgentSpan onReceiveFrameStart( + final HandlerContext.Receiver handlerContext, final Object data, boolean partialDelivery) { + handlerContext.recordChunkData(data, partialDelivery); + return onFrameStart( + WEBSOCKET_RECEIVE, SPAN_KIND_CONSUMER, handlerContext, SPAN_ATTRIBUTES_RECEIVE, true); + } + + @Nonnull + public AgentSpan onSessionCloseIssued( + final HandlerContext.Sender handlerContext, CharSequence closeReason, int closeCode) { + return onFrameStart( + WEBSOCKET_CLOSE, SPAN_KIND_PRODUCER, handlerContext, SPAN_ATTRIBUTES_SEND, false) + .setTag(WEBSOCKET_CLOSE_CODE, closeCode) + .setTag(WEBSOCKET_CLOSE_REASON, closeReason); + } + + @Nonnull + public AgentSpan onSessionCloseReceived( + final HandlerContext.Receiver handlerContext, CharSequence closeReason, int closeCode) { + return onFrameStart( + WEBSOCKET_CLOSE, SPAN_KIND_CONSUMER, handlerContext, SPAN_ATTRIBUTES_RECEIVE, true) + .setTag(WEBSOCKET_CLOSE_CODE, closeCode) + .setTag(WEBSOCKET_CLOSE_REASON, closeReason); + } + + @Nonnull + public AgentSpan onSendFrameStart( + final HandlerContext.Sender handlerContext, final CharSequence msgType, final int msgSize) { + handlerContext.recordChunkData(msgType, msgSize); + return onFrameStart( + WEBSOCKET_SEND, SPAN_KIND_PRODUCER, handlerContext, SPAN_ATTRIBUTES_SEND, false); + } + + public void onFrameEnd(final HandlerContext handlerContext) { + if (handlerContext == null) { + return; + } + final AgentSpan wsSpan = handlerContext.getWebsocketSpan(); + try { + final long startTime = handlerContext.getFirstFrameTimestamp(); + if (startTime > 0) { + wsSpan.setTag( + WEBSOCKET_MESSAGE_RECEIVE_TIME, + SystemTimeSource.INSTANCE.getCurrentTimeNanos() - startTime); + } + final long chunks = handlerContext.getMsgChunks(); + if (chunks > 0) { + wsSpan.setTag(WEBSOCKET_MESSAGE_FRAMES, chunks); + wsSpan.setTag(WEBSOCKET_MESSAGE_LENGTH, handlerContext.getMsgSize()); + wsSpan.setTag(WEBSOCKET_MESSAGE_TYPE, handlerContext.getMessageType()); + } + (beforeFinish(wsSpan)).finish(); + } finally { + handlerContext.reset(); + } + } + + private AgentSpan onFrameStart( + final CharSequence operationName, + final CharSequence spanKind, + final HandlerContext handlerContext, + final SpanAttributes linkAttributes, + boolean traceStarter) { + AgentSpan wsSpan = handlerContext.getWebsocketSpan(); + if (wsSpan == null) { + final Config config = Config.get(); + final AgentSpan handshakeSpan = handlerContext.getHandshakeSpan(); + boolean inheritSampling = config.isWebsocketMessagesInheritSampling(); + boolean useDedicatedTraces = config.isWebsocketMessagesSeparateTraces(); + if (traceStarter) { + if (useDedicatedTraces) { + wsSpan = startSpan(WEBSOCKET.toString(), operationName, null); + if (inheritSampling) { + wsSpan.setTag(DECISION_MAKER_INHERITED, 1); + wsSpan.setTag(DECISION_MAKER_SERVICE, handshakeSpan.getServiceName()); + wsSpan.setTag(DECISION_MAKER_RESOURCE, handshakeSpan.getResourceName()); + } + } else { + wsSpan = startSpan(WEBSOCKET.toString(), operationName, handshakeSpan.context()); + } + } else { + wsSpan = startSpan(WEBSOCKET.toString(), operationName); + } + handlerContext.setWebsocketSpan(wsSpan); + afterStart(wsSpan); + wsSpan.setTag(SPAN_KIND, spanKind); + wsSpan.setResourceName(handlerContext.getWsResourceName()); + // carry over peer information for inferred services + final String handshakePeerAddress = (String) handshakeSpan.getTag(Tags.PEER_HOSTNAME); + if (handshakePeerAddress != null) { + wsSpan.setTag(Tags.PEER_HOSTNAME, handshakePeerAddress); + } + if (config.isWebsocketTagSessionId()) { + wsSpan.setTag(WEBSOCKET_SESSION_ID, handlerContext.getSessionId()); + } + if (useDedicatedTraces || !traceStarter) { + // the link is not added if the user wants to have receive frames on the same trace as the + // handshake + wsSpan.addLink( + SpanLink.from(handshakeSpan.context(), (byte) 0, "", linkAttributes, inheritSampling)); + } + } + return wsSpan; + } +} diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/HandlerContext.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/HandlerContext.java new file mode 100644 index 00000000000..d0f1894be23 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/HandlerContext.java @@ -0,0 +1,127 @@ +package datadog.trace.bootstrap.instrumentation.websocket; + +import datadog.trace.api.time.SystemTimeSource; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; + +public abstract class HandlerContext { + + private final AgentSpan handshakeSpan; + private AgentSpan websocketSpan; + private final String sessionId; + protected long msgChunks = 0; + protected long msgSize = 0; + private final CharSequence wsResourceName; + protected long firstFrameTimestamp; + + public HandlerContext(AgentSpan handshakeSpan, String sessionId) { + this.handshakeSpan = handshakeSpan; + this.sessionId = sessionId; + wsResourceName = ResourceNameExtractor.extractResourceName(handshakeSpan.getResourceName()); + } + + public AgentSpan getHandshakeSpan() { + return handshakeSpan; + } + + public AgentSpan getWebsocketSpan() { + return websocketSpan; + } + + public void setWebsocketSpan(AgentSpan websocketSpan) { + this.websocketSpan = websocketSpan; + } + + public String getSessionId() { + return sessionId; + } + + public long getMsgChunks() { + return msgChunks; + } + + public long getMsgSize() { + return msgSize; + } + + public CharSequence getWsResourceName() { + return wsResourceName; + } + + public long getFirstFrameTimestamp() { + return firstFrameTimestamp; + } + + public abstract CharSequence getMessageType(); + + public void reset() { + msgChunks = 0; + websocketSpan = null; + msgSize = 0; + firstFrameTimestamp = 0; + } + + public static class Receiver extends HandlerContext { + private boolean msgSizeExtractorInitialized = false; + private HandlersExtractor.SizeCalculator msgSizeCalculator; + + public Receiver(AgentSpan handshakeSpan, String sessionId) { + super(handshakeSpan, sessionId); + } + + @Override + public CharSequence getMessageType() { + return msgSizeCalculator != null ? msgSizeCalculator.getFormat() : null; + } + + public void recordChunkData(Object data, boolean partialDelivery) { + if (msgChunks++ == 0 && partialDelivery) { + firstFrameTimestamp = SystemTimeSource.INSTANCE.getCurrentTimeNanos(); + } + if (data == null) { + return; + } + if (!msgSizeExtractorInitialized) { + msgSizeExtractorInitialized = true; + msgSizeCalculator = HandlersExtractor.getSizeCalculator(data); + } + + if (msgSizeCalculator != null) { + try { + int sz = msgSizeCalculator.getLengthFunction().applyAsInt(data); + msgSize += sz; + if (partialDelivery && sz == 0) { + msgChunks--; // if we receive an empty frame with the fin bit don't count it as a chunk + } + } catch (Throwable ignored) { + } + } + } + } + + public static class Sender extends HandlerContext { + private CharSequence msgType; + + public Sender(AgentSpan handshakeSpan, String sessionId) { + super(handshakeSpan, sessionId); + } + + @Override + public CharSequence getMessageType() { + return msgType; + } + + @Override + public void reset() { + super.reset(); + msgType = null; + } + + public void recordChunkData(CharSequence type, int size) { + msgChunks++; + if (msgType == null) { + msgType = type; + } + msgSize += size; + } + } +} diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/HandlersExtractor.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/HandlersExtractor.java new file mode 100644 index 00000000000..9d38b66f2ac --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/HandlersExtractor.java @@ -0,0 +1,91 @@ +package datadog.trace.bootstrap.instrumentation.websocket; + +import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; +import java.io.InputStream; +import java.io.Reader; +import java.nio.ByteBuffer; +import java.util.function.ToIntFunction; +import javax.annotation.Nonnull; + +public class HandlersExtractor { + public static final CharSequence MESSAGE_TYPE_TEXT = UTF8BytesString.create("text"); + public static final CharSequence MESSAGE_TYPE_BINARY = UTF8BytesString.create("binary"); + + static class CharSequenceLenFunction implements ToIntFunction { + @Override + public int applyAsInt(CharSequence value) { + return value != null ? value.length() : 0; + } + } + + static class ByteArrayLenFunction implements ToIntFunction { + @Override + public int applyAsInt(byte[] value) { + return value != null ? value.length : 0; + } + } + + static class ByteBufferLenFunction implements ToIntFunction { + @Override + public int applyAsInt(ByteBuffer value) { + return value != null ? value.remaining() : 0; + } + } + + static class NoopLenFunction implements ToIntFunction { + @Override + public int applyAsInt(Object ignored) { + return 0; + } + } + + public static class SizeCalculator { + @Nonnull private final ToIntFunction lengthFunction; + @Nonnull private final CharSequence format; + + SizeCalculator(@Nonnull ToIntFunction lengthFunction, @Nonnull CharSequence format) { + this.lengthFunction = lengthFunction; + this.format = format; + } + + @Nonnull + public ToIntFunction getLengthFunction() { + return lengthFunction; + } + + @Nonnull + public CharSequence getFormat() { + return format; + } + } + + public static final SizeCalculator CHAR_SEQUENCE_SIZE_CALCULATOR = + new SizeCalculator<>(new CharSequenceLenFunction(), MESSAGE_TYPE_TEXT); + public static final SizeCalculator BYTES_SIZE_CALCULATOR = + new SizeCalculator<>(new ByteArrayLenFunction(), MESSAGE_TYPE_BINARY); + public static final SizeCalculator BYTE_BUFFER_SIZE_CALCULATOR = + new SizeCalculator<>(new ByteBufferLenFunction(), MESSAGE_TYPE_BINARY); + public static final SizeCalculator TEXT_STREAM_SIZE_CALCULATOR = + new SizeCalculator<>(new NoopLenFunction(), MESSAGE_TYPE_TEXT); + public static final SizeCalculator BYTE_STREAM_SIZE_CALCULATOR = + new SizeCalculator<>(new NoopLenFunction(), MESSAGE_TYPE_BINARY); + + public static SizeCalculator getSizeCalculator(Object data) { + // we only extract "safely" the message size from byte[], ByteBuffer and String. + // Other types will contain streaming data (i.e. InputStream, Reader) + if (data instanceof CharSequence) { + return CHAR_SEQUENCE_SIZE_CALCULATOR; + } else if (data instanceof byte[]) { + return BYTES_SIZE_CALCULATOR; + } else if (data instanceof ByteBuffer) { + return BYTE_BUFFER_SIZE_CALCULATOR; + } else if (data instanceof Reader) { + return TEXT_STREAM_SIZE_CALCULATOR; + } else if (data instanceof InputStream) { + return BYTE_STREAM_SIZE_CALCULATOR; + } + return null; + } + + private HandlersExtractor() {} +} diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/ResourceNameExtractor.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/ResourceNameExtractor.java new file mode 100644 index 00000000000..b9fe9c179c0 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/ResourceNameExtractor.java @@ -0,0 +1,36 @@ +package datadog.trace.bootstrap.instrumentation.websocket; + +import datadog.trace.api.Functions; +import datadog.trace.api.cache.DDCache; +import datadog.trace.api.cache.DDCaches; +import datadog.trace.api.normalize.HttpResourceNames; +import java.util.function.Function; + +public class ResourceNameExtractor { + private static final String SPACE = " "; + private static final String WEBSOCKET_SPACE = "websocket" + SPACE; + private static final DDCache CACHE = DDCaches.newFixedSizeCache(128); + private static final Function EXTRACTOR = + s -> { + if (s == null || s.length() == 0) { + return HttpResourceNames.DEFAULT_RESOURCE_NAME; + } + int idx = s.toString().indexOf(SPACE); + if (idx < 0 || idx == s.length() - 1) { + return s; + } + final CharSequence ret = s.subSequence(idx + 1, s.length()); + if (ret.length() == 0) { + return HttpResourceNames.DEFAULT_RESOURCE_NAME; + } + return ret; + }; + private static final Function ADDER = + EXTRACTOR.andThen(new Functions.Prefix(WEBSOCKET_SPACE)); + + private ResourceNameExtractor() {} + + public static CharSequence extractResourceName(CharSequence handshakeResourceName) { + return CACHE.computeIfAbsent(handshakeResourceName, ADDER); + } +} diff --git a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie index bce733ae3d7..8fad583b62f 100644 --- a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie +++ b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie @@ -81,6 +81,7 @@ 1 org.groovy.* 1 org.jinspired.* 1 org.springframework.context.support.ContextTypeMatchClassLoader +1 org.springframework.context.support.DefaultLifecycleProcessor* 1 org.springframework.core.DecoratingClassLoader 1 org.springframework.core.OverridingClassLoader 1 org.springframework.core.$Proxy* @@ -349,6 +350,7 @@ 0 org.springframework.web.context.support.XmlWebApplicationContext 0 org.springframework.web.reactive.* 0 org.springframework.web.servlet.* +0 org.springframework.web.socket.* # Need for IAST so propagation of tainted objects is complete in spring 2.7.5 0 org.springframework.util.StreamUtils$NonClosingInputStream # Need for IAST gson propagation diff --git a/dd-java-agent/instrumentation/armeria-jetty/src/test/jetty9/groovy/ArmeriaJetty9ServerTest.groovy b/dd-java-agent/instrumentation/armeria-jetty/src/test/jetty9/groovy/ArmeriaJetty9ServerTest.groovy index e12f4e2aa4a..81074ab7086 100644 --- a/dd-java-agent/instrumentation/armeria-jetty/src/test/jetty9/groovy/ArmeriaJetty9ServerTest.groovy +++ b/dd-java-agent/instrumentation/armeria-jetty/src/test/jetty9/groovy/ArmeriaJetty9ServerTest.groovy @@ -2,8 +2,8 @@ import com.linecorp.armeria.server.Server import com.linecorp.armeria.server.jetty.JettyService import datadog.trace.agent.test.base.HttpServer import datadog.trace.agent.test.base.HttpServerTest -import datadog.trace.instrumentation.jetty9.JettyServer -import datadog.trace.instrumentation.jetty9.TestHandler +import test.JettyServer +import test.TestHandler import java.util.concurrent.TimeoutException diff --git a/dd-java-agent/instrumentation/jetty-11/build.gradle b/dd-java-agent/instrumentation/jetty-11/build.gradle index c4714e10de3..e4497de5156 100644 --- a/dd-java-agent/instrumentation/jetty-11/build.gradle +++ b/dd-java-agent/instrumentation/jetty-11/build.gradle @@ -38,6 +38,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(project(':dd-java-agent:testing')) { exclude group: 'org.eclipse.jetty', module: 'jetty-server' @@ -52,11 +54,16 @@ dependencies { testImplementation "org.eclipse.jetty:jetty-servlet:11.0.0", { exclude group: 'org.slf4j', module: 'slf4j-api' } + testImplementation ("org.eclipse.jetty.websocket:websocket-jakarta-server:11.0.0") { + exclude group: 'org.slf4j', module: 'slf4j-api' + } testImplementation(project(':dd-java-agent:instrumentation:jetty-appsec-9.3')) testImplementation testFixtures(project(':dd-java-agent:instrumentation:servlet:request-5')) testImplementation testFixtures(project(':dd-java-agent:appsec')) testRuntimeOnly project(':dd-java-agent:instrumentation:jetty-9') 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') latestDepTestImplementation group: 'org.eclipse.jetty', name: 'jetty-server', version: '11.+', { exclude group: 'org.slf4j', module: 'slf4j-api' @@ -64,6 +71,9 @@ dependencies { latestDepTestImplementation group: 'org.eclipse.jetty', name: 'jetty-servlet', version: '11.+', { exclude group: 'org.slf4j', module: 'slf4j-api' } + latestDepTestImplementation group: 'org.eclipse.jetty.websocket', name: 'websocket-jakarta-server', version: '11.+', { + exclude group: 'org.slf4j', module: 'slf4j-api' + } // just to mix things up, see if there's no conflict latestDepTestRuntimeOnly group: 'javax.servlet', name: 'javax.servlet-api', version: '4.0.1' diff --git a/dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy b/dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy index 01dd60631ca..d557b97a7bf 100644 --- a/dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy +++ b/dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy @@ -1,4 +1,4 @@ -import datadog.trace.agent.test.base.HttpServer +import datadog.trace.agent.test.base.WebsocketServer import jakarta.servlet.Filter import jakarta.servlet.FilterChain import jakarta.servlet.MultipartConfigElement @@ -7,21 +7,33 @@ import jakarta.servlet.ServletException import jakarta.servlet.ServletRequest import jakarta.servlet.ServletResponse import jakarta.servlet.http.HttpServletRequest -import org.eclipse.jetty.server.Handler +import jakarta.websocket.CloseReason +import jakarta.websocket.Endpoint +import jakarta.websocket.EndpointConfig +import jakarta.websocket.MessageHandler +import jakarta.websocket.Session +import jakarta.websocket.server.ServerEndpointConfig import org.eclipse.jetty.server.Server -import org.eclipse.jetty.server.handler.AbstractHandler import org.eclipse.jetty.server.handler.ErrorHandler import org.eclipse.jetty.server.session.SessionHandler import org.eclipse.jetty.servlet.FilterMapping import org.eclipse.jetty.servlet.ServletContextHandler +import org.eclipse.jetty.websocket.jakarta.server.config.JakartaWebSocketServletContainerInitializer -class JettyServer implements HttpServer { +import java.nio.ByteBuffer + +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace + +class JettyServer implements WebsocketServer { def port = 0 final server = new Server(0) // select random open port - JettyServer(Handler handler) { + JettyServer(ServletContextHandler handler) { server.handler = handler server.addBean(errorHandler) + JakartaWebSocketServletContainerInitializer.configure(handler, (servletContext, container) -> { + container.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.class, "/websocket").build()) + }) } @Override @@ -46,7 +58,7 @@ class JettyServer implements HttpServer { return this.class.name } - static AbstractHandler servletHandler(Class servlet) { + static ServletContextHandler servletHandler(Class servlet) { ServletContextHandler handler = new ServletContextHandler(null, "/context-path") final sessionHandler = new SessionHandler() handler.sessionHandler = sessionHandler @@ -58,6 +70,7 @@ class JettyServer implements HttpServer { static class EnableMultipartFilter implements Filter { static final MultipartConfigElement MULTIPART_CONFIG_ELEMENT = new MultipartConfigElement(System.getProperty('java.io.tmpdir')) + @Override void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException { request.setAttribute('org.eclipse.jetty.multipartConfig', MULTIPART_CONFIG_ELEMENT) @@ -75,4 +88,87 @@ class JettyServer implements HttpServer { } } } + + @Override + void serverSendText(String[] messages) { + if (messages.length == 1) { + WsEndpoint.activeSession.getBasicRemote().sendText(messages[0]) + } else { + def remoteEndpoint = WsEndpoint.activeSession.getBasicRemote() + for (int i = 0; i < messages.length; i++) { + remoteEndpoint.sendText(messages[i], i == messages.length - 1) + } + } + } + + @Override + void serverSendBinary(byte[][] binaries) { + if (binaries.length == 1) { + WsEndpoint.activeSession.getBasicRemote().sendBinary(ByteBuffer.wrap(binaries[0])) + } else { + try (def stream = WsEndpoint.activeSession.getBasicRemote().getSendStream()) { + binaries.each { stream.write(it) } + } + } + } + + @Override + synchronized void awaitConnected() { + synchronized (WsEndpoint) { + try { + while (WsEndpoint.activeSession == null) { + WsEndpoint.wait() + } + } catch (InterruptedException ie) { + Thread.currentThread().interrupt() + } + } + } + + @Override + void serverClose() { + WsEndpoint.activeSession?.close() + WsEndpoint.activeSession = null + } + + @Override + void setMaxPayloadSize(int size) { + WsEndpoint.activeSession?.setMaxTextMessageBufferSize(size) + WsEndpoint.activeSession?.setMaxBinaryMessageBufferSize(size) + } + + @Override + boolean canSplitLargeWebsocketPayloads() { + false + } + + static class WsEndpoint extends Endpoint { + static Session activeSession + + @Override + void onOpen(Session session, EndpointConfig endpointConfig) { + session.addMessageHandler(new MessageHandler.Partial() { + @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 + runUnderTrace("onRead", {}) + } + }) + session.addMessageHandler(new MessageHandler.Partial() { + @Override + void onMessage(byte[] b, boolean last) { + runUnderTrace("onRead", {}) + } + }) + activeSession = session + synchronized (WsEndpoint) { + WsEndpoint.notifyAll() + } + } + + @Override + void onClose(Session session, CloseReason closeReason) { + activeSession = null + } + } } diff --git a/dd-java-agent/instrumentation/jetty-12/build.gradle b/dd-java-agent/instrumentation/jetty-12/build.gradle index b5b7ff50cad..a1a1aa600bd 100644 --- a/dd-java-agent/instrumentation/jetty-12/build.gradle +++ b/dd-java-agent/instrumentation/jetty-12/build.gradle @@ -68,6 +68,8 @@ dependencies { testImplementation(project(':dd-java-agent:instrumentation:jetty-appsec-9.3')) testImplementation testFixtures(project(':dd-java-agent:instrumentation:servlet:request-5')) 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') testImplementation testFixtures(project(':dd-java-agent:appsec')) testRuntimeOnly project(':dd-java-agent:instrumentation:jetty-9') testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-5') @@ -84,21 +86,40 @@ dependencies { ee8TestImplementation group: 'org.eclipse.jetty.ee8', name: 'jetty-ee8-servlet', version: '12.0.0' , { exclude group: 'org.slf4j', module: 'slf4j-api' } + ee8TestImplementation group: 'org.eclipse.jetty.ee8.websocket', name: 'jetty-ee8-websocket-javax-server', version: '12.0.0' , { + exclude group: 'org.slf4j', module: 'slf4j-api' + } + ee8LatestDepTestImplementation group: 'org.eclipse.jetty.ee8', name: 'jetty-ee8-servlet', version: '12.+' , { exclude group: 'org.slf4j', module: 'slf4j-api' } + ee8LatestDepTestImplementation group: 'org.eclipse.jetty.ee8.websocket', name: 'jetty-ee8-websocket-javax-server', version: '12.+' , { + exclude group: 'org.slf4j', module: 'slf4j-api' + } ee9TestImplementation group: 'org.eclipse.jetty.ee9', name: 'jetty-ee9-servlet', version: '12.0.0' , { exclude group: 'org.slf4j', module: 'slf4j-api' } + ee9TestImplementation group: 'org.eclipse.jetty.ee9.websocket', name: 'jetty-ee9-websocket-jakarta-server', version: '12.0.0' , { + exclude group: 'org.slf4j', module: 'slf4j-api' + } ee9LatestDepTestImplementation group: 'org.eclipse.jetty.ee9', name: 'jetty-ee9-servlet', version: '12.+' , { exclude group: 'org.slf4j', module: 'slf4j-api' } + ee9LatestDepTestImplementation group: 'org.eclipse.jetty.ee9.websocket', name: 'jetty-ee9-websocket-jakarta-server', version: '12.+' , { + exclude group: 'org.slf4j', module: 'slf4j-api' + } ee10TestImplementation group: 'org.eclipse.jetty.ee10', name: 'jetty-ee10-servlet', version: '12.0.0' , { exclude group: 'org.slf4j', module: 'slf4j-api' } + ee10TestImplementation group: 'org.eclipse.jetty.ee10.websocket', name: 'jetty-ee10-websocket-jakarta-server', version: '12.0.0' , { + exclude group: 'org.slf4j', module: 'slf4j-api' + } ee10LatestDepTestImplementation group: 'org.eclipse.jetty.ee10', name: 'jetty-ee10-servlet', version: '12.+' , { exclude group: 'org.slf4j', module: 'slf4j-api' } + ee10LatestDepTestImplementation group: 'org.eclipse.jetty.ee10.websocket', name: 'jetty-ee10-websocket-jakarta-server', version: '12.+' , { + exclude group: 'org.slf4j', module: 'slf4j-api' + } } diff --git a/dd-java-agent/instrumentation/jetty-12/src/test/ee10/groovy/JettyServer.groovy b/dd-java-agent/instrumentation/jetty-12/src/test/ee10/groovy/JettyServer.groovy index 27bc9ee3aa1..7f523f8e1bd 100644 --- a/dd-java-agent/instrumentation/jetty-12/src/test/ee10/groovy/JettyServer.groovy +++ b/dd-java-agent/instrumentation/jetty-12/src/test/ee10/groovy/JettyServer.groovy @@ -1,23 +1,36 @@ -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.UNKNOWN - -import datadog.trace.agent.test.base.HttpServer import datadog.trace.agent.test.base.HttpServerTest +import datadog.trace.agent.test.base.WebsocketServer import jakarta.servlet.Servlet import jakarta.servlet.ServletException import jakarta.servlet.http.HttpServletRequest +import jakarta.websocket.CloseReason +import jakarta.websocket.Endpoint +import jakarta.websocket.EndpointConfig +import jakarta.websocket.MessageHandler +import jakarta.websocket.Session +import jakarta.websocket.server.ServerEndpointConfig import org.eclipse.jetty.ee10.servlet.ErrorHandler import org.eclipse.jetty.ee10.servlet.ServletContextHandler +import org.eclipse.jetty.ee10.websocket.jakarta.server.config.JakartaWebSocketServletContainerInitializer import org.eclipse.jetty.server.Handler import org.eclipse.jetty.server.Server -class JettyServer implements HttpServer { +import java.nio.ByteBuffer + +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.UNKNOWN +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace + +class JettyServer implements WebsocketServer { def port = 0 final server = new Server(0) // select random open port - JettyServer(Handler handler) { + JettyServer(ServletContextHandler handler) { server.handler = handler server.addBean(errorHandler) + JakartaWebSocketServletContainerInitializer.configure(handler, (servletContext, container) -> { + container.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.class, "/websocket").build()) + }) } @Override @@ -65,4 +78,87 @@ class JettyServer implements HttpServer { } } } + + @Override + void serverSendText(String[] messages) { + if (messages.length == 1) { + WsEndpoint.activeSession.getBasicRemote().sendText(messages[0]) + } else { + def remoteEndpoint = WsEndpoint.activeSession.getBasicRemote() + for (int i = 0; i < messages.length; i++) { + remoteEndpoint.sendText(messages[i], i == messages.length - 1) + } + } + } + + @Override + void serverSendBinary(byte[][] binaries) { + if (binaries.length == 1) { + WsEndpoint.activeSession.getBasicRemote().sendBinary(ByteBuffer.wrap(binaries[0])) + } else { + try (def stream = WsEndpoint.activeSession.getBasicRemote().getSendStream()) { + binaries.each { stream.write(it) } + } + } + } + + @Override + synchronized void awaitConnected() { + synchronized (WsEndpoint) { + try { + while (WsEndpoint.activeSession == null) { + WsEndpoint.wait() + } + } catch (InterruptedException ie) { + Thread.currentThread().interrupt() + } + } + } + + @Override + void serverClose() { + WsEndpoint.activeSession?.close() + WsEndpoint.activeSession = null + } + + @Override + void setMaxPayloadSize(int size) { + WsEndpoint.activeSession?.setMaxTextMessageBufferSize(size) + WsEndpoint.activeSession?.setMaxBinaryMessageBufferSize(size) + } + + @Override + boolean canSplitLargeWebsocketPayloads() { + false + } + + static class WsEndpoint extends Endpoint { + static Session activeSession + + @Override + void onOpen(Session session, EndpointConfig endpointConfig) { + session.addMessageHandler(new MessageHandler.Partial() { + @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 + runUnderTrace("onRead", {}) + } + }) + session.addMessageHandler(new MessageHandler.Partial() { + @Override + void onMessage(byte[] b, boolean last) { + runUnderTrace("onRead", {}) + } + }) + activeSession = session + synchronized (WsEndpoint) { + WsEndpoint.notifyAll() + } + } + + @Override + void onClose(Session session, CloseReason closeReason) { + activeSession = null + } + } } diff --git a/dd-java-agent/instrumentation/jetty-12/src/test/ee8/groovy/JettyServer.groovy b/dd-java-agent/instrumentation/jetty-12/src/test/ee8/groovy/JettyServer.groovy index 38f2f7a5328..9170a6d606c 100644 --- a/dd-java-agent/instrumentation/jetty-12/src/test/ee8/groovy/JettyServer.groovy +++ b/dd-java-agent/instrumentation/jetty-12/src/test/ee8/groovy/JettyServer.groovy @@ -1,30 +1,40 @@ -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.UNKNOWN - -import datadog.trace.agent.test.base.HttpServer import datadog.trace.agent.test.base.HttpServerTest +import datadog.trace.agent.test.base.WebsocketServer import org.eclipse.jetty.ee8.nested.ErrorHandler -import org.eclipse.jetty.ee8.servlet.ServletContextHandler import org.eclipse.jetty.ee8.nested.SessionHandler as EE8SessionHandler -import org.eclipse.jetty.server.Handler +import org.eclipse.jetty.ee8.servlet.ServletContextHandler +import org.eclipse.jetty.ee8.websocket.javax.server.config.JavaxWebSocketServletContainerInitializer import org.eclipse.jetty.server.Server import org.eclipse.jetty.session.SessionHandler - import javax.servlet.Servlet import javax.servlet.ServletException import javax.servlet.http.HttpServletRequest +import javax.websocket.CloseReason +import javax.websocket.Endpoint +import javax.websocket.EndpointConfig +import javax.websocket.MessageHandler +import javax.websocket.Session +import javax.websocket.server.ServerEndpointConfig +import java.nio.ByteBuffer -class JettyServer implements HttpServer { +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.UNKNOWN +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace + +class JettyServer implements WebsocketServer { def port = 0 final server = new Server(0) // select random open port - JettyServer(Handler handler) { + JettyServer(ServletContextHandler handler) { final sessionHandler = new SessionHandler() sessionHandler.handler = handler server.handler = sessionHandler server.addBean(errorHandler) server.addBean(sessionHandler.sessionIdManager, true) + JavaxWebSocketServletContainerInitializer.configure(handler, (servletContext, container) -> { + container.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.class, "/websocket").build()) + }) } @Override @@ -49,7 +59,7 @@ class JettyServer implements HttpServer { return this.class.name } - static Handler servletHandler(Class servlet) { + static ServletContextHandler servletHandler(Class servlet) { // defaults to jakarta servlet ServletContextHandler handler = new ServletContextHandler(null, "/context-path") handler.sessionHandler = new EE8SessionHandler() @@ -59,7 +69,7 @@ class JettyServer implements HttpServer { .each { handler.servletHandler.addServletWithMapping(servlet, it.path) } - handler.get() + handler } static errorHandler = new ErrorHandler() { @@ -73,4 +83,87 @@ class JettyServer implements HttpServer { } } } + + @Override + void serverSendText(String[] messages) { + if (messages.length == 1) { + WsEndpoint.activeSession.getBasicRemote().sendText(messages[0]) + } else { + def remoteEndpoint = WsEndpoint.activeSession.getBasicRemote() + for (int i = 0; i < messages.length; i++) { + remoteEndpoint.sendText(messages[i], i == messages.length - 1) + } + } + } + + @Override + void serverSendBinary(byte[][] binaries) { + if (binaries.length == 1) { + WsEndpoint.activeSession.getBasicRemote().sendBinary(ByteBuffer.wrap(binaries[0])) + } else { + try (def stream = WsEndpoint.activeSession.getBasicRemote().getSendStream()) { + binaries.each { stream.write(it) } + } + } + } + + @Override + synchronized void awaitConnected() { + synchronized (WsEndpoint) { + try { + while (WsEndpoint.activeSession == null) { + WsEndpoint.wait() + } + } catch (InterruptedException ie) { + Thread.currentThread().interrupt() + } + } + } + + @Override + void serverClose() { + WsEndpoint.activeSession?.close() + WsEndpoint.activeSession = null + } + + @Override + void setMaxPayloadSize(int size) { + WsEndpoint.activeSession?.setMaxTextMessageBufferSize(size) + WsEndpoint.activeSession?.setMaxBinaryMessageBufferSize(size) + } + + @Override + boolean canSplitLargeWebsocketPayloads() { + false + } + + static class WsEndpoint extends Endpoint { + static Session activeSession + + @Override + void onOpen(Session session, EndpointConfig endpointConfig) { + session.addMessageHandler(new MessageHandler.Partial() { + @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 + runUnderTrace("onRead", {}) + } + }) + session.addMessageHandler(new MessageHandler.Partial() { + @Override + void onMessage(byte[] b, boolean last) { + runUnderTrace("onRead", {}) + } + }) + activeSession = session + synchronized (WsEndpoint) { + WsEndpoint.notifyAll() + } + } + + @Override + void onClose(Session session, CloseReason closeReason) { + activeSession = null + } + } } diff --git a/dd-java-agent/instrumentation/jetty-12/src/test/ee9/groovy/JettyServer.groovy b/dd-java-agent/instrumentation/jetty-12/src/test/ee9/groovy/JettyServer.groovy index 246542bb1fc..5baad3db396 100644 --- a/dd-java-agent/instrumentation/jetty-12/src/test/ee9/groovy/JettyServer.groovy +++ b/dd-java-agent/instrumentation/jetty-12/src/test/ee9/groovy/JettyServer.groovy @@ -1,23 +1,35 @@ -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.UNKNOWN - -import datadog.trace.agent.test.base.HttpServer import datadog.trace.agent.test.base.HttpServerTest +import datadog.trace.agent.test.base.WebsocketServer import jakarta.servlet.Servlet import jakarta.servlet.ServletException import jakarta.servlet.http.HttpServletRequest +import jakarta.websocket.CloseReason +import jakarta.websocket.Endpoint +import jakarta.websocket.EndpointConfig +import jakarta.websocket.MessageHandler +import jakarta.websocket.Session +import jakarta.websocket.server.ServerEndpointConfig import org.eclipse.jetty.ee9.nested.ErrorHandler import org.eclipse.jetty.ee9.servlet.ServletContextHandler -import org.eclipse.jetty.server.Handler +import org.eclipse.jetty.ee9.websocket.jakarta.server.config.JakartaWebSocketServletContainerInitializer import org.eclipse.jetty.server.Server -class JettyServer implements HttpServer { +import java.nio.ByteBuffer + +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.UNKNOWN +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace + +class JettyServer implements WebsocketServer { def port = 0 final server = new Server(0) // select random open port - JettyServer(Handler handler) { + JettyServer(ServletContextHandler handler) { server.handler = handler server.addBean(errorHandler) + JakartaWebSocketServletContainerInitializer.configure(handler, (servletContext, container) -> { + container.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.class, "/websocket").build()) + }) } @Override @@ -42,7 +54,7 @@ class JettyServer implements HttpServer { return this.class.name } - static Handler servletHandler(Class servlet) { + static ServletContextHandler servletHandler(Class servlet) { // defaults to jakarta servlet ServletContextHandler handler = new ServletContextHandler(null, "/context-path") handler.errorHandler = errorHandler @@ -51,7 +63,7 @@ class JettyServer implements HttpServer { .each { handler.servletHandler.addServletWithMapping(servlet, it.path) } - handler.get() + handler } static errorHandler = new ErrorHandler() { @@ -65,4 +77,87 @@ class JettyServer implements HttpServer { } } } + + @Override + void serverSendText(String[] messages) { + if (messages.length == 1) { + WsEndpoint.activeSession.getBasicRemote().sendText(messages[0]) + } else { + def remoteEndpoint = WsEndpoint.activeSession.getBasicRemote() + for (int i = 0; i < messages.length; i++) { + remoteEndpoint.sendText(messages[i], i == messages.length - 1) + } + } + } + + @Override + void serverSendBinary(byte[][] binaries) { + if (binaries.length == 1) { + WsEndpoint.activeSession.getBasicRemote().sendBinary(ByteBuffer.wrap(binaries[0])) + } else { + try (def stream = WsEndpoint.activeSession.getBasicRemote().getSendStream()) { + binaries.each { stream.write(it) } + } + } + } + + @Override + synchronized void awaitConnected() { + synchronized (WsEndpoint) { + try { + while (WsEndpoint.activeSession == null) { + WsEndpoint.wait() + } + } catch (InterruptedException ie) { + Thread.currentThread().interrupt() + } + } + } + + @Override + void serverClose() { + WsEndpoint.activeSession?.close() + WsEndpoint.activeSession = null + } + + @Override + void setMaxPayloadSize(int size) { + WsEndpoint.activeSession?.setMaxTextMessageBufferSize(size) + WsEndpoint.activeSession?.setMaxBinaryMessageBufferSize(size) + } + + @Override + boolean canSplitLargeWebsocketPayloads() { + false + } + + static class WsEndpoint extends Endpoint { + static Session activeSession + + @Override + void onOpen(Session session, EndpointConfig endpointConfig) { + session.addMessageHandler(new MessageHandler.Partial() { + @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 + runUnderTrace("onRead", {}) + } + }) + session.addMessageHandler(new MessageHandler.Partial() { + @Override + void onMessage(byte[] b, boolean last) { + runUnderTrace("onRead", {}) + } + }) + activeSession = session + synchronized (WsEndpoint) { + WsEndpoint.notifyAll() + } + } + + @Override + void onClose(Session session, CloseReason closeReason) { + activeSession = null + } + } } diff --git a/dd-java-agent/instrumentation/jetty-9/build.gradle b/dd-java-agent/instrumentation/jetty-9/build.gradle index 314bc926a9d..4114dc46eae 100644 --- a/dd-java-agent/instrumentation/jetty-9/build.gradle +++ b/dd-java-agent/instrumentation/jetty-9/build.gradle @@ -123,6 +123,7 @@ tasks.named("compileLatestDepForkedTestGroovy").configure { dependencies { compileOnly group: 'org.eclipse.jetty', name: 'jetty-server', version: '9.0.0.v20130308' + implementation project(':dd-java-agent:instrumentation:jetty-common') main_jetty904CompileOnly group: 'org.eclipse.jetty', name: 'jetty-server', version: '9.0.4.v20130625' @@ -173,12 +174,15 @@ dependencies { String jetty9Version = '9.0.0.v20130308' testFixturesCompileOnly group: 'org.eclipse.jetty', name: 'jetty-server', version: jetty9Version testFixturesCompileOnly group: 'org.eclipse.jetty', name: 'jetty-servlet', version: jetty9Version + testFixturesImplementation group: 'javax.websocket', name: 'javax.websocket-api', version: '1.0' testImplementation group: 'org.eclipse.jetty', name: 'jetty-server', version: jetty9Version testImplementation group: 'org.eclipse.jetty', name: 'jetty-servlet', version: jetty9Version testImplementation group: 'org.eclipse.jetty', name: 'jetty-continuation', version: jetty9Version testImplementation project(':dd-java-agent:instrumentation:jetty-appsec-7') testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-2') testRuntimeOnly project(':dd-java-agent:instrumentation:jetty-appsec-8.1.3') + testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:javax-websocket-1.0') + testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:jakarta-websocket-2.0') testImplementation testFixtures(project(':dd-java-agent:instrumentation:servlet:request-3')) testFixturesImplementation testFixtures(project(':dd-java-agent:instrumentation:servlet:request-3')) testImplementation testFixtures(project(':dd-java-agent:appsec')) @@ -186,23 +190,28 @@ dependencies { jetty92TestImplementation group: 'org.eclipse.jetty', name: 'jetty-server', version: '9.2.30.v20200428' jetty92TestImplementation group: 'org.eclipse.jetty', name: 'jetty-servlet', version: '9.2.30.v20200428' jetty92TestImplementation group: 'org.eclipse.jetty', name: 'jetty-continuation', version: '9.2.30.v20200428' + jetty92TestImplementation group: 'org.eclipse.jetty.websocket', name: 'javax-websocket-server-impl', version: '9.2.30.v20200428' jetty92TestImplementation project(':dd-java-agent:instrumentation:jetty-appsec-9.2') jetty92TestImplementation testFixtures(project(':dd-java-agent:instrumentation:servlet:request-3')) jetty94TestImplementation group: 'org.eclipse.jetty', name: 'jetty-server', version: '9.4.15.v20190215' jetty94TestImplementation group: 'org.eclipse.jetty', name: 'jetty-servlet', version: '9.4.15.v20190215' jetty94TestImplementation group: 'org.eclipse.jetty', name: 'jetty-continuation', version: '9.4.15.v20190215' + jetty94TestImplementation group: 'org.eclipse.jetty.websocket', name: 'javax-websocket-server-impl', version: '9.4.15.v20190215' jetty94TestImplementation project(':dd-java-agent:instrumentation:jetty-appsec-9.3') jetty94TestImplementation testFixtures(project(':dd-java-agent:instrumentation:servlet:request-3')) latestDepJetty9TestImplementation group: 'org.eclipse.jetty', name: 'jetty-server', version: '9.+' latestDepJetty9TestImplementation group: 'org.eclipse.jetty', name: 'jetty-servlet', version: '9.+' latestDepJetty9TestImplementation group: 'org.eclipse.jetty', name: 'jetty-continuation', version: '9.+' + latestDepJetty9TestImplementation group: 'org.eclipse.jetty.websocket', name: 'javax-websocket-server-impl', version: '9.+' + latestDepJetty9TestImplementation project(':dd-java-agent:instrumentation:jetty-appsec-9.3') latestDepJetty9TestImplementation testFixtures(project(':dd-java-agent:instrumentation:servlet:request-3')) latestDepTestImplementation group: 'org.eclipse.jetty', name: 'jetty-server', version: '10.+' latestDepTestImplementation group: 'org.eclipse.jetty', name: 'jetty-servlet', version: '10.+' + latestDepTestImplementation group: 'org.eclipse.jetty.websocket', name: 'websocket-javax-server', version: '10.+' latestDepTestImplementation project(':dd-java-agent:instrumentation:jetty-appsec-9.3') latestDepTestImplementation testFixtures(project(':dd-java-agent:instrumentation:servlet:request-3')) } diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty10/JettyServerInstrumentation.java b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty10/JettyServerInstrumentation.java index 874a618c0db..f16673cfd9e 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty10/JettyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty10/JettyServerInstrumentation.java @@ -99,6 +99,7 @@ public void typeAdvice(TypeTransformer transformer) { public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice(takesNoArguments().and(named("handle")), packageName + ".HandleAdvice"); transformer.applyAdvice(named("recycle").and(takesNoArguments()), packageName + ".ResetAdvice"); + transformer.applyAdvice(named("onCompleted"), packageName + ".OnCompletedAdvice"); if (appSecNotFullyDisabled) { transformer.applyAdvice( diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/WebSocketSessionInstrumentation.java b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/WebSocketSessionInstrumentation.java new file mode 100644 index 00000000000..44024246627 --- /dev/null +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/WebSocketSessionInstrumentation.java @@ -0,0 +1,88 @@ +package datadog.trace.instrumentation.jetty9; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE; +import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; +import java.util.Collections; +import java.util.Map; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.matcher.ElementMatcher; + +@AutoService(InstrumenterModule.class) +public class WebSocketSessionInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + + public WebSocketSessionInstrumentation() { + super("jetty", "jetty-websocket"); + } + + @Override + public String instrumentedType() { + return "org.eclipse.jetty.websocket.common.WebSocketSession"; + } + + @Override + public Map contextStore() { + return Collections.singletonMap( + "javax.websocket.Session", + "datadog.trace.bootstrap.instrumentation.websocket.HandlerContext$Sender"); + } + + @Override + public String muzzleDirective() { + return "jetty-websocket"; + } + + @Override + public ElementMatcher.Junction classLoaderMatcher() { + return hasClassNamed("org.eclipse.jetty.websocket.jsr356.JsrSession"); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("close").and(takesNoArguments()), getClass().getName() + "$CloseAdvice"); + } + + public static class CloseAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope before( + @Advice.This final Object session, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext) { + + // that class is not implementing javax.websocket.Session but hides the close() method + // hence we need an ad hoc advice + handlerContext = + (HandlerContext.Sender) + InstrumentationContext.get( + "javax.websocket.Session", + "datadog.trace.bootstrap.instrumentation.websocket.HandlerContext$Sender") + .remove(session); + if (handlerContext == null) { + return null; + } + return activateSpan(DECORATE.onSessionCloseIssued(handlerContext, null, 1000)); + } + + @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) + public static void after( + @Advice.Enter final AgentScope scope, + @Advice.Thrown final Throwable thrown, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext) { + if (scope != null) { + DECORATE.onError(scope, thrown); + DECORATE.onFrameEnd(handlerContext); + scope.close(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/OnCompletedAdvice.java b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/OnCompletedAdvice.java new file mode 100644 index 00000000000..9881cb5be76 --- /dev/null +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/OnCompletedAdvice.java @@ -0,0 +1,29 @@ +package datadog.trace.instrumentation.jetty10; + +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; + +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import net.bytebuddy.asm.Advice; +import org.eclipse.jetty.server.HttpChannel; +import org.eclipse.jetty.server.Request; + +public class OnCompletedAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope enter(@Advice.This final HttpChannel channel) { + Request req = channel.getRequest(); + Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); + if (existingSpan instanceof AgentSpan) { + // return activateSpan((AgentSpan) existingSpan); + } + return null; + } + + @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) + public static void exit( + @Advice.This final HttpChannel channel, @Advice.Enter final AgentScope scope) { + if (scope != null) { + scope.close(); + } + } +} diff --git a/dd-java-agent/instrumentation/jetty-9/src/test/groovy/datadog/trace/instrumentation/jetty9/Jetty9InactiveAppSecTest.groovy b/dd-java-agent/instrumentation/jetty-9/src/test/groovy/datadog/trace/instrumentation/jetty9/Jetty9InactiveAppSecTest.groovy index f50269f55bf..b54000a00ec 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/test/groovy/datadog/trace/instrumentation/jetty9/Jetty9InactiveAppSecTest.groovy +++ b/dd-java-agent/instrumentation/jetty-9/src/test/groovy/datadog/trace/instrumentation/jetty9/Jetty9InactiveAppSecTest.groovy @@ -2,6 +2,8 @@ package datadog.trace.instrumentation.jetty9 import com.datadog.appsec.AppSecInactiveHttpServerTest import datadog.trace.agent.test.base.HttpServer +import test.JettyServer +import test.TestHandler class Jetty9InactiveAppSecTest extends AppSecInactiveHttpServerTest { HttpServer server() { diff --git a/dd-java-agent/instrumentation/jetty-9/src/test/groovy/datadog/trace/instrumentation/jetty9/Jetty9Test.groovy b/dd-java-agent/instrumentation/jetty-9/src/test/groovy/datadog/trace/instrumentation/jetty9/Jetty9Test.groovy index c556f23f912..46cb709d200 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/test/groovy/datadog/trace/instrumentation/jetty9/Jetty9Test.groovy +++ b/dd-java-agent/instrumentation/jetty-9/src/test/groovy/datadog/trace/instrumentation/jetty9/Jetty9Test.groovy @@ -5,15 +5,14 @@ import datadog.trace.agent.test.base.HttpServerTest import datadog.trace.agent.test.naming.TestingGenericHttpNamingConventions import org.eclipse.jetty.server.Server import org.eclipse.jetty.server.handler.AbstractHandler -import org.eclipse.jetty.server.session.SessionHandler +import test.JettyServer +import test.TestHandler abstract class Jetty9Test extends HttpServerTest { @Override HttpServer server() { - final sessionHandler = new SessionHandler() - sessionHandler.handler = handler() - new JettyServer(sessionHandler) + new JettyServer(handler()) } AbstractHandler handler() { @@ -85,6 +84,11 @@ abstract class Jetty9Test extends HttpServerTest { boolean testSessionId() { true } + + @Override + boolean testWebsockets() { + return super.testWebsockets() && (getServer() as JettyServer).websocketAvailable + } } class Jetty9V0ForkedTest extends Jetty9Test implements TestingGenericHttpNamingConventions.ServerV0 { diff --git a/dd-java-agent/instrumentation/jetty-9/src/test/groovy/datadog/trace/instrumentation/jetty9/JettyContinuationHandlerTest.groovy b/dd-java-agent/instrumentation/jetty-9/src/test/groovy/datadog/trace/instrumentation/jetty9/JettyContinuationHandlerTest.groovy index 5be31bff28d..a9999d16b8b 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/test/groovy/datadog/trace/instrumentation/jetty9/JettyContinuationHandlerTest.groovy +++ b/dd-java-agent/instrumentation/jetty-9/src/test/groovy/datadog/trace/instrumentation/jetty9/JettyContinuationHandlerTest.groovy @@ -8,6 +8,7 @@ import org.eclipse.jetty.continuation.Continuation import org.eclipse.jetty.continuation.ContinuationSupport import org.eclipse.jetty.server.Request import org.eclipse.jetty.server.handler.AbstractHandler +import org.eclipse.jetty.server.session.SessionHandler import javax.servlet.MultipartConfigElement import javax.servlet.ServletException @@ -16,7 +17,7 @@ import javax.servlet.http.HttpServletResponse import java.util.concurrent.ExecutorService import java.util.concurrent.Executors -import static TestHandler.handleRequest +import static test.TestHandler.handleRequest import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.TIMEOUT import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.TIMEOUT_ERROR @@ -26,7 +27,9 @@ abstract class JettyContinuationHandlerTest extends Jetty9Test { @Override AbstractHandler handler() { - ContinuationTestHandler.INSTANCE + def ret = new SessionHandler() + ret.handler = ContinuationTestHandler.INSTANCE + ret } static class ContinuationTestHandler extends AbstractHandler { diff --git a/dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/datadog/trace/instrumentation/jetty9/JettyServer.groovy b/dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/datadog/trace/instrumentation/jetty9/JettyServer.groovy deleted file mode 100644 index b18c1aebe7b..00000000000 --- a/dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/datadog/trace/instrumentation/jetty9/JettyServer.groovy +++ /dev/null @@ -1,37 +0,0 @@ -package datadog.trace.instrumentation.jetty9 - -import datadog.trace.agent.test.base.HttpServer -import org.eclipse.jetty.server.Server -import org.eclipse.jetty.server.handler.AbstractHandler - -class JettyServer implements HttpServer { - def port = 0 - final server = new Server(0) // select random open port - - JettyServer(AbstractHandler handler) { - server.setHandler(handler) - server.addBean(TestHandler.errorHandler) - } - - @Override - void start() { - server.start() - port = server.connectors[0].localPort - assert port > 0 - } - - @Override - void stop() { - server.stop() - } - - @Override - URI address() { - new URI("http://localhost:$port/") - } - - @Override - String toString() { - this.class.name - } -} diff --git a/dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/test/JettyServer.groovy b/dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/test/JettyServer.groovy new file mode 100644 index 00000000000..21034f45255 --- /dev/null +++ b/dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/test/JettyServer.groovy @@ -0,0 +1,144 @@ +package test + +import datadog.trace.agent.test.base.WebsocketServer +import org.eclipse.jetty.server.Server +import org.eclipse.jetty.server.handler.AbstractHandler + +import javax.websocket.CloseReason +import javax.websocket.Endpoint +import javax.websocket.EndpointConfig +import javax.websocket.MessageHandler +import javax.websocket.Session +import javax.websocket.server.ServerEndpointConfig +import java.nio.ByteBuffer + +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace + +class JettyServer implements WebsocketServer { + def port = 0 + final server = new Server(0) // select random open port + def websocketAvailable = true + + JettyServer(AbstractHandler handler) { + server.setHandler(handler) + try { + def container = ("org.eclipse.jetty.websocket.jsr356.server.deploy.WebSocketServerContainerInitializer" as Class)."configureContext"(handler) + container."addEndpoint"(ServerEndpointConfig.Builder.create(WsEndpoint.class, "/websocket").build()) + + } catch (Throwable t) { + try { + ("org.eclipse.jetty.websocket.javax.server.config.JavaxWebSocketServletContainerInitializer" as Class)."configure"(handler, { servletContext, container -> + container.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.class, "/websocket").build()) + }) + + } catch (Throwable tt) { + websocketAvailable = false + } + } + } + + @Override + void start() { + server.start() + port = server.connectors[0].localPort + assert port > 0 + } + + @Override + void stop() { + server.stop() + } + + @Override + URI address() { + new URI("http://localhost:$port/") + } + + @Override + String toString() { + this.class.name + } + + @Override + void serverSendText(String[] messages) { + if (messages.length == 1) { + WsEndpoint.activeSession.getBasicRemote().sendText(messages[0]) + } else { + def remoteEndpoint = WsEndpoint.activeSession.getBasicRemote() + for (int i = 0; i < messages.length; i++) { + remoteEndpoint.sendText(messages[i], i == messages.length - 1) + } + } + } + + @Override + void serverSendBinary(byte[][] binaries) { + if (binaries.length == 1) { + WsEndpoint.activeSession.getBasicRemote().sendBinary(ByteBuffer.wrap(binaries[0])) + } else { + try (def stream = WsEndpoint.activeSession.getBasicRemote().getSendStream()) { + binaries.each { stream.write(it) } + } + } + } + + @Override + synchronized void awaitConnected() { + synchronized (WsEndpoint) { + try { + while (WsEndpoint.activeSession == null) { + WsEndpoint.wait() + } + } catch (InterruptedException ie) { + Thread.currentThread().interrupt() + } + } + } + + @Override + void serverClose() { + WsEndpoint.activeSession?.close() + WsEndpoint.activeSession = null + } + + @Override + void setMaxPayloadSize(int size) { + WsEndpoint.activeSession?.setMaxTextMessageBufferSize(size) + WsEndpoint.activeSession?.setMaxBinaryMessageBufferSize(size) + } + + @Override + boolean canSplitLargeWebsocketPayloads() { + false + } + + static class WsEndpoint extends Endpoint { + static Session activeSession + + @Override + void onOpen(Session session, EndpointConfig endpointConfig) { + session.addMessageHandler(new MessageHandler.Partial() { + @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 + runUnderTrace("onRead", {}) + } + }) + session.addMessageHandler(new MessageHandler.Partial() { + @Override + void onMessage(byte[] b, boolean last) { + runUnderTrace("onRead", {}) + } + }) + activeSession = session + synchronized (WsEndpoint) { + WsEndpoint.notifyAll() + } + } + + @Override + void onClose(Session session, CloseReason closeReason) { + activeSession = null + } + } +} diff --git a/dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/datadog/trace/instrumentation/jetty9/TestHandler.groovy b/dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/test/TestHandler.groovy similarity index 50% rename from dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/datadog/trace/instrumentation/jetty9/TestHandler.groovy rename to dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/test/TestHandler.groovy index fff5654ade0..1c726d92a29 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/datadog/trace/instrumentation/jetty9/TestHandler.groovy +++ b/dd-java-agent/instrumentation/jetty-9/src/testFixtures/groovy/test/TestHandler.groovy @@ -1,23 +1,63 @@ -package datadog.trace.instrumentation.jetty9 +package test import datadog.trace.agent.test.base.HttpServerTest import datadog.trace.instrumentation.servlet3.TestServlet3 import groovy.transform.CompileStatic import org.eclipse.jetty.server.Request -import org.eclipse.jetty.server.handler.AbstractHandler import org.eclipse.jetty.server.handler.ErrorHandler +import org.eclipse.jetty.server.session.SessionHandler +import org.eclipse.jetty.servlet.ServletContextHandler import javax.servlet.DispatcherType +import javax.servlet.Filter +import javax.servlet.FilterChain +import javax.servlet.FilterConfig import javax.servlet.MultipartConfigElement import javax.servlet.ServletException +import javax.servlet.ServletRequest +import javax.servlet.ServletResponse import javax.servlet.http.HttpServletRequest import javax.servlet.http.HttpServletResponse -class TestHandler extends AbstractHandler { +class TestHandler extends ServletContextHandler { static final TestHandler INSTANCE = new TestHandler() private static final MultipartConfigElement MULTIPART_CONFIG_ELEMENT = new MultipartConfigElement(System.getProperty('java.io.tmpdir')) - final TestServlet3.Sync testServlet3 = new TestServlet3.Sync() { + TestHandler() { + setSessionHandler(new SessionHandler()) + setErrorHandler(new ErrorHandler() { + @Override + protected void handleErrorPage(HttpServletRequest request, Writer writer, int code, String message) throws IOException { + Throwable th = (Throwable) request.getAttribute("javax.servlet.error.exception") + message = th ? th.message : message + if (message) { + writer.write(message) + } + } + }) + addFilter(MultipartFilter, "/*", EnumSet.of(DispatcherType.REQUEST)) + addServlet(TestServlet, "/*") + } + + static class MultipartFilter implements Filter { + + @Override + void init(FilterConfig filterConfig) throws ServletException { + } + + @Override + void doFilter(ServletRequest request, ServletResponse response, FilterChain filterChain) throws IOException, ServletException { + request.setAttribute('org.eclipse.jetty.multipartConfig', MULTIPART_CONFIG_ELEMENT) + request.setAttribute('org.eclipse.multipartConfig', MULTIPART_CONFIG_ELEMENT) + filterChain.doFilter(request, response) + } + + @Override + void destroy() { + } + } + + static class TestServlet extends TestServlet3.Sync { @Override HttpServerTest.ServerEndpoint determineEndpoint(HttpServletRequest req) { HttpServerTest.ServerEndpoint.forPath(req.requestURI) @@ -41,30 +81,6 @@ class TestHandler extends AbstractHandler { @CompileStatic static void handleRequest(Request request, HttpServletResponse response) { - INSTANCE.testServlet3.service(request, response) - } - - @Override - void handle(String target, Request baseRequest, HttpServletRequest request, HttpServletResponse response) throws IOException, ServletException { - // name of attribute varies depending on the jetty version - request.setAttribute('org.eclipse.jetty.multipartConfig', MULTIPART_CONFIG_ELEMENT) - request.setAttribute('org.eclipse.multipartConfig', MULTIPART_CONFIG_ELEMENT) - if (baseRequest.dispatcherType != DispatcherType.ERROR) { - handleRequest(baseRequest, response) - baseRequest.handled = true - } else { - errorHandler.handle(target, baseRequest, response, response) - } - } - - static errorHandler = new ErrorHandler() { - @Override - protected void handleErrorPage(HttpServletRequest request, Writer writer, int code, String message) throws IOException { - Throwable th = (Throwable) request.getAttribute("javax.servlet.error.exception") - message = th ? th.message : message - if (message) { - writer.write(message) - } - } + new TestServlet().service(request, response) } } diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/build.gradle b/dd-java-agent/instrumentation/spring-webmvc-3.1/build.gradle index 8c2193e0c4b..6e29753b618 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/build.gradle +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/build.gradle @@ -45,7 +45,7 @@ dependencies { // compileOnly group: 'javax.servlet', name: 'servlet-api', version: '2.4' testImplementation(project(':dd-java-agent:testing')) { - exclude(module: 'jetty-server') // incompatable servlet api + exclude(module: 'jetty-server') // incompatible servlet api } testImplementation testFixtures(project(':dd-java-agent:appsec')) @@ -62,6 +62,8 @@ dependencies { testRuntimeOnly project(':dd-java-agent:instrumentation:tomcat-5.5') testRuntimeOnly project(':dd-java-agent:instrumentation:tomcat-appsec-6') testRuntimeOnly project(':dd-java-agent:instrumentation:tomcat-appsec-7') + testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:javax-websocket-1.0') + testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:jakarta-websocket-2.0') testImplementation group: 'javax.validation', name: 'validation-api', version: '1.1.0.Final' testImplementation group: 'org.hibernate', name: 'hibernate-validator', version: '5.4.2.Final' @@ -70,6 +72,7 @@ dependencies { testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-test', version: '1.5.17.RELEASE' testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-web', version: '1.5.17.RELEASE' + testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-websocket', version: '1.5.17.RELEASE' testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-security', version: '1.5.17.RELEASE' testImplementation group: 'org.springframework.security.oauth', name: 'spring-security-oauth2', version: '2.0.16.RELEASE' diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/CustomBeanClassloaderTest.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/CustomBeanClassloaderTest.groovy index 0eb2ffabaf5..ac681cbbcc3 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/CustomBeanClassloaderTest.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/CustomBeanClassloaderTest.groovy @@ -6,6 +6,6 @@ class CustomBeanClassloaderTest extends SpringBootBasedTest { @Override SpringApplication application() { - return new SpringApplication(AppConfig, SecurityConfig, AuthServerConfig, CustomClassloaderConfig, TestController) + return new SpringApplication(AppConfig, SecurityConfig, AuthServerConfig, CustomClassloaderConfig, TestController, WebsocketConfig) } } diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy index 694f5a660a9..18a2fce745c 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy @@ -34,6 +34,7 @@ import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_HE import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.PATH_PARAM import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.REDIRECT import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.WEBSOCKET class SpringBootBasedTest extends HttpServerTest { @@ -46,7 +47,7 @@ class SpringBootBasedTest extends HttpServerTest Map extraServerTags = [:] SpringApplication application() { - new SpringApplication(AppConfig, SecurityConfig, AuthServerConfig, TestController) + new SpringApplication(AppConfig, SecurityConfig, AuthServerConfig, TestController, WebsocketConfig) } def setupSpec() { @@ -239,7 +240,7 @@ class SpringBootBasedTest extends HttpServerTest when: def response = client.newCall(request).execute() TEST_WRITER.waitForTraces(2) - DDSpan span = TEST_WRITER.flatten().find {it.operationName =='appsec-span' } + DDSpan span = TEST_WRITER.flatten().find { it.operationName == 'appsec-span' } then: response.code() == PATH_PARAM.status @@ -274,12 +275,12 @@ class SpringBootBasedTest extends HttpServerTest when: def response = client.newCall(request).execute() TEST_WRITER.waitForTraces(2) - DDSpan span = TEST_WRITER.flatten().find {it.operationName =='appsec-span' } + DDSpan span = TEST_WRITER.flatten().find { it.operationName == 'appsec-span' } then: response.code() == MATRIX_PARAM.status response.body().string() == MATRIX_PARAM.body - span.getTag(IG_PATH_PARAMS_TAG) == [var:['a=x,y;a=z', [a:['x', 'y', 'z']]]] + span.getTag(IG_PATH_PARAMS_TAG) == [var: ['a=x,y;a=z', [a: ['x', 'y', 'z']]]] } void 'tainting on matrix var'() { @@ -322,7 +323,7 @@ class SpringBootBasedTest extends HttpServerTest when: client.newCall(request).execute() TEST_WRITER.waitForTraces(1) - DDSpan span = TEST_WRITER.flatten().find {"servlet.request".contentEquals(it.operationName)} + DDSpan span = TEST_WRITER.flatten().find { "servlet.request".contentEquals(it.operationName) } then: span.getResourceName().toString() == "GET " + testPathParam() @@ -391,7 +392,9 @@ class SpringBootBasedTest extends HttpServerTest serviceName expectedServiceName() operationName "spring.handler" resourceName { - it == "TestController.${endpoint.name().toLowerCase()}" || endpoint == NOT_FOUND && it == "ResourceHttpRequestHandler.handleRequest" + it == "TestController.${endpoint.name().toLowerCase()}" + || endpoint == NOT_FOUND && it == "ResourceHttpRequestHandler.handleRequest" + || endpoint == WEBSOCKET && it == "WebSocketHttpRequestHandler.handleRequest" } spanType DDSpanTypes.HTTP_SERVER errored endpoint == EXCEPTION diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootServer.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootServer.groovy index d011c2a9814..bc875bb70d5 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootServer.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootServer.groovy @@ -1,14 +1,20 @@ package test.boot -import datadog.trace.agent.test.base.HttpServer + +import datadog.trace.agent.test.base.WebsocketServer +import org.springframework.beans.BeansException import org.springframework.boot.SpringApplication import org.springframework.boot.context.embedded.EmbeddedWebApplicationContext +import org.springframework.web.socket.BinaryMessage +import org.springframework.web.socket.TextMessage +import org.springframework.web.socket.WebSocketSession -class SpringBootServer implements HttpServer { +class SpringBootServer implements WebsocketServer { def port = 0 final SpringApplication app EmbeddedWebApplicationContext context String servletContext + WebsocketEndpoint endpoint SpringBootServer(SpringApplication app, String servletContext) { this.app = app @@ -20,6 +26,11 @@ class SpringBootServer implements HttpServer { app.setDefaultProperties(["server.port": 0, "server.context-path": "/$servletContext"]) context = app.run() as EmbeddedWebApplicationContext port = context.embeddedServletContainer.port + try { + endpoint = context.getBean(WebsocketEndpoint) + } catch (BeansException ignored) { + // silently ignore since not all the tests are deploying this endpoint + } assert port > 0 } @@ -37,4 +48,56 @@ class SpringBootServer implements HttpServer { String toString() { this.class.name } + + @Override + void serverSendText(String[] messages) { + WebSocketSession session = endpoint?.activeSession + if (session != null) { + if (messages.size() == 1) { + session.sendMessage(new TextMessage(messages[0])) + } else { + for (def i = 0; i < messages.size(); i++) { + session.sendMessage(new TextMessage(messages[i], i == messages.size() - 1)) + } + } + } + } + + @Override + void serverSendBinary(byte[][] binaries) { + WebSocketSession session = endpoint?.activeSession + if (session != null) { + if (binaries.length == 1) { + session.sendMessage(new BinaryMessage(binaries[0])) + } else { + for (def i = 0; i < binaries.length; i++) { + session.sendMessage(new BinaryMessage(binaries[i], i == binaries.length - 1)) + } + } + } + } + + @Override + void serverClose() { + endpoint?.activeSession?.close() + } + + @Override + synchronized void awaitConnected() { + synchronized (WebsocketEndpoint) { + try { + while (endpoint?.activeSession == null) { + WebsocketEndpoint.wait() + } + } catch (InterruptedException ie) { + Thread.currentThread().interrupt() + } + } + } + + @Override + void setMaxPayloadSize(int size) { + endpoint?.activeSession?.setBinaryMessageSizeLimit(size) + endpoint?.activeSession?.setTextMessageSizeLimit(size) + } } diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/WebsocketConfig.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/WebsocketConfig.groovy new file mode 100644 index 00000000000..cdd568e3023 --- /dev/null +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/WebsocketConfig.groovy @@ -0,0 +1,21 @@ +package test.boot + +import org.springframework.context.annotation.Bean +import org.springframework.context.annotation.Configuration +import org.springframework.web.socket.config.annotation.EnableWebSocket +import org.springframework.web.socket.config.annotation.WebSocketConfigurer +import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry + +@Configuration +@EnableWebSocket +class WebsocketConfig implements WebSocketConfigurer { + @Bean + WebsocketEndpoint websocketEndpoint() { + new WebsocketEndpoint() + } + + @Override + void registerWebSocketHandlers(WebSocketHandlerRegistry webSocketHandlerRegistry) { + webSocketHandlerRegistry.addHandler(websocketEndpoint(), "/websocket") + } +} diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/WebsocketEndpoint.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/WebsocketEndpoint.groovy new file mode 100644 index 00000000000..63307a6d635 --- /dev/null +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/WebsocketEndpoint.groovy @@ -0,0 +1,42 @@ +package test.boot + + +import org.springframework.web.socket.BinaryMessage +import org.springframework.web.socket.CloseStatus +import org.springframework.web.socket.TextMessage +import org.springframework.web.socket.WebSocketSession +import org.springframework.web.socket.handler.AbstractWebSocketHandler + +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace + +class WebsocketEndpoint extends AbstractWebSocketHandler { + volatile WebSocketSession activeSession + + @Override + boolean supportsPartialMessages() { + true + } + + @Override + void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception { + activeSession = null + } + + @Override + void afterConnectionEstablished(WebSocketSession session) throws Exception { + activeSession = session + synchronized (WebsocketEndpoint) { + WebsocketEndpoint.notifyAll() + } + } + + @Override + protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception { + runUnderTrace("onRead", {}) + } + + @Override + protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception { + runUnderTrace("onRead", {}) + } +} diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/build.gradle b/dd-java-agent/instrumentation/spring-webmvc-6.0/build.gradle index 99e2e6b2155..574ba859b3b 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-6.0/build.gradle +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/build.gradle @@ -45,6 +45,8 @@ dependencies { testRuntimeOnly project(':dd-java-agent:instrumentation:tomcat-5.5') testRuntimeOnly project(':dd-java-agent:instrumentation:tomcat-appsec-6') testRuntimeOnly project(':dd-java-agent:instrumentation:tomcat-appsec-7') + testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:javax-websocket-1.0') + testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:jakarta-websocket-2.0') testImplementation project(':dd-java-agent:instrumentation:servlet:request-5') testImplementation group: 'org.spockframework', name: 'spock-spring', version: libs.versions.spock.get() @@ -52,11 +54,13 @@ dependencies { testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-test', version: '3.0.0' testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-web', version: '3.0.0' testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-security', version: '3.0.0' + testImplementation group: 'org.springframework', name: 'spring-websocket', version: '6.0.2' testImplementation group: 'com.google.code.gson', name: 'gson', version: '+' latestDepTestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-test', version: '3.+' latestDepTestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-web', version: '3.+' latestDepTestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-security', version: '3.+' + latestDepTestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-websocket', version: '3.+' } diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/SpringBootBasedTest.groovy b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/SpringBootBasedTest.groovy index d3b2c80236b..102ce0497bf 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/SpringBootBasedTest.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/SpringBootBasedTest.groovy @@ -1,18 +1,9 @@ package datadog.trace.instrumentation.springweb6.boot -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.LOGIN -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.MATRIX_PARAM -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_HERE -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.PATH_PARAM -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.REDIRECT -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SECURE_SUCCESS -import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS - import datadog.trace.agent.test.asserts.TraceAssert import datadog.trace.agent.test.base.HttpServer import datadog.trace.agent.test.base.HttpServerTest +import datadog.trace.agent.test.base.WebsocketServer import datadog.trace.api.ConfigDefaults import datadog.trace.api.DDSpanTypes import datadog.trace.api.DDTags @@ -31,12 +22,27 @@ import okhttp3.FormBody import okhttp3.Request import okhttp3.RequestBody import okhttp3.Response +import org.springframework.beans.BeansException import org.springframework.boot.SpringApplication import org.springframework.boot.web.servlet.context.ServletWebServerApplicationContext import org.springframework.context.ConfigurableApplicationContext import org.springframework.web.servlet.HandlerInterceptor +import org.springframework.web.socket.BinaryMessage +import org.springframework.web.socket.TextMessage +import org.springframework.web.socket.WebSocketSession import spock.lang.Shared +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.LOGIN +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.MATRIX_PARAM +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_HERE +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.PATH_PARAM +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.REDIRECT +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SECURE_SUCCESS +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.WEBSOCKET + class SpringBootBasedTest extends HttpServerTest { @Override @@ -51,20 +57,27 @@ class SpringBootBasedTest extends HttpServerTest Map extraServerTags = [:] SpringApplication application() { - return new SpringApplication(SecurityConfig, TestController, AppConfig) + return new SpringApplication(SecurityConfig, TestController, AppConfig, WebsocketConfig) } - class SpringBootServer implements HttpServer { + class SpringBootServer implements WebsocketServer { def port = 0 final app = application() + WebsocketEndpoint endpoint + @Override void start() { - app.setDefaultProperties(["server.port": 0, "server.context-path": "/$servletContext", + app.setDefaultProperties(["server.port" : 0, "server.context-path": "/$servletContext", "spring.mvc.throw-exception-if-no-handler-found": false, - "spring.web.resources.add-mappings" : false]) + "spring.web.resources.add-mappings" : false]) context = app.run() port = (context as ServletWebServerApplicationContext).webServer.port + try { + endpoint = context.getBean(WebsocketEndpoint) + } catch (BeansException ignored) { + // silently ignore since not all the tests are deploying this endpoint + } assert port > 0 } @@ -82,6 +95,58 @@ class SpringBootBasedTest extends HttpServerTest String toString() { return this.class.name } + + @Override + void serverSendText(String[] messages) { + WebSocketSession session = endpoint?.activeSession + if (session != null) { + if (messages.size() == 1) { + session.sendMessage(new TextMessage(messages[0])) + } else { + for (def i = 0; i < messages.size(); i++) { + session.sendMessage(new TextMessage(messages[i], i == messages.size() - 1)) + } + } + } + } + + @Override + void serverSendBinary(byte[][] binaries) { + WebSocketSession session = endpoint?.activeSession + if (session != null) { + if (binaries.length == 1) { + session.sendMessage(new BinaryMessage(binaries[0])) + } else { + for (def i = 0; i < binaries.length; i++) { + session.sendMessage(new BinaryMessage(binaries[i], i == binaries.length - 1)) + } + } + } + } + + @Override + void serverClose() { + endpoint?.activeSession?.close() + } + + @Override + synchronized void awaitConnected() { + synchronized (WebsocketEndpoint) { + try { + while (endpoint?.activeSession == null) { + WebsocketEndpoint.wait() + } + } catch (InterruptedException ie) { + Thread.currentThread().interrupt() + } + } + } + + @Override + void setMaxPayloadSize(int size) { + endpoint?.activeSession?.setBinaryMessageSizeLimit(size) + endpoint?.activeSession?.setTextMessageSizeLimit(size) + } } def setupSpec() { @@ -155,12 +220,11 @@ class SpringBootBasedTest extends HttpServerTest @Override Map expectedExtraErrorInformation(ServerEndpoint endpoint) { - System.err.println("CALLED") // latest DispatcherServlet throws if no handlers have been found if (endpoint == NOT_FOUND && isLatestDepTest) { return [(DDTags.ERROR_STACK): { String }, - (DDTags.ERROR_MSG) : {String}, - (DDTags.ERROR_TYPE) :{String}] + (DDTags.ERROR_MSG) : { String }, + (DDTags.ERROR_TYPE) : { String }] } return super.expectedExtraErrorInformation(endpoint) } @@ -297,7 +361,7 @@ class SpringBootBasedTest extends HttpServerTest when: def response = client.newCall(request).execute() TEST_WRITER.waitForTraces(2) - DDSpan span = TEST_WRITER.flatten().find {it.operationName =='appsec-span' } + DDSpan span = TEST_WRITER.flatten().find { it.operationName == 'appsec-span' } then: response.code() == PATH_PARAM.status @@ -333,14 +397,14 @@ class SpringBootBasedTest extends HttpServerTest when: def response = client.newCall(request).execute() TEST_WRITER.waitForTraces(2) - DDSpan span = TEST_WRITER.flatten().find {it.operationName =='appsec-span' } + DDSpan span = TEST_WRITER.flatten().find { it.operationName == 'appsec-span' } then: response.code() == MATRIX_PARAM.status response.body().string() == MATRIX_PARAM.body //FIXME: tomcat seems removing the part after the ';' on the decodedUri // should be [var:['a=x,y;a=z' - span.getTag(IG_PATH_PARAMS_TAG) == [var:['a=x,y', [a:['x', 'y', 'z']]]] + span.getTag(IG_PATH_PARAMS_TAG) == [var: ['a=x,y', [a: ['x', 'y', 'z']]]] } void 'tainting on matrix var'() { @@ -384,7 +448,7 @@ class SpringBootBasedTest extends HttpServerTest when: client.newCall(request).execute() TEST_WRITER.waitForTraces(1) - DDSpan span = TEST_WRITER.flatten().find {"servlet.request".contentEquals(it.operationName)} + DDSpan span = TEST_WRITER.flatten().find { "servlet.request".contentEquals(it.operationName) } then: span.getResourceName().toString() == "GET " + testPathParam() @@ -407,7 +471,6 @@ class SpringBootBasedTest extends HttpServerTest } - @Override void handlerSpan(TraceAssert trace, ServerEndpoint endpoint = SUCCESS) { if (endpoint == NOT_FOUND) { @@ -417,7 +480,9 @@ class SpringBootBasedTest extends HttpServerTest serviceName expectedServiceName() operationName "spring.handler" resourceName { - it == "TestController.${endpoint.name().toLowerCase()}" || endpoint == NOT_FOUND && it == "ResourceHttpRequestHandler.handleRequest" + it == "TestController.${endpoint.name().toLowerCase()}" + || endpoint == NOT_FOUND && it == "ResourceHttpRequestHandler.handleRequest" + || endpoint == WEBSOCKET && it == "WebSocketHttpRequestHandler.handleRequest" } spanType DDSpanTypes.HTTP_SERVER errored endpoint == EXCEPTION diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/WebsocketConfig.groovy b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/WebsocketConfig.groovy new file mode 100644 index 00000000000..2a27cc2ffe5 --- /dev/null +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/WebsocketConfig.groovy @@ -0,0 +1,21 @@ +package datadog.trace.instrumentation.springweb6.boot + +import org.springframework.context.annotation.Bean +import org.springframework.context.annotation.Configuration +import org.springframework.web.socket.config.annotation.EnableWebSocket +import org.springframework.web.socket.config.annotation.WebSocketConfigurer +import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry + +@Configuration +@EnableWebSocket +class WebsocketConfig implements WebSocketConfigurer { + @Bean + WebsocketEndpoint websocketEndpoint() { + new WebsocketEndpoint() + } + + @Override + void registerWebSocketHandlers(WebSocketHandlerRegistry webSocketHandlerRegistry) { + webSocketHandlerRegistry.addHandler(websocketEndpoint(), "/websocket") + } +} diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/WebsocketEndpoint.groovy b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/WebsocketEndpoint.groovy new file mode 100644 index 00000000000..f707dc6617c --- /dev/null +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/WebsocketEndpoint.groovy @@ -0,0 +1,42 @@ +package datadog.trace.instrumentation.springweb6.boot + + +import org.springframework.web.socket.BinaryMessage +import org.springframework.web.socket.CloseStatus +import org.springframework.web.socket.TextMessage +import org.springframework.web.socket.WebSocketSession +import org.springframework.web.socket.handler.AbstractWebSocketHandler + +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace + +class WebsocketEndpoint extends AbstractWebSocketHandler { + volatile WebSocketSession activeSession + + @Override + boolean supportsPartialMessages() { + true + } + + @Override + void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception { + activeSession = null + } + + @Override + void afterConnectionEstablished(WebSocketSession session) throws Exception { + activeSession = session + synchronized (WebsocketEndpoint) { + WebsocketEndpoint.notifyAll() + } + } + + @Override + protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception { + runUnderTrace("onRead", {}) + } + + @Override + protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception { + runUnderTrace("onRead", {}) + } +} diff --git a/dd-java-agent/instrumentation/tomcat-5.5/build.gradle b/dd-java-agent/instrumentation/tomcat-5.5/build.gradle index 8ee94a412e3..a7fb77b990b 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/build.gradle +++ b/dd-java-agent/instrumentation/tomcat-5.5/build.gradle @@ -43,6 +43,8 @@ apply from: "$rootDir/gradle/java.gradle" addTestSuite('latestDepTest') addTestSuiteExtendingForDir('latestDepForkedTest', 'latestDepTest', 'latestDepTest') +addTestSuite("tomcat9Test") + addTestSuiteForDir('latest10Test', 'latestDepTest') addTestSuiteExtendingForDir('latest10ForkedTest', 'latest10Test', 'latestDepTest') @@ -56,12 +58,20 @@ configurations.all { } } +[compileLatest10TestGroovy, compileLatest10ForkedTestGroovy].each { + it.javaLauncher = getJavaLauncherFor(11) +} +[compileLatestDepTestGroovy, compileLatestDepForkedTestGroovy].each { + it.javaLauncher = getJavaLauncherFor(17) +} dependencies { compileOnly group: 'tomcat', name: 'catalina', version: tomcatVersion compileOnly group: 'tomcat', name: 'tomcat-coyote', version: tomcatVersion compileOnly group: 'tomcat', name: 'tomcat-util', version: tomcatVersion + compileOnly group: 'org.apache.tomcat', name: 'tomcat-websocket', version: '8.0.1', transitive: false + // Version that corresponds with Tomcat 5.5 // https://tomcat.apache.org/whichversion.html @@ -90,9 +100,16 @@ dependencies { testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-2') testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-3') 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') + testImplementation testFixtures(project(':dd-java-agent:instrumentation:servlet:request-5')) + tomcat9TestImplementation group: 'org.apache.tomcat.embed', name: 'tomcat-embed-core', version: '9.+' + tomcat9TestImplementation group: 'org.apache.tomcat.embed', name: 'tomcat-embed-websocket', version: '9.+' + latest10TestImplementation group: 'org.apache.tomcat.embed', name: 'tomcat-embed-core', version: '10.+' + latest10TestImplementation group: 'org.apache.tomcat.embed', name: 'tomcat-embed-websocket', version: '10.+' latest10TestImplementation group: 'org.apache.tomcat', name: 'jakartaee-migration', version: '1.+' latest10TestImplementation(testFixtures(project(":dd-java-agent:appsec"))) @@ -101,6 +118,7 @@ dependencies { latest10TestImplementation(project(':dd-java-agent:instrumentation:tomcat-classloading-9')) latestDepTestImplementation group: 'org.apache.tomcat.embed', name: 'tomcat-embed-core', version: '11.+' + latestDepTestImplementation group: 'org.apache.tomcat.embed', name: 'tomcat-embed-websocket', version: '11.+' latestDepTestImplementation group: 'org.apache.tomcat', name: 'jakartaee-migration', version: '1.+' latestDepTestImplementation(testFixtures(project(":dd-java-agent:appsec"))) @@ -109,16 +127,23 @@ dependencies { latestDepTestRuntimeOnly(project(':dd-java-agent:instrumentation:tomcat-classloading-9')) } -// Exclude all the dependencies from test for latestDepTest since the names are completely different. -configurations { - latestDepTestImplementation { - exclude group: 'tomcat', module: 'catalina' - exclude group: 'tomcat', module: 'tomcat-coyote' - exclude group: 'tomcat', module: 'tomcat-util' - exclude group: 'tomcat', module: 'tomcat-http' - exclude group: 'tomcat', module: 'naming-resources' - exclude group: 'tomcat', module: 'naming-factory' - exclude group: 'commons-modeler', module: 'commons-modeler' - exclude group: 'javax.servlet', module: 'servlet-api' +project.afterEvaluate { + tasks.withType(Test).configureEach { + if (javaLauncher.get().metadata.languageVersion.asInt() >= 16) { + jvmArgs += ['--add-opens', 'java.base/java.io=ALL-UNNAMED'] + } } } + +// Exclude all the dependencies from test for latestDepTest since the names are completely different. +configurations.findAll { it.name in ["latestDepTestImplementation", "tomcat9TestImplementation", "latest10TestImplementation"] } +.each { + it.exclude group: 'tomcat', module: 'catalina' + it.exclude group: 'tomcat', module: 'tomcat-coyote' + it.exclude group: 'tomcat', module: 'tomcat-util' + it.exclude group: 'tomcat', module: 'tomcat-http' + it.exclude group: 'tomcat', module: 'naming-resources' + it.exclude group: 'tomcat', module: 'naming-factory' + it.exclude group: 'commons-modeler', module: 'commons-modeler' + it.exclude group: 'javax.servlet', module: 'servlet-api' +} diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatAppSecInactiveHttpServerTest.groovy b/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatAppSecInactiveHttpServerForkedTest.groovy similarity index 90% rename from dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatAppSecInactiveHttpServerTest.groovy rename to dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatAppSecInactiveHttpServerForkedTest.groovy index a58f7bcf9be..e5de6e44ea1 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatAppSecInactiveHttpServerTest.groovy +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatAppSecInactiveHttpServerForkedTest.groovy @@ -8,7 +8,7 @@ import org.apache.catalina.Wrapper import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.UNKNOWN -class TomcatAppSecInactiveHttpServerTest extends AppSecInactiveHttpServerTest { +class TomcatAppSecInactiveHttpServerForkedTest extends AppSecInactiveHttpServerTest { HttpServer server() { new TomcatServer('tomcat-context', false, { Context ctx -> HttpServerTest.ServerEndpoint.values().findAll { @@ -21,6 +21,6 @@ class TomcatAppSecInactiveHttpServerTest extends AppSecInactiveHttpServerTest { ctx.addChild(wrapper) ctx.addServletMappingDecoded(it.path, wrapper.name) } - }) + }, {}) } } diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatNoPropagationForkedTest.groovy b/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatNoPropagationForkedTest.groovy index 66370856159..f9928fb95ea 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatNoPropagationForkedTest.groovy +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatNoPropagationForkedTest.groovy @@ -21,7 +21,7 @@ class TomcatNoPropagationForkedTest extends AgentTestRunner { wrapper.asyncSupported = true ctx.addChild(wrapper) ctx.addServletMappingDecoded("/*", wrapper.name) - }) + }, {}) @Override protected void configurePreAgent() { @@ -39,6 +39,7 @@ class TomcatNoPropagationForkedTest extends AgentTestRunner { server.stop() } + def "should not extract propagated context but collect headers"() { setup: def client = OkHttpUtils.client() diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatServer.groovy b/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatServer.groovy index 0edb4a2f4a4..2ac2491d37d 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatServer.groovy +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatServer.groovy @@ -1,20 +1,31 @@ import com.google.common.io.Files import datadog.trace.agent.test.base.HttpServer +import datadog.trace.agent.test.base.WebsocketServer +import jakarta.servlet.ServletContextEvent +import jakarta.servlet.ServletContextListener +import jakarta.websocket.Session +import jakarta.websocket.server.ServerContainer import org.apache.catalina.Context import org.apache.catalina.core.StandardHost import org.apache.catalina.startup.Tomcat import org.apache.tomcat.JarScanFilter import org.apache.tomcat.JarScanType +import org.apache.tomcat.websocket.server.WsSci -class TomcatServer implements HttpServer { +import java.nio.ByteBuffer + +class TomcatServer implements WebsocketServer { def port = 0 final Tomcat server final String context final boolean dispatch + volatile Session activeSession + final boolean wsAsyncSend - TomcatServer(String context, boolean dispatch, Closure setupServlets) { + TomcatServer(String context, boolean dispatch, Closure setupServlets, Closure setupWebsockets, boolean wsAsyncSend = false) { this.context = context this.dispatch = dispatch + this.wsAsyncSend = wsAsyncSend server = new Tomcat() def baseDir = Files.createTempDir() @@ -39,7 +50,10 @@ class TomcatServer implements HttpServer { } } setupServlets(servletContext) - + servletContext.addServletContainerInitializer(new WsSci(), null) + def listeners = new ArrayList(Arrays.asList(servletContext.getApplicationLifecycleListeners())) + listeners.add(new EndpointDeployer(setupWebsockets)) + servletContext.setApplicationLifecycleListeners(listeners.toArray()) (server.host as StandardHost).errorReportValveClass = TomcatServletTest.ErrorHandlerValve.name } @@ -75,4 +89,77 @@ class TomcatServer implements HttpServer { String toString() { return this.class.name } + + @Override + void serverSendText(String[] messages) { + if (wsAsyncSend && messages.length == 1) { // async does not support partial write + WsEndpoint.activeSession.getAsyncRemote().sendText(messages[0]) + } else { + if (messages.length == 1) { + WsEndpoint.activeSession.getBasicRemote().sendText(messages[0]) + } else { + def remoteEndpoint = WsEndpoint.activeSession.getBasicRemote() + for (int i = 0; i < messages.length; i++) { + remoteEndpoint.sendText(messages[i], i == messages.length - 1) + } + } + } + } + + @Override + void serverSendBinary(byte[][] binaries) { + if (wsAsyncSend && binaries.length == 1) { // async does not support partial write + WsEndpoint.activeSession.getAsyncRemote().sendBinary(ByteBuffer.wrap(binaries[0])) + } else { + if (binaries.length == 1) { + WsEndpoint.activeSession.getBasicRemote().sendBinary(ByteBuffer.wrap(binaries[0])) + } else { + try (def stream = WsEndpoint.activeSession.getBasicRemote().getSendStream()) { + binaries.each { stream.write(it) } + } + } + } + } + + @Override + synchronized void awaitConnected() { + synchronized (WsEndpoint) { + try { + while (WsEndpoint.activeSession == null) { + WsEndpoint.wait() + } + } catch (InterruptedException ie) { + Thread.currentThread().interrupt() + } + } + } + + @Override + void serverClose() { + WsEndpoint.activeSession.close() + WsEndpoint.activeSession = null + } + + @Override + void setMaxPayloadSize(int size) { + WsEndpoint.activeSession.setMaxTextMessageBufferSize(size) + WsEndpoint.activeSession.setMaxBinaryMessageBufferSize(size) + } + + static class EndpointDeployer implements ServletContextListener { + final Closure wsDeployCallback + + EndpointDeployer(Closure wsDeployCallback) { + this.wsDeployCallback = wsDeployCallback + } + + @Override + void contextInitialized(ServletContextEvent sce) { + wsDeployCallback.call((ServerContainer) sce.getServletContext().getAttribute(ServerContainer.class.getName())) + } + + @Override + void contextDestroyed(ServletContextEvent servletContextEvent) { + } + } } diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatServletTest.groovy b/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatServletTest.groovy index 5b3f05d15db..3d6b9db9fe8 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatServletTest.groovy +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/TomcatServletTest.groovy @@ -1,8 +1,12 @@ +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.WEBSOCKET + import datadog.trace.agent.test.base.HttpServer import datadog.trace.instrumentation.servlet5.TestServlet5 import jakarta.servlet.Filter import jakarta.servlet.Servlet import jakarta.servlet.ServletException +import jakarta.websocket.server.ServerContainer +import jakarta.websocket.server.ServerEndpointConfig import org.apache.catalina.Context import org.apache.catalina.Wrapper import org.apache.catalina.connector.Request @@ -67,8 +71,8 @@ class TomcatServletTest extends AbstractServletTest { if (endpoint.throwsException) { // Exception classes get wrapped in ServletException ["error.message": { endpoint == EXCEPTION ? "Servlet execution threw an exception" : it == endpoint.body }, - "error.type": { it == ServletException.name || it == InputMismatchException.name }, - "error.stack": String] + "error.type" : { it == ServletException.name || it == InputMismatchException.name }, + "error.stack" : String] } else { Collections.emptyMap() } @@ -95,7 +99,7 @@ class TomcatServletTest extends AbstractServletTest { @Override HttpServer server() { - new TomcatServer(context, dispatch, this.&setupServlets) + new TomcatServer(context, dispatch, this.&setupServlets, this.&setupWebsockets, isWebsocketAsyncSend()) } @Override @@ -108,7 +112,7 @@ class TomcatServletTest extends AbstractServletTest { Wrapper wrapper = servletContext.createWrapper() wrapper.name = UUID.randomUUID() wrapper.servletClass = servlet.name - wrapper.asyncSupported =true + wrapper.asyncSupported = true servletContext.addChild(wrapper) servletContext.addServletMappingDecoded(path, wrapper.name) } @@ -131,6 +135,15 @@ class TomcatServletTest extends AbstractServletTest { TestServlet5 } + protected boolean isWebsocketAsyncSend() { + false + } + + protected void setupWebsockets(ServerContainer serverContainer) { + serverContainer.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.StdPartialEndpoint, + WEBSOCKET.path).build()) + } + def "test exception with custom status"() { setup: assumeTrue(testException()) @@ -212,16 +225,27 @@ class TomcatServletClassloaderNamingForkedTest extends TomcatServletTest { // will not set the service name according to the servlet context value injectSysConfig("trace.experimental.jee.split-by-deployment", "true") } + + @Override + protected boolean isWebsocketAsyncSend() { + true + } + + @Override + protected void setupWebsockets(ServerContainer serverContainer) { + serverContainer.addEndpoint(WsEndpoint.PojoEndpoint) + } } class TomcatServletEnvEntriesTagTest extends TomcatServletTest { - def addEntry (context, name, value) { + def addEntry(context, name, value) { def envEntry = new ContextEnvironment() envEntry.setName(name) envEntry.setValue(value) envEntry.setType("java.lang.String") context.getNamingResources().addEnvironment(envEntry) } + @Override protected void setupServlets(Context context) { super.setupServlets(context) @@ -238,6 +262,11 @@ class TomcatServletEnvEntriesTagTest extends TomcatServletTest { Map expectedExtraServerTags(ServerEndpoint endpoint) { super.expectedExtraServerTags(endpoint) + ["custom-tag": "custom-value"] as Map } + + @Override + boolean testWebsockets() { + false + } } diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/WsEndpoint.groovy b/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/WsEndpoint.groovy new file mode 100644 index 00000000000..3500c03d1b3 --- /dev/null +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/latestDepTest/groovy/WsEndpoint.groovy @@ -0,0 +1,76 @@ +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 java.nio.ByteBuffer + +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace + +class WsEndpoint { + static volatile Session activeSession = null + + static void logReadSpan() { + runUnderTrace("onRead", {}) + } + + + static class StdPartialEndpoint extends Endpoint { + @Override + void onOpen(Session session, EndpointConfig endpointConfig) { + session.addMessageHandler(new MessageHandler.Partial() { + @Override + void onMessage(String s, boolean last) { + logReadSpan() + } + }) + session.addMessageHandler(new MessageHandler.Partial() { + @Override + void onMessage(byte[] buffer, boolean last) { + logReadSpan() + } + }) + activeSession = session + synchronized (WsEndpoint) { + WsEndpoint.notifyAll() + } + } + + @Override + void onClose(Session session, CloseReason closeReason) { + activeSession = null + } + } + + @ServerEndpoint("/websocket") + static class PojoEndpoint { + @OnOpen + void onOpen(Session session, EndpointConfig endpointConfig) { + activeSession = session + synchronized (WsEndpoint) { + WsEndpoint.notifyAll() + } + } + + @OnClose + void onClose(Session session, CloseReason closeReason) { + activeSession = null + } + + @OnMessage + void onTextMessage(String message, boolean last, Session session) { + logReadSpan() + } + + @OnMessage + void onBinaryMessage(ByteBuffer message, boolean last, Session session) { + logReadSpan() + } + } +} + diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java new file mode 100644 index 00000000000..7ef56a13942 --- /dev/null +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java @@ -0,0 +1,52 @@ +package datadog.trace.instrumentation.tomcat; + +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static net.bytebuddy.matcher.ElementMatchers.isConstructor; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.ClassloaderConfigurationOverrides; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.util.Collections; +import java.util.Map; +import net.bytebuddy.asm.Advice; +import org.apache.tomcat.websocket.server.WsHandshakeRequest; + +@AutoService(InstrumenterModule.class) +public class WsHandshakeRequestInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + public WsHandshakeRequestInstrumentation() { + super("tomcat", "tomcat-websocket", "websocket"); + } + + @Override + public String instrumentedType() { + return "org.apache.tomcat.websocket.server.WsHandshakeRequest"; + } + + @Override + public Map contextStore() { + return Collections.singletonMap(instrumentedType(), AgentSpan.class.getName()); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice(isConstructor(), getClass().getName() + "$CaptureHandshakeSpanAdvice"); + } + + public static class CaptureHandshakeSpanAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void captureHandshakeSpan(@Advice.This final WsHandshakeRequest self) { + final AgentSpan span = activeSpan(); + if (span != null) { + // apply jee configuration overrides if any since the servlet instrumentation won't kick in + // for this span. + ClassloaderConfigurationOverrides.maybeEnrichSpan(span); + InstrumentationContext.get(WsHandshakeRequest.class, AgentSpan.class) + .putIfAbsent(self, span); + } + } + } +} diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHttpUpgradeHandlerInstrumentation.java b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHttpUpgradeHandlerInstrumentation.java new file mode 100644 index 00000000000..2374c18099e --- /dev/null +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHttpUpgradeHandlerInstrumentation.java @@ -0,0 +1,56 @@ +package datadog.trace.instrumentation.tomcat; + +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static net.bytebuddy.matcher.ElementMatchers.named; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.util.Collections; +import java.util.Map; +import net.bytebuddy.asm.Advice; +import org.apache.tomcat.websocket.server.WsHandshakeRequest; + +@AutoService(InstrumenterModule.class) +public class WsHttpUpgradeHandlerInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + public WsHttpUpgradeHandlerInstrumentation() { + super("tomcat", "tomcat-websocket", "websocket"); + } + + @Override + public String instrumentedType() { + return "org.apache.tomcat.websocket.server.WsHttpUpgradeHandler"; + } + + @Override + public Map contextStore() { + return Collections.singletonMap( + "org.apache.tomcat.websocket.server.WsHandshakeRequest", AgentSpan.class.getName()); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice(named("init"), getClass().getName() + "$CaptureHandshakeSpanAdvice"); + } + + public static class CaptureHandshakeSpanAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope before( + @Advice.FieldValue("handshakeRequest") final WsHandshakeRequest request) { + final AgentSpan span = + InstrumentationContext.get(WsHandshakeRequest.class, AgentSpan.class).get(request); + return span != null ? activateSpan(span) : null; + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void after(@Advice.Enter final AgentScope scope) { + if (scope != null) { + scope.close(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/test/groovy/TomcatServletTest.groovy b/dd-java-agent/instrumentation/tomcat-5.5/src/test/groovy/TomcatServletTest.groovy index c7f80b84a67..629b41cf1e8 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/test/groovy/TomcatServletTest.groovy +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/test/groovy/TomcatServletTest.groovy @@ -310,7 +310,7 @@ abstract class TomcatServletTest extends AbstractServletTest } } -class TomcatServletV0ForkedTest extends TomcatServletTest implements TestingGenericHttpNamingConventions.ServerV0 { +class TomcatServletV0Test extends TomcatServletTest implements TestingGenericHttpNamingConventions.ServerV0 { } diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/tomcat9Test/groovy/TomcatServer.groovy b/dd-java-agent/instrumentation/tomcat-5.5/src/tomcat9Test/groovy/TomcatServer.groovy new file mode 100644 index 00000000000..99118a4d222 --- /dev/null +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/tomcat9Test/groovy/TomcatServer.groovy @@ -0,0 +1,161 @@ +import com.google.common.io.Files +import datadog.trace.agent.test.base.WebsocketServer +import org.apache.catalina.Context +import org.apache.catalina.core.StandardHost +import org.apache.catalina.startup.Tomcat +import org.apache.tomcat.JarScanFilter +import org.apache.tomcat.JarScanType +import org.apache.tomcat.websocket.server.WsSci + +import javax.servlet.ServletContextEvent +import javax.servlet.ServletContextListener +import javax.websocket.server.ServerContainer +import java.nio.ByteBuffer +import java.util.concurrent.CountDownLatch + +class TomcatServer implements WebsocketServer { + + def port = 0 + final Tomcat server + final String context + final boolean wsAsyncSend + + TomcatServer(String context, Closure setupServlets, Closure setupWebsockets, boolean wsAsyncSend = false) { + this.context = context + this.wsAsyncSend = wsAsyncSend + server = new Tomcat() + + def baseDir = Files.createTempDir() + baseDir.deleteOnExit() + server.basedir = baseDir.absolutePath + + server.port = 0 // select random open port + server.connector.enableLookups = true // get localhost instead of 127.0.0.1 + + final File applicationDir = new File(baseDir, "/webapps/ROOT") + if (!applicationDir.exists()) { + applicationDir.mkdirs() + applicationDir.deleteOnExit() + } + Context servletContext = server.addWebapp("/$context", applicationDir.getAbsolutePath()) + servletContext.allowCasualMultipartParsing = true + // Speed up startup by disabling jar scanning: + servletContext.jarScanner.jarScanFilter = new JarScanFilter() { + @Override + boolean check(JarScanType jarScanType, String jarName) { + return false + } + } + setupServlets(servletContext) + servletContext.addServletContainerInitializer(new WsSci(), null) + def listeners = new ArrayList(Arrays.asList(servletContext.getApplicationLifecycleListeners())) + listeners.add(new EndpointDeployer(setupWebsockets)) + servletContext.setApplicationLifecycleListeners(EndpointDeployer.class.getName()) + servletContext.setApplicationLifecycleListeners(listeners.toArray()) + + (server.host as StandardHost).errorReportValveClass = TomcatWebsocketTest.ErrorHandlerValve.name + } + + @Override + void start() { + server.start() + port = server.service.findConnectors()[0].localPort + assert port > 0 + } + + @Override + void stop() { + Thread.start { + sleep 50 + // tomcat doesn't seem to interrupt accept() on stop() + // so connect to force the loop to continue + def sock = new Socket('localhost', port) + sock.close() + } + server.stop() + server.destroy() + } + + @Override + URI address() { + return new URI("http://localhost:$port/$context/") + } + + @Override + String toString() { + return this.class.name + } + + @Override + synchronized void awaitConnected() { + synchronized (WsEndpoint) { + try { + while (WsEndpoint.activeSession == null) { + WsEndpoint.wait() + } + } catch (InterruptedException ie) { + Thread.currentThread().interrupt() + } + } + } + + @Override + void serverSendText(String[] messages) { + if (wsAsyncSend && messages.length == 1) { // async does not support partial write + WsEndpoint.activeSession.getAsyncRemote().sendText(messages[0]) + } else { + if (messages.length == 1) { + WsEndpoint.activeSession.getBasicRemote().sendText(messages[0]) + } else { + def remoteEndpoint = WsEndpoint.activeSession.getBasicRemote() + for (int i = 0; i < messages.length; i++) { + remoteEndpoint.sendText(messages[i], i == messages.length - 1) + } + } + } + } + + @Override + void serverSendBinary(byte[][] binaries) { + if (wsAsyncSend && binaries.length == 1) { // async does not support partial write + WsEndpoint.activeSession.getAsyncRemote().sendBinary(ByteBuffer.wrap(binaries[0])) + } else { + if (binaries.length == 1) { + WsEndpoint.activeSession.getBasicRemote().sendBinary(ByteBuffer.wrap(binaries[0])) + } else { + try (def stream = WsEndpoint.activeSession.getBasicRemote().getSendStream()) { + binaries.each { stream.write(it) } + } + } + } + } + + @Override + void serverClose() { + WsEndpoint.activeSession.close() + WsEndpoint.activeSession = null + } + + @Override + void setMaxPayloadSize(int size) { + WsEndpoint.activeSession.setMaxTextMessageBufferSize(size) + WsEndpoint.activeSession.setMaxBinaryMessageBufferSize(size) + } + + static class EndpointDeployer implements ServletContextListener { + final Closure wsDeployCallback + + EndpointDeployer(Closure wsDeployCallback) { + this.wsDeployCallback = wsDeployCallback + } + + @Override + void contextInitialized(ServletContextEvent sce) { + wsDeployCallback.call((ServerContainer) sce.getServletContext().getAttribute(ServerContainer.class.getName())) + } + + @Override + void contextDestroyed(ServletContextEvent servletContextEvent) { + } + } +} diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/tomcat9Test/groovy/TomcatWebsocketTest.groovy b/dd-java-agent/instrumentation/tomcat-5.5/src/tomcat9Test/groovy/TomcatWebsocketTest.groovy new file mode 100644 index 00000000000..9846d02df25 --- /dev/null +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/tomcat9Test/groovy/TomcatWebsocketTest.groovy @@ -0,0 +1,374 @@ +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.CUSTOM_EXCEPTION +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ERROR +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.REDIRECT +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.TIMEOUT_ERROR +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.UNKNOWN +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.WEBSOCKET +import static org.junit.Assume.assumeTrue + +import datadog.trace.agent.test.asserts.TraceAssert +import datadog.trace.agent.test.base.HttpServer +import datadog.trace.agent.test.base.HttpServerTest +import datadog.trace.api.DDTags +import datadog.trace.instrumentation.servlet3.TestServlet3 +import datadog.trace.instrumentation.tomcat.TomcatDecorator +import okhttp3.Interceptor +import okhttp3.OkHttpClient +import org.apache.catalina.Context +import org.apache.catalina.Wrapper +import org.apache.catalina.connector.Request +import org.apache.catalina.connector.Response +import org.apache.catalina.startup.Tomcat +import org.apache.catalina.valves.ErrorReportValve +import org.apache.tomcat.util.descriptor.web.FilterDef +import org.apache.tomcat.util.descriptor.web.FilterMap + +import javax.servlet.Filter +import javax.servlet.FilterChain +import javax.servlet.FilterConfig +import javax.servlet.Servlet +import javax.servlet.ServletException +import javax.servlet.ServletRequest +import javax.servlet.ServletResponse +import javax.servlet.http.HttpServletRequest +import javax.servlet.http.HttpServletRequestWrapper +import javax.websocket.server.ServerContainer +import javax.websocket.server.ServerEndpointConfig +import java.security.Principal + +abstract class TomcatWebsocketTest extends HttpServerTest { + + @Override + boolean hasExtraErrorInformation() { + true + } + + @Override + boolean testBodyUrlencoded() { + true + } + + @Override + boolean testBlocking() { + false + } + + @Override + boolean testRequestBody() { + true + } + + @Override + boolean testRequestBodyISVariant() { + true + } + + @Override + boolean testBodyMultipart() { + false + } + + @Override + boolean testBlockingOnResponse() { + false + } + + @Override + boolean testSessionId() { + true + } + + @Override + boolean testEncodedPath() { + false + } + + @Override + boolean testEncodedQuery() { + false + } + + boolean hasResponseSpan(ServerEndpoint endpoint) { + def responseSpans = [REDIRECT, NOT_FOUND, ERROR] + return responseSpans.contains(endpoint) + } + + @Override + void responseSpan(TraceAssert trace, ServerEndpoint endpoint) { + switch (endpoint) { + case REDIRECT: + trace.span { + operationName "servlet.response" + resourceName "HttpServletResponse.sendRedirect" + childOfPrevious() + tags { + "component" "java-web-servlet-response" + if ({ isDataStreamsEnabled() }) { + "$DDTags.PATHWAY_HASH" { String } + } + defaultTags() + } + } + break + case ERROR: + case NOT_FOUND: + trace.span { + operationName "servlet.response" + resourceName "HttpServletResponse.sendError" + childOfPrevious() + tags { + "component" "java-web-servlet-response" + defaultTags() + } + } + break + default: + throw new UnsupportedOperationException("responseSpan not implemented for " + endpoint) + } + } + + @Override + URI buildAddress(int port) { + return new URI("http://localhost:$port/$context/") + } + + @Override + String component() { + return TomcatDecorator.DECORATE.component() + } + + @Override + String expectedServiceName() { + context + } + + @Override + String expectedOperationName() { + return "servlet.request" + } + + boolean hasHandlerSpan() { + return false + } + + + @Override + OkHttpClient getClient() { + return super.getClient().newBuilder() + .addInterceptor(new Interceptor() { + @Override + okhttp3.Response intercept(Interceptor.Chain chain) throws IOException { + return chain.proceed(chain.request().newBuilder() + .header("Cookie", "somethingcouldbreaktherequest=" + UUID.randomUUID()) + + .build()) + } + }).build() + } + + static class SecurityFilter implements Filter { + + @Override + void init(FilterConfig filterConfig) throws ServletException { + } + + @Override + void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException { + chain.doFilter(new HttpServletRequestWrapper((HttpServletRequest) request) { + @Override + Principal getUserPrincipal() { + return new Principal() { + @Override + String getName() { + return "superadmin" + } + } + } + }, response) + } + + @Override + void destroy() { + } + } + + protected void setupServlets(Context context) { + ServerEndpoint.values().findAll { it != NOT_FOUND && it != UNKNOWN }.each { + addServlet(context, it.path, TestServlet3.Sync) + } + addFilter(context, "/*", SecurityFilter) + } + + /** + * Use async send + * @return + */ + protected boolean isWebsocketAsyncSend() { + false + } + + protected abstract void setupWebsockets(ServerContainer serverContainer) + + + @Override + Map expectedExtraErrorInformation(ServerEndpoint endpoint) { + if (endpoint.throwsException) { + // Exception classes get wrapped in ServletException + ["error.message": { endpoint == EXCEPTION ? "Servlet execution threw an exception" : it == endpoint.body }, + "error.type" : { it == ServletException.name || it == InputMismatchException.name }, + "error.stack" : String] + } else { + Collections.emptyMap() + } + } + + @Override + Map expectedExtraServerTags(ServerEndpoint endpoint) { + Map map = ["servlet.path": endpoint.path] + if (context) { + map.put("servlet.context", "/$context") + } + map + } + + @Override + boolean expectedErrored(ServerEndpoint endpoint) { + (endpoint.errored && bubblesResponse()) || [EXCEPTION, CUSTOM_EXCEPTION, TIMEOUT_ERROR].contains(endpoint) + } + + @Override + Serializable expectedStatus(ServerEndpoint endpoint) { + return { !bubblesResponse() || it == endpoint.status } + } + + @Override + HttpServer server() { + new TomcatServer(context, this.&setupServlets, this.&setupWebsockets, isWebsocketAsyncSend()) + } + + static String getContext() { + return "tomcat-context" + } + + static void addServlet(Context servletContext, String path, Class servlet) { + Wrapper wrapper = servletContext.createWrapper() + wrapper.name = UUID.randomUUID() + wrapper.servletClass = servlet.name + wrapper.asyncSupported = true + servletContext.addChild(wrapper) + servletContext.addServletMappingDecoded(path, wrapper.name) + } + + static void addFilter(Context context, String path, Class filter) { + def filterDef = new FilterDef() + def filterMap = new FilterMap() + filterDef.filterClass = filter.getName() + filterDef.asyncSupported = true + filterDef.filterName = UUID.randomUUID() + filterMap.filterName = filterDef.filterName + filterMap.addURLPattern(path) + context.addFilterDef(filterDef) + context.addFilterMap(filterMap) + } + + def "test exception with custom status"() { + setup: + assumeTrue(testException()) + def request = request(CUSTOM_EXCEPTION, method, body).build() + def response = client.newCall(request).execute() + + expect: + response.code() == CUSTOM_EXCEPTION.status + if (testExceptionBody()) { + assert response.body().string() == CUSTOM_EXCEPTION.body + } + + and: + assertTraces(1) { + trace(spanCount(CUSTOM_EXCEPTION)) { + sortSpansByStart() + serverSpan(it, null, null, method, CUSTOM_EXCEPTION) + if (hasHandlerSpan()) { + handlerSpan(it, CUSTOM_EXCEPTION) + } + controllerSpan(it, CUSTOM_EXCEPTION) + if (hasResponseSpan(CUSTOM_EXCEPTION)) { + responseSpan(it, CUSTOM_EXCEPTION) + } + } + } + + where: + method = "GET" + body = null + } + + def "test user principal extracted"() { + setup: + injectSysConfig("trace.servlet.principal.enabled", "true") + when: + def request = request(SUCCESS, "GET", null).build() + def response = client.newCall(request).execute() + then: + response.code() == 200 + and: + assertTraces(1) { + trace(2) { + serverSpan(it, null, null, "GET", SUCCESS, ["user.principal": "superadmin"]) + controllerSpan(it) + } + } + } + + static class ErrorHandlerValve extends ErrorReportValve { + @Override + protected void report(Request request, Response response, Throwable t) { + if (!response.error) { + return + } + try { + if (t) { + if (t instanceof ServletException) { + t = t.rootCause + } + if (t instanceof InputMismatchException) { + response.status = CUSTOM_EXCEPTION.status + } + response.reporter.write(t.message) + } else if (response.message) { + response.reporter.write(response.message) + } + } catch (IOException e) { + e.printStackTrace() + } + } + } +} + + +class TomcatWebsocketSyncPartialTest extends TomcatWebsocketTest { + @Override + protected void setupWebsockets(ServerContainer serverContainer) { + serverContainer.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.StdPartialEndpoint, WEBSOCKET.path).build()) + } +} + +class TomcatWebsocketAsyncPojoTest extends TomcatWebsocketTest { + @Override + protected boolean isWebsocketAsyncSend() { + true + } + + @Override + protected void setupWebsockets(ServerContainer serverContainer) { + serverContainer.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.PojoEndpoint, WEBSOCKET.path).build()) + } +} + + + + + diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/tomcat9Test/groovy/WsEndpoint.groovy b/dd-java-agent/instrumentation/tomcat-5.5/src/tomcat9Test/groovy/WsEndpoint.groovy new file mode 100644 index 00000000000..f81a83a1ced --- /dev/null +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/tomcat9Test/groovy/WsEndpoint.groovy @@ -0,0 +1,79 @@ +import javax.websocket.CloseReason +import javax.websocket.Endpoint +import javax.websocket.EndpointConfig +import javax.websocket.MessageHandler +import javax.websocket.OnClose +import javax.websocket.OnMessage +import javax.websocket.OnOpen +import javax.websocket.Session +import java.nio.ByteBuffer + +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace + +class WsEndpoint { + static volatile Session activeSession = null + + static void logReadSpan() { + runUnderTrace("onRead", {}) + } + + abstract static class StdEndpoint extends Endpoint { + @Override + void onOpen(Session session, EndpointConfig endpointConfig) { + activeSession = session + synchronized (WsEndpoint) { + WsEndpoint.notifyAll() + } + } + + @Override + void onClose(Session session, CloseReason closeReason) { + activeSession = null + } + } + + static class StdPartialEndpoint extends StdEndpoint { + @Override + void onOpen(Session session, EndpointConfig endpointConfig) { + session.addMessageHandler(new MessageHandler.Partial() { + @Override + void onMessage(String s, boolean last) { + logReadSpan() + } + }) + session.addMessageHandler(new MessageHandler.Partial() { + @Override + void onMessage(byte[] buffer, boolean last) { + logReadSpan() + } + }) + super.onOpen(session, endpointConfig) + } + } + + static class PojoEndpoint { + @OnOpen + void onOpen(Session session, EndpointConfig endpointConfig) { + activeSession = session + synchronized (WsEndpoint) { + WsEndpoint.notifyAll() + } + } + + @OnClose + void onClose(Session session, CloseReason closeReason) { + activeSession = null + } + + @OnMessage + void onTextMessage(String message, boolean last, Session session) { + logReadSpan() + } + + @OnMessage + void onBinaryMessage(ByteBuffer message, boolean last, Session session) { + logReadSpan() + } + } +} + diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/build.gradle b/dd-java-agent/instrumentation/undertow/undertow-2.0/build.gradle index 142b5519e99..be782a4591a 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/build.gradle +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/build.gradle @@ -35,6 +35,8 @@ dependencies { implementation project(':dd-java-agent:instrumentation:undertow') testImplementation group: 'io.undertow', name: 'undertow-servlet', version: '2.0.0.Final' + testImplementation group: 'io.undertow', name: 'undertow-websockets-jsr', version: '2.0.0.Final' + testRuntimeOnly project(':dd-java-agent:instrumentation:servlet') testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-2') @@ -42,4 +44,5 @@ dependencies { testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-5') latestDepTestImplementation group: 'io.undertow', name: 'undertow-servlet', version: '2.2.+' + latestDepTestImplementation group: 'io.undertow', name: 'undertow-websockets-jsr', version: '2.2.+' } diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/TestEndpoint.groovy b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/TestEndpoint.groovy new file mode 100644 index 00000000000..91c3c4278b6 --- /dev/null +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/TestEndpoint.groovy @@ -0,0 +1,33 @@ +import javax.websocket.OnClose +import javax.websocket.OnMessage +import javax.websocket.OnOpen +import javax.websocket.Session +import javax.websocket.server.ServerEndpoint + +@ServerEndpoint("/websocket") +class TestEndpoint { + static volatile Session activeSession + + @OnOpen + void onOpen(Session session) { + activeSession = session + synchronized (TestEndpoint) { + TestEndpoint.notifyAll() + } + } + + @OnClose + void onClose() { + activeSession = null + } + + @OnMessage + void onMessage(String s, boolean last) { + datadog.trace.agent.test.utils.TraceUtils.runUnderTrace("onRead", {}) + } + + @OnMessage + void onMessage(byte[] b, boolean last) { + datadog.trace.agent.test.utils.TraceUtils.runUnderTrace("onRead", {}) + } +} diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy index b4745e081dd..4fa8a555dbc 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy @@ -1,6 +1,6 @@ import datadog.trace.agent.test.asserts.TraceAssert -import datadog.trace.agent.test.base.HttpServer import datadog.trace.agent.test.base.HttpServerTest +import datadog.trace.agent.test.base.WebsocketServer import datadog.trace.agent.test.naming.TestingGenericHttpNamingConventions import datadog.trace.bootstrap.instrumentation.api.Tags import io.undertow.Handlers @@ -10,9 +10,11 @@ import io.undertow.servlet.api.DeploymentInfo import io.undertow.servlet.api.DeploymentManager import io.undertow.servlet.api.ServletContainer import io.undertow.servlet.api.ServletInfo +import io.undertow.websockets.jsr.WebSocketDeploymentInfo import spock.lang.IgnoreIf import javax.servlet.MultipartConfigElement +import java.nio.ByteBuffer import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_MULTIPART import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_URLENCODED @@ -35,7 +37,7 @@ import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.USER_B abstract class UndertowServletTest extends HttpServerTest { private static final CONTEXT = "ctx" - class UndertowServer implements HttpServer { + class UndertowServer implements WebsocketServer { def port = 0 Undertow undertowServer @@ -62,6 +64,7 @@ abstract class UndertowServletTest extends HttpServerTest { .addServlet(new ServletInfo("CreatedISServlet", CreatedISServlet).addMapping(CREATED_IS.path)) .addServlet(new ServletInfo("BodyUrlEncodedServlet", BodyUrlEncodedServlet).addMapping(BODY_URLENCODED.path)) .addServlet(new ServletInfo("BodyMultipartServlet", BodyMultipartServlet).addMapping(BODY_MULTIPART.path)) + .addServletContextAttribute(WebSocketDeploymentInfo.ATTRIBUTE_NAME, new WebSocketDeploymentInfo().addEndpoint(TestEndpoint)) DeploymentManager manager = container.addDeployment(builder) manager.deploy() @@ -90,6 +93,56 @@ abstract class UndertowServletTest extends HttpServerTest { URI address() { return new URI("http://localhost:$port/$CONTEXT/") } + + @Override + void awaitConnected() { + while (TestEndpoint.activeSession == null) { + synchronized (TestEndpoint) { + TestEndpoint.wait() + } + } + } + + @Override + void serverSendText(String[] messages) { + if (messages.length == 1) { + TestEndpoint.activeSession.getBasicRemote().sendText(messages[0]) + } else { + def remoteEndpoint = TestEndpoint.activeSession.getBasicRemote() + for (int i = 0; i < messages.length; i++) { + remoteEndpoint.sendText(messages[i], i == messages.length - 1) + } + } + } + + @Override + boolean canSplitLargeWebsocketPayloads() { + false + } + + @Override + void serverSendBinary(byte[][] binaries) { + if (binaries.length == 1) { + TestEndpoint.activeSession.getBasicRemote().sendBinary(ByteBuffer.wrap(binaries[0])) + } else { + try (def stream = TestEndpoint.activeSession.getBasicRemote().getSendStream()) { + binaries.each { + stream.write(it) + } + } + } + } + + @Override + void serverClose() { + TestEndpoint.activeSession?.close() + } + + @Override + void setMaxPayloadSize(int size) { + TestEndpoint.activeSession?.setMaxTextMessageBufferSize(size) + TestEndpoint.activeSession?.setMaxBinaryMessageBufferSize(size) + } } @Override @@ -214,11 +267,11 @@ abstract class UndertowServletTest extends HttpServerTest { switch (endpoint) { case LOGIN: case NOT_FOUND: - return null + return null case PATH_PARAM: - return testPathParam() + return testPathParam() default: - return endpoint.path + return endpoint.path } } @@ -271,6 +324,7 @@ abstract class UndertowServletTest extends HttpServerTest { class UndertowServletV0Test extends UndertowServletTest implements TestingGenericHttpNamingConventions.ServerV0 { } + class UndertowServletNoHttpRouteForkedTest extends UndertowServletTest implements TestingGenericHttpNamingConventions.ServerV0 { @Override def generateHttpRoute() { diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.2/build.gradle b/dd-java-agent/instrumentation/undertow/undertow-2.2/build.gradle index 18eac6b59ad..9375c69892b 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.2/build.gradle +++ b/dd-java-agent/instrumentation/undertow/undertow-2.2/build.gradle @@ -29,15 +29,22 @@ dependencies { implementation project(':dd-java-agent:instrumentation:undertow') testImplementation group: 'io.undertow', name: 'undertow-servlet-jakarta', version: '2.2.14.Final' + testImplementation group: 'io.undertow', name: 'undertow-websockets-jsr-jakarta', version: '2.2.14.Final' testRuntimeOnly project(':dd-java-agent:instrumentation:undertow:undertow-2.0') testRuntimeOnly project(':dd-java-agent:instrumentation:servlet') testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-2') testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-3') 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') latest22TestImplementation group: 'io.undertow', name: 'undertow-servlet-jakarta', version: '2.2.+' + latest22TestImplementation group: 'io.undertow', name: 'undertow-websockets-jsr-jakarta', version: '2.2.+' latestDepTestImplementation group: 'io.undertow', name: 'undertow-servlet', version: '+', { exclude group: 'io.undertow', module: 'undertow-servlet-jakarta' } + latestDepTestImplementation group: 'io.undertow', name: 'undertow-websockets-jsr', version: '+', { + exclude group: 'io.undertow', module: 'undertow-websocket-jsr-jakarta' + } } diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.2/src/test/groovy/TestEndpoint.groovy b/dd-java-agent/instrumentation/undertow/undertow-2.2/src/test/groovy/TestEndpoint.groovy new file mode 100644 index 00000000000..e858a2e4f5a --- /dev/null +++ b/dd-java-agent/instrumentation/undertow/undertow-2.2/src/test/groovy/TestEndpoint.groovy @@ -0,0 +1,35 @@ +import jakarta.websocket.OnClose +import jakarta.websocket.OnMessage +import jakarta.websocket.OnOpen +import jakarta.websocket.Session +import jakarta.websocket.server.ServerEndpoint + +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace + +@ServerEndpoint("/websocket") +class TestEndpoint { + static volatile Session activeSession + + @OnOpen + void onOpen(Session session) { + activeSession = session + synchronized (TestEndpoint) { + TestEndpoint.notifyAll() + } + } + + @OnClose + void onClose() { + activeSession = null + } + + @OnMessage + void onMessage(String s, boolean last) { + runUnderTrace("onRead", {}) + } + + @OnMessage + void onMessage(byte[] b, boolean last) { + runUnderTrace("onRead", {}) + } +} diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.2/src/test/groovy/UndertowServletTest.groovy b/dd-java-agent/instrumentation/undertow/undertow-2.2/src/test/groovy/UndertowServletTest.groovy index dd3769209b8..a620835433e 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.2/src/test/groovy/UndertowServletTest.groovy +++ b/dd-java-agent/instrumentation/undertow/undertow-2.2/src/test/groovy/UndertowServletTest.groovy @@ -1,6 +1,6 @@ import datadog.trace.agent.test.asserts.TraceAssert -import datadog.trace.agent.test.base.HttpServer import datadog.trace.agent.test.base.HttpServerTest +import datadog.trace.agent.test.base.WebsocketServer import io.undertow.Handlers import io.undertow.Undertow import io.undertow.UndertowOptions @@ -8,8 +8,11 @@ import io.undertow.servlet.api.DeploymentInfo import io.undertow.servlet.api.DeploymentManager import io.undertow.servlet.api.ServletContainer import io.undertow.servlet.api.ServletInfo +import io.undertow.websockets.jsr.WebSocketDeploymentInfo import jakarta.servlet.MultipartConfigElement +import java.nio.ByteBuffer + import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_MULTIPART import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_URLENCODED import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.CREATED @@ -28,7 +31,7 @@ import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.USER_B class UndertowServletTest extends HttpServerTest { private static final CONTEXT = "ctx" - class UndertowServer implements HttpServer { + class UndertowServer implements WebsocketServer { def port = 0 Undertow undertowServer @@ -54,6 +57,7 @@ class UndertowServletTest extends HttpServerTest { .addServlet(new ServletInfo("CreatedISServlet", CreatedISServlet).addMapping(CREATED_IS.path)) .addServlet(new ServletInfo("BodyUrlEncodedServlet", BodyUrlEncodedServlet).addMapping(BODY_URLENCODED.path)) .addServlet(new ServletInfo("BodyMultipartServlet", BodyMultipartServlet).addMapping(BODY_MULTIPART.path)) + .addServletContextAttribute(WebSocketDeploymentInfo.ATTRIBUTE_NAME, new WebSocketDeploymentInfo().addEndpoint(TestEndpoint)) DeploymentManager manager = container.addDeployment(builder) manager.deploy() @@ -62,7 +66,7 @@ class UndertowServletTest extends HttpServerTest { undertowServer = Undertow.builder() .addHttpListener(port, "localhost") .setServerOption(UndertowOptions.DECODE_URL, true) - .setHandler(Handlers.httpContinueRead (root)) + .setHandler(Handlers.httpContinueRead(root)) .build() } @@ -82,6 +86,56 @@ class UndertowServletTest extends HttpServerTest { URI address() { return new URI("http://localhost:$port/$CONTEXT/") } + + @Override + void awaitConnected() { + while (TestEndpoint.activeSession == null) { + synchronized (TestEndpoint) { + TestEndpoint.wait() + } + } + } + + @Override + void serverSendText(String[] messages) { + if (messages.length == 1) { + TestEndpoint.activeSession.getBasicRemote().sendText(messages[0]) + } else { + def remoteEndpoint = TestEndpoint.activeSession.getBasicRemote() + for (int i = 0; i < messages.length; i++) { + remoteEndpoint.sendText(messages[i], i == messages.length - 1) + } + } + } + + @Override + boolean canSplitLargeWebsocketPayloads() { + false + } + + @Override + void serverSendBinary(byte[][] binaries) { + if (binaries.length == 1) { + TestEndpoint.activeSession.getBasicRemote().sendBinary(ByteBuffer.wrap(binaries[0])) + } else { + try (def stream = TestEndpoint.activeSession.getBasicRemote().getSendStream()) { + binaries.each { + stream.write(it) + } + } + } + } + + @Override + void serverClose() { + TestEndpoint.activeSession?.close() + } + + @Override + void setMaxPayloadSize(int size) { + TestEndpoint.activeSession?.setMaxTextMessageBufferSize(size) + TestEndpoint.activeSession?.setMaxBinaryMessageBufferSize(size) + } } @Override diff --git a/dd-java-agent/instrumentation/websocket/build.gradle b/dd-java-agent/instrumentation/websocket/build.gradle new file mode 100644 index 00000000000..fe680a977ab --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/build.gradle @@ -0,0 +1,4 @@ +apply from: "$rootDir/gradle/java.gradle" + +dependencies { +} diff --git a/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/build.gradle b/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/build.gradle new file mode 100644 index 00000000000..5bc84784fb8 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/build.gradle @@ -0,0 +1,25 @@ +ext { + minJavaVersionForTests = JavaVersion.VERSION_11 + minJavaVersionForLatestDepTests = JavaVersion.VERSION_17 +} +muzzle { + pass { + name = "jakarta-websocket" + group = "jakarta.websocket" + module = "jakarta.websocket-client-api" + versions = "[2.0.0,)" + javaVersion = "17" + } +} + +apply from: "$rootDir/gradle/java.gradle" +addTestSuiteForDir("latestDepTest", "test") + +dependencies { + compileOnly group: 'jakarta.websocket', name: 'jakarta.websocket-client-api', version: '2.0.0' + implementation project(":dd-java-agent:instrumentation:websocket:javax-websocket-1.0") + + testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:javax-websocket-1.0") + testImplementation group: 'org.glassfish.tyrus', name: 'tyrus-container-inmemory', version: '2.0.0' + latestDepTestImplementation group: 'org.glassfish.tyrus', name: 'tyrus-container-inmemory', version: '+' +} diff --git a/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/JakartaWebsocketModule.java b/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/JakartaWebsocketModule.java new file mode 100644 index 00000000000..02b864544a2 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/JakartaWebsocketModule.java @@ -0,0 +1,25 @@ +package datadog.trace.instrumentation.websocket.jsr256; + +import static java.util.Collections.singletonMap; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.InstrumenterModule; +import java.util.Map; + +@AutoService(InstrumenterModule.class) +public class JakartaWebsocketModule extends JavaxWebsocketModule { + + public JakartaWebsocketModule() { + super("jakarta", "jakarta-websocket", "websocket"); + } + + @Override + public Map adviceShading() { + return singletonMap("javax", "jakarta"); + } + + @Override + public String muzzleDirective() { + return "jakarta-websocket"; + } +} diff --git a/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/EndpointWrapper.groovy b/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/EndpointWrapper.groovy new file mode 100644 index 00000000000..37c1bd1ebce --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/EndpointWrapper.groovy @@ -0,0 +1,39 @@ +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan + +import jakarta.websocket.CloseReason +import jakarta.websocket.Endpoint +import jakarta.websocket.EndpointConfig +import jakarta.websocket.Session + +class EndpointWrapper extends Endpoint { + + EndpointWrapper() { + } + + @Override + void onOpen(Session session, EndpointConfig endpointConfig) { + def span = endpointConfig.getUserProperties().get(AgentSpan.class.getName()) as AgentSpan + def endpoint = endpointConfig.getUserProperties().get(Endpoint.class.getName()) as Endpoint + assert endpoint != null + session.getUserProperties().put(Endpoint.class.getName(), endpoint) + try (def ignored = span != null ? activateSpan(span) : null) { + endpoint.onOpen(session, endpointConfig) + } + } + + @Override + void onClose(Session session, CloseReason closeReason) { + def endpoint = session.getUserProperties().get(Endpoint.class.getName()) as Endpoint + assert endpoint != null + endpoint.onClose(session, closeReason) + } + + @Override + void onError(Session session, Throwable thr) { + def endpoint = session.getUserProperties().get(Endpoint.class.getName()) as Endpoint + assert endpoint != null + endpoint.onError(session, thr) + } +} diff --git a/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/Endpoints.groovy b/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/Endpoints.groovy new file mode 100644 index 00000000000..dc49cb7b987 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/Endpoints.groovy @@ -0,0 +1,130 @@ +import org.apache.commons.io.IOUtils + +import jakarta.websocket.DecodeException +import jakarta.websocket.Decoder +import jakarta.websocket.EncodeException +import jakarta.websocket.Encoder +import jakarta.websocket.Endpoint +import jakarta.websocket.EndpointConfig +import jakarta.websocket.MessageHandler +import jakarta.websocket.Session +import java.nio.ByteBuffer + +class Endpoints { + static class ClientTestEndpoint extends Endpoint { + @Override + void onOpen(Session session, EndpointConfig endpointConfig) { + } + } + + static class TestEndpoint extends Endpoint { + final MessageHandler handler + + TestEndpoint(final MessageHandler handler) { + this.handler = handler + } + + @Override + void onOpen(Session session, EndpointConfig endpointConfig) { + session.addMessageHandler(handler) + } + } + + // Full handlers + static class FullStringHandler implements MessageHandler.Whole { + @Override + void onMessage(String s) { + } + } + + static class FullReaderHandler implements MessageHandler.Whole { + @Override + void onMessage(Reader reader) { + IOUtils.toString(reader) + } + } + + static class FullBytesHandler implements MessageHandler.Whole { + @Override + void onMessage(byte[] b) { + } + } + + static class FullByteBufferHandler implements MessageHandler.Whole { + @Override + void onMessage(ByteBuffer b) { + } + } + + static class CustomMessage { + } + + static class CustomMessageEncoder implements Encoder.Text { + + @Override + void init(EndpointConfig endpointConfig) { + } + + @Override + void destroy() { + } + + @Override + String encode(CustomMessage customMessage) throws EncodeException { + return "CustomMessage" + } + } + + static class CustomMessageDecoder implements Decoder.Text { + + @Override + void init(EndpointConfig endpointConfig) { + } + + @Override + void destroy() { + } + + @Override + CustomMessage decode(String s) throws DecodeException { + return new CustomMessage() + } + + @Override + boolean willDecode(String s) { + return s == "CustomMessage" + } + } + + static class FullObjectHandler implements MessageHandler.Whole { + @Override + void onMessage(CustomMessage o) { + } + } + + static class FullStreamHandler implements MessageHandler.Whole { + @Override + void onMessage(InputStream is) { + IOUtils.toByteArray(is) + } + } + // Partial Handlers + static class PartialStringHandler implements MessageHandler.Partial { + @Override + void onMessage(String s, boolean last) { + } + } + + static class PartialBytesHandler implements MessageHandler.Partial { + @Override + void onMessage(byte[] b, boolean last) { + + } + } + + static class PartialByteBufferHandler implements MessageHandler.Partial { + @Override + void onMessage(ByteBuffer b, boolean last) { + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/WebsocketTest.groovy b/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/WebsocketTest.groovy new file mode 100644 index 00000000000..a2644fab0d5 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/WebsocketTest.groovy @@ -0,0 +1,577 @@ +import static datadog.trace.agent.test.base.HttpServerTest.someBytes +import static datadog.trace.agent.test.base.HttpServerTest.websocketCloseSpan +import static datadog.trace.agent.test.base.HttpServerTest.websocketReceiveSpan +import static datadog.trace.agent.test.base.HttpServerTest.websocketSendSpan +import static datadog.trace.agent.test.utils.TraceUtils.basicSpan +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_CLASSES_EXCLUDE +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_ENABLED +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_INHERIT_SAMPLING +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_SEPARATE_TRACES +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_TAG_SESSION_ID +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan + +import datadog.trace.agent.test.AgentTestRunner +import datadog.trace.api.DDTags +import datadog.trace.bootstrap.instrumentation.api.AgentSpan +import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags +import datadog.trace.bootstrap.instrumentation.api.Tags +import datadog.trace.core.DDSpan +import net.bytebuddy.utility.RandomString +import org.glassfish.tyrus.container.inmemory.InMemoryClientContainer +import org.glassfish.tyrus.server.TyrusServerConfiguration + +import jakarta.websocket.ClientEndpointConfig +import jakarta.websocket.CloseReason +import jakarta.websocket.ContainerProvider +import jakarta.websocket.Endpoint +import jakarta.websocket.server.ServerApplicationConfig +import jakarta.websocket.server.ServerEndpointConfig +import java.nio.ByteBuffer + +class WebsocketTest extends AgentTestRunner { + + @Override + protected void configurePreAgent() { + super.configurePreAgent() + injectSysConfig(TRACE_WEBSOCKET_MESSAGES_ENABLED, "true") + injectSysConfig(TRACE_CLASSES_EXCLUDE, "EndpointWrapper") + } + + def createHandshakeSpan(String spanName, String url) { + def span = TEST_TRACER.startSpan("test", spanName, null) + handshakeTags(url).each { span.setTag(it.key, it.value) } + span.finish() + span + } + + def deployEndpointAndConnect(Endpoint endpoint, Object handshakeClientSpan, Object handshakeServerSpan, String url) { + def webSocketContainer = ContainerProvider.getWebSocketContainer() + def sec = ServerEndpointConfig.Builder.create(EndpointWrapper, "/test") + .encoders([Endpoints.CustomMessageEncoder]) + .decoders([Endpoints.CustomMessageDecoder]) + .build() + + sec.getUserProperties().put(Endpoint.class.getName(), endpoint) + sec.getUserProperties().put(AgentSpan.class.getName(), handshakeServerSpan) + + final ServerApplicationConfig serverConfig = + new TyrusServerConfiguration(Collections.singleton(EndpointWrapper.class), + Collections.singleton(sec)) + + ClientEndpointConfig cec = ClientEndpointConfig.Builder.create() + .encoders([Endpoints.CustomMessageEncoder]) + .decoders([Endpoints.CustomMessageDecoder]) + .build() + cec.getUserProperties().put(InMemoryClientContainer.SERVER_CONFIG, serverConfig) + + try (def ignored = handshakeClientSpan != null ? activateSpan(handshakeClientSpan as AgentSpan) : null) { + def session = webSocketContainer.connectToServer(new Endpoints.ClientTestEndpoint(), cec, URI.create(url)) + session + } + } + + def handshakeTags(url) { + [(Tags.HTTP_METHOD): "GET", (Tags.HTTP_URL): url] + } + + def "test full sync send and receive for endpoint #endpoint.class with #sendSize len #msgType message"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(endpoint), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + if (msgType == "text") { + session.getBasicRemote().sendText(message as String) + } else { + session.getBasicRemote().sendBinary(ByteBuffer.wrap(message as byte[])) + } + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, msgType, sendSize, 1, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + + trace(1) { + websocketReceiveSpan(it, serverHandshake, msgType, rcvSize, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + where: + endpoint | message | msgType | sendSize | rcvSize + // text full + new Endpoints.FullStringHandler() | RandomString.make(10) | "text" | 10 | 10 + new Endpoints.FullReaderHandler() | RandomString.make(20) | "text" | 20 | 0 + // binary full + new Endpoints.FullByteBufferHandler() | someBytes(10) | "binary" | 10 | 10 + new Endpoints.FullBytesHandler() | someBytes(25) | "binary" | 25 | 25 + new Endpoints.FullStreamHandler() | someBytes(10) | "binary" | 10 | 0 + } + + def "test partial sync send and receive for endpoint #endpoint.class with #chunks chunks and #size len #msgType message"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(endpoint), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + def remote = session.getBasicRemote() + if (msgType == "text") { + for (int i = 0; i < message.size(); i++) { + remote.sendText(message[i] as String, (i == message.size() - 1)) + } + } else { + for (int i = 0; i < message.size(); i++) { + remote.sendBinary(ByteBuffer.wrap(message[i] as byte[]), i == (message.size() - 1)) + } + } + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, msgType, size, chunks, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + trace(1) { + websocketReceiveSpan(it, serverHandshake, msgType, size, chunks) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + where: + endpoint | message | msgType | chunks | size + // text partial + new Endpoints.PartialStringHandler() | [RandomString.make(10)] | "text" | 1 | 10 + new Endpoints.PartialStringHandler() | [RandomString.make(10), RandomString.make(10)] | "text" | 2 | 20 + // binary Partial + new Endpoints.PartialByteBufferHandler() | [someBytes(10)] | "binary" | 1 | 10 + new Endpoints.PartialByteBufferHandler() | [someBytes(10), someBytes(15)] | "binary" | 2 | 25 + new Endpoints.PartialBytesHandler() | [someBytes(10)] | "binary" | 1 | 10 + new Endpoints.PartialBytesHandler() | [someBytes(10), someBytes(15)] | "binary" | 2 | 25 + } + + def "test stream sync send and receive for endpoint #endpoint.class with #chunks chunks and #size len #msgType message"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(endpoint), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + if (msgType == "text") { + // tyrus text writer send the flushes when closing so the fin bit is sent along with the last message + try (def writer = session.getBasicRemote().getSendWriter()) { + message.each { + writer.write(it as String) + } + } + } else { + // tyrus binary writer send the fin bit with an empty frame when close is called + try (def os = session.getBasicRemote().getSendStream()) { + message.each { + os.write(it as byte[]) + } + } + } + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, msgType, size, chunks, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + trace(1) { + websocketReceiveSpan(it, serverHandshake, msgType, size, chunks) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + where: + endpoint | message | msgType | chunks | size + // text partial + new Endpoints.PartialStringHandler() | [RandomString.make(10)] | "text" | 1 | 10 + new Endpoints.PartialStringHandler() | [RandomString.make(10), RandomString.make(15)] | "text" | 2 | 25 + // binary Partial + new Endpoints.PartialByteBufferHandler() | [someBytes(10)] | "binary" | 1 | 10 + new Endpoints.PartialByteBufferHandler() | [someBytes(10), someBytes(15)] | "binary" | 2 | 25 + new Endpoints.PartialBytesHandler() | [someBytes(10)] | "binary" | 1 | 10 + new Endpoints.PartialBytesHandler() | [someBytes(10), someBytes(15)] | "binary" | 2 | 25 + } + + def "test full async (future) send and receive for endpoint #endpoint.class with #sendSize len #msgType message"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(endpoint), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + if (msgType == "text") { + session.getAsyncRemote().sendText(message as String) + } else { + session.getAsyncRemote().sendBinary(ByteBuffer.wrap(message as byte[])) + } + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, msgType, sendSize, 1, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + + trace(1) { + websocketReceiveSpan(it, serverHandshake, msgType, rcvSize, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + where: + endpoint | message | msgType | sendSize | rcvSize + // text full + new Endpoints.FullStringHandler() | RandomString.make(10) | "text" | 10 | 10 + new Endpoints.FullReaderHandler() | RandomString.make(20) | "text" | 20 | 0 + // binary full + new Endpoints.FullByteBufferHandler() | someBytes(10) | "binary" | 10 | 10 + new Endpoints.FullBytesHandler() | someBytes(25) | "binary" | 25 | 25 + new Endpoints.FullStreamHandler() | someBytes(10) | "binary" | 10 | 0 + } + + def "test full async (SendHandler) send and receive for endpoint #endpoint.class with #sendSize len #msgType message"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(endpoint), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + if (msgType == "text") { + session.getAsyncRemote().sendText(message as String, { assert it.OK }) + } else { + session.getAsyncRemote().sendBinary(ByteBuffer.wrap(message as byte[]), { assert it.OK }) + } + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, msgType, sendSize, 1, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + + trace(1) { + websocketReceiveSpan(it, serverHandshake, msgType, rcvSize, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + where: + endpoint | message | msgType | sendSize | rcvSize + // text full + new Endpoints.FullStringHandler() | RandomString.make(10) | "text" | 10 | 10 + new Endpoints.FullReaderHandler() | RandomString.make(20) | "text" | 20 | 0 + // binary full + new Endpoints.FullByteBufferHandler() | someBytes(10) | "binary" | 10 | 10 + new Endpoints.FullBytesHandler() | someBytes(25) | "binary" | 25 | 25 + new Endpoints.FullStreamHandler() | someBytes(10) | "binary" | 10 | 0 + } + + def "test session close code #code and reason #reason"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullStringHandler()), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + session.close(new CloseReason(CloseReason.CloseCodes.getCloseCode(code), reason)) + } + then: + assertTraces(4, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(2) { + sortSpansByStart() + basicSpan(it, "parent") + websocketCloseSpan(it, clientHandshake, true, code, reason, span(0)) + } + + trace(1) { + websocketCloseSpan(it, serverHandshake, false, code, reason) + } + }) + where: + code | reason + 1000 | null + 1000 | "bye" + 1001 | "see you" + } + + def "test session id logged as tag"() { + setup: + injectSysConfig(TRACE_WEBSOCKET_TAG_SESSION_ID, "true") + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullStringHandler()), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + session.getBasicRemote().sendText("Hello") + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, "text", 5, 1, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + trace(1) { + websocketReceiveSpan(it, serverHandshake, "text", 5, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + //it's normally already tested by websocket(Receive|Send|Close)Span but we enforce that check + TEST_WRITER.flatten().findAll { span -> (span as DDSpan).getSpanType() == "websocket" }.each { + assert (it as DDSpan).getTag(InstrumentationTags.WEBSOCKET_SESSION_ID) != null + } + } + + def "test close and receive on same handshake trace"() { + setup: + injectSysConfig(TRACE_WEBSOCKET_MESSAGES_SEPARATE_TRACES, "false") + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullStringHandler()), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + session.getBasicRemote().sendText("Hello") + session.close() + } + then: + // in reality we have 3 traces but since the handshake finishes soon, the trace structure writer is collecting 5 chunks + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(1) { + websocketReceiveSpan(it, serverHandshake, "text", 5, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }, serverHandshake) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, "text", 5, 1, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + }) + } + + def "test sampling not inherited"() { + setup: + injectSysConfig(TRACE_WEBSOCKET_MESSAGES_INHERIT_SAMPLING, "false") + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullStringHandler()), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + session.getBasicRemote().sendText("Hello") + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, "text", 5, 1, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + trace(1) { + websocketReceiveSpan(it, serverHandshake, "text", 5, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + //it's normally already tested by websocket(Receive|Send|Close)Span but we enforce that check + TEST_WRITER.flatten().findAll { span -> (span as DDSpan).getSpanType() == "websocket" }.each { + assert (it as DDSpan).getTag(DDTags.DECISION_MAKER_INHERITED) == null + } + } + + def "if handshake is not captured traces are not generated"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullStringHandler()), + null, null, url) + + runUnderTrace("parent") { + session.getBasicRemote().sendText("Hello") + session.close() + } + then: + assertTraces(1, { + trace(1) { + basicSpan(it, "parent") + } + }) + } + + def "test send and receive object"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullObjectHandler()), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + session.getBasicRemote().sendObject(new Endpoints.CustomMessage()) + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, null, 0, 1, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + trace(1) { + websocketReceiveSpan(it, serverHandshake, null, 0, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + } +} diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/build.gradle b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/build.gradle new file mode 100644 index 00000000000..ed48fb2bb34 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/build.gradle @@ -0,0 +1,19 @@ +muzzle { + pass { + name = "javax-websocket" + group = "javax.websocket" + module = "javax.websocket-api" + versions = "[1.0,)" + } +} + +apply from: "$rootDir/gradle/java.gradle" + +addTestSuiteForDir("latestDepTest", "test") + +dependencies { + compileOnly group: 'javax.websocket', name: 'javax.websocket-api', version: '1.0-rc1' + testImplementation group: 'org.glassfish.tyrus', name: 'tyrus-container-inmemory', version: '1.3.1' + latestDepTestImplementation group: 'org.glassfish.tyrus', name: 'tyrus-container-inmemory', version: '1.+' +} + diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/AsyncRemoteEndpointInstrumentation.java b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/AsyncRemoteEndpointInstrumentation.java new file mode 100644 index 00000000000..47b760ae5cd --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/AsyncRemoteEndpointInstrumentation.java @@ -0,0 +1,245 @@ +package datadog.trace.instrumentation.websocket.jsr256; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE; +import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.BYTE_BUFFER_SIZE_CALCULATOR; +import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.CHAR_SEQUENCE_SIZE_CALCULATOR; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.bootstrap.CallDepthThreadLocalMap; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; +import java.nio.ByteBuffer; +import java.util.concurrent.Future; +import javax.websocket.RemoteEndpoint; +import javax.websocket.SendHandler; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.implementation.bytecode.assign.Assigner; +import net.bytebuddy.matcher.ElementMatcher; + +public class AsyncRemoteEndpointInstrumentation + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + private final String namespace; + + public AsyncRemoteEndpointInstrumentation(String namespace) { + this.namespace = namespace; + } + + @Override + public String hierarchyMarkerType() { + return namespace + ".websocket.RemoteEndpoint$Async"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named(hierarchyMarkerType())); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isPublic() + .and(named("sendText")) + .and( + takesArguments(1) + .or( + takesArguments(2) + .and(takesArgument(1, named(namespace + ".websocket.SendHandler"))))) + .and(takesArgument(0, named("java.lang.String"))), + getClass().getName() + "$SendTextAdvice"); + transformer.applyAdvice( + isPublic() + .and(named("sendBinary")) + .and( + takesArguments(1) + .or( + takesArguments(2) + .and(takesArgument(1, named(namespace + ".websocket.SendHandler"))))) + .and(takesArgument(0, named("java.nio.ByteBuffer"))), + getClass().getName() + "$SendBinaryAdvice"); + transformer.applyAdvice( + isPublic() + .and(named("sendObject")) + .and( + takesArguments(1) + .or( + takesArguments(2) + .and(takesArgument(1, named(namespace + ".websocket.SendHandler"))))), + getClass().getName() + "$SendObjectAdvice"); + } + + public static class SendTextAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope before( + @Advice.This final RemoteEndpoint.Async self, + @Advice.Argument(0) String text, + @Advice.Argument( + value = 1, + optional = true, + readOnly = false, + typing = Assigner.Typing.DYNAMIC) + SendHandler sendHandler, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext) { + handlerContext = + InstrumentationContext.get(RemoteEndpoint.class, HandlerContext.Sender.class).get(self); + if (handlerContext == null + || CallDepthThreadLocalMap.incrementCallDepth(RemoteEndpoint.class) > 0) { + return null; + } + + final AgentSpan wsSpan = + DECORATE.onSendFrameStart( + handlerContext, + CHAR_SEQUENCE_SIZE_CALCULATOR.getFormat(), + CHAR_SEQUENCE_SIZE_CALCULATOR.getLengthFunction().applyAsInt(text)); + if (sendHandler != null) { + sendHandler = new TracingSendHandler(sendHandler, handlerContext); + } + return activateSpan(wsSpan); + } + + @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) + public static void after( + @Advice.Enter final AgentScope scope, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext, + @Advice.Thrown final Throwable throwable, + @Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) Future future) { + CallDepthThreadLocalMap.decrementCallDepth(RemoteEndpoint.class); + if (scope == null) { + return; + } + try { + if (throwable != null) { + DECORATE.onError(scope, throwable); + DECORATE.onFrameEnd(handlerContext); + } else if (future != null) { + // FIXME: Knowing when the future really completes would imply instrumenting all the + // possible implementations. + // In this case we will just finish the span to have a trace of this send even if the + // duration is not exact + DECORATE.onFrameEnd(handlerContext); + } + } finally { + scope.close(); + } + } + } + + public static class SendBinaryAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope before( + @Advice.This final RemoteEndpoint.Async self, + @Advice.Argument(0) ByteBuffer buffer, + @Advice.Argument( + value = 1, + optional = true, + readOnly = false, + typing = Assigner.Typing.DYNAMIC) + SendHandler sendHandler, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext) { + handlerContext = + InstrumentationContext.get(RemoteEndpoint.class, HandlerContext.Sender.class).get(self); + if (handlerContext == null + || CallDepthThreadLocalMap.incrementCallDepth(RemoteEndpoint.class) > 0) { + return null; + } + + final AgentSpan wsSpan = + DECORATE.onSendFrameStart( + handlerContext, + BYTE_BUFFER_SIZE_CALCULATOR.getFormat(), + BYTE_BUFFER_SIZE_CALCULATOR.getLengthFunction().applyAsInt(buffer)); + if (sendHandler != null) { + sendHandler = new TracingSendHandler(sendHandler, handlerContext); + } + return activateSpan(wsSpan); + } + + @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) + public static void after( + @Advice.Enter final AgentScope scope, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext, + @Advice.Thrown final Throwable throwable, + @Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) Future future) { + CallDepthThreadLocalMap.decrementCallDepth(RemoteEndpoint.class); + if (scope == null) { + return; + } + try { + if (throwable != null) { + DECORATE.onError(scope, throwable); + DECORATE.onFrameEnd(handlerContext); + } else if (future != null) { + // FIXME: Knowing when the future really completes would imply instrumenting all the + // possible implementations. + // In this case we will just finish the span to have a trace of this send even if the + // duration is not exact + DECORATE.onFrameEnd(handlerContext); + } + } finally { + scope.close(); + } + } + } + + public static class SendObjectAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope before( + @Advice.This final RemoteEndpoint.Async self, + @Advice.Argument( + value = 1, + optional = true, + readOnly = false, + typing = Assigner.Typing.DYNAMIC) + SendHandler sendHandler, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext) { + handlerContext = + InstrumentationContext.get(RemoteEndpoint.class, HandlerContext.Sender.class).get(self); + if (handlerContext == null + || CallDepthThreadLocalMap.incrementCallDepth(RemoteEndpoint.class) > 0) { + return null; + } + + final AgentSpan wsSpan = + DECORATE.onSendFrameStart(handlerContext, BYTE_BUFFER_SIZE_CALCULATOR.getFormat(), 0); + if (sendHandler != null) { + sendHandler = new TracingSendHandler(sendHandler, handlerContext); + } + return activateSpan(wsSpan); + } + + @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) + public static void after( + @Advice.Enter final AgentScope scope, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext, + @Advice.Thrown final Throwable throwable, + @Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) Future future) { + CallDepthThreadLocalMap.decrementCallDepth(RemoteEndpoint.class); + if (scope == null) { + return; + } + try { + if (throwable != null) { + DECORATE.onError(scope, throwable); + DECORATE.onFrameEnd(handlerContext); + } else if (future != null) { + // FIXME: Knowing when the future really completes would imply instrumenting all the + // possible implementations. + // In this case we will just finish the span to have a trace of this send even if the + // duration is not exact + DECORATE.onFrameEnd(handlerContext); + } + } finally { + scope.close(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/BasicRemoteEndpointInstrumentation.java b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/BasicRemoteEndpointInstrumentation.java new file mode 100644 index 00000000000..40850478789 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/BasicRemoteEndpointInstrumentation.java @@ -0,0 +1,238 @@ +package datadog.trace.instrumentation.websocket.jsr256; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE; +import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.BYTE_BUFFER_SIZE_CALCULATOR; +import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.CHAR_SEQUENCE_SIZE_CALCULATOR; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.returns; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.bootstrap.CallDepthThreadLocalMap; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; +import java.io.OutputStream; +import java.io.Writer; +import java.nio.ByteBuffer; +import javax.websocket.RemoteEndpoint; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +public class BasicRemoteEndpointInstrumentation + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + private final String namespace; + + public BasicRemoteEndpointInstrumentation(String namespace) { + this.namespace = namespace; + } + + @Override + public String hierarchyMarkerType() { + return namespace + ".websocket.RemoteEndpoint$Basic"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named(hierarchyMarkerType())); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isPublic() + .and(named("sendText")) + .and(takesArguments(1).or(takesArguments(2).and(takesArgument(1, boolean.class)))) + .and(takesArgument(0, named("java.lang.String"))) + .and(returns(void.class)), + getClass().getName() + "$SendTextAdvice"); + transformer.applyAdvice( + isPublic() + .and(named("sendBinary")) + .and(takesArguments(1).or(takesArguments(2).and(takesArgument(1, boolean.class)))) + .and(takesArgument(0, named("java.nio.ByteBuffer"))) + .and(returns(void.class)), + getClass().getName() + "$SendBinaryAdvice"); + transformer.applyAdvice( + isPublic().and(named("sendObject")).and(takesArguments(1)).and(returns(void.class)), + getClass().getName() + "$SendObjectAdvice"); + transformer.applyAdvice( + isPublic().and(named("getSendStream")).and(takesNoArguments()), + getClass().getName() + "$WrapStreamAdvice"); + transformer.applyAdvice( + isPublic().and(named("getSendWriter")).and(takesNoArguments()), + getClass().getName() + "$WrapWriterAdvice"); + } + + public static class SendTextAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope before( + @Advice.This final RemoteEndpoint.Basic self, + @Advice.Argument(0) String text, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext) { + handlerContext = + InstrumentationContext.get(RemoteEndpoint.class, HandlerContext.Sender.class).get(self); + if (handlerContext == null + || CallDepthThreadLocalMap.incrementCallDepth(RemoteEndpoint.class) > 0) { + return null; + } + + final AgentSpan wsSpan = + DECORATE.onSendFrameStart( + handlerContext, + CHAR_SEQUENCE_SIZE_CALCULATOR.getFormat(), + CHAR_SEQUENCE_SIZE_CALCULATOR.getLengthFunction().applyAsInt(text)); + return activateSpan(wsSpan); + } + + @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) + public static void after( + @Advice.Enter final AgentScope scope, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext, + @Advice.Thrown final Throwable throwable, + @Advice.Argument(value = 1, optional = true) final Boolean last) { + CallDepthThreadLocalMap.decrementCallDepth(RemoteEndpoint.class); + + if (scope == null) { + return; + } + try { + boolean finishSpan = last == null || last; + if (throwable != null) { + finishSpan = true; + DECORATE.onError(scope, throwable); + } + if (finishSpan) { + DECORATE.onFrameEnd(handlerContext); + } + } finally { + scope.close(); + } + } + } + + public static class SendBinaryAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope before( + @Advice.This final RemoteEndpoint.Basic self, + @Advice.Argument(0) ByteBuffer buffer, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext) { + handlerContext = + InstrumentationContext.get(RemoteEndpoint.class, HandlerContext.Sender.class).get(self); + if (handlerContext == null + || CallDepthThreadLocalMap.incrementCallDepth(RemoteEndpoint.class) > 0) { + return null; + } + + final AgentSpan wsSpan = + DECORATE.onSendFrameStart( + handlerContext, + BYTE_BUFFER_SIZE_CALCULATOR.getFormat(), + BYTE_BUFFER_SIZE_CALCULATOR.getLengthFunction().applyAsInt(buffer)); + return activateSpan(wsSpan); + } + + @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) + public static void after( + @Advice.Enter final AgentScope scope, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext, + @Advice.Thrown final Throwable throwable, + @Advice.Argument(value = 1, optional = true) final Boolean last) { + CallDepthThreadLocalMap.decrementCallDepth(RemoteEndpoint.class); + if (scope == null) { + return; + } + try { + boolean finishSpan = last == null || last; + if (throwable != null) { + finishSpan = true; + DECORATE.onError(scope, throwable); + } + if (finishSpan) { + DECORATE.onFrameEnd(handlerContext); + } + } finally { + scope.close(); + } + } + } + + public static class SendObjectAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope before( + @Advice.This final RemoteEndpoint.Basic self, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext) { + handlerContext = + InstrumentationContext.get(RemoteEndpoint.class, HandlerContext.Sender.class).get(self); + if (handlerContext == null + || CallDepthThreadLocalMap.incrementCallDepth(RemoteEndpoint.class) > 0) { + return null; + } + + // we actually cannot know the size and the type since this the conversion is done by + // encoders/decoders. + // we can anyway instrument also the Encoders but that would add much more complexity. + // right now this is not in scope + final AgentSpan wsSpan = DECORATE.onSendFrameStart(handlerContext, null, 0); + return activateSpan(wsSpan); + } + + @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) + public static void after( + @Advice.Enter final AgentScope scope, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext, + @Advice.Thrown final Throwable throwable) { + CallDepthThreadLocalMap.decrementCallDepth(RemoteEndpoint.class); + if (scope == null) { + return; + } + try { + if (throwable != null) { + DECORATE.onError(scope, throwable); + } + DECORATE.onFrameEnd(handlerContext); + } finally { + scope.close(); + } + } + } + + public static class WrapWriterAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void after( + @Advice.This final RemoteEndpoint.Basic self, + @Advice.Return(readOnly = false) Writer writer) { + if (writer instanceof TracingWriter) { + return; + } + final HandlerContext.Sender handlerContext = + InstrumentationContext.get(RemoteEndpoint.class, HandlerContext.Sender.class).get(self); + if (handlerContext != null) { + writer = new TracingWriter(writer, handlerContext); + } + } + } + + public static class WrapStreamAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void after( + @Advice.This final RemoteEndpoint.Basic self, + @Advice.Return(readOnly = false) OutputStream outputStream) { + if (outputStream instanceof TracingOutputStream) { + return; + } + final HandlerContext.Sender handlerContext = + InstrumentationContext.get(RemoteEndpoint.class, HandlerContext.Sender.class).get(self); + if (handlerContext != null) { + outputStream = new TracingOutputStream(outputStream, handlerContext); + } + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/EndpointInstrumentation.java b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/EndpointInstrumentation.java new file mode 100644 index 00000000000..21ec05b00f7 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/EndpointInstrumentation.java @@ -0,0 +1,109 @@ +package datadog.trace.instrumentation.websocket.jsr256; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.api.Config; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; +import javax.websocket.CloseReason; +import javax.websocket.Session; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +public class EndpointInstrumentation + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + private final String namespace; + + public EndpointInstrumentation(String namespace) { + this.namespace = namespace; + } + + @Override + public String hierarchyMarkerType() { + return namespace + ".websocket.Endpoint"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return extendsClass(named(hierarchyMarkerType())); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isPublic() + .and( + named("onOpen") + .and(takesArguments(2)) + .and(takesArgument(0, named(namespace + ".websocket.Session")))), + getClass().getName() + "$CaptureHandshakeSpanAdvice"); + transformer.applyAdvice( + isPublic() + .and( + named("onClose") + .and(takesArguments(2)) + .and(takesArgument(0, named(namespace + ".websocket.Session"))) + .and(takesArgument(1, named(namespace + ".websocket.CloseReason")))), + getClass().getName() + "$SessionCloseAdvice"); + } + + public static class CaptureHandshakeSpanAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter(@Advice.Argument(0) final Session session) { + final AgentSpan current = AgentTracer.get().activeSpan(); + if (current != null) { + // we need to force the sampling decision in case the span is linked + if (Config.get().isWebsocketMessagesInheritSampling()) { + current.forceSamplingDecision(); + } + InstrumentationContext.get(Session.class, HandlerContext.Sender.class) + .putIfAbsent( + session, new HandlerContext.Sender(current.getLocalRootSpan(), session.getId())); + } + } + } + + public static class SessionCloseAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope onEnter( + @Advice.Local("handlerContext") HandlerContext.Receiver handlerContext, + @Advice.Argument(0) final Session session, + @Advice.Argument(1) final CloseReason closeReason) { + final HandlerContext.Sender sessionState = + InstrumentationContext.get(Session.class, HandlerContext.Sender.class).remove(session); + if (sessionState == null) { + return null; + } + handlerContext = + new HandlerContext.Receiver(sessionState.getHandshakeSpan(), session.getId()); + + return activateSpan( + DECORATE.onSessionCloseReceived( + handlerContext, closeReason.getReasonPhrase(), closeReason.getCloseCode().getCode())); + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void onExit( + @Advice.Enter final AgentScope scope, + @Advice.Local("handlerContext") HandlerContext.Receiver handlerContext, + @Advice.Thrown final Throwable thrown) { + if (scope != null) { + final AgentSpan span = scope.span(); + DECORATE.onError(span, thrown); + DECORATE.onFrameEnd(handlerContext); + scope.close(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/JavaxWebsocketModule.java b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/JavaxWebsocketModule.java new file mode 100644 index 00000000000..51014538f6e --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/JavaxWebsocketModule.java @@ -0,0 +1,64 @@ +package datadog.trace.instrumentation.websocket.jsr256; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.InstrumenterConfig; +import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +@AutoService(InstrumenterModule.class) +public class JavaxWebsocketModule extends InstrumenterModule.Tracing { + private final String namespace; + + public JavaxWebsocketModule() { + this("javax", "javax-websocket", "websocket"); + } + + protected JavaxWebsocketModule( + String namespace, String instrumentationName, String... additionalNames) { + super(instrumentationName, additionalNames); + this.namespace = namespace; + } + + @Override + public Map contextStore() { + final Map map = new HashMap<>(); + map.put(namespace + ".websocket.Session", HandlerContext.Sender.class.getName()); + map.put(namespace + ".websocket.RemoteEndpoint", HandlerContext.Sender.class.getName()); + map.put(namespace + ".websocket.MessageHandler", HandlerContext.Receiver.class.getName()); + return map; + } + + @Override + public String[] helperClassNames() { + return new String[] { + packageName + ".TracingOutputStream", + packageName + ".TracingWriter", + packageName + ".TracingSendHandler", + }; + } + + @Override + protected boolean defaultEnabled() { + return InstrumenterConfig.get().isWebsocketTracingEnabled(); + } + + @Override + public String muzzleDirective() { + return "javax-websocket"; + } + + @Override + public List typeInstrumentations() { + return Arrays.asList( + new EndpointInstrumentation(namespace), + new SessionInstrumentation(namespace), + new MessageHandlerInstrumentation(namespace), + new BasicRemoteEndpointInstrumentation(namespace), + new AsyncRemoteEndpointInstrumentation(namespace)); + } +} diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/MessageHandlerInstrumentation.java b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/MessageHandlerInstrumentation.java new file mode 100644 index 00000000000..6d1c93d4cf6 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/MessageHandlerInstrumentation.java @@ -0,0 +1,98 @@ +package datadog.trace.instrumentation.websocket.jsr256; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.bootstrap.CallDepthThreadLocalMap; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; +import javax.websocket.MessageHandler; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.implementation.bytecode.assign.Assigner; +import net.bytebuddy.matcher.ElementMatcher; + +public class MessageHandlerInstrumentation + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + private final String namespace; + + public MessageHandlerInstrumentation(String namespace) { + this.namespace = namespace; + } + + @Override + public String hierarchyMarkerType() { + return namespace + ".websocket.MessageHandler"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named(hierarchyMarkerType())); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isPublic() + .and(named("onMessage")) + .and( + takesArguments(1) // whole + .or(takesArguments(2).and(takesArgument(1, boolean.class)))), // partial + getClass().getName() + "$OnMessageAdvice"); + } + + public static class OnMessageAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope onEnter( + @Advice.This final MessageHandler handler, + @Advice.Argument(value = 0, typing = Assigner.Typing.DYNAMIC) final Object data, + @Advice.Argument(value = 1, optional = true) final Boolean last, + @Advice.Local("handlerContext") HandlerContext.Receiver handlerContext) { + handlerContext = + InstrumentationContext.get(MessageHandler.class, HandlerContext.Receiver.class) + .get(handler); + if (handlerContext == null) { + return null; + } + if (CallDepthThreadLocalMap.incrementCallDepth(MessageHandler.class) > 0) { + return null; + } + + final AgentSpan wsSpan = + DECORATE.onReceiveFrameStart(handlerContext, data, last != null && last); + return activateSpan(wsSpan); + } + + @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) + public static void onExit( + @Advice.Enter final AgentScope scope, + @Advice.Local("handlerContext") HandlerContext.Receiver handlerContext, + @Advice.Thrown final Throwable throwable, + @Advice.Argument(value = 1, optional = true) final Boolean last) { + if (scope == null) { + return; + } + CallDepthThreadLocalMap.reset(MessageHandler.class); + try { + boolean finishSpan = last == null || last; + if (throwable != null) { + finishSpan = true; + DECORATE.onError(scope, throwable); + } + if (finishSpan) { + DECORATE.onFrameEnd(handlerContext); + } + } finally { + scope.close(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/SessionInstrumentation.java b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/SessionInstrumentation.java new file mode 100644 index 00000000000..4c1c5bb8f41 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/SessionInstrumentation.java @@ -0,0 +1,203 @@ +package datadog.trace.instrumentation.websocket.jsr256; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; +import javax.websocket.CloseReason; +import javax.websocket.MessageHandler; +import javax.websocket.RemoteEndpoint; +import javax.websocket.Session; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +public class SessionInstrumentation + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + private final String namespace; + + public SessionInstrumentation(String namespace) { + this.namespace = namespace; + } + + @Override + public String hierarchyMarkerType() { + return namespace + ".websocket.Session"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named(hierarchyMarkerType())); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isPublic() + .and( + named("addMessageHandler") + .and(takesArgument(0, named(namespace + ".websocket.MessageHandler")))), + getClass().getName() + "$LinkReceiverSessionArg0Advice"); + + transformer.applyAdvice( + isPublic() + .and( + named("addMessageHandler") + .and( + takesArgument( + 1, + namedOneOf( + namespace + ".websocket.MessageHandler$Whole", + namespace + ".websocket.MessageHandler$Partial")))), + getClass().getName() + "$LinkReceiverSessionArg1Advice"); + + transformer.applyAdvice( + isPublic().and(namedOneOf("getBasicRemote", "getAsyncRemote")).and(takesNoArguments()), + getClass().getName() + "$LinkSenderSessionAdvice"); + + transformer.applyAdvice( + isPublic() + .and(named("close")) + .and( + takesArguments(1) + .and(takesArgument(0, named(namespace + ".websocket.CloseReason")))), + getClass().getName() + "$SessionCloseAdvice"); + + transformer.applyAdvice( + named("close").and(takesNoArguments()), + getClass().getName() + "$DefaultSessionCloseAdvice"); + } + + public static class LinkReceiverSessionArg0Advice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void onExit( + @Advice.This final Session session, @Advice.Argument(0) final MessageHandler handler) { + if (handler != null) { + final HandlerContext.Sender sessionState = + InstrumentationContext.get(Session.class, HandlerContext.Sender.class).get(session); + if (sessionState != null) { + // If the user is adding singletons that is not going to work. However, in this case there + // is no chance to have the session linked in any way (even if wrapping the + // messagehandler). + // Now hopefully this is not the usual habit and implementations of that API are not doing + // that as well when creating proxies for the annotated pojos. + // Also, adding a field here is preferred than wrapping the message handler since the + // implementations are introspecting it to understand the type of message handled and, + // with + // erasures, it won't work. + InstrumentationContext.get(MessageHandler.class, HandlerContext.Receiver.class) + .put( + handler, + new HandlerContext.Receiver(sessionState.getHandshakeSpan(), session.getId())); + } + } + } + } + + public static class LinkReceiverSessionArg1Advice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void onExit( + @Advice.This final Session session, @Advice.Argument(1) final MessageHandler handler) { + if (handler != null) { + final HandlerContext.Sender sessionState = + InstrumentationContext.get(Session.class, HandlerContext.Sender.class).get(session); + if (sessionState != null) { + // If the user is adding singletons that is not going to work. However, in this case there + // is no chance to have the session linked in any way (even if wrapping the + // messagehandler). + // Now hopefully this is not the usual habit and implementations of that API are not doing + // that as well when creating proxies for the annotated pojos. + // Also, adding a field here is preferred than wrapping the message handler since the + // implementations are introspecting it to understand the type of message handled and, + // with + // erasures, it won't work. + InstrumentationContext.get(MessageHandler.class, HandlerContext.Receiver.class) + .put( + handler, + new HandlerContext.Receiver(sessionState.getHandshakeSpan(), session.getId())); + } + } + } + } + + public static class LinkSenderSessionAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void onExit( + @Advice.This final Session session, @Advice.Return final RemoteEndpoint remoteEndpoint) { + if (remoteEndpoint != null) { + final HandlerContext.Sender sessionState = + InstrumentationContext.get(Session.class, HandlerContext.Sender.class).get(session); + if (sessionState != null) { + InstrumentationContext.get(RemoteEndpoint.class, HandlerContext.Sender.class) + .put(remoteEndpoint, sessionState); + } + } + } + } + + public static class SessionCloseAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope before( + @Advice.This final Session session, + @Advice.Argument(0) final CloseReason reason, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext) { + handlerContext = + InstrumentationContext.get(Session.class, HandlerContext.Sender.class).remove(session); + if (handlerContext == null) { + return null; + } + return activateSpan( + DECORATE.onSessionCloseIssued( + handlerContext, reason.getReasonPhrase(), reason.getCloseCode().getCode())); + } + + @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) + public static void after( + @Advice.Enter final AgentScope scope, + @Advice.Thrown final Throwable thrown, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext) { + if (scope != null) { + DECORATE.onError(scope, thrown); + DECORATE.onFrameEnd(handlerContext); + scope.close(); + } + } + } + + public static class DefaultSessionCloseAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope before( + @Advice.This final Session session, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext) { + + handlerContext = + InstrumentationContext.get(Session.class, HandlerContext.Sender.class).remove(session); + if (handlerContext == null) { + return null; + } + return activateSpan(DECORATE.onSessionCloseIssued(handlerContext, null, 1000)); + } + + @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) + public static void after( + @Advice.Enter final AgentScope scope, + @Advice.Thrown final Throwable thrown, + @Advice.Local("handlerContext") HandlerContext.Sender handlerContext) { + if (scope != null) { + DECORATE.onError(scope, thrown); + DECORATE.onFrameEnd(handlerContext); + scope.close(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/TracingOutputStream.java b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/TracingOutputStream.java new file mode 100644 index 00000000000..6aa3d2f01eb --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/TracingOutputStream.java @@ -0,0 +1,68 @@ +package datadog.trace.instrumentation.websocket.jsr256; + +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE; +import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_BINARY; + +import datadog.trace.bootstrap.CallDepthThreadLocalMap; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; +import java.io.IOException; +import java.io.OutputStream; + +public class TracingOutputStream extends OutputStream { + private final OutputStream delegate; + private final HandlerContext.Sender handlerContext; + + public TracingOutputStream(OutputStream delegate, HandlerContext.Sender handlerContext) { + super(); + this.delegate = delegate; + this.handlerContext = handlerContext; + } + + @Override + public void write(int b) throws IOException { + final boolean doTrace = CallDepthThreadLocalMap.incrementCallDepth(HandlerContext.class) == 0; + if (doTrace) { + DECORATE.onSendFrameStart(handlerContext, MESSAGE_TYPE_BINARY, 1); + } + try (final AgentScope ignored = activateSpan(handlerContext.getWebsocketSpan())) { + delegate.write(b); + } finally { + if (doTrace) { + CallDepthThreadLocalMap.reset(HandlerContext.class); + } + } + } + + @Override + public void write(byte[] b, int off, int len) throws IOException { + final boolean doTrace = CallDepthThreadLocalMap.incrementCallDepth(HandlerContext.class) == 0; + if (doTrace) { + DECORATE.onSendFrameStart(handlerContext, MESSAGE_TYPE_BINARY, len); + } + try (final AgentScope ignored = activateSpan(handlerContext.getWebsocketSpan())) { + delegate.write(b, off, len); + } finally { + if (doTrace) { + CallDepthThreadLocalMap.reset(HandlerContext.class); + } + } + } + + @Override + public void close() throws IOException { + final boolean doTrace = CallDepthThreadLocalMap.incrementCallDepth(HandlerContext.class) == 0; + try (final AgentScope ignored = + handlerContext.getWebsocketSpan() != null + ? activateSpan(handlerContext.getWebsocketSpan()) + : null) { + delegate.close(); + } finally { + if (doTrace) { + CallDepthThreadLocalMap.reset(HandlerContext.class); + DECORATE.onFrameEnd(handlerContext); + } + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/TracingSendHandler.java b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/TracingSendHandler.java new file mode 100644 index 00000000000..1401ae0f502 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/TracingSendHandler.java @@ -0,0 +1,33 @@ +package datadog.trace.instrumentation.websocket.jsr256; + +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE; + +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; +import javax.websocket.SendHandler; +import javax.websocket.SendResult; + +public class TracingSendHandler implements SendHandler { + private final SendHandler delegate; + private final HandlerContext handlerContext; + + public TracingSendHandler(SendHandler delegate, HandlerContext handlerContext) { + this.delegate = delegate; + this.handlerContext = handlerContext; + } + + @Override + public void onResult(SendResult sendResult) { + final AgentSpan wsSpan = handlerContext.getWebsocketSpan(); + try (final AgentScope ignored = activateSpan(wsSpan)) { + delegate.onResult(sendResult); + } finally { + if (sendResult.getException() != null) { + DECORATE.onError(wsSpan, sendResult.getException()); + } + DECORATE.onFrameEnd(handlerContext); + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/TracingWriter.java b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/TracingWriter.java new file mode 100644 index 00000000000..4d37d1ba307 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/TracingWriter.java @@ -0,0 +1,58 @@ +package datadog.trace.instrumentation.websocket.jsr256; + +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE; +import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_TEXT; + +import datadog.trace.bootstrap.CallDepthThreadLocalMap; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; +import java.io.IOException; +import java.io.Writer; + +public class TracingWriter extends Writer { + private final Writer delegate; + private final HandlerContext.Sender handlerContext; + + public TracingWriter(Writer delegate, HandlerContext.Sender handlerContext) { + super(); + this.delegate = delegate; + this.handlerContext = handlerContext; + } + + @Override + public void write(char[] cbuf, int off, int len) throws IOException { + final boolean doTrace = CallDepthThreadLocalMap.incrementCallDepth(HandlerContext.class) == 0; + if (doTrace) { + DECORATE.onSendFrameStart(handlerContext, MESSAGE_TYPE_TEXT, len); + } + try (final AgentScope ignored = activateSpan(handlerContext.getWebsocketSpan())) { + delegate.write(cbuf, off, len); + } finally { + if (doTrace) { + CallDepthThreadLocalMap.reset(HandlerContext.class); + } + } + } + + @Override + public void flush() throws IOException { + delegate.flush(); + } + + @Override + public void close() throws IOException { + final boolean doTrace = CallDepthThreadLocalMap.incrementCallDepth(HandlerContext.class) == 0; + try (final AgentScope ignored = + handlerContext.getWebsocketSpan() != null + ? activateSpan(handlerContext.getWebsocketSpan()) + : null) { + delegate.close(); + } finally { + if (doTrace) { + CallDepthThreadLocalMap.reset(HandlerContext.class); + DECORATE.onFrameEnd(handlerContext); + } + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/EndpointWrapper.groovy b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/EndpointWrapper.groovy new file mode 100644 index 00000000000..3889b0e0d2a --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/EndpointWrapper.groovy @@ -0,0 +1,39 @@ +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan + +import javax.websocket.CloseReason +import javax.websocket.Endpoint +import javax.websocket.EndpointConfig +import javax.websocket.Session + +class EndpointWrapper extends Endpoint { + + EndpointWrapper() { + } + + @Override + void onOpen(Session session, EndpointConfig endpointConfig) { + def span = endpointConfig.getUserProperties().get(AgentSpan.class.getName()) as AgentSpan + def endpoint = endpointConfig.getUserProperties().get(Endpoint.class.getName()) as Endpoint + assert endpoint != null + session.getUserProperties().put(Endpoint.class.getName(), endpoint) + try (def ignored = span != null ? activateSpan(span) : null) { + endpoint.onOpen(session, endpointConfig) + } + } + + @Override + void onClose(Session session, CloseReason closeReason) { + def endpoint = session.getUserProperties().get(Endpoint.class.getName()) as Endpoint + assert endpoint != null + endpoint.onClose(session, closeReason) + } + + @Override + void onError(Session session, Throwable thr) { + def endpoint = session.getUserProperties().get(Endpoint.class.getName()) as Endpoint + assert endpoint != null + endpoint.onError(session, thr) + } +} diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/Endpoints.groovy b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/Endpoints.groovy new file mode 100644 index 00000000000..232c7ee3849 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/Endpoints.groovy @@ -0,0 +1,130 @@ +import org.apache.commons.io.IOUtils + +import javax.websocket.DecodeException +import javax.websocket.Decoder +import javax.websocket.EncodeException +import javax.websocket.Encoder +import javax.websocket.Endpoint +import javax.websocket.EndpointConfig +import javax.websocket.MessageHandler +import javax.websocket.Session +import java.nio.ByteBuffer + +class Endpoints { + static class ClientTestEndpoint extends Endpoint { + @Override + void onOpen(Session session, EndpointConfig endpointConfig) { + } + } + + static class TestEndpoint extends Endpoint { + final MessageHandler handler + + TestEndpoint(final MessageHandler handler) { + this.handler = handler + } + + @Override + void onOpen(Session session, EndpointConfig endpointConfig) { + session.addMessageHandler(handler) + } + } + + // Full handlers + static class FullStringHandler implements MessageHandler.Whole { + @Override + void onMessage(String s) { + } + } + + static class FullReaderHandler implements MessageHandler.Whole { + @Override + void onMessage(Reader reader) { + IOUtils.toString(reader) + } + } + + static class FullBytesHandler implements MessageHandler.Whole { + @Override + void onMessage(byte[] b) { + } + } + + static class FullByteBufferHandler implements MessageHandler.Whole { + @Override + void onMessage(ByteBuffer b) { + } + } + + static class CustomMessage { + } + + static class CustomMessageEncoder implements Encoder.Text { + + @Override + void init(EndpointConfig endpointConfig) { + } + + @Override + void destroy() { + } + + @Override + String encode(CustomMessage customMessage) throws EncodeException { + return "CustomMessage" + } + } + + static class CustomMessageDecoder implements Decoder.Text { + + @Override + void init(EndpointConfig endpointConfig) { + } + + @Override + void destroy() { + } + + @Override + CustomMessage decode(String s) throws DecodeException { + return new CustomMessage() + } + + @Override + boolean willDecode(String s) { + return s == "CustomMessage" + } + } + + static class FullObjectHandler implements MessageHandler.Whole { + @Override + void onMessage(CustomMessage o) { + } + } + + static class FullStreamHandler implements MessageHandler.Whole { + @Override + void onMessage(InputStream is) { + IOUtils.toByteArray(is) + } + } + // Partial Handlers + static class PartialStringHandler implements MessageHandler.Partial { + @Override + void onMessage(String s, boolean last) { + } + } + + static class PartialBytesHandler implements MessageHandler.Partial { + @Override + void onMessage(byte[] b, boolean last) { + + } + } + + static class PartialByteBufferHandler implements MessageHandler.Partial { + @Override + void onMessage(ByteBuffer b, boolean last) { + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/WebsocketTest.groovy b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/WebsocketTest.groovy new file mode 100644 index 00000000000..948ba636ce2 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/WebsocketTest.groovy @@ -0,0 +1,577 @@ +import static datadog.trace.agent.test.base.HttpServerTest.someBytes +import static datadog.trace.agent.test.base.HttpServerTest.websocketCloseSpan +import static datadog.trace.agent.test.base.HttpServerTest.websocketReceiveSpan +import static datadog.trace.agent.test.base.HttpServerTest.websocketSendSpan +import static datadog.trace.agent.test.utils.TraceUtils.basicSpan +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_CLASSES_EXCLUDE +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_ENABLED +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_INHERIT_SAMPLING +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_SEPARATE_TRACES +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_TAG_SESSION_ID +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan + +import datadog.trace.agent.test.AgentTestRunner +import datadog.trace.api.DDTags +import datadog.trace.bootstrap.instrumentation.api.AgentSpan +import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags +import datadog.trace.bootstrap.instrumentation.api.Tags +import datadog.trace.core.DDSpan +import net.bytebuddy.utility.RandomString +import org.glassfish.tyrus.container.inmemory.InMemoryClientContainer +import org.glassfish.tyrus.server.TyrusServerConfiguration + +import javax.websocket.ClientEndpointConfig +import javax.websocket.CloseReason +import javax.websocket.ContainerProvider +import javax.websocket.Endpoint +import javax.websocket.server.ServerApplicationConfig +import javax.websocket.server.ServerEndpointConfig +import java.nio.ByteBuffer + +class WebsocketTest extends AgentTestRunner { + + @Override + protected void configurePreAgent() { + super.configurePreAgent() + injectSysConfig(TRACE_WEBSOCKET_MESSAGES_ENABLED, "true") + injectSysConfig(TRACE_CLASSES_EXCLUDE, "EndpointWrapper") + } + + def createHandshakeSpan(String spanName, String url) { + def span = TEST_TRACER.startSpan("test", spanName, null) + handshakeTags(url).each { span.setTag(it.key, it.value) } + span.finish() + span + } + + def deployEndpointAndConnect(Endpoint endpoint, Object handshakeClientSpan, Object handshakeServerSpan, String url) { + def webSocketContainer = ContainerProvider.getWebSocketContainer() + def sec = ServerEndpointConfig.Builder.create(EndpointWrapper, "/test") + .encoders([Endpoints.CustomMessageEncoder]) + .decoders([Endpoints.CustomMessageDecoder]) + .build() + + sec.getUserProperties().put(Endpoint.class.getName(), endpoint) + sec.getUserProperties().put(AgentSpan.class.getName(), handshakeServerSpan) + + final ServerApplicationConfig serverConfig = + new TyrusServerConfiguration(Collections.singleton(EndpointWrapper.class), + Collections.singleton(sec)) + + ClientEndpointConfig cec = ClientEndpointConfig.Builder.create() + .encoders([Endpoints.CustomMessageEncoder]) + .decoders([Endpoints.CustomMessageDecoder]) + .build() + cec.getUserProperties().put(InMemoryClientContainer.SERVER_CONFIG, serverConfig) + + try (def ignored = handshakeClientSpan != null ? activateSpan(handshakeClientSpan as AgentSpan) : null) { + def session = webSocketContainer.connectToServer(new Endpoints.ClientTestEndpoint(), cec, URI.create(url)) + session + } + } + + def handshakeTags(url) { + [(Tags.HTTP_METHOD): "GET", (Tags.HTTP_URL): url] + } + + def "test full sync send and receive for endpoint #endpoint.class with #sendSize len #msgType message"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(endpoint), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + if (msgType == "text") { + session.getBasicRemote().sendText(message as String) + } else { + session.getBasicRemote().sendBinary(ByteBuffer.wrap(message as byte[])) + } + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, msgType, sendSize, 1, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + + trace(1) { + websocketReceiveSpan(it, serverHandshake, msgType, rcvSize, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + where: + endpoint | message | msgType | sendSize | rcvSize + // text full + new Endpoints.FullStringHandler() | RandomString.make(10) | "text" | 10 | 10 + new Endpoints.FullReaderHandler() | RandomString.make(20) | "text" | 20 | 0 + // binary full + new Endpoints.FullByteBufferHandler() | someBytes(10) | "binary" | 10 | 10 + new Endpoints.FullBytesHandler() | someBytes(25) | "binary" | 25 | 25 + new Endpoints.FullStreamHandler() | someBytes(10) | "binary" | 10 | 0 + } + + def "test partial sync send and receive for endpoint #endpoint.class with #chunks chunks and #size len #msgType message"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(endpoint), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + def remote = session.getBasicRemote() + if (msgType == "text") { + for (int i = 0; i < message.size(); i++) { + remote.sendText(message[i] as String, (i == message.size() - 1)) + } + } else { + for (int i = 0; i < message.size(); i++) { + remote.sendBinary(ByteBuffer.wrap(message[i] as byte[]), i == (message.size() - 1)) + } + } + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, msgType, size, chunks, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + trace(1) { + websocketReceiveSpan(it, serverHandshake, msgType, size, chunks) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + where: + endpoint | message | msgType | chunks | size + // text partial + new Endpoints.PartialStringHandler() | [RandomString.make(10)] | "text" | 1 | 10 + new Endpoints.PartialStringHandler() | [RandomString.make(10), RandomString.make(10)] | "text" | 2 | 20 + // binary Partial + new Endpoints.PartialByteBufferHandler() | [someBytes(10)] | "binary" | 1 | 10 + new Endpoints.PartialByteBufferHandler() | [someBytes(10), someBytes(15)] | "binary" | 2 | 25 + new Endpoints.PartialBytesHandler() | [someBytes(10)] | "binary" | 1 | 10 + new Endpoints.PartialBytesHandler() | [someBytes(10), someBytes(15)] | "binary" | 2 | 25 + } + + def "test stream sync send and receive for endpoint #endpoint.class with #chunks chunks and #size len #msgType message"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(endpoint), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + if (msgType == "text") { + // tyrus text writer send the flushes when closing so the fin bit is sent along with the last message + try (def writer = session.getBasicRemote().getSendWriter()) { + message.each { + writer.write(it as String) + } + } + } else { + // tyrus binary writer send the fin bit with an empty frame when close is called + try (def os = session.getBasicRemote().getSendStream()) { + message.each { + os.write(it as byte[]) + } + } + } + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, msgType, size, chunks, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + trace(1) { + websocketReceiveSpan(it, serverHandshake, msgType, size, chunks) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + where: + endpoint | message | msgType | chunks | size + // text partial + new Endpoints.PartialStringHandler() | [RandomString.make(10)] | "text" | 1 | 10 + new Endpoints.PartialStringHandler() | [RandomString.make(10), RandomString.make(15)] | "text" | 2 | 25 + // binary Partial + new Endpoints.PartialByteBufferHandler() | [someBytes(10)] | "binary" | 1 | 10 + new Endpoints.PartialByteBufferHandler() | [someBytes(10), someBytes(15)] | "binary" | 2 | 25 + new Endpoints.PartialBytesHandler() | [someBytes(10)] | "binary" | 1 | 10 + new Endpoints.PartialBytesHandler() | [someBytes(10), someBytes(15)] | "binary" | 2 | 25 + } + + def "test full async (future) send and receive for endpoint #endpoint.class with #sendSize len #msgType message"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(endpoint), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + if (msgType == "text") { + session.getAsyncRemote().sendText(message as String) + } else { + session.getAsyncRemote().sendBinary(ByteBuffer.wrap(message as byte[])) + } + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, msgType, sendSize, 1, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + + trace(1) { + websocketReceiveSpan(it, serverHandshake, msgType, rcvSize, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + where: + endpoint | message | msgType | sendSize | rcvSize + // text full + new Endpoints.FullStringHandler() | RandomString.make(10) | "text" | 10 | 10 + new Endpoints.FullReaderHandler() | RandomString.make(20) | "text" | 20 | 0 + // binary full + new Endpoints.FullByteBufferHandler() | someBytes(10) | "binary" | 10 | 10 + new Endpoints.FullBytesHandler() | someBytes(25) | "binary" | 25 | 25 + new Endpoints.FullStreamHandler() | someBytes(10) | "binary" | 10 | 0 + } + + def "test full async (SendHandler) send and receive for endpoint #endpoint.class with #sendSize len #msgType message"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(endpoint), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + if (msgType == "text") { + session.getAsyncRemote().sendText(message as String, { assert it.OK }) + } else { + session.getAsyncRemote().sendBinary(ByteBuffer.wrap(message as byte[]), { assert it.OK }) + } + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, msgType, sendSize, 1, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + + trace(1) { + websocketReceiveSpan(it, serverHandshake, msgType, rcvSize, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + where: + endpoint | message | msgType | sendSize | rcvSize + // text full + new Endpoints.FullStringHandler() | RandomString.make(10) | "text" | 10 | 10 + new Endpoints.FullReaderHandler() | RandomString.make(20) | "text" | 20 | 0 + // binary full + new Endpoints.FullByteBufferHandler() | someBytes(10) | "binary" | 10 | 10 + new Endpoints.FullBytesHandler() | someBytes(25) | "binary" | 25 | 25 + new Endpoints.FullStreamHandler() | someBytes(10) | "binary" | 10 | 0 + } + + def "test session close code #code and reason #reason"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullStringHandler()), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + session.close(new CloseReason(CloseReason.CloseCodes.getCloseCode(code), reason)) + } + then: + assertTraces(4, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(2) { + sortSpansByStart() + basicSpan(it, "parent") + websocketCloseSpan(it, clientHandshake, true, code, reason, span(0)) + } + + trace(1) { + websocketCloseSpan(it, serverHandshake, false, code, reason) + } + }) + where: + code | reason + 1000 | null + 1000 | "bye" + 1001 | "see you" + } + + def "test session id logged as tag"() { + setup: + injectSysConfig(TRACE_WEBSOCKET_TAG_SESSION_ID, "true") + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullStringHandler()), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + session.getBasicRemote().sendText("Hello") + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, "text", 5, 1, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + trace(1) { + websocketReceiveSpan(it, serverHandshake, "text", 5, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + //it's normally already tested by websocket(Receive|Send|Close)Span but we enforce that check + TEST_WRITER.flatten().findAll { span -> (span as DDSpan).getSpanType() == "websocket" }.each { + assert (it as DDSpan).getTag(InstrumentationTags.WEBSOCKET_SESSION_ID) != null + } + } + + def "test close and receive on same handshake trace"() { + setup: + injectSysConfig(TRACE_WEBSOCKET_MESSAGES_SEPARATE_TRACES, "false") + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullStringHandler()), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + session.getBasicRemote().sendText("Hello") + session.close() + } + then: + // in reality we have 3 traces but since the handshake finishes soon, the trace structure writer is collecting 5 chunks + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(1) { + websocketReceiveSpan(it, serverHandshake, "text", 5, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }, serverHandshake) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, "text", 5, 1, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + }) + } + + def "test sampling not inherited"() { + setup: + injectSysConfig(TRACE_WEBSOCKET_MESSAGES_INHERIT_SAMPLING, "false") + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullStringHandler()), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + session.getBasicRemote().sendText("Hello") + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, "text", 5, 1, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + trace(1) { + websocketReceiveSpan(it, serverHandshake, "text", 5, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + //it's normally already tested by websocket(Receive|Send|Close)Span but we enforce that check + TEST_WRITER.flatten().findAll { span -> (span as DDSpan).getSpanType() == "websocket" }.each { + assert (it as DDSpan).getTag(DDTags.DECISION_MAKER_INHERITED) == null + } + } + + def "if handshake is not captured traces are not generated"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullStringHandler()), + null, null, url) + + runUnderTrace("parent") { + session.getBasicRemote().sendText("Hello") + session.close() + } + then: + assertTraces(1, { + trace(1) { + basicSpan(it, "parent") + } + }) + } + + def "test send and receive object"() { + when: + String url = "ws://inmemory/test" + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullObjectHandler()), + createHandshakeSpan("http.request", url), //simulate client span + createHandshakeSpan("servlet.request", url), // simulate server span + url) + + runUnderTrace("parent") { + session.getBasicRemote().sendObject(new Endpoints.CustomMessage()) + session.close() + } + then: + assertTraces(5, { + DDSpan serverHandshake, clientHandshake + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) + clientHandshake = span(0) + } + trace(1) { + basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) + serverHandshake = span(0) + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake, null, 0, 1, span(0)) + websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) + } + trace(1) { + websocketReceiveSpan(it, serverHandshake, null, 0, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }) + } + }) + } +} diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy new file mode 100644 index 00000000000..97ef890f266 --- /dev/null +++ b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy @@ -0,0 +1,50 @@ +package datadog.trace.agent.test.asserts + + +import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink +import datadog.trace.bootstrap.instrumentation.api.SpanAttributes +import datadog.trace.bootstrap.instrumentation.api.SpanLink +import datadog.trace.core.DDSpan +import groovy.transform.stc.ClosureParams +import groovy.transform.stc.SimpleType + +class LinksAssert { + private final List links + private final Set assertedLinks = [] + + private LinksAssert(DDSpan span) { + this.links = span.links // this is class protected but for the moment groovy can access it + } + + static void assertLinks(DDSpan span, + @ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.LinksAssert']) + @DelegatesTo(value = LinksAssert, strategy = Closure.DELEGATE_FIRST) Closure spec, + boolean checkAllLinks = true) { + def asserter = new LinksAssert(span) + def clone = (Closure) spec.clone() + clone.delegate = asserter + clone.resolveStrategy = Closure.DELEGATE_FIRST + clone(asserter) + if (checkAllLinks) { + asserter.assertLinksAllVerified() + } + } + + def link(DDSpan linked, byte flags = SpanLink.DEFAULT_FLAGS, SpanAttributes attributes = SpanAttributes.EMPTY, String traceState = '') { + def found = links.find { + it.spanId() == linked.spanId && + it.traceId() == linked.traceId + } + assert found != null + assert found.traceFlags() == flags + assert found.attributes() == attributes + assert found.traceState() == traceState + assertedLinks.add(found) + } + + void assertLinksAllVerified() { + def list = new ArrayList(links) + list.removeAll(assertedLinks) + assert list.isEmpty() + } +} diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy index 6dfc6933e18..c4be54bf7c8 100644 --- a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy +++ b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy @@ -1,5 +1,8 @@ package datadog.trace.agent.test.asserts +import static TagsAssert.assertTags +import static datadog.trace.agent.test.asserts.LinksAssert.assertLinks + import datadog.trace.api.DDSpanId import datadog.trace.api.DDTraceId import datadog.trace.core.DDSpan @@ -8,8 +11,6 @@ import groovy.transform.stc.SimpleType import java.util.regex.Pattern -import static TagsAssert.assertTags - class SpanAssert { private final DDSpan span private final DDSpan previous @@ -167,6 +168,9 @@ class SpanAssert { if (!checked.errored) { errored(false) } + if (!checked.links) { + assert span.tags['_dd.span_links'] == null + } hasServiceName() } @@ -181,6 +185,19 @@ class SpanAssert { assertTags(span, spec, checkAllTags) } + void links(@ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.LinksAssert']) + @DelegatesTo(value = LinksAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) { + checked.links = true + assertLinks(span, spec) + } + + void links(boolean checkAllLinks, + @ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.LinksAssert']) + @DelegatesTo(value = LinksAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) { + checked.links = true + assertLinks(span, spec, checkAllLinks) + } + DDSpan getSpan() { return span } diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/TagsAssert.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/TagsAssert.groovy index cb361270cec..bbd7e8593a5 100644 --- a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/TagsAssert.groovy +++ b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/TagsAssert.groovy @@ -91,6 +91,7 @@ class TagsAssert { assertedTags.add(DDTags.DSM_ENABLED) assertedTags.add(DDTags.DJM_ENABLED) assertedTags.add(DDTags.PARENT_ID) + assertedTags.add(DDTags.SPAN_LINKS) // this is checked by LinksAsserter assert tags["thread.name"] != null assert tags["thread.id"] != null diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy index 6dbc2ac8510..0edf3993aaa 100644 --- a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy +++ b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy @@ -29,6 +29,8 @@ import datadog.trace.api.normalize.SimpleHttpPathNormalizer import datadog.trace.bootstrap.blocking.BlockingActionHelper import datadog.trace.bootstrap.instrumentation.api.AgentTracer import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags +import datadog.trace.bootstrap.instrumentation.api.SpanAttributes +import datadog.trace.bootstrap.instrumentation.api.SpanLink import datadog.trace.bootstrap.instrumentation.api.Tags import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter import datadog.trace.bootstrap.instrumentation.api.URIUtils @@ -37,12 +39,15 @@ import datadog.trace.core.datastreams.StatsGroup import datadog.trace.test.util.Flaky import groovy.transform.Canonical import groovy.transform.CompileStatic +import net.bytebuddy.utility.RandomString import okhttp3.HttpUrl import okhttp3.MediaType import okhttp3.MultipartBody import okhttp3.Request import okhttp3.RequestBody import okhttp3.Response +import okhttp3.WebSocketListener +import okio.ByteString import org.slf4j.Logger import org.slf4j.LoggerFactory @@ -74,12 +79,14 @@ import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.TIMEOU import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.TIMEOUT_ERROR import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.UNKNOWN import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.USER_BLOCK +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.WEBSOCKET import static datadog.trace.agent.test.utils.TraceUtils.basicSpan import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_SERVER_RAW_QUERY_STRING import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_SERVER_RAW_RESOURCE import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_SERVER_TAG_QUERY_STRING import static datadog.trace.api.config.TraceInstrumentationConfig.SERVLET_ASYNC_TIMEOUT_ERROR +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_ENABLED import static datadog.trace.api.config.TracerConfig.HEADER_TAGS import static datadog.trace.api.config.TracerConfig.REQUEST_HEADER_TAGS import static datadog.trace.api.config.TracerConfig.RESPONSE_HEADER_TAGS @@ -153,6 +160,7 @@ abstract class HttpServerTest extends WithHttpServer { injectSysConfig(HEADER_TAGS, 'x-datadog-test-both-header:both_header_tag') injectSysConfig(REQUEST_HEADER_TAGS, 'x-datadog-test-request-header:request_header_tag') // We don't inject a matching response header tag here since it would be always on and show up in all the tests + injectSysConfig(TRACE_WEBSOCKET_MESSAGES_ENABLED, "true") } // used in blocking tests to check if the handler was skipped @@ -211,9 +219,9 @@ abstract class HttpServerTest extends WithHttpServer { // Only used if hasExtraErrorInformation is true Map expectedExtraErrorInformation(ServerEndpoint endpoint) { if (endpoint.errored) { - ["error.message" : { it == null || it == EXCEPTION.body }, - "error.type" : { it == null || it == Exception.name }, - "error.stack": { it == null || it instanceof String }] + ["error.message": { it == null || it == EXCEPTION.body }, + "error.type" : { it == null || it == Exception.name }, + "error.stack" : { it == null || it instanceof String }] } else { Collections.emptyMap() } @@ -392,6 +400,10 @@ abstract class HttpServerTest extends WithHttpServer { true } + boolean testWebsockets() { + server instanceof WebsocketServer + } + @Override int version() { return 0 @@ -444,6 +456,7 @@ abstract class HttpServerTest extends WithHttpServer { SECURE_SUCCESS("secure/success", 200, null), SESSION_ID("session", 200, null), + WEBSOCKET("websocket", 101, null) private final String path private final String rawPath @@ -503,8 +516,8 @@ abstract class HttpServerTest extends WithHttpServer { } private static final Map PATH_MAP = { - Map map = values().collectEntries { [it.path, it]} - map.putAll(values().collectEntries { [it.rawPath, it]}) + Map map = values().collectEntries { [it.path, it] } + map.putAll(values().collectEntries { [it.rawPath, it] }) map }.call() @@ -550,7 +563,7 @@ abstract class HttpServerTest extends WithHttpServer { } responses = (1..count).collect { completionService.take().get() } } else { - responses = (1..count).collect {client.newCall(request).execute()} + responses = (1..count).collect { client.newCall(request).execute() } } if (isDataStreamsEnabled()) { @@ -723,9 +736,9 @@ abstract class HttpServerTest extends WithHttpServer { } where: - method | body | header | value | tags - 'GET' | null | 'x-datadog-test-both-header' | 'foo' | [ 'both_header_tag': 'foo' ] - 'GET' | null | 'x-datadog-test-request-header' | 'bar' | [ 'request_header_tag': 'bar' ] + method | body | header | value | tags + 'GET' | null | 'x-datadog-test-both-header' | 'foo' | ['both_header_tag': 'foo'] + 'GET' | null | 'x-datadog-test-request-header' | 'bar' | ['request_header_tag': 'bar'] } @Flaky(value = "https://github.com/DataDog/dd-trace-java/issues/4690", suites = ["MuleHttpServerForkedTest"]) @@ -737,7 +750,7 @@ abstract class HttpServerTest extends WithHttpServer { def body = null def header = IG_RESPONSE_HEADER def mapping = 'mapped_response_header_tag' - def tags = ['mapped_response_header_tag': "$IG_RESPONSE_HEADER_VALUE" ] + def tags = ['mapped_response_header_tag': "$IG_RESPONSE_HEADER_VALUE"] injectSysConfig(HTTP_SERVER_TAG_QUERY_STRING, "true") injectSysConfig(RESPONSE_HEADER_TAGS, "$header:$mapping") @@ -820,13 +833,13 @@ abstract class HttpServerTest extends WithHttpServer { } where: - rawQuery | endpoint | encoded - true | SUCCESS | false - true | QUERY_PARAM | false - true | QUERY_ENCODED_QUERY | true - false | SUCCESS | false - false | QUERY_PARAM | false - false | QUERY_ENCODED_QUERY | true + rawQuery | endpoint | encoded + true | SUCCESS | false + true | QUERY_PARAM | false + true | QUERY_ENCODED_QUERY | true + false | SUCCESS | false + false | QUERY_PARAM | false + false | QUERY_ENCODED_QUERY | true method = "GET" body = null @@ -939,7 +952,7 @@ abstract class HttpServerTest extends WithHttpServer { } then: - DDSpan span = TEST_WRITER.flatten().find {it.operationName =='appsec-span' } + DDSpan span = TEST_WRITER.flatten().find { it.operationName == 'appsec-span' } span.getTag(IG_PATH_PARAMS_TAG) == expectedIGPathParams() and: @@ -1633,7 +1646,7 @@ abstract class HttpServerTest extends WithHttpServer { then: TEST_WRITER.waitForTraces(1) def trace = TEST_WRITER.get(0) - assert trace.find {it.isError() } == null + assert trace.find { it.isError() } == null } def 'test blocking of request for path parameters'() { @@ -1715,7 +1728,8 @@ abstract class HttpServerTest extends WithHttpServer { if (testBlockingErrorTypeSet()) { spans.find { it.error && - it.tags['error.type'] == BlockingException.name } != null + it.tags['error.type'] == BlockingException.name + } != null } and: @@ -1897,13 +1911,195 @@ abstract class HttpServerTest extends WithHttpServer { if (isDataStreamsEnabled()) { TEST_DATA_STREAMS_WRITER.waitForGroups(1) } - DDSpan span = TEST_WRITER.flatten().find {it.operationName =='appsec-span' } + DDSpan span = TEST_WRITER.flatten().find { it.operationName == 'appsec-span' } span != null final sessionId = span.tags[IG_SESSION_ID_TAG] sessionId != null secondResponse.body().string().contains(sessionId as String) } + + def 'test websocket server send #msgType message of size #size and #chunks chunks'() { + setup: + assumeTrue(testWebsockets()) + def wsServer = getServer() as WebsocketServer + + when: + def request = new Request.Builder().url(HttpUrl.get(WEBSOCKET.resolve(address))) + .get().build() + + client.newWebSocket(request, new WebSocketListener() {}) + wsServer.awaitConnected() + runUnderTrace("parent", { + if (messages[0] instanceof String) { + wsServer.serverSendText(messages as String[]) + } else { + wsServer.serverSendBinary(messages as byte[][]) + } + wsServer.serverClose() + }) + + then: + assertTraces(2, { + DDSpan handshake + trace(hasHandlerSpan() ? 2 : 1, { + handshake = span(0) + serverSpan(it, null, null, "GET", WEBSOCKET) + if (hasHandlerSpan()) { + handlerSpan(it, WEBSOCKET) + } + }) + trace(3, { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, handshake, msgType, size, chunks, span(0)) + websocketCloseSpan(it, handshake, true, 1000, null, span(0)) + }) + }) + + where: + + messages | msgType | chunks | size + [RandomString.make(10)] | "text" | 1 | 10 + [someBytes(20)] | "binary" | 1 | 20 + [RandomString.make(10), RandomString.make(5)] | "text" | 2 | 10 + 5 + [someBytes(10), someBytes(15), someBytes(30)] | "binary" | 3 | 10 + 15 + 30 + } + + def 'test websocket server receive #msgType message of size #size and #chunks chunks'() { + setup: + assumeTrue(testWebsockets()) + def wsServer = getServer() as WebsocketServer + assumeTrue(chunks == 1 || wsServer.canSplitLargeWebsocketPayloads()) + + when: + def request = new Request.Builder().url(HttpUrl.get(WEBSOCKET.resolve(address))) + .get().build() + + def ws = client.newWebSocket(request, new WebSocketListener() {}) + wsServer.awaitConnected() + wsServer.setMaxPayloadSize(10) + if (message instanceof String) { + ws.send(message as String) + } else { + ws.send(ByteString.of(message as byte[])) + } + ws.close(1000, "goodbye") + + then: + assertTraces(3, { + DDSpan handshake + trace(hasHandlerSpan() ? 2 : 1, { + handshake = span(0) + serverSpan(it, null, null, "GET", WEBSOCKET) + if (hasHandlerSpan()) { + handlerSpan(it, WEBSOCKET) + } + }) + trace(1 + chunks, { + websocketReceiveSpan(it, handshake, msgType, size, chunks) + for (int i = 0; i < chunks; i++) { + basicSpan(it, "onRead", span(0)) + } + }) + trace(1, { + websocketCloseSpan(it, handshake, false, 1000, "goodbye") + }) + }) + where: + + message | msgType | chunks | size + RandomString.make(10) | "text" | 1 | 10 + someBytes(10) | "binary" | 1 | 10 + RandomString.make(20) | "text" | 2 | 20 + someBytes(30) | "binary" | 3 | 30 + } + + static def someBytes(nb) { + def b = new byte[nb] + new Random().nextBytes(b) + b + } + + static void websocketSendSpan(TraceAssert trace, DDSpan handshake, String messageType, int messageLength, + int nbOfChunks = 1, DDSpan parentSpan = null, Map extraTags = [:]) { + websocketSpan(trace, handshake, "websocket.send", messageType, messageLength, nbOfChunks, + false, parentSpan, extraTags) + } + + static void websocketReceiveSpan(TraceAssert trace, DDSpan handshake, String messageType, int messageLength, int nbOfChunks = 1, Map extraTags = [:]) { + websocketSpan(trace, handshake, "websocket.receive", messageType, messageLength, nbOfChunks, true, + Config.get().isWebsocketMessagesSeparateTraces() ? null : handshake, + extraTags + [(InstrumentationTags.WEBSOCKET_MESSAGE_RECEIVE_TIME): { Number }]) + } + + static void websocketCloseSpan(TraceAssert trace, DDSpan handshake, boolean closeStarter, int closeCode, closeReason = null, + DDSpan parentSpan = null, Map extraTags = [:]) { + Map tags = new HashMap(extraTags) + tags.put(InstrumentationTags.WEBSOCKET_CLOSE_REASON, closeReason) + tags.put(InstrumentationTags.WEBSOCKET_CLOSE_CODE, closeCode) + websocketSpan(trace, handshake, "websocket.close", null, null, null, !closeStarter, parentSpan, tags) + } + + static void websocketSpan(TraceAssert trace, DDSpan handshake, String operation, + String messageType, Integer messageLength, Integer nbOfChunks, + boolean traceStarter, + DDSpan parentSpan, + Map extraTags = [:]) { + byte linkFlags = SpanLink.DEFAULT_FLAGS + if (handshake.getSamplingPriority() > 0 && Config.get().isWebsocketMessagesInheritSampling()) { + linkFlags |= SpanLink.SAMPLED_FLAG + } + + def linkAttributes = SpanAttributes.builder() + .put("dd.kind", traceStarter ? "executed_from" : "resuming") + .build() + + trace.span { + operationName operation + if (handshake.getTag(Tags.HTTP_ROUTE) != null) { + resourceName "websocket ${handshake.getTag(Tags.HTTP_ROUTE) as String}" + } else { + resourceName "websocket ${URI.create(handshake.getTag(Tags.HTTP_URL) as String).path}" + } + if (traceStarter && Config.get().isWebsocketMessagesSeparateTraces()) { + parent() + } else { + if (parentSpan != null) { + childOf(parentSpan) + } else { + childOfPrevious() + } + } + spanType(DDSpanTypes.WEBSOCKET) + if (Config.get().isWebsocketMessagesSeparateTraces() || !traceStarter) { + links { + link(handshake, linkFlags, linkAttributes) + } + } + tags { + tag(Tags.SPAN_KIND, traceStarter ? Tags.SPAN_KIND_CONSUMER : Tags.SPAN_KIND_PRODUCER) + tag(Tags.COMPONENT, "websocket") + if (traceStarter && Config.get().isWebsocketMessagesSeparateTraces()) { + if (Config.get().isWebsocketMessagesInheritSampling()) { + tag(DDTags.DECISION_MAKER_INHERITED, 1) + tag(DDTags.DECISION_MAKER_SERVICE, handshake.getServiceName()) + tag(DDTags.DECISION_MAKER_RESOURCE, handshake.getResourceName()) + } + } + tag(InstrumentationTags.WEBSOCKET_MESSAGE_LENGTH, messageLength) + tag(InstrumentationTags.WEBSOCKET_MESSAGE_TYPE, messageType) + tag(InstrumentationTags.WEBSOCKET_MESSAGE_FRAMES, nbOfChunks) + tag(Tags.PEER_HOSTNAME, handshake.getTag(Tags.PEER_HOSTNAME)) + if (Config.get().isWebsocketTagSessionId()) { + tag(InstrumentationTags.WEBSOCKET_SESSION_ID, { it != null }) // it can be an incremental thing + } + extraTags.each { tag(it.key, it.value) } + defaultTagsNoPeerService() + } + } + } + void controllerSpan(TraceAssert trace, ServerEndpoint endpoint = null) { def exception = endpoint == CUSTOM_EXCEPTION ? expectedCustomExceptionType() : expectedExceptionType() def errorMessage = endpoint?.body diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/WebsocketServer.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/WebsocketServer.groovy new file mode 100644 index 00000000000..1a318347be8 --- /dev/null +++ b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/WebsocketServer.groovy @@ -0,0 +1,39 @@ +package datadog.trace.agent.test.base + +/** + * Implement this interface if the tested server supports websockets. + * The server is supposed to handle both binary and text messages. + * Each time a message is received, a span called `onRead` should be logged. + */ +interface WebsocketServer extends HttpServer { + /** + * Blocks until connected. + */ + void awaitConnected() + /** + * Send text fragments from the current server active session. + * @param messages + */ + void serverSendText(String[] messages) + /** + * Send binary fragments from the current server active session. + * @param messages + */ + void serverSendBinary(byte[][] binaries) + /** + * Close the active server session. + */ + void serverClose() + /** + * Set the max size for both text and binary payloads on the active session. + * @param size + */ + void setMaxPayloadSize(int size) + /** + * If false, receiver tests with multiple chunks will be skipped. + * @return + */ + default boolean canSplitLargeWebsocketPayloads() { + true + } +} diff --git a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java index f3b58a88f71..6a4158cdf6a 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java @@ -258,6 +258,10 @@ public final class ConfigDefaults { static final boolean DEFAULT_SPARK_APP_NAME_AS_SERVICE = false; static final boolean DEFAULT_JAX_RS_EXCEPTION_AS_ERROR_ENABLED = true; static final boolean DEFAULT_TELEMETRY_DEBUG_REQUESTS_ENABLED = false; + static final boolean DEFAULT_WEBSOCKET_MESSAGES_ENABLED = false; + static final boolean DEFAULT_WEBSOCKET_MESSAGES_INHERIT_SAMPLING = true; + static final boolean DEFAULT_WEBSOCKET_MESSAGES_SEPARATE_TRACES = true; + static final boolean DEFAULT_WEBSOCKET_TAG_SESSION_ID = false; static final Set DEFAULT_TRACE_CLOUD_PAYLOAD_TAGGING_SERVICES = new HashSet<>( diff --git a/dd-trace-api/src/main/java/datadog/trace/api/DDSpanTypes.java b/dd-trace-api/src/main/java/datadog/trace/api/DDSpanTypes.java index c380b6078c9..335dde0effe 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/DDSpanTypes.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/DDSpanTypes.java @@ -37,6 +37,7 @@ public class DDSpanTypes { public static final String MULE = "mule"; public static final String VALKEY = "valkey"; + public static final String WEBSOCKET = "websocket"; public static final String SERVERLESS = "serverless"; } diff --git a/dd-trace-api/src/main/java/datadog/trace/api/DDTags.java b/dd-trace-api/src/main/java/datadog/trace/api/DDTags.java index 11bb60129fb..d8e459e3eee 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/DDTags.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/DDTags.java @@ -77,4 +77,7 @@ public class DDTags { public static final String BASE_SERVICE = "_dd.base_service"; public static final String PARENT_ID = "_dd.parent_id"; public static final String APM_ENABLED = "_dd.apm.enabled"; + public static final String DECISION_MAKER_INHERITED = "_dd.dm.inherited"; + public static final String DECISION_MAKER_SERVICE = "_dd.dm.service"; + public static final String DECISION_MAKER_RESOURCE = "_dd.dm.resource"; } diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java index b124e03131f..687f68aa494 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java @@ -160,6 +160,13 @@ public final class TraceInstrumentationConfig { public static final String SPARK_TASK_HISTOGRAM_ENABLED = "spark.task-histogram.enabled"; public static final String SPARK_APP_NAME_AS_SERVICE = "spark.app-name-as-service"; + public static final String TRACE_WEBSOCKET_MESSAGES_ENABLED = "trace.websocket.messages.enabled"; + public static final String TRACE_WEBSOCKET_MESSAGES_INHERIT_SAMPLING = + "trace.websocket.messages.inherit.sampling"; + public static final String TRACE_WEBSOCKET_MESSAGES_SEPARATE_TRACES = + "trace.websocket.messages.separate.traces"; + public static final String TRACE_WEBSOCKET_TAG_SESSION_ID = "trace.websocket.tag.session-id"; + public static final String JAX_RS_EXCEPTION_AS_ERROR_ENABLED = "trace.jax-rs.exception-as-error.enabled"; public static final String JAX_RS_ADDITIONAL_ANNOTATIONS = "trace.jax-rs.additional.annotations"; diff --git a/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java b/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java index 0b576654936..5fafa0f29a8 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java @@ -107,7 +107,7 @@ static DDSpan create( */ private volatile int longRunningVersion = 0; - private final List links; + protected final List links; /** * Spans should be constructed using the builder, not by calling the constructor directly. diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java index 3dba6f5b8f4..17e776f7980 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -537,7 +537,9 @@ public static String getHostName() { private final boolean sparkTaskHistogramEnabled; private final boolean sparkAppNameAsService; private final boolean jaxRsExceptionAsErrorsEnabled; - + private final boolean websocketMessagesInheritSampling; + private final boolean websocketMessagesSeparateTraces; + private final boolean websocketTagSessionId; private final boolean axisPromoteResourceName; private final float traceFlushIntervalSeconds; private final long tracePostProcessingTimeout; @@ -1847,6 +1849,15 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment()) axisPromoteResourceName = configProvider.getBoolean(AXIS_PROMOTE_RESOURCE_NAME, false); + websocketMessagesInheritSampling = + configProvider.getBoolean( + TRACE_WEBSOCKET_MESSAGES_INHERIT_SAMPLING, DEFAULT_WEBSOCKET_MESSAGES_INHERIT_SAMPLING); + websocketMessagesSeparateTraces = + configProvider.getBoolean( + TRACE_WEBSOCKET_MESSAGES_SEPARATE_TRACES, DEFAULT_WEBSOCKET_MESSAGES_SEPARATE_TRACES); + websocketTagSessionId = + configProvider.getBoolean(TRACE_WEBSOCKET_TAG_SESSION_ID, DEFAULT_WEBSOCKET_TAG_SESSION_ID); + this.traceFlushIntervalSeconds = configProvider.getFloat( TracerConfig.TRACE_FLUSH_INTERVAL, ConfigDefaults.DEFAULT_TRACE_FLUSH_INTERVAL); @@ -3486,6 +3497,18 @@ public boolean isAxisPromoteResourceName() { return axisPromoteResourceName; } + public boolean isWebsocketMessagesInheritSampling() { + return websocketMessagesInheritSampling; + } + + public boolean isWebsocketMessagesSeparateTraces() { + return websocketMessagesSeparateTraces; + } + + public boolean isWebsocketTagSessionId() { + return websocketTagSessionId; + } + public boolean isDataJobsEnabled() { return dataJobsEnabled; } diff --git a/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java b/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java index 2732e02746b..28375709574 100644 --- a/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java +++ b/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java @@ -19,6 +19,7 @@ import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_METHODS; import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_OTEL_ENABLED; import static datadog.trace.api.ConfigDefaults.DEFAULT_USM_ENABLED; +import static datadog.trace.api.ConfigDefaults.DEFAULT_WEBSOCKET_MESSAGES_ENABLED; import static datadog.trace.api.config.AppSecConfig.APPSEC_ENABLED; import static datadog.trace.api.config.CiVisibilityConfig.CIVISIBILITY_ENABLED; import static datadog.trace.api.config.GeneralConfig.INTERNAL_EXIT_ON_FAILURE; @@ -66,6 +67,7 @@ import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_METHODS; import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_OTEL_ENABLED; import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_THREAD_POOL_EXECUTORS_EXCLUDE; +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_ENABLED; import static datadog.trace.api.config.UsmConfig.USM_ENABLED; import static datadog.trace.util.CollectionUtils.tryMakeImmutableList; import static datadog.trace.util.CollectionUtils.tryMakeImmutableSet; @@ -126,6 +128,7 @@ public class InstrumenterConfig { private final String httpURLConnectionClassName; private final String axisTransportClassName; + private final boolean websocketTracingEnabled; private final boolean directAllocationProfilingEnabled; @@ -277,6 +280,9 @@ private InstrumenterConfig() { this.additionalJaxRsAnnotations = tryMakeImmutableSet(configProvider.getList(JAX_RS_ADDITIONAL_ANNOTATIONS)); + this.websocketTracingEnabled = + configProvider.getBoolean( + TRACE_WEBSOCKET_MESSAGES_ENABLED, DEFAULT_WEBSOCKET_MESSAGES_ENABLED); } public boolean isCodeOriginEnabled() { @@ -504,6 +510,10 @@ public Collection getAdditionalJaxRsAnnotations() { return additionalJaxRsAnnotations; } + public boolean isWebsocketTracingEnabled() { + return websocketTracingEnabled; + } + /** * Check whether asynchronous result types are supported with @Trace annotation. * @@ -636,6 +646,8 @@ public String toString() { + internalExitOnFailure + ", additionalJaxRsAnnotations=" + additionalJaxRsAnnotations + + ", websocketTracingEnabled=" + + websocketTracingEnabled + '}'; } } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InstrumentationTags.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InstrumentationTags.java index 899bc3f0b8d..44d496e7582 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InstrumentationTags.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InstrumentationTags.java @@ -128,4 +128,12 @@ public class InstrumentationTags { public static final String MULE_CORRELATION_ID = "mule.correlation_id"; public static final String MULE_LOCATION = "mule.location"; + + public static final String WEBSOCKET_SESSION_ID = "websocket.session.id"; + public static final String WEBSOCKET_MESSAGE_TYPE = "websocket.message.type"; + public static final String WEBSOCKET_MESSAGE_LENGTH = "websocket.message.length"; + public static final String WEBSOCKET_MESSAGE_FRAMES = "websocket.message.frames"; + public static final String WEBSOCKET_MESSAGE_RECEIVE_TIME = "websocket.message.receive_time"; + public static final String WEBSOCKET_CLOSE_CODE = "websocket.close.code"; + public static final String WEBSOCKET_CLOSE_REASON = "websocket.close.reason"; } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InternalSpanTypes.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InternalSpanTypes.java index 72e47c2a131..c67cfe6b6e4 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InternalSpanTypes.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InternalSpanTypes.java @@ -49,6 +49,7 @@ public class InternalSpanTypes { public static final UTF8BytesString TIBCO_BW = UTF8BytesString.create("tibco_bw"); public static final UTF8BytesString MULE = UTF8BytesString.create(DDSpanTypes.MULE); public static final CharSequence VALKEY = UTF8BytesString.create(DDSpanTypes.VALKEY); + public static final UTF8BytesString WEBSOCKET = UTF8BytesString.create(DDSpanTypes.WEBSOCKET); public static final CharSequence SERVERLESS = UTF8BytesString.create(DDSpanTypes.SERVERLESS); } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanAttributes.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanAttributes.java index 6e329545c55..169e01da374 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanAttributes.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanAttributes.java @@ -6,6 +6,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Objects; public class SpanAttributes { /** Represent an empty attributes. */ @@ -125,4 +126,16 @@ public SpanAttributes build() { return new SpanAttributes(this.attributes); } } + + @Override + public boolean equals(Object o) { + if (!(o instanceof SpanAttributes)) return false; + SpanAttributes that = (SpanAttributes) o; + return Objects.equals(attributes, that.attributes); + } + + @Override + public int hashCode() { + return Objects.hashCode(attributes); + } } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanLink.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanLink.java index 11351d2d594..e34e8af6438 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanLink.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanLink.java @@ -48,7 +48,28 @@ public static SpanLink from(AgentSpanContext context) { */ public static SpanLink from( AgentSpanContext context, byte traceFlags, String traceState, SpanAttributes attributes) { - if (context.getSamplingPriority() > 0) { + return from(context, traceFlags, traceState, attributes, true); + } + + /** + * Creates a span link from a span context with W3C trace state and custom attributes. Gathers the + * trace and span identifiers from the given instance. + * + * @param context The context of the span to get the link to. + * @param traceFlags The W3C formatted trace flags. + * @param traceState The W3C formatted trace state. + * @param attributes The link attributes. + * @param inheritSampling if true the sampling flag will be set based on the context's sampling + * priority. + * @return A span link to the given context. + */ + public static SpanLink from( + AgentSpanContext context, + byte traceFlags, + String traceState, + SpanAttributes attributes, + boolean inheritSampling) { + if (inheritSampling && context.getSamplingPriority() > 0) { traceFlags = (byte) (traceFlags | SAMPLED_FLAG); } return new SpanLink( diff --git a/settings.gradle b/settings.gradle index 2c1aee4607d..3ba7be61fe5 100644 --- a/settings.gradle +++ b/settings.gradle @@ -519,6 +519,9 @@ include ':dd-java-agent:instrumentation:redisson:redisson-2.0.0' include ':dd-java-agent:instrumentation:redisson:redisson-2.3.0' include ':dd-java-agent:instrumentation:redisson:redisson-3.10.3' include ':dd-java-agent:instrumentation:weaver' +include ':dd-java-agent:instrumentation:websocket' +include ':dd-java-agent:instrumentation:websocket:jakarta-websocket-2.0' +include ':dd-java-agent:instrumentation:websocket:javax-websocket-1.0' include ':dd-java-agent:instrumentation:websphere-jmx' include ':dd-java-agent:instrumentation:wildfly-9' include ':dd-java-agent:instrumentation:zio' From fef4b8208c39d4cba7f79c0e4dfc5cdeb98e854a Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Mon, 3 Mar 2025 14:37:24 +0100 Subject: [PATCH 02/15] Port span link related tests --- .../groovy/ArmeriaJetty11ServerTest.groovy | 5 ++ .../src/test/groovy/S3ClientTest.groovy | 76 +++++++++---------- .../testFixtures/groovy/JettyServer.groovy | 9 ++- .../test/groovy/OpenTelemetry14Test.groovy | 4 + ...bstractSparkStructuredStreamingTest.groovy | 7 +- .../undertow/undertow-2.0/build.gradle | 2 + .../test/groovy/UndertowServletTest.groovy | 10 +++ .../agent/test/asserts/LinksAssert.groovy | 9 ++- .../agent/test/asserts/SpanAssert.groovy | 13 +++- .../tagprocessor/SpanPointersProcessor.java | 8 +- 10 files changed, 88 insertions(+), 55 deletions(-) diff --git a/dd-java-agent/instrumentation/armeria-jetty/src/test/jetty11/groovy/ArmeriaJetty11ServerTest.groovy b/dd-java-agent/instrumentation/armeria-jetty/src/test/jetty11/groovy/ArmeriaJetty11ServerTest.groovy index 3891a6d74d3..9ddecf3c6ae 100644 --- a/dd-java-agent/instrumentation/armeria-jetty/src/test/jetty11/groovy/ArmeriaJetty11ServerTest.groovy +++ b/dd-java-agent/instrumentation/armeria-jetty/src/test/jetty11/groovy/ArmeriaJetty11ServerTest.groovy @@ -119,4 +119,9 @@ class ArmeriaJetty11ServerTest extends HttpServerTest { boolean testSessionId() { true } + + @Override + boolean testWebsockets() { + false + } } diff --git a/dd-java-agent/instrumentation/aws-java-s3-2.0/src/test/groovy/S3ClientTest.groovy b/dd-java-agent/instrumentation/aws-java-s3-2.0/src/test/groovy/S3ClientTest.groovy index fe9decb6f11..f5e2c2f9352 100644 --- a/dd-java-agent/instrumentation/aws-java-s3-2.0/src/test/groovy/S3ClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-s3-2.0/src/test/groovy/S3ClientTest.groovy @@ -1,7 +1,9 @@ import datadog.trace.agent.test.AgentTestRunner import datadog.trace.api.DDSpanTypes +import datadog.trace.api.DDTraceId +import datadog.trace.bootstrap.instrumentation.api.SpanAttributes +import datadog.trace.bootstrap.instrumentation.api.SpanLink import datadog.trace.core.tagprocessor.SpanPointersProcessor -import groovy.json.JsonSlurper import org.testcontainers.containers.GenericContainer import org.testcontainers.utility.DockerImageName import software.amazon.awssdk.auth.credentials.AwsBasicCredentials @@ -70,6 +72,14 @@ class S3ClientTest extends AgentTestRunner { operationName "aws.http" resourceName "S3.PutObject" spanType DDSpanTypes.HTTP_CLIENT + links { + link(DDTraceId.ZERO, (long)0, SpanLink.DEFAULT_FLAGS, + SpanAttributes.builder() + .put("ptr.kind", SpanPointersProcessor.S3_PTR_KIND) + .put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION) + .put("ptr.hash","6d1a2fe194c6579187408f827f942be3") + .put("link.kind",SpanPointersProcessor.LINK_KIND).build()) + } tags { defaultTags() tag "component", "java-aws-sdk" @@ -87,19 +97,6 @@ class S3ClientTest extends AgentTestRunner { tag "peer.port", { it instanceof Integer } tag "span.kind", "client" tag "aws.requestId", { it != null } - tag "_dd.span_links", { it != null } - - // Assert the span links - def spanLinks = tags["_dd.span_links"] - assert spanLinks != null - def links = new JsonSlurper().parseText(spanLinks) - assert links.size() == 1 - def link = links[0] - assert link["attributes"] != null - assert link["attributes"]["ptr.kind"] == SpanPointersProcessor.S3_PTR_KIND - assert link["attributes"]["ptr.dir"] == SpanPointersProcessor.DOWN_DIRECTION - assert link["attributes"]["ptr.hash"] == "6d1a2fe194c6579187408f827f942be3" - assert link["attributes"]["link.kind"] == SpanPointersProcessor.LINK_KIND } } } @@ -134,6 +131,14 @@ class S3ClientTest extends AgentTestRunner { operationName "aws.http" resourceName "S3.PutObject" spanType DDSpanTypes.HTTP_CLIENT + links { + link(DDTraceId.ZERO, (long)0, SpanLink.DEFAULT_FLAGS, + SpanAttributes.builder() + .put("ptr.kind", SpanPointersProcessor.S3_PTR_KIND) + .put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION) + .put("ptr.hash","6d1a2fe194c6579187408f827f942be3") + .put("link.kind",SpanPointersProcessor.LINK_KIND).build()) + } tags { defaultTags() tag "component", "java-aws-sdk" @@ -151,7 +156,6 @@ class S3ClientTest extends AgentTestRunner { tag "peer.port", { it instanceof Integer } tag "span.kind", "client" tag "aws.requestId", { it != null } - tag "_dd.span_links", { it != null } } } } @@ -161,6 +165,14 @@ class S3ClientTest extends AgentTestRunner { operationName "aws.http" resourceName "S3.CopyObject" spanType DDSpanTypes.HTTP_CLIENT + links { + link(DDTraceId.ZERO, (long)0, SpanLink.DEFAULT_FLAGS, + SpanAttributes.builder() + .put("ptr.kind", SpanPointersProcessor.S3_PTR_KIND) + .put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION) + .put("ptr.hash","1542053ce6d393c424b1374bac1fc0c5") + .put("link.kind",SpanPointersProcessor.LINK_KIND).build()) + } tags { defaultTags() tag "component", "java-aws-sdk" @@ -178,19 +190,6 @@ class S3ClientTest extends AgentTestRunner { tag "peer.port", { it instanceof Integer } tag "span.kind", "client" tag "aws.requestId", { it != null } - tag "_dd.span_links", { it != null } - - // Assert the span links - def spanLinks = tags["_dd.span_links"] - assert spanLinks != null - def links = new JsonSlurper().parseText(spanLinks) - assert links.size() == 1 - def link = links[0] - assert link["attributes"] != null - assert link["attributes"]["ptr.kind"] == SpanPointersProcessor.S3_PTR_KIND - assert link["attributes"]["ptr.dir"] == SpanPointersProcessor.DOWN_DIRECTION - assert link["attributes"]["ptr.hash"] == "1542053ce6d393c424b1374bac1fc0c5" - assert link["attributes"]["link.kind"] == SpanPointersProcessor.LINK_KIND } } } @@ -346,6 +345,14 @@ class S3ClientTest extends AgentTestRunner { operationName "aws.http" resourceName "S3.CompleteMultipartUpload" spanType DDSpanTypes.HTTP_CLIENT + links { + link(DDTraceId.ZERO, (long)0, SpanLink.DEFAULT_FLAGS, + SpanAttributes.builder() + .put("ptr.kind", SpanPointersProcessor.S3_PTR_KIND) + .put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION) + .put("ptr.hash","422412aa6b472a7194f3e24f4b12b4a6") + .put("link.kind",SpanPointersProcessor.LINK_KIND).build()) + } tags { defaultTags() tag "component", "java-aws-sdk" @@ -364,19 +371,6 @@ class S3ClientTest extends AgentTestRunner { tag "span.kind", "client" tag "aws.requestId", { it != null } tag "http.query.string", { it != null } - tag "_dd.span_links", { it != null } - - // Assert the span links - def spanLinks = tags["_dd.span_links"] - assert spanLinks != null - def links = new JsonSlurper().parseText(spanLinks) - assert links.size() == 1 - def link = links[0] - assert link["attributes"] != null - assert link["attributes"]["ptr.kind"] == SpanPointersProcessor.S3_PTR_KIND - assert link["attributes"]["ptr.dir"] == SpanPointersProcessor.DOWN_DIRECTION - assert link["attributes"]["ptr.hash"] == "422412aa6b472a7194f3e24f4b12b4a6" - assert link["attributes"]["link.kind"] == SpanPointersProcessor.LINK_KIND } } } diff --git a/dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy b/dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy index d557b97a7bf..47ba2866d0b 100644 --- a/dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy +++ b/dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy @@ -31,9 +31,12 @@ class JettyServer implements WebsocketServer { JettyServer(ServletContextHandler handler) { server.handler = handler server.addBean(errorHandler) - JakartaWebSocketServletContainerInitializer.configure(handler, (servletContext, container) -> { - container.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.class, "/websocket").build()) - }) + try { + JakartaWebSocketServletContainerInitializer.configure(handler, (servletContext, container) -> { + container.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.class, "/websocket").build()) + }) + } catch (Throwable ignored) { + } } @Override diff --git a/dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/OpenTelemetry14Test.groovy b/dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/OpenTelemetry14Test.groovy index 532b2b44fc7..3f52a661c80 100644 --- a/dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/OpenTelemetry14Test.groovy +++ b/dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/OpenTelemetry14Test.groovy @@ -311,6 +311,7 @@ class OpenTelemetry14Test extends AgentTestRunner { assertTraces(1) { trace(1) { span { + ignoreSpanLinks() // check is done on the content of the tag below tags { defaultTags() "$SPAN_KIND" "$SPAN_KIND_INTERNAL" @@ -345,6 +346,7 @@ class OpenTelemetry14Test extends AgentTestRunner { assertTraces(1) { trace(1) { span { + ignoreSpanLinks() // check is done on the content of the tag below tags { defaultTags() "$SPAN_KIND" "$SPAN_KIND_INTERNAL" @@ -380,6 +382,7 @@ class OpenTelemetry14Test extends AgentTestRunner { assertTraces(1) { trace(1) { span { + ignoreSpanLinks() // check is done on the content of the tag below tags { defaultTags() "$SPAN_KIND" "$SPAN_KIND_INTERNAL" @@ -421,6 +424,7 @@ class OpenTelemetry14Test extends AgentTestRunner { assertTraces(1) { trace(1) { span { + ignoreSpanLinks() // check is done on the content of the tag below tags { defaultTags() "$SPAN_KIND" "$SPAN_KIND_INTERNAL" diff --git a/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkStructuredStreamingTest.groovy b/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkStructuredStreamingTest.groovy index 46332f80236..94984225e14 100644 --- a/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkStructuredStreamingTest.groovy +++ b/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkStructuredStreamingTest.groovy @@ -2,9 +2,11 @@ package datadog.trace.instrumentation.spark import datadog.trace.agent.test.AgentTestRunner import datadog.trace.api.DDTags +import datadog.trace.api.DDTraceId import datadog.trace.api.Platform import datadog.trace.api.sampling.PrioritySampling import datadog.trace.api.sampling.SamplingMechanism +import datadog.trace.core.DDSpan import org.apache.spark.sql.Dataset import org.apache.spark.sql.Encoders import org.apache.spark.sql.execution.streaming.MemoryStream @@ -300,8 +302,9 @@ class AbstractSparkStructuredStreamingTest extends AgentTestRunner { resourceName "test-query" spanType "spark" parent() - assert span.tags.containsKey(DDTags.SPAN_LINKS) - assert span.tags[DDTags.SPAN_LINKS] != null + links({ + link(DDTraceId.from((long)12052652441736835200), (long)-6394091631972716416) + }) } span { operationName "spark.sql" diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/build.gradle b/dd-java-agent/instrumentation/undertow/undertow-2.0/build.gradle index be782a4591a..874a4b1e006 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/build.gradle +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/build.gradle @@ -42,6 +42,8 @@ dependencies { testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-2') testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-3') 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') latestDepTestImplementation group: 'io.undertow', name: 'undertow-servlet', version: '2.2.+' latestDepTestImplementation group: 'io.undertow', name: 'undertow-websockets-jsr', version: '2.2.+' diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy index 4fa8a555dbc..3683768588b 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy @@ -323,6 +323,10 @@ abstract class UndertowServletTest extends HttpServerTest { } class UndertowServletV0Test extends UndertowServletTest implements TestingGenericHttpNamingConventions.ServerV0 { + @Override + boolean testWebsockets() { + false + } } class UndertowServletNoHttpRouteForkedTest extends UndertowServletTest implements TestingGenericHttpNamingConventions.ServerV0 { @@ -331,6 +335,12 @@ class UndertowServletNoHttpRouteForkedTest extends UndertowServletTest implement false } + @Override + boolean testWebsockets() { + // only test websocket in this test class not to have the same done too many times + true + } + @Override protected void configurePreAgent() { super.configurePreAgent() diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy index 97ef890f266..27b33a9995c 100644 --- a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy +++ b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy @@ -1,6 +1,6 @@ package datadog.trace.agent.test.asserts - +import datadog.trace.api.DDTraceId import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink import datadog.trace.bootstrap.instrumentation.api.SpanAttributes import datadog.trace.bootstrap.instrumentation.api.SpanLink @@ -31,9 +31,12 @@ class LinksAssert { } def link(DDSpan linked, byte flags = SpanLink.DEFAULT_FLAGS, SpanAttributes attributes = SpanAttributes.EMPTY, String traceState = '') { + link(linked.traceId, linked.spanId, flags, attributes, traceState) + } + + def link(DDTraceId traceId, def spanId, byte flags = SpanLink.DEFAULT_FLAGS, SpanAttributes attributes = SpanAttributes.EMPTY, String traceState = '') { def found = links.find { - it.spanId() == linked.spanId && - it.traceId() == linked.traceId + it.spanId() == spanId && it.traceId() == traceId } assert found != null assert found.traceFlags() == flags diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy index c4be54bf7c8..1e0c88ec3fb 100644 --- a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy +++ b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy @@ -11,9 +11,12 @@ import groovy.transform.stc.SimpleType import java.util.regex.Pattern +import static datadog.trace.agent.test.asserts.TagsAssert.assertTags + class SpanAssert { private final DDSpan span private final DDSpan previous + private boolean checkLinks = true private final checked = [:] private SpanAssert(span, DDSpan previous) { @@ -168,8 +171,10 @@ class SpanAssert { if (!checked.errored) { errored(false) } - if (!checked.links) { - assert span.tags['_dd.span_links'] == null + if (checkLinks) { + if (!checked.links) { + assert span.tags['_dd.span_links'] == null + } } hasServiceName() } @@ -185,6 +190,10 @@ class SpanAssert { assertTags(span, spec, checkAllTags) } + void ignoreSpanLinks() { + this.checkLinks = false + } + void links(@ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.LinksAssert']) @DelegatesTo(value = LinksAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) { checked.links = true diff --git a/dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/SpanPointersProcessor.java b/dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/SpanPointersProcessor.java index ea86ae2e76f..c3932b87785 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/SpanPointersProcessor.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/SpanPointersProcessor.java @@ -30,10 +30,10 @@ public class SpanPointersProcessor implements TagsPostProcessor { // The pointer direction will always be down. The serverless agent handles cases where the // direction is up. - static final String DOWN_DIRECTION = "d"; - static final String DYNAMODB_PTR_KIND = "aws.dynamodb.item"; - static final String S3_PTR_KIND = "aws.s3.object"; - static final String LINK_KIND = "span-pointer"; + public static final String DOWN_DIRECTION = "d"; + public static final String DYNAMODB_PTR_KIND = "aws.dynamodb.item"; + public static final String S3_PTR_KIND = "aws.s3.object"; + public static final String LINK_KIND = "span-pointer"; @Override public Map processTags( From aa9501d7305462699d47a87d611d23df91863a6f Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 4 Mar 2025 09:56:53 +0100 Subject: [PATCH 03/15] Handle route for undertow upgrades --- .../undertow/ExchangeEndSpanListener.java | 7 +++++++ .../src/test/groovy/UndertowServletTest.groovy | 10 ---------- 2 files changed, 7 insertions(+), 10 deletions(-) diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java index 76cb6392f4c..2999e37d29a 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.undertow; +import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR; import static datadog.trace.instrumentation.undertow.UndertowDecorator.DD_UNDERTOW_CONTINUATION; import static datadog.trace.instrumentation.undertow.UndertowDecorator.DECORATE; @@ -27,6 +28,12 @@ public void exchangeEvent(HttpServerExchange exchange, NextListener nextListener if (throwable != null) { DECORATE.onError(span, throwable); } + if (exchange.isUpgrade() && UndertowDecorator.UNDERTOW_LEGACY_TRACING) { + // make sure the resource is set correctly when upgrade is performed. + // we set it in the ServletInstrumentation but it's too early + HTTP_RESOURCE_DECORATOR.withRoute( + span, exchange.getRequestMethod().toString(), exchange.getRelativePath(), false); + } DECORATE.onResponse(span, exchange); DECORATE.beforeFinish(span); diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy index 3683768588b..4fa8a555dbc 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy @@ -323,10 +323,6 @@ abstract class UndertowServletTest extends HttpServerTest { } class UndertowServletV0Test extends UndertowServletTest implements TestingGenericHttpNamingConventions.ServerV0 { - @Override - boolean testWebsockets() { - false - } } class UndertowServletNoHttpRouteForkedTest extends UndertowServletTest implements TestingGenericHttpNamingConventions.ServerV0 { @@ -335,12 +331,6 @@ class UndertowServletNoHttpRouteForkedTest extends UndertowServletTest implement false } - @Override - boolean testWebsockets() { - // only test websocket in this test class not to have the same done too many times - true - } - @Override protected void configurePreAgent() { super.configurePreAgent() From dba58969fb54aa32d17607ea2d344188820ed82d Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 4 Mar 2025 10:07:15 +0100 Subject: [PATCH 04/15] fix jetty11 test --- .../src/test/groovy/JettyAsyncHandlerTest.groovy | 5 +++++ .../src/testFixtures/groovy/JettyServer.groovy | 15 +++++++++------ 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/dd-java-agent/instrumentation/jetty-11/src/test/groovy/JettyAsyncHandlerTest.groovy b/dd-java-agent/instrumentation/jetty-11/src/test/groovy/JettyAsyncHandlerTest.groovy index 8cc1ebb30cf..38f5b1449ab 100644 --- a/dd-java-agent/instrumentation/jetty-11/src/test/groovy/JettyAsyncHandlerTest.groovy +++ b/dd-java-agent/instrumentation/jetty-11/src/test/groovy/JettyAsyncHandlerTest.groovy @@ -20,6 +20,11 @@ class JettyAsyncHandlerTest extends Jetty11Test implements TestingGenericHttpNam false // continuation test handler not working with sessions } + @Override + boolean testWebsockets() { + false + } + static class ContinuationTestHandler implements Handler { @Delegate private final Handler delegate diff --git a/dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy b/dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy index 47ba2866d0b..67f29825fac 100644 --- a/dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy +++ b/dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy @@ -13,6 +13,7 @@ import jakarta.websocket.EndpointConfig import jakarta.websocket.MessageHandler import jakarta.websocket.Session import jakarta.websocket.server.ServerEndpointConfig +import org.eclipse.jetty.server.Handler import org.eclipse.jetty.server.Server import org.eclipse.jetty.server.handler.ErrorHandler import org.eclipse.jetty.server.session.SessionHandler @@ -28,14 +29,16 @@ class JettyServer implements WebsocketServer { def port = 0 final server = new Server(0) // select random open port - JettyServer(ServletContextHandler handler) { + JettyServer(Handler handler) { server.handler = handler server.addBean(errorHandler) - try { - JakartaWebSocketServletContainerInitializer.configure(handler, (servletContext, container) -> { - container.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.class, "/websocket").build()) - }) - } catch (Throwable ignored) { + if (handler instanceof ServletContextHandler) { + try { + JakartaWebSocketServletContainerInitializer.configure(handler, (servletContext, container) -> { + container.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.class, "/websocket").build()) + }) + } catch (Throwable ignored) { + } } } From 36fbdbaf98c33a866e8f962603cfc30fae0503ef Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 4 Mar 2025 10:34:13 +0100 Subject: [PATCH 05/15] make sure instrumentation is applied in spring boot test --- ...groovy => HandlerMappingResourceNameFilterForkedTest.groovy} | 2 +- ...Test.groovy => ResteasySpringBeanProcessorForkedTest.groovy} | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/{HandlerMappingResourceNameFilterTest.groovy => HandlerMappingResourceNameFilterForkedTest.groovy} (96%) rename dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/{ResteasySpringBeanProcessorTest.groovy => ResteasySpringBeanProcessorForkedTest.groovy} (97%) diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterTest.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy similarity index 96% rename from dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterTest.groovy rename to dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy index ec941fbbab7..17188ac34d6 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterTest.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy @@ -21,7 +21,7 @@ import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecora // https://mvnrepository.com/artifact/org.springframework/spring-test-mvc?repo=springio-plugins-release // https://github.com/spring-attic/spring-test-mvc @ContextConfiguration(classes = TestConfiguration, loader = AnnotationConfigWebContextLoader) -class HandlerMappingResourceNameFilterTest extends AgentTestRunner { +class HandlerMappingResourceNameFilterForkedTest extends AgentTestRunner { @EnableWebMvc @Configuration @ComponentScan(basePackages = "datadog.trace.instrumentation.springweb") diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/ResteasySpringBeanProcessorTest.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/ResteasySpringBeanProcessorForkedTest.groovy similarity index 97% rename from dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/ResteasySpringBeanProcessorTest.groovy rename to dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/ResteasySpringBeanProcessorForkedTest.groovy index a8c082fef6c..66c8c66ff44 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/ResteasySpringBeanProcessorTest.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/ResteasySpringBeanProcessorForkedTest.groovy @@ -7,7 +7,7 @@ import org.springframework.beans.factory.support.BeanDefinitionRegistry import org.springframework.beans.factory.support.GenericBeanDefinition import test.boot.TestProvider -class ResteasySpringBeanProcessorTest extends AgentTestRunner { +class ResteasySpringBeanProcessorForkedTest extends AgentTestRunner { interface TestBeanFactory extends ConfigurableListableBeanFactory, BeanDefinitionRegistry {} From e6d51a27f251c417d2a9ae0ac4861a9e611f6196 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 4 Mar 2025 10:44:40 +0100 Subject: [PATCH 06/15] add coverage for spanlinkattributes --- .../instrumentation/api/SpanLinkTest.groovy | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/internal-api/src/test/groovy/datadog/trace/bootstrap/instrumentation/api/SpanLinkTest.groovy b/internal-api/src/test/groovy/datadog/trace/bootstrap/instrumentation/api/SpanLinkTest.groovy index a1e9b93ca38..a8227194fc6 100644 --- a/internal-api/src/test/groovy/datadog/trace/bootstrap/instrumentation/api/SpanLinkTest.groovy +++ b/internal-api/src/test/groovy/datadog/trace/bootstrap/instrumentation/api/SpanLinkTest.groovy @@ -142,4 +142,18 @@ class SpanLinkTest extends Specification { then: notThrown(NullPointerException) } + + def "test span link attributes equals and hashcode"() { + when: + def a = SpanAttributes.builder().put("test", "value").build() + def b = SpanAttributes.builder().put("test", "value").build() + def c = SpanAttributes.builder().build() + + then: + assert a == b + assert a != c + assert !c.equals(null) + assert a.hashCode() == b.hashCode() + assert a.hashCode() != c.hashCode() + } } From 54da9d3d2c0aea2a2e7de05e6564fc417a61b4d8 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 4 Mar 2025 11:57:43 +0100 Subject: [PATCH 07/15] spotless --- .../datadog/trace/agent/test/asserts/LinksAssert.groovy | 7 ++++++- .../datadog/trace/agent/test/asserts/SpanAssert.groovy | 4 +--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy index 27b33a9995c..390cab56dff 100644 --- a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy +++ b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy @@ -1,6 +1,7 @@ package datadog.trace.agent.test.asserts import datadog.trace.api.DDTraceId +import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink import datadog.trace.bootstrap.instrumentation.api.SpanAttributes import datadog.trace.bootstrap.instrumentation.api.SpanLink @@ -31,7 +32,11 @@ class LinksAssert { } def link(DDSpan linked, byte flags = SpanLink.DEFAULT_FLAGS, SpanAttributes attributes = SpanAttributes.EMPTY, String traceState = '') { - link(linked.traceId, linked.spanId, flags, attributes, traceState) + link(linked.context(), flags, attributes, traceState) + } + + def link(AgentSpanContext context, byte flags = SpanLink.DEFAULT_FLAGS, SpanAttributes attributes = SpanAttributes.EMPTY, String traceState = '') { + link(context.traceId, context.spanId, flags, attributes, traceState) } def link(DDTraceId traceId, def spanId, byte flags = SpanLink.DEFAULT_FLAGS, SpanAttributes attributes = SpanAttributes.EMPTY, String traceState = '') { diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy index 1e0c88ec3fb..92097c2c96a 100644 --- a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy +++ b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy @@ -1,8 +1,5 @@ package datadog.trace.agent.test.asserts -import static TagsAssert.assertTags -import static datadog.trace.agent.test.asserts.LinksAssert.assertLinks - import datadog.trace.api.DDSpanId import datadog.trace.api.DDTraceId import datadog.trace.core.DDSpan @@ -11,6 +8,7 @@ import groovy.transform.stc.SimpleType import java.util.regex.Pattern +import static datadog.trace.agent.test.asserts.LinksAssert.assertLinks import static datadog.trace.agent.test.asserts.TagsAssert.assertTags class SpanAssert { From 9f7a27ed60d820f2cbda0ef191b816c18f23750e Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 4 Mar 2025 14:05:01 +0100 Subject: [PATCH 08/15] fix undertow test --- .../instrumentation/undertow/ExchangeEndSpanListener.java | 7 ------- .../undertow-2.0/src/test/groovy/TestEndpoint.groovy | 6 ++++-- .../src/test/groovy/UndertowServletTest.groovy | 5 +++++ 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java index 2999e37d29a..76cb6392f4c 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java @@ -1,6 +1,5 @@ package datadog.trace.instrumentation.undertow; -import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR; import static datadog.trace.instrumentation.undertow.UndertowDecorator.DD_UNDERTOW_CONTINUATION; import static datadog.trace.instrumentation.undertow.UndertowDecorator.DECORATE; @@ -28,12 +27,6 @@ public void exchangeEvent(HttpServerExchange exchange, NextListener nextListener if (throwable != null) { DECORATE.onError(span, throwable); } - if (exchange.isUpgrade() && UndertowDecorator.UNDERTOW_LEGACY_TRACING) { - // make sure the resource is set correctly when upgrade is performed. - // we set it in the ServletInstrumentation but it's too early - HTTP_RESOURCE_DECORATOR.withRoute( - span, exchange.getRequestMethod().toString(), exchange.getRelativePath(), false); - } DECORATE.onResponse(span, exchange); DECORATE.beforeFinish(span); diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/TestEndpoint.groovy b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/TestEndpoint.groovy index 91c3c4278b6..99c7ec50746 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/TestEndpoint.groovy +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/TestEndpoint.groovy @@ -4,6 +4,8 @@ import javax.websocket.OnOpen import javax.websocket.Session import javax.websocket.server.ServerEndpoint +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace + @ServerEndpoint("/websocket") class TestEndpoint { static volatile Session activeSession @@ -23,11 +25,11 @@ class TestEndpoint { @OnMessage void onMessage(String s, boolean last) { - datadog.trace.agent.test.utils.TraceUtils.runUnderTrace("onRead", {}) + runUnderTrace("onRead", {}) } @OnMessage void onMessage(byte[] b, boolean last) { - datadog.trace.agent.test.utils.TraceUtils.runUnderTrace("onRead", {}) + runUnderTrace("onRead", {}) } } diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy index 4fa8a555dbc..f61d85630b1 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy @@ -33,6 +33,7 @@ import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.QUERY_ import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.REDIRECT import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.USER_BLOCK +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.WEBSOCKET abstract class UndertowServletTest extends HttpServerTest { private static final CONTEXT = "ctx" @@ -218,6 +219,9 @@ abstract class UndertowServletTest extends HttpServerTest { return "404" } else if (endpoint.hasPathParam) { return "$method ${testPathParam()}" + } else if (endpoint == WEBSOCKET) { + // the route is not set on websocket handlers + return "$method ${endpoint.resolve(address).path}" } def base = endpoint == LOGIN ? address : address.resolve("/") return "$method ${endpoint.resolve(base).path}" @@ -267,6 +271,7 @@ abstract class UndertowServletTest extends HttpServerTest { switch (endpoint) { case LOGIN: case NOT_FOUND: + case WEBSOCKET: return null case PATH_PARAM: return testPathParam() From 046ca28f36d9abf3d08d6ff26d4d1935e40e1172 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 4 Mar 2025 17:49:07 +0100 Subject: [PATCH 09/15] fix muzzle for tomcat --- dd-java-agent/instrumentation/tomcat-5.5/build.gradle | 8 ++++++++ .../tomcat/WsHandshakeRequestInstrumentation.java | 5 +++++ .../tomcat/WsHttpUpgradeHandlerInstrumentation.java | 5 +++++ 3 files changed, 18 insertions(+) diff --git a/dd-java-agent/instrumentation/tomcat-5.5/build.gradle b/dd-java-agent/instrumentation/tomcat-5.5/build.gradle index a7fb77b990b..8811a084214 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/build.gradle +++ b/dd-java-agent/instrumentation/tomcat-5.5/build.gradle @@ -35,6 +35,14 @@ muzzle { extraDependency 'tomcat:tomcat-util:5.5.23' assertInverse = true } + pass { + name = "tomcat-websocket" + group = "org.apache.tomcat" + module = 'tomcat-websocket' + versions = "[8.0.1,]" + extraDependency 'org.apache.tomcat:tomcat-catalina-ha:8.0.1' + assertInverse = true + } // org.apache.catalina.connector.CoyoteAdapter introduced in Catalina 5.5 } diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java index 7ef56a13942..41de9e9f2d5 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java @@ -36,6 +36,11 @@ public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice(isConstructor(), getClass().getName() + "$CaptureHandshakeSpanAdvice"); } + @Override + public String muzzleDirective() { + return "tomcat-websocket"; + } + public static class CaptureHandshakeSpanAdvice { @Advice.OnMethodExit(suppress = Throwable.class) public static void captureHandshakeSpan(@Advice.This final WsHandshakeRequest self) { diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHttpUpgradeHandlerInstrumentation.java b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHttpUpgradeHandlerInstrumentation.java index 2374c18099e..e25179bb2ca 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHttpUpgradeHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHttpUpgradeHandlerInstrumentation.java @@ -32,6 +32,11 @@ public Map contextStore() { "org.apache.tomcat.websocket.server.WsHandshakeRequest", AgentSpan.class.getName()); } + @Override + public String muzzleDirective() { + return "tomcat-websocket"; + } + @Override public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice(named("init"), getClass().getName() + "$CaptureHandshakeSpanAdvice"); From 2421a36067d3bb9f6e1b3a302aa69979a1cb98db Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 4 Mar 2025 17:51:53 +0100 Subject: [PATCH 10/15] codenarc --- .../spark/AbstractSparkStructuredStreamingTest.groovy | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkStructuredStreamingTest.groovy b/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkStructuredStreamingTest.groovy index 94984225e14..38ea89e0ca0 100644 --- a/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkStructuredStreamingTest.groovy +++ b/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkStructuredStreamingTest.groovy @@ -6,11 +6,10 @@ import datadog.trace.api.DDTraceId import datadog.trace.api.Platform import datadog.trace.api.sampling.PrioritySampling import datadog.trace.api.sampling.SamplingMechanism -import datadog.trace.core.DDSpan import org.apache.spark.sql.Dataset import org.apache.spark.sql.Encoders -import org.apache.spark.sql.execution.streaming.MemoryStream import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.execution.streaming.MemoryStream import org.apache.spark.sql.streaming.StreamingQuery import scala.Option import scala.collection.JavaConverters From 2614843b24a346ddd23fbe2b365867214d1ca3cc Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Wed, 5 Mar 2025 09:07:51 +0100 Subject: [PATCH 11/15] fix forbidden api --- dd-java-agent/instrumentation/tomcat-5.5/build.gradle | 1 + .../tomcat/WsHttpUpgradeHandlerInstrumentation.java | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/dd-java-agent/instrumentation/tomcat-5.5/build.gradle b/dd-java-agent/instrumentation/tomcat-5.5/build.gradle index 8811a084214..2d9cc2a1102 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/build.gradle +++ b/dd-java-agent/instrumentation/tomcat-5.5/build.gradle @@ -79,6 +79,7 @@ dependencies { compileOnly group: 'tomcat', name: 'tomcat-coyote', version: tomcatVersion compileOnly group: 'tomcat', name: 'tomcat-util', version: tomcatVersion compileOnly group: 'org.apache.tomcat', name: 'tomcat-websocket', version: '8.0.1', transitive: false + compileOnly group: 'org.apache.tomcat', name: 'tomcat-websocket-api', version: '8.0.1' // Version that corresponds with Tomcat 5.5 diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHttpUpgradeHandlerInstrumentation.java b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHttpUpgradeHandlerInstrumentation.java index e25179bb2ca..bc178cba700 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHttpUpgradeHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHttpUpgradeHandlerInstrumentation.java @@ -1,7 +1,7 @@ package datadog.trace.instrumentation.tomcat; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; -import static net.bytebuddy.matcher.ElementMatchers.named; import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; From 5be57c078cff240708f7deab180c1988a6ff10bb Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Wed, 5 Mar 2025 09:31:49 +0100 Subject: [PATCH 12/15] apply review --- .../decorator/WebsocketDecorator.java | 3 ++ .../websocket/HandlerContext.java | 9 +++++- .../jetty10/JettyServerInstrumentation.java | 1 - .../jetty10/OnCompletedAdvice.java | 29 ------------------- .../boot/SpringBootBasedTest.groovy | 2 +- .../instrumentation/tomcat-5.5/build.gradle | 1 + .../instrumentation/websocket/build.gradle | 3 -- 7 files changed, 13 insertions(+), 35 deletions(-) delete mode 100644 dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/OnCompletedAdvice.java diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java index 63a2100bc62..125d0a6f684 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java @@ -99,6 +99,9 @@ public void onFrameEnd(final HandlerContext handlerContext) { return; } final AgentSpan wsSpan = handlerContext.getWebsocketSpan(); + if (wsSpan == null) { + return; + } try { final long startTime = handlerContext.getFirstFrameTimestamp(); if (startTime > 0) { diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/HandlerContext.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/HandlerContext.java index d0f1894be23..39425d1102e 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/HandlerContext.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/websocket/HandlerContext.java @@ -2,8 +2,11 @@ import datadog.trace.api.time.SystemTimeSource; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public abstract class HandlerContext { + final Logger LOGGER = LoggerFactory.getLogger(HandlerContext.class); private final AgentSpan handshakeSpan; private AgentSpan websocketSpan; @@ -92,7 +95,11 @@ public void recordChunkData(Object data, boolean partialDelivery) { if (partialDelivery && sz == 0) { msgChunks--; // if we receive an empty frame with the fin bit don't count it as a chunk } - } catch (Throwable ignored) { + } catch (Throwable t) { + LOGGER.debug( + "Unable to calculate websocket message size for data type {}", + data.getClass().getName(), + t); } } } diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty10/JettyServerInstrumentation.java b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty10/JettyServerInstrumentation.java index f16673cfd9e..874a618c0db 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty10/JettyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty10/JettyServerInstrumentation.java @@ -99,7 +99,6 @@ public void typeAdvice(TypeTransformer transformer) { public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice(takesNoArguments().and(named("handle")), packageName + ".HandleAdvice"); transformer.applyAdvice(named("recycle").and(takesNoArguments()), packageName + ".ResetAdvice"); - transformer.applyAdvice(named("onCompleted"), packageName + ".OnCompletedAdvice"); if (appSecNotFullyDisabled) { transformer.applyAdvice( diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/OnCompletedAdvice.java b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/OnCompletedAdvice.java deleted file mode 100644 index 9881cb5be76..00000000000 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/OnCompletedAdvice.java +++ /dev/null @@ -1,29 +0,0 @@ -package datadog.trace.instrumentation.jetty10; - -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; - -import datadog.trace.bootstrap.instrumentation.api.AgentScope; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; -import net.bytebuddy.asm.Advice; -import org.eclipse.jetty.server.HttpChannel; -import org.eclipse.jetty.server.Request; - -public class OnCompletedAdvice { - @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope enter(@Advice.This final HttpChannel channel) { - Request req = channel.getRequest(); - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - // return activateSpan((AgentSpan) existingSpan); - } - return null; - } - - @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) - public static void exit( - @Advice.This final HttpChannel channel, @Advice.Enter final AgentScope scope) { - if (scope != null) { - scope.close(); - } - } -} diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/SpringBootBasedTest.groovy b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/SpringBootBasedTest.groovy index 102ce0497bf..8d6a2a81b72 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/SpringBootBasedTest.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/SpringBootBasedTest.groovy @@ -68,7 +68,7 @@ class SpringBootBasedTest extends HttpServerTest @Override void start() { - app.setDefaultProperties(["server.port" : 0, "server.context-path": "/$servletContext", + app.setDefaultProperties(["server.port": 0, "server.context-path": "/$servletContext", "spring.mvc.throw-exception-if-no-handler-found": false, "spring.web.resources.add-mappings" : false]) context = app.run() diff --git a/dd-java-agent/instrumentation/tomcat-5.5/build.gradle b/dd-java-agent/instrumentation/tomcat-5.5/build.gradle index 2d9cc2a1102..a1fa18fdb31 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/build.gradle +++ b/dd-java-agent/instrumentation/tomcat-5.5/build.gradle @@ -139,6 +139,7 @@ dependencies { project.afterEvaluate { tasks.withType(Test).configureEach { if (javaLauncher.get().metadata.languageVersion.asInt() >= 16) { + // to avoid java.lang.IllegalAccessException: class org.apache.tomcat.util.compat.JreCompat cannot access a member of class java.io.FileSystem (in module java.base) with modifiers "static final" jvmArgs += ['--add-opens', 'java.base/java.io=ALL-UNNAMED'] } } diff --git a/dd-java-agent/instrumentation/websocket/build.gradle b/dd-java-agent/instrumentation/websocket/build.gradle index fe680a977ab..5e69c67bd78 100644 --- a/dd-java-agent/instrumentation/websocket/build.gradle +++ b/dd-java-agent/instrumentation/websocket/build.gradle @@ -1,4 +1 @@ apply from: "$rootDir/gradle/java.gradle" - -dependencies { -} From f760e30226901d81309729611759d5e0aa722981 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Wed, 5 Mar 2025 11:19:17 +0100 Subject: [PATCH 13/15] change linkassert signature --- .../groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy index 390cab56dff..3d583398cac 100644 --- a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy +++ b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy @@ -39,7 +39,7 @@ class LinksAssert { link(context.traceId, context.spanId, flags, attributes, traceState) } - def link(DDTraceId traceId, def spanId, byte flags = SpanLink.DEFAULT_FLAGS, SpanAttributes attributes = SpanAttributes.EMPTY, String traceState = '') { + def link(DDTraceId traceId, long spanId, byte flags = SpanLink.DEFAULT_FLAGS, SpanAttributes attributes = SpanAttributes.EMPTY, String traceState = '') { def found = links.find { it.spanId() == spanId && it.traceId() == traceId } From 55238f7645f5ef5822c0b5adb061436e2c8632f8 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Wed, 5 Mar 2025 11:40:29 +0100 Subject: [PATCH 14/15] Add suggestions --- .../decorator/WebsocketDecorator.java | 9 +++- internal-api/build.gradle | 1 + .../api/NotSampledSpanContext.java | 45 +++++++++++++++++++ .../instrumentation/api/SpanLink.java | 23 +--------- 4 files changed, 55 insertions(+), 23 deletions(-) create mode 100644 internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/NotSampledSpanContext.java diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java index 125d0a6f684..cdc758829d9 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java @@ -19,6 +19,7 @@ import datadog.trace.api.time.SystemTimeSource; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.NotSampledSpanContext; import datadog.trace.bootstrap.instrumentation.api.SpanAttributes; import datadog.trace.bootstrap.instrumentation.api.SpanLink; import datadog.trace.bootstrap.instrumentation.api.Tags; @@ -163,7 +164,13 @@ private AgentSpan onFrameStart( // the link is not added if the user wants to have receive frames on the same trace as the // handshake wsSpan.addLink( - SpanLink.from(handshakeSpan.context(), (byte) 0, "", linkAttributes, inheritSampling)); + SpanLink.from( + inheritSampling + ? handshakeSpan.context() + : new NotSampledSpanContext(handshakeSpan.context()), + SpanLink.DEFAULT_FLAGS, + "", + linkAttributes)); } } return wsSpan; diff --git a/internal-api/build.gradle b/internal-api/build.gradle index f9eef196cf6..4731d37a255 100644 --- a/internal-api/build.gradle +++ b/internal-api/build.gradle @@ -96,6 +96,7 @@ excludedClassesCoverage += [ "datadog.trace.bootstrap.instrumentation.api.NoopScope", "datadog.trace.bootstrap.instrumentation.api.NoopSpan", "datadog.trace.bootstrap.instrumentation.api.NoopSpanContext", + "datadog.trace.bootstrap.instrumentation.api.NotSampledSpanContext", "datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities", "datadog.trace.bootstrap.instrumentation.api.Schema", "datadog.trace.bootstrap.instrumentation.api.ScopeSource", diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/NotSampledSpanContext.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/NotSampledSpanContext.java new file mode 100644 index 00000000000..db84c70cf16 --- /dev/null +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/NotSampledSpanContext.java @@ -0,0 +1,45 @@ +package datadog.trace.bootstrap.instrumentation.api; + +import datadog.trace.api.DDTraceId; +import datadog.trace.api.datastreams.PathwayContext; +import datadog.trace.api.sampling.PrioritySampling; +import java.util.Map; + +/** An {@link AgentSpanContext} that hides the sampling priority. */ +public final class NotSampledSpanContext implements AgentSpanContext { + private final AgentSpanContext delegate; + + public NotSampledSpanContext(AgentSpanContext delegate) { + this.delegate = delegate; + } + + @Override + public DDTraceId getTraceId() { + return delegate.getTraceId(); + } + + @Override + public long getSpanId() { + return delegate.getSpanId(); + } + + @Override + public AgentTraceCollector getTraceCollector() { + return delegate.getTraceCollector(); + } + + @Override + public int getSamplingPriority() { + return PrioritySampling.UNSET; + } + + @Override + public Iterable> baggageItems() { + return delegate.baggageItems(); + } + + @Override + public PathwayContext getPathwayContext() { + return delegate.getPathwayContext(); + } +} diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanLink.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanLink.java index e34e8af6438..11351d2d594 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanLink.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanLink.java @@ -48,28 +48,7 @@ public static SpanLink from(AgentSpanContext context) { */ public static SpanLink from( AgentSpanContext context, byte traceFlags, String traceState, SpanAttributes attributes) { - return from(context, traceFlags, traceState, attributes, true); - } - - /** - * Creates a span link from a span context with W3C trace state and custom attributes. Gathers the - * trace and span identifiers from the given instance. - * - * @param context The context of the span to get the link to. - * @param traceFlags The W3C formatted trace flags. - * @param traceState The W3C formatted trace state. - * @param attributes The link attributes. - * @param inheritSampling if true the sampling flag will be set based on the context's sampling - * priority. - * @return A span link to the given context. - */ - public static SpanLink from( - AgentSpanContext context, - byte traceFlags, - String traceState, - SpanAttributes attributes, - boolean inheritSampling) { - if (inheritSampling && context.getSamplingPriority() > 0) { + if (context.getSamplingPriority() > 0) { traceFlags = (byte) (traceFlags | SAMPLED_FLAG); } return new SpanLink( From 978afb6bc805f6e2c99008decac34487b7f78414 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Mon, 10 Mar 2025 11:20:10 +0100 Subject: [PATCH 15/15] port dynamodb span link tests --- .../src/test/groovy/DynamoDbClientTest.groovy | 89 +++++++++---------- 1 file changed, 40 insertions(+), 49 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/test/groovy/DynamoDbClientTest.groovy b/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/test/groovy/DynamoDbClientTest.groovy index 7e839fcdb3a..60ca0ad91f9 100644 --- a/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/test/groovy/DynamoDbClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/test/groovy/DynamoDbClientTest.groovy @@ -1,7 +1,9 @@ import datadog.trace.agent.test.AgentTestRunner import datadog.trace.api.DDSpanTypes +import datadog.trace.api.DDTraceId +import datadog.trace.bootstrap.instrumentation.api.SpanAttributes +import datadog.trace.bootstrap.instrumentation.api.SpanLink import datadog.trace.core.tagprocessor.SpanPointersProcessor -import groovy.json.JsonSlurper import org.testcontainers.containers.GenericContainer import org.testcontainers.utility.DockerImageName import software.amazon.awssdk.auth.credentials.AwsBasicCredentials @@ -23,6 +25,7 @@ import software.amazon.awssdk.services.dynamodb.model.PutItemRequest import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest import spock.lang.Shared + import java.time.Duration class DynamoDbClientTest extends AgentTestRunner { @@ -188,6 +191,15 @@ class DynamoDbClientTest extends AgentTestRunner { operationName "aws.http" resourceName "DynamoDb.UpdateItem" spanType DDSpanTypes.HTTP_CLIENT + links { + link(DDTraceId.ZERO, 0, SpanLink.DEFAULT_FLAGS, SpanAttributes.builder() + .put("ptr.kind", SpanPointersProcessor.DYNAMODB_PTR_KIND) + .put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION) + // First 32 chars of SHA256("dynamodb-one-key-table|id|test-id-1||") + .put("ptr.hash", "ca8daaa857b00545ed5186a915cf1ab5") + .put("link.kind", SpanPointersProcessor.LINK_KIND) + .build()) + } tags { defaultTags() tag "component", "java-aws-sdk" @@ -205,18 +217,6 @@ class DynamoDbClientTest extends AgentTestRunner { tag "span.kind", "client" tag "aws.requestId", { it != null } tag "_dd.span_links", { it != null } - // Assert the span links - def spanLinks = tags["_dd.span_links"] - assert spanLinks != null - def links = new JsonSlurper().parseText(spanLinks) - assert links.size() == 1 - def link = links[0] - assert link["attributes"] != null - assert link["attributes"]["ptr.kind"] == SpanPointersProcessor.DYNAMODB_PTR_KIND - assert link["attributes"]["ptr.dir"] == SpanPointersProcessor.DOWN_DIRECTION - // First 32 chars of SHA256("dynamodb-one-key-table|id|test-id-1||") - assert link["attributes"]["ptr.hash"] == "ca8daaa857b00545ed5186a915cf1ab5" - assert link["attributes"]["link.kind"] == SpanPointersProcessor.LINK_KIND } } } @@ -293,6 +293,15 @@ class DynamoDbClientTest extends AgentTestRunner { operationName "aws.http" resourceName "DynamoDb.UpdateItem" spanType DDSpanTypes.HTTP_CLIENT + links { + link(DDTraceId.ZERO, 0, SpanLink.DEFAULT_FLAGS, SpanAttributes.builder() + .put("ptr.kind", SpanPointersProcessor.DYNAMODB_PTR_KIND) + .put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION) + // First 32 chars of SHA256("dynamodb-two-key-table|primaryKey|customer-123|sortKey|order-456") + .put("ptr.hash", "90922c7899a82ea34406fdcdfb95161e") + .put("link.kind", SpanPointersProcessor.LINK_KIND) + .build()) + } tags { defaultTags() tag "component", "java-aws-sdk" @@ -310,18 +319,6 @@ class DynamoDbClientTest extends AgentTestRunner { tag "span.kind", "client" tag "aws.requestId", { it != null } tag "_dd.span_links", { it != null } - // Assert the span links - def spanLinks = tags["_dd.span_links"] - assert spanLinks != null - def links = new JsonSlurper().parseText(spanLinks) - assert links.size() == 1 - def link = links[0] - assert link["attributes"] != null - assert link["attributes"]["ptr.kind"] == SpanPointersProcessor.DYNAMODB_PTR_KIND - assert link["attributes"]["ptr.dir"] == SpanPointersProcessor.DOWN_DIRECTION - // First 32 chars of SHA256("dynamodb-two-key-table|primaryKey|customer-123|sortKey|order-456") - assert link["attributes"]["ptr.hash"] == "90922c7899a82ea34406fdcdfb95161e" - assert link["attributes"]["link.kind"] == SpanPointersProcessor.LINK_KIND } } } @@ -384,6 +381,15 @@ class DynamoDbClientTest extends AgentTestRunner { operationName "aws.http" resourceName "DynamoDb.DeleteItem" spanType DDSpanTypes.HTTP_CLIENT + links { + link(DDTraceId.ZERO, 0, SpanLink.DEFAULT_FLAGS, SpanAttributes.builder() + .put("ptr.kind", SpanPointersProcessor.DYNAMODB_PTR_KIND) + .put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION) + // First 32 chars of SHA256("dynamodb-one-key-table|id|delete-test-id||") + .put("ptr.hash", "65031164be5e929fddd274a02cba3f9f") + .put("link.kind", SpanPointersProcessor.LINK_KIND) + .build()) + } tags { defaultTags() tag "component", "java-aws-sdk" @@ -401,18 +407,6 @@ class DynamoDbClientTest extends AgentTestRunner { tag "span.kind", "client" tag "aws.requestId", { it != null } tag "_dd.span_links", { it != null } - // Assert the span links - def spanLinks = tags["_dd.span_links"] - assert spanLinks != null - def links = new JsonSlurper().parseText(spanLinks) - assert links.size() == 1 - def link = links[0] - assert link["attributes"] != null - assert link["attributes"]["ptr.kind"] == SpanPointersProcessor.DYNAMODB_PTR_KIND - assert link["attributes"]["ptr.dir"] == SpanPointersProcessor.DOWN_DIRECTION - // First 32 chars of SHA256("dynamodb-one-key-table|id|delete-test-id||") - assert link["attributes"]["ptr.hash"] == "65031164be5e929fddd274a02cba3f9f" - assert link["attributes"]["link.kind"] == SpanPointersProcessor.LINK_KIND } } } @@ -486,6 +480,15 @@ class DynamoDbClientTest extends AgentTestRunner { operationName "aws.http" resourceName "DynamoDb.DeleteItem" spanType DDSpanTypes.HTTP_CLIENT + links { + link(DDTraceId.ZERO, 0, SpanLink.DEFAULT_FLAGS, SpanAttributes.builder() + .put("ptr.kind", SpanPointersProcessor.DYNAMODB_PTR_KIND) + .put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION) + // First 32 chars of SHA256("dynamodb-two-key-table|primaryKey|user-789|sortKey|profile") + .put("ptr.hash", "e5ce1148208c6f88041c73ceb9bbbf3a") + .put("link.kind", SpanPointersProcessor.LINK_KIND) + .build()) + } tags { defaultTags() tag "component", "java-aws-sdk" @@ -503,18 +506,6 @@ class DynamoDbClientTest extends AgentTestRunner { tag "span.kind", "client" tag "aws.requestId", { it != null } tag "_dd.span_links", { it != null } - // Assert the span links - def spanLinks = tags["_dd.span_links"] - assert spanLinks != null - def links = new JsonSlurper().parseText(spanLinks) - assert links.size() == 1 - def link = links[0] - assert link["attributes"] != null - assert link["attributes"]["ptr.kind"] == SpanPointersProcessor.DYNAMODB_PTR_KIND - assert link["attributes"]["ptr.dir"] == SpanPointersProcessor.DOWN_DIRECTION - // First 32 chars of SHA256("dynamodb-two-key-table|primaryKey|user-789|sortKey|profile") - assert link["attributes"]["ptr.hash"] == "e5ce1148208c6f88041c73ceb9bbbf3a" - assert link["attributes"]["link.kind"] == SpanPointersProcessor.LINK_KIND } } }