Add Java-WebSocket instrumentation

dd-java-agent/instrumentation/websocket/java-websocket/build.gradle

@@ -0,0 +1,17 @@
+muzzle {
+  pass {
+    name = "Java-WebSocket"
+    group = "org.java-websocket"
+    module = "Java-WebSocket"
+    versions = "[1.4.1,)"
+  }
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+
+
+dependencies {
+  compileOnly group: 'org.java-websocket', name: 'Java-WebSocket', version: '1.5.3'
+
+  testImplementation group: 'org.java-websocket', name: 'Java-WebSocket', version: '1.5.3'
+}

dd-java-agent/instrumentation/websocket/java-websocket/src/main/java/datadog/trace/instrumentation/websocket/org/OrgWebsocketModule.java

@@ -0,0 +1,63 @@
+package datadog.trace.instrumentation.websocket.org;
+
+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.api.AgentSpan;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+@AutoService(InstrumenterModule.class)
+public class OrgWebsocketModule extends InstrumenterModule.Tracing {
+
+  public OrgWebsocketModule() {
+    this("java-websocket", "websocket");
+  }
+
+  protected OrgWebsocketModule(String instrumentationName, String... additionalNames) {
+    super(instrumentationName, additionalNames);
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    Map<String, String> contextStores = new HashMap<>();
+    contextStores.put("org.java_websocket.WebSocket", WebsocketAgentSpanContext.class.getName());
+    contextStores.put("org.java_websocket.client.WebSocketClient", AgentSpan.class.getName());
+    return contextStores;
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".WebSocketDecorator",
+      packageName + ".WebSocketClientDecorator",
+      packageName + ".WebSocketServerDecorator",
+      packageName + ".WebsocketExtractAdapter",
+      packageName + ".TraceDraft_6455",
+      packageName + ".WebsocketHeaderInjector",
+      packageName + ".WebsocketHeaderExtract",
+      packageName + ".WebsocketAgentSpanContext"
+    };
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isWebsocketTracingEnabled();
+  }
+
+  @Override
+  public String muzzleDirective() {
+    return "Java-WebSocket";
+  }
+
+  @Override
+  public List<Instrumenter> typeInstrumentations() {
+    return Arrays.asList(
+        new WebsocketClientInstrumentation(),
+        new WebsocketServerInstrumentation(),
+        new WebSocketSendInstrumentation());
+  }
+}

dd-java-agent/instrumentation/websocket/java-websocket/src/main/java/datadog/trace/instrumentation/websocket/org/TraceDraft_6455.java

@@ -0,0 +1,68 @@
+package datadog.trace.instrumentation.websocket.org;
+
+import java.util.ArrayList;
+import java.util.List;
+import org.java_websocket.drafts.Draft;
+import org.java_websocket.drafts.Draft_6455;
+import org.java_websocket.extensions.IExtension;
+import org.java_websocket.protocols.IProtocol;
+
+public class TraceDraft_6455 extends Draft_6455 {
+
+  public TraceDraft_6455() {
+    super();
+  }
+
+  public TraceDraft_6455(List<IExtension> inputExtensions) {
+    super(inputExtensions);
+  }
+
+  public TraceDraft_6455(List<IExtension> inputExtensions, List<IProtocol> inputProtocols) {
+    super(inputExtensions, inputProtocols);
+  }
+
+  public TraceDraft_6455(
+      List<IExtension> inputExtensions, List<IProtocol> inputProtocols, int maxFrameSize) {
+    super(inputExtensions, inputProtocols, maxFrameSize);
+  }
+
+  public static TraceDraft_6455 fromDraft6455(Draft_6455 original) {
+    List<IExtension> extensions = new ArrayList<>();
+    for (IExtension ext : original.getKnownExtensions()) {
+      extensions.add(ext);
+    }
+
+    List<IProtocol> protocols = new ArrayList<>();
+    for (IProtocol proto : original.getKnownProtocols()) {
+      protocols.add(proto);
+    }
+
+    return new TraceDraft_6455(extensions, protocols, original.getMaxFrameSize());
+  }
+
+  @Override
+  public Draft copyInstance() {
+    ArrayList<IExtension> newExtensions = new ArrayList<>();
+    for (IExtension knownExtension : this.getKnownExtensions()) {
+      newExtensions.add(knownExtension.copyInstance());
+    }
+
+    ArrayList<IProtocol> newProtocols = new ArrayList<>();
+
+    for (IProtocol knownProtocol : this.getKnownProtocols()) {
+      newProtocols.add(knownProtocol.copyInstance());
+    }
+
+    return new TraceDraft_6455(newExtensions, newProtocols, getMaxFrameSize());
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    return this == o || super.equals(o);
+  }
+
+  @Override
+  public int hashCode() {
+    return super.hashCode();
+  }
+}

