Capture buffers in netty 4.0, small cleanup (#227)

Signed-off-by: Pavol Loffay <[email protected]>

Author: pavolloffay

instrumentation/netty/netty-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/netty/v4_0/server/DataCaptureUtils.java

@@ -16,6 +16,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.hypertrace.netty.v4_0.server;
 
+import io.netty.buffer.ByteBuf;
 import io.netty.channel.Channel;
 import io.netty.handler.codec.http.HttpContent;
 import io.netty.handler.codec.http.HttpMessage;
@@ -35,7 +36,7 @@ public static void captureBody(
       Span span,
       Channel channel,
       AttributeKey<BoundedByteArrayOutputStream> attributeKey,
-      HttpContent httpContent) {
+      Object httpContentOrBuffer) {
 
     Attribute<BoundedByteArrayOutputStream> bufferAttr = channel.attr(attributeKey);
     BoundedByteArrayOutputStream buffer = bufferAttr.get();
@@ -44,16 +45,17 @@ public static void captureBody(
       return;
     }
 
-    final ByteArrayOutputStream finalBuffer = buffer;
-    httpContent
-        .content()
-        .forEachByte(
-            value -> {
-              finalBuffer.write(value);
-              return true;
-            });
+    ByteBuf content = castToBuf(httpContentOrBuffer);
+    if (content != null && content.isReadable()) {
+      final ByteArrayOutputStream finalBuffer = buffer;
+      content.forEachByte(
+          value -> {
+            finalBuffer.write(value);
+            return true;
+          });
+    }
 
-    if (httpContent instanceof LastHttpContent) {
+    if (httpContentOrBuffer instanceof LastHttpContent) {
       bufferAttr.remove();
       try {
         span.setAttribute(attributeKey.name(), buffer.toStringWithSuppliedCharset());
@@ -63,6 +65,16 @@ public static void captureBody(
     }
   }
 
+  private static ByteBuf castToBuf(Object msg) {
+    if (msg instanceof ByteBuf) {
+      return (ByteBuf) msg;
+    } else if (msg instanceof HttpContent) {
+      HttpContent httpContent = (HttpContent) msg;
+      return httpContent.content();
+    }
+    return null;
+  }
+
   // see io.netty.handler.codec.http.HttpUtil
   public static CharSequence getContentType(HttpMessage message) {
     return message.headers().get("content-type");

instrumentation/netty/netty-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/netty/v4_0/server/HttpServerRequestTracingHandler.java

@@ -85,8 +85,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) {
 
     if (msg instanceof HttpContent
         && agentConfig.getDataCapture().getHttpBody().getRequest().getValue()) {
-      DataCaptureUtils.captureBody(
-          span, channel, AttributeKeys.REQUEST_BODY_BUFFER, (HttpContent) msg);
+      DataCaptureUtils.captureBody(span, channel, AttributeKeys.REQUEST_BODY_BUFFER, msg);
     }
 
     ctx.fireChannelRead(msg);

instrumentation/netty/netty-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/netty/v4_0/server/HttpServerResponseTracingHandler.java

@@ -85,8 +85,7 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) {
 
     if (msg instanceof HttpContent
         && agentConfig.getDataCapture().getHttpBody().getResponse().getValue()) {
-      DataCaptureUtils.captureBody(
-          span, ctx.channel(), AttributeKeys.RESPONSE_BODY_BUFFER, (HttpContent) msg);
+      DataCaptureUtils.captureBody(span, ctx.channel(), AttributeKeys.RESPONSE_BODY_BUFFER, msg);
     }
 
     try (Scope ignored = context.makeCurrent()) {

instrumentation/netty/netty-4.1/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/netty/v4_1/server/DataCaptureUtils.java

@@ -19,6 +19,7 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.channel.Channel;
 import io.netty.handler.codec.http.HttpContent;
+import io.netty.handler.codec.http.HttpHeaderNames;
 import io.netty.handler.codec.http.HttpMessage;
 import io.netty.handler.codec.http.LastHttpContent;
 import io.netty.util.Attribute;
@@ -77,10 +78,10 @@ private static ByteBuf castToBuf(Object msg) {
 
   // see io.netty.handler.codec.http.HttpUtil
   public static CharSequence getContentType(HttpMessage message) {
-    return message.headers().get("content-type");
+    return message.headers().get(HttpHeaderNames.CONTENT_TYPE);
   }
 
   public static CharSequence getContentLength(HttpMessage message) {
-    return message.headers().get("content-length");
+    return message.headers().get(HttpHeaderNames.CONTENT_LENGTH);
   }
 }