Fix Netty 4.1 server span completion for streaming responses

dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java

@@ -15,10 +15,13 @@
 import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.handler.codec.http.HttpHeaders;
 import io.netty.handler.codec.http.HttpRequest;
+import java.util.Deque;
 
 @ChannelHandler.Sharable
 public class HttpServerRequestTracingHandler extends ChannelInboundHandlerAdapter {
   public static HttpServerRequestTracingHandler INSTANCE = new HttpServerRequestTracingHandler();
+  private static final String INCOMPLETE_RESPONSE_MESSAGE =
+      "Channel closed before response completed";
 
   @Override
   public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
@@ -54,7 +57,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
       DECORATE.onRequest(span, channel, request, parentContext);
 
       final ServerRequestContext serverContext =
-          ServerRequestContext.add(channel, context, request.headers());
+          ServerRequestContext.add(channel, context, request);
 
       Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction();
       if (rba != null) {
@@ -89,9 +92,47 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception {
       super.channelInactive(ctx);
     } finally {
       try {
-        ServerRequestContext.closeAll(ctx.channel());
+        final Deque<ServerRequestContext> storedContexts =
+            ServerRequestContext.removeAll(ctx.channel());
+        if (storedContexts != null) {
+          ServerRequestContext storedContext;
+          while ((storedContext = storedContexts.pollFirst()) != null) {
+            if (storedContext.isResponseStarted()) {
+              finishSpanOnChannelClose(storedContext);
+            } else {
+              publishSpanOnChannelClose(storedContext.tracingContext());
+            }
+          }
+        }
       } catch (final Throwable ignored) {
       }
     }
   }
+
+  private static void finishSpanOnChannelClose(final ServerRequestContext serverContext) {
+    final Context storedContext = serverContext.tracingContext();
+    final AgentSpan span = AgentSpan.fromContext(storedContext);
+    if (span == null) {
+      return;
+    }
+    if (!serverContext.isResponseCloseDelimited()) {
+      // The response declared a body via Content-Length or chunked encoding, but the channel closed
+      // before that body completed. Close-delimited responses, in contrast, end normally when the
+      // connection closes.
+      DECORATE.onError(span, new IllegalStateException(INCOMPLETE_RESPONSE_MESSAGE));
+    }
+    if (!serverContext.isBeforeFinishCalled()) {
+      serverContext.markBeforeFinishCalled();
+      DECORATE.beforeFinish(storedContext);
+    }
+    span.finish();
+  }
+
+  private static void publishSpanOnChannelClose(final Context storedContext) {
+    final AgentSpan span = AgentSpan.fromContext(storedContext);
+    if (span != null && span.phasedFinish()) {
+      // At this point we can just publish this span to avoid losing the rest of the trace.
+      span.publish();
+    }
+  }
 }

dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java

@@ -13,20 +13,19 @@
 import io.netty.channel.ChannelOutboundHandlerAdapter;
 import io.netty.channel.ChannelPromise;
 import io.netty.handler.codec.http.HttpHeaderNames;
+import io.netty.handler.codec.http.HttpHeaderValues;
 import io.netty.handler.codec.http.HttpResponse;
 import io.netty.handler.codec.http.HttpResponseStatus;
+import io.netty.handler.codec.http.HttpUtil;
+import io.netty.handler.codec.http.LastHttpContent;
+import io.netty.util.concurrent.Future;
 
 @ChannelHandler.Sharable
 public class HttpServerResponseTracingHandler extends ChannelOutboundHandlerAdapter {
   public static HttpServerResponseTracingHandler INSTANCE = new HttpServerResponseTracingHandler();
 
   @Override
   public void write(final ChannelHandlerContext ctx, final Object msg, final ChannelPromise prm) {
-    if (!(msg instanceof HttpResponse)) {
-      ctx.write(msg, prm);
-      return;
-    }
-
     final ServerRequestContext serverContext = ServerRequestContext.nextResponse(ctx.channel());
     final Context storedContext = serverContext == null ? null : serverContext.tracingContext();
     final AgentSpan span = AgentSpan.fromContext(storedContext);
@@ -36,33 +35,136 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann
       return;
     }
 
-    try (final ContextScope scope = storedContext.attach()) {
-      final HttpResponse response = (HttpResponse) msg;
+    if (!(msg instanceof HttpResponse) && !serverContext.isResponseStarted()) {
+      ctx.write(msg, prm);
+      return;
+    }
 
+    try (final ContextScope ignored = storedContext.attach()) {
+      final boolean terminalResponse = isTerminalResponse(ctx, span, serverContext, msg);
+      final ChannelPromise writePromise = terminalResponse && prm.isVoid() ? ctx.newPromise() : prm;
       try {
-        ctx.write(msg, prm);
+        if (terminalResponse) {
+          ServerRequestContext.remove(ctx.channel(), serverContext);
+          writePromise.addListener(
+              future -> finishSpan(serverContext, storedContext, span, future));
+        }
+        ctx.write(msg, writePromise);
+        if (terminalResponse) {
+          // Run the request-ended callbacks (AppSec/IAST) here, while the context is attached
+          beforeFinish(serverContext, storedContext);
+        }
       } catch (final Throwable throwable) {
         DECORATE.onError(span, throwable);
         span.setHttpStatusCode(500);
-        span.finish(); // Finish the span manually since finishSpanOnClose was false
-        ServerRequestContext.remove(ctx.channel(), serverContext);
+        if (!terminalResponse) {
+          ServerRequestContext.remove(ctx.channel(), serverContext);
+        }
+        finishSpan(serverContext, storedContext, span);
         throw throwable;
       }
-      final boolean isWebsocketUpgrade =
-          response.status() == HttpResponseStatus.SWITCHING_PROTOCOLS
-              && "websocket".equals(response.headers().get(HttpHeaderNames.UPGRADE));
+    }
+  }
+
+  private static boolean isTerminalResponse(
+      final ChannelHandlerContext ctx,
+      final AgentSpan span,
+      final ServerRequestContext serverContext,
+      final Object msg) {
+    if (msg instanceof HttpResponse) {
+      final HttpResponse response = (HttpResponse) msg;
+
+      final boolean isWebsocketUpgrade = isWebsocketUpgrade(response);
+      if (isInformationalResponse(response) && !isWebsocketUpgrade) {
+        return false;
+      }
+
       if (isWebsocketUpgrade) {
         ctx.channel()
             .attr(WEBSOCKET_SENDER_HANDLER_CONTEXT)
             .set(new HandlerContext.Sender(span, ctx.channel().id().asShortText()));
       }
-      if (response.status() != HttpResponseStatus.CONTINUE
-          && (response.status() != HttpResponseStatus.SWITCHING_PROTOCOLS || isWebsocketUpgrade)) {
-        DECORATE.onResponse(span, response);
-        DECORATE.beforeFinish(scope.context());
-        span.finish(); // Finish the span manually since finishSpanOnClose was false
-        ServerRequestContext.remove(ctx.channel(), serverContext);
+      DECORATE.onResponse(span, response);
+      serverContext.markResponseStarted();
+      if (msg instanceof LastHttpContent
+          || isBodylessResponse(serverContext, response)
+          || isWebsocketUpgrade) {
+        return true;
+      }
+      // A response with neither a Content-Length nor chunked transfer-encoding is delimited by the
+      // connection closing, so a later channel close is the normal end of this response rather than
+      // an incomplete one.
+      if (!hasKnownBodyLength(response)) {
+        serverContext.markResponseCloseDelimited();
       }
+      return false;
+    }
+    return serverContext.isResponseStarted() && msg instanceof LastHttpContent;
+  }
+
+  private static boolean isInformationalResponse(final HttpResponse response) {
+    final int statusCode = response.status().code();
+    return statusCode >= 100 && statusCode < 200;
+  }
+
+  private static boolean isWebsocketUpgrade(final HttpResponse response) {
+    return response.status() == HttpResponseStatus.SWITCHING_PROTOCOLS
+        && response
+            .headers()
+            .containsValue(HttpHeaderNames.UPGRADE, HttpHeaderValues.WEBSOCKET, true);
+  }
+
+  private static boolean isBodylessResponse(
+      final ServerRequestContext serverContext, final HttpResponse response) {
+    final int statusCode = response.status().code();
+    return serverContext.isHeadRequest()
+        || statusCode == 204
+        || statusCode == 205
+        || statusCode == 304
+        || (serverContext.isConnectRequest() && statusCode >= 200 && statusCode < 300)
+        || (contentLength(response) == 0 && !HttpUtil.isTransferEncodingChunked(response));
+  }
+
+  private static boolean hasKnownBodyLength(final HttpResponse response) {
+    return contentLength(response) >= 0 || HttpUtil.isTransferEncodingChunked(response);
+  }
+
+  /**
+   * Returns the response {@code Content-Length}, or {@code -1} when it is absent or malformed. A
+   * malformed value is left for Netty's encoder to reject rather than failing the write from the
+   * tracing handler.
+   */
+  private static long contentLength(final HttpResponse response) {
+    try {
+      return HttpUtil.getContentLength(response, -1L);
+    } catch (final NumberFormatException e) {
+      return -1L;
+    }
+  }
+
+  private static void finishSpan(
+      final ServerRequestContext serverContext,
+      final Context storedContext,
+      final AgentSpan span,
+      final Future<?> future) {
+    if (!future.isSuccess()) {
+      DECORATE.onError(span, future.cause());
+      span.setHttpStatusCode(500);
+    }
+    finishSpan(serverContext, storedContext, span);
+  }
+
+  private static void finishSpan(
+      final ServerRequestContext serverContext, final Context storedContext, final AgentSpan span) {
+    beforeFinish(serverContext, storedContext);
+    span.finish(); // Finish the span manually since finishSpanOnClose was false
+  }
+
+  private static void beforeFinish(
+      final ServerRequestContext serverContext, final Context storedContext) {
+    if (!serverContext.isBeforeFinishCalled()) {
+      serverContext.markBeforeFinishCalled();
+      DECORATE.beforeFinish(storedContext);
     }
   }
 }

dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java

@@ -62,7 +62,14 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) thr
       return;
     }
     HttpResponse origResponse = (HttpResponse) msg;
-    if (origResponse.status().code() == HttpResponseStatus.CONTINUE.code()) {
+    int statusCode = origResponse.status().code();
+    // Interim 1xx responses (e.g. 100 Continue, 103 Early Hints) precede the final response.
+    // Analyzing one here would consume the one-shot response analysis before the final response is
+    // written, so its status and headers would never be inspected. Switching Protocols (101) is
+    // terminal, so it is still analyzed.
+    if (statusCode >= 100
+        && statusCode < 200
+        && statusCode != HttpResponseStatus.SWITCHING_PROTOCOLS.code()) {
       super.write(ctx, msg, prm);
       return;
     }