dd-java-agent/instrumentation/websocket/java-websocket/src/main/java/datadog/trace/instrumentation/websocket/org/WebSocketClientDecorator.java

@@ -0,0 +1,39 @@
+package datadog.trace.instrumentation.websocket.org;
+
+import static datadog.context.propagation.Propagators.defaultPropagator;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_CLIENT;
+import static datadog.trace.instrumentation.websocket.org.WebsocketHeaderInjector.SETTER;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import org.java_websocket.client.WebSocketClient;
+
+public class WebSocketClientDecorator extends WebSocketDecorator {
+  public static final WebSocketClientDecorator CLIENT_DECORATE = new WebSocketClientDecorator();
+
+  @Override
+  protected CharSequence spanKind() {
+    return SPAN_KIND_CLIENT;
+  }
+
+  private static final String OPERATION_NAME = "websocket.handshake";
+
+  public AgentScope startHandshakeSpan(WebSocketClient client) {
+    String uri = client.getURI().toString();
+    AgentSpan span = startSpan("websocket.handshake", OPERATION_NAME);
+    span.setTag(Tags.HTTP_URL, uri);
+    span.setTag(Tags.HTTP_METHOD, "GET");
+    AgentScope scope = activateSpan(span);
+    defaultPropagator().inject(scope.context(), client, SETTER);
+    afterStart(span);
+    return scope;
+  }
+
+  public void onHandshakeSuccess(AgentSpan span, int httpStatus) {
+    span.setTag(Tags.HTTP_STATUS, httpStatus);
+    span.setTag("websocket.handshake.success", true);
+  }
+}

dd-java-agent/instrumentation/websocket/java-websocket/src/main/java/datadog/trace/instrumentation/websocket/org/WebSocketDecorator.java

@@ -0,0 +1,107 @@
+package datadog.trace.instrumentation.websocket.org;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_BINARY;
+import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_TEXT;
+import static datadog.trace.instrumentation.websocket.org.WebsocketExtractAdapter.GETTER;
+
+import datadog.trace.api.DDSpanTypes;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
+import datadog.trace.bootstrap.instrumentation.decorator.BaseDecorator;
+import java.net.InetSocketAddress;
+import java.nio.ByteBuffer;
+import org.java_websocket.WebSocket;
+import org.java_websocket.handshake.Handshakedata;
+
+public abstract class WebSocketDecorator extends BaseDecorator {
+  public static final CharSequence WEBSOCKET = UTF8BytesString.create("org-java-websocket");
+  public static final CharSequence WEBSOCKET_OPEN = UTF8BytesString.create("websocket.open");
+  public static final CharSequence WEBSOCKET_RECEIVE = UTF8BytesString.create("websocket.receive");
+  public static final CharSequence WEBSOCKET_SEND = UTF8BytesString.create("websocket.send");
+  public static final CharSequence WEBSOCKET_CLOSE = UTF8BytesString.create("websocket.close");
+  private static final String[] INSTRUMENTATION_NAMES = {WEBSOCKET.toString()};
+  private static final String REMOTE_ADDRESS = "remote.address";
+  private static final String MESSAGE_TYPE = "message.type";
+  private static final String MESSAGE_SIZE = "message.size";
+  private static final String MESSAGE_REMOTE = "message.remote";
+  private static final String MESSAGE_CODE = "message.code";
+  private static final String MESSAGE_REASON = "message.reason";
+
+  @Override
+  protected String[] instrumentationNames() {
+    return INSTRUMENTATION_NAMES;
+  }
+
+  @Override
+  protected CharSequence spanType() {
+    return DDSpanTypes.WEBSOCKET;
+  }
+
+  @Override
+  protected CharSequence component() {
+    return WEBSOCKET;
+  }
+
+  protected abstract CharSequence spanKind();
+
+  public AgentSpan open(WebSocket conn, Handshakedata handshake) {
+    AgentSpanContext parentContext = extractContextAndGetSpanContext(handshake, GETTER);
+    AgentSpan span = startSpan(instrumentationNames()[0], WEBSOCKET_OPEN, parentContext);
+    InetSocketAddress remoteAddress = conn.getRemoteSocketAddress();
+    if (remoteAddress != null) {
+      span.setTag(REMOTE_ADDRESS, remoteAddress.toString());
+    }
+    afterStart(span);
+    return span;
+  }
+
+  public AgentSpan onMessage(Object message, AgentSpanContext spanContext) {
+    AgentSpan span = startSpan(instrumentationNames()[0], WEBSOCKET_RECEIVE, spanContext);
+    CharSequence messageType;
+    int msgSize;
+    if (message instanceof ByteBuffer) {
+      messageType = MESSAGE_TYPE_BINARY;
+      msgSize = ((ByteBuffer) message).remaining();
+    } else {
+      messageType = MESSAGE_TYPE_TEXT;
+      msgSize = message == null ? 0 : ((CharSequence) message).length();
+    }
+    span.setTag(MESSAGE_SIZE, msgSize);
+    span.setTag(MESSAGE_TYPE, messageType);
+    afterStart(span);
+    return span;
+  }
+
+  public AgentSpan send(WebsocketAgentSpanContext spanContext) {
+    AgentSpan span;
+    if (spanContext == null) {
+      span = startSpan(instrumentationNames()[0], WEBSOCKET_SEND);
+    } else {
+      span = startSpan(instrumentationNames()[0], WEBSOCKET_SEND, spanContext.getSpanContext());
+    }
+    afterStart(span);
+    if (spanContext != null) {
+      span.setTag(Tags.SPAN_KIND, spanContext.getSpanKind());
+    }
+    return span;
+  }
+
+  public AgentSpan onClose(int code, String reason, boolean remote, AgentSpanContext spanContext) {
+    AgentSpan span = startSpan(instrumentationNames()[0], WEBSOCKET_CLOSE, spanContext);
+    span.setTag(MESSAGE_REASON, reason);
+    span.setTag(MESSAGE_CODE, code);
+    span.setTag(MESSAGE_REMOTE, remote);
+    afterStart(span);
+    return span;
+  }
+
+  @Override
+  public AgentSpan afterStart(AgentSpan span) {
+    span.setTag(Tags.SPAN_KIND, spanKind());
+    return super.afterStart(span);
+  }
+}

dd-java-agent/instrumentation/websocket/java-websocket/src/main/java/datadog/trace/instrumentation/websocket/org/WebSocketSendInstrumentation.java

@@ -0,0 +1,81 @@
+package datadog.trace.instrumentation.websocket.org;
+
+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.api.AgentTracer.noopSpan;
+import static datadog.trace.instrumentation.websocket.org.WebSocketServerDecorator.SERVER_DECORATOR;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPrivate;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import java.util.Collection;
+import net.bytebuddy.asm.Advice;
+import org.java_websocket.WebSocket;
+import org.java_websocket.WebSocketImpl;
+import org.java_websocket.framing.CloseFrame;
+import org.java_websocket.framing.Framedata;
+import org.java_websocket.framing.PingFrame;
+import org.java_websocket.framing.PongFrame;
+
+public class WebSocketSendInstrumentation
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public WebSocketSendInstrumentation() {}
+
+  @Override
+  public String instrumentedType() {
+    return "org.java_websocket.WebSocketImpl";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(isPrivate())
+            .and(named("send"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, Collection.class)),
+        getClass().getName() + "$OnSendAdvice");
+  }
+
+  public static class OnSendAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope onEnter(
+        @Advice.This WebSocketImpl impl, @Advice.Argument(value = 0) Collection<Framedata> frames) {
+      WebsocketAgentSpanContext context =
+          InstrumentationContext.get(WebSocket.class, WebsocketAgentSpanContext.class).get(impl);
+      if (context == null) {
+        // close after send
+        final AgentSpan wsSpan = SERVER_DECORATOR.send(null);
+        return activateSpan(wsSpan);
+      }
+      // ignore ping/pong/close
+      for (Framedata frame : frames) {
+        if (frame instanceof PingFrame
+            || frame instanceof PongFrame
+            || frame instanceof CloseFrame) {
+          return activateSpan(noopSpan());
+        }
+      }
+      final AgentSpan wsSpan = SERVER_DECORATOR.send(context);
+      return activateSpan(wsSpan);
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void onExit(
+        @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) {
+      if (scope == null) {
+        return;
+      }
+      SERVER_DECORATOR.onError(scope.span(), throwable);
+      SERVER_DECORATOR.beforeFinish(scope.span());
+      scope.close();
+      scope.span().finish();
+    }
+  }
+}

dd-java-agent/instrumentation/websocket/java-websocket/src/main/java/datadog/trace/instrumentation/websocket/org/WebSocketServerDecorator.java

@@ -0,0 +1,12 @@
+package datadog.trace.instrumentation.websocket.org;
+
+import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_SERVER;
+
+public class WebSocketServerDecorator extends WebSocketDecorator {
+  public static final WebSocketServerDecorator SERVER_DECORATOR = new WebSocketServerDecorator();
+
+  @Override
+  protected CharSequence spanKind() {
+    return SPAN_KIND_SERVER;
+  }
+}