WebSocket support for Netty

dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelPipelineAdviceUtil.java

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.netty38;
 
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.instrumentation.netty38.client.HttpClientRequestTracingHandler;
@@ -9,6 +10,9 @@
 import datadog.trace.instrumentation.netty38.server.HttpServerResponseTracingHandler;
 import datadog.trace.instrumentation.netty38.server.HttpServerTracingHandler;
 import datadog.trace.instrumentation.netty38.server.MaybeBlockResponseHandler;
+import datadog.trace.instrumentation.netty38.server.websocket.WebSocketServerRequestTracingHandler;
+import datadog.trace.instrumentation.netty38.server.websocket.WebSocketServerResponseTracingHandler;
+import datadog.trace.instrumentation.netty38.server.websocket.WebSocketServerTracingHandler;
 import org.jboss.netty.channel.Channel;
 import org.jboss.netty.channel.ChannelHandler;
 import org.jboss.netty.channel.ChannelPipeline;
@@ -18,6 +22,9 @@
 import org.jboss.netty.handler.codec.http.HttpResponseDecoder;
 import org.jboss.netty.handler.codec.http.HttpResponseEncoder;
 import org.jboss.netty.handler.codec.http.HttpServerCodec;
+import org.jboss.netty.handler.codec.http.websocketx.WebSocket13FrameDecoder;
+import org.jboss.netty.handler.codec.http.websocketx.WebSocket13FrameEncoder;
+import org.jboss.netty.handler.codec.http.websocketx.WebSocketServerProtocolHandler;
 
 /**
  * When certain handlers are added to the pipeline, we want to add our corresponding tracing
@@ -46,6 +53,33 @@ public static void wrapHandler(
             new HttpServerResponseTracingHandler(contextStore));
         pipeline.addLast(
             MaybeBlockResponseHandler.class.getName(), new MaybeBlockResponseHandler(contextStore));
+      } else if (handler instanceof WebSocketServerProtocolHandler) {
+        if (InstrumenterConfig.get().isWebsocketTracingEnabled()) {
+          if (pipeline.get(HttpServerTracingHandler.class) != null) {
+            addHandlerAfter(
+                pipeline,
+                "datadog.trace.instrumentation.netty38.server.HttpServerTracingHandler",
+                new WebSocketServerTracingHandler(contextStore));
+          }
+        }
+      } else if (handler instanceof WebSocket13FrameEncoder) {
+        if (InstrumenterConfig.get().isWebsocketTracingEnabled()) {
+          if (pipeline.get(HttpServerRequestTracingHandler.class) != null) {
+            addHandlerAfter(
+                pipeline,
+                "datadog.trace.instrumentation.netty38.server.HttpServerRequestTracingHandler",
+                new WebSocketServerRequestTracingHandler(contextStore));
+          }
+        }
+      } else if (handler instanceof WebSocket13FrameDecoder) {
+        if (InstrumenterConfig.get().isWebsocketTracingEnabled()) {
+          if (pipeline.get(HttpServerResponseTracingHandler.class) != null) {
+            addHandlerAfter(
+                pipeline,
+                "datadog.trace.instrumentation.netty38.server.HttpServerResponseTracingHandler",
+                new WebSocketServerResponseTracingHandler(contextStore));
+          }
+        }
       } else
       // Client pipeline handlers
       if (handler instanceof HttpClientCodec) {
@@ -64,4 +98,13 @@ public static void wrapHandler(
       CallDepthThreadLocalMap.reset(ChannelPipeline.class);
     }
   }
+
+  private static void addHandlerAfter(
+      final ChannelPipeline pipeline, final String name, final ChannelHandler handler) {
+    ChannelHandler existing = pipeline.get(handler.getClass());
+    if (existing != null) {
+      pipeline.remove(existing);
+    }
+    pipeline.addAfter(name, handler.getClass().getName(), handler);
+  }
 }

dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelTraceContext.java

@@ -3,6 +3,7 @@
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext;
 import org.jboss.netty.handler.codec.http.HttpHeaders;
 
 public class ChannelTraceContext {
@@ -23,6 +24,9 @@ public ChannelTraceContext create() {
   boolean analyzedResponse;
   boolean blockedResponse;
 
+  HandlerContext.Sender senderHandlerContext;
+  HandlerContext.Receiver receiverHandlerContext;
+
   public void reset() {
     this.connectionContinuation = null;
     this.serverSpan = null;
@@ -88,4 +92,20 @@ public void setClientSpan(AgentSpan clientSpan) {
   public void setClientParentSpan(AgentSpan clientParentSpan) {
     this.clientParentSpan = clientParentSpan;
   }
+
+  public HandlerContext.Sender getSenderHandlerContext() {
+    return senderHandlerContext;
+  }
+
+  public void setSenderHandlerContext(HandlerContext.Sender senderHandlerContext) {
+    this.senderHandlerContext = senderHandlerContext;
+  }
+
+  public HandlerContext.Receiver getReceiverHandlerContext() {
+    return receiverHandlerContext;
+  }
+
+  public void setReceiverHandlerContext(HandlerContext.Receiver receiverHandlerContext) {
+    this.receiverHandlerContext = receiverHandlerContext;
+  }
 }

dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/NettyChannelPipelineInstrumentation.java

@@ -68,6 +68,9 @@ public String[] helperClassNames() {
       packageName + ".server.HttpServerResponseTracingHandler",
       packageName + ".server.HttpServerTracingHandler",
       packageName + ".server.MaybeBlockResponseHandler",
+      packageName + ".server.websocket.WebSocketServerTracingHandler",
+      packageName + ".server.websocket.WebSocketServerRequestTracingHandler",
+      packageName + ".server.websocket.WebSocketServerResponseTracingHandler",
     };
   }
 

dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java

@@ -6,6 +6,7 @@
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext;
 import datadog.trace.instrumentation.netty38.ChannelTraceContext;
 import org.jboss.netty.channel.Channel;
 import org.jboss.netty.channel.ChannelHandlerContext;
@@ -17,6 +18,7 @@
 public class HttpServerResponseTracingHandler extends SimpleChannelDownstreamHandler {
 
   private final ContextStore<Channel, ChannelTraceContext> contextStore;
+  private static final String UPGRADE_HEADER = "upgrade";
 
   public HttpServerResponseTracingHandler(
       final ContextStore<Channel, ChannelTraceContext> contextStore) {
@@ -45,7 +47,16 @@ public void writeRequested(final ChannelHandlerContext ctx, final MessageEvent m
         span.finish(); // Finish the span manually since finishSpanOnClose was false
         throw throwable;
       }
-      if (response.getStatus() != HttpResponseStatus.CONTINUE) {
+      final boolean isWebsocketUpgrade =
+          response.getStatus() == HttpResponseStatus.SWITCHING_PROTOCOLS
+              && "websocket".equals(response.headers().get(UPGRADE_HEADER));
+      if (isWebsocketUpgrade) {
+        String channelId = ctx.getChannel().getId().toString();
+        channelTraceContext.setSenderHandlerContext(new HandlerContext.Sender(span, channelId));
+      }
+      if (response.getStatus() != HttpResponseStatus.CONTINUE
+          && (response.getStatus() != HttpResponseStatus.SWITCHING_PROTOCOLS
+              || isWebsocketUpgrade)) {
         DECORATE.onResponse(span, response);
         DECORATE.beforeFinish(span);
         span.finish(); // Finish the span manually since finishSpanOnClose was false

dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/websocket/WebSocketServerRequestTracingHandler.java

@@ -0,0 +1,135 @@
+package datadog.trace.instrumentation.netty38.server.websocket;
+
+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.ContextStore;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext;
+import datadog.trace.instrumentation.netty38.ChannelTraceContext;
+import org.jboss.netty.channel.Channel;
+import org.jboss.netty.channel.ChannelHandlerContext;
+import org.jboss.netty.channel.MessageEvent;
+import org.jboss.netty.channel.SimpleChannelUpstreamHandler;
+import org.jboss.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
+import org.jboss.netty.handler.codec.http.websocketx.CloseWebSocketFrame;
+import org.jboss.netty.handler.codec.http.websocketx.ContinuationWebSocketFrame;
+import org.jboss.netty.handler.codec.http.websocketx.TextWebSocketFrame;
+import org.jboss.netty.handler.codec.http.websocketx.WebSocketFrame;
+
+public class WebSocketServerRequestTracingHandler extends SimpleChannelUpstreamHandler {
+
+  private final ContextStore<Channel, ChannelTraceContext> contextStore;
+
+  public WebSocketServerRequestTracingHandler(
+      final ContextStore<Channel, ChannelTraceContext> contextStore) {
+    this.contextStore = contextStore;
+  }
+
+  @Override
+  public void messageReceived(ChannelHandlerContext ctx, MessageEvent event) throws Exception {
+    Object frame = event.getMessage();
+    if (frame instanceof WebSocketFrame) {
+      Channel channel = ctx.getChannel();
+
+      ChannelTraceContext traceContext = this.contextStore.get(channel);
+      if (traceContext != null) {
+
+        HandlerContext.Receiver receiverContext = traceContext.getReceiverHandlerContext();
+        if (receiverContext == null) {
+          HandlerContext.Sender sessionState = traceContext.getSenderHandlerContext();
+          if (sessionState != null) {
+            receiverContext =
+                new HandlerContext.Receiver(
+                    sessionState.getHandshakeSpan(), channel.getId().toString());
+            traceContext.setReceiverHandlerContext(receiverContext);
+          }
+        }
+        if (receiverContext != null) {
+          if (frame instanceof TextWebSocketFrame) {
+            // WebSocket Read Text Start
+            TextWebSocketFrame textFrame = (TextWebSocketFrame) frame;
+
+            final AgentSpan span =
+                DECORATE.onReceiveFrameStart(
+                    receiverContext, textFrame.getText(), textFrame.isFinalFragment());
+            try (final AgentScope scope = activateSpan(span)) {
+              ctx.sendUpstream(event);
+              // WebSocket Read Text Start
+            } finally {
+              if (textFrame.isFinalFragment()) {
+                traceContext.setReceiverHandlerContext(null);
+                DECORATE.onFrameEnd(receiverContext);
+              }
+            }
+            return;
+          }
+
+          if (frame instanceof BinaryWebSocketFrame) {
+            // WebSocket Read Binary Start
+            BinaryWebSocketFrame binaryFrame = (BinaryWebSocketFrame) frame;
+            final AgentSpan span =
+                DECORATE.onReceiveFrameStart(
+                    receiverContext,
+                    binaryFrame.getBinaryData().array(),
+                    binaryFrame.isFinalFragment());
+            try (final AgentScope scope = activateSpan(span)) {
+              ctx.sendUpstream(event);
+            } finally {
+              // WebSocket Read Binary End
+              if (binaryFrame.isFinalFragment()) {
+                traceContext.setReceiverHandlerContext(null);
+                DECORATE.onFrameEnd(receiverContext);
+              }
+            }
+
+            return;
+          }
+
+          if (frame instanceof ContinuationWebSocketFrame) {
+            ContinuationWebSocketFrame continuationWebSocketFrame =
+                (ContinuationWebSocketFrame) frame;
+            final AgentSpan span =
+                DECORATE.onReceiveFrameStart(
+                    receiverContext,
+                    MESSAGE_TYPE_TEXT.equals(receiverContext.getMessageType())
+                        ? continuationWebSocketFrame.getText()
+                        : continuationWebSocketFrame.getBinaryData().array(),
+                    continuationWebSocketFrame.isFinalFragment());
+            try (final AgentScope scope = activateSpan(span)) {
+              ctx.sendUpstream(event);
+            } finally {
+              if (continuationWebSocketFrame.isFinalFragment()) {
+                traceContext.setReceiverHandlerContext(null);
+                DECORATE.onFrameEnd(receiverContext);
+              }
+            }
+            return;
+          }
+
+          if (frame instanceof CloseWebSocketFrame) {
+            // WebSocket Closed by client
+            CloseWebSocketFrame closeFrame = (CloseWebSocketFrame) frame;
+            int statusCode = closeFrame.getStatusCode();
+            String reasonText = closeFrame.getReasonText();
+            traceContext.setSenderHandlerContext(null);
+            traceContext.setReceiverHandlerContext(null);
+            final AgentSpan span =
+                DECORATE.onSessionCloseReceived(receiverContext, reasonText, statusCode);
+            try (final AgentScope scope = activateSpan(span)) {
+              ctx.sendUpstream(event);
+              if (closeFrame.isFinalFragment()) {
+                DECORATE.onFrameEnd(receiverContext);
+              }
+            }
+            return;
+          }
+        }
+      }
+    }
+
+    ctx.sendUpstream(event); // superclass does not throw
+  }
+}