ENG-17140 Ensure requests can be filtered because of method body

instrumentation/servlet/servlet-3.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/servlet/v3_0/nowrapping/Servlet30AndFilterInstrumentation.java

@@ -154,61 +154,63 @@ public static void exit(
       HttpServletRequest httpRequest = (HttpServletRequest) request;
       InstrumentationConfig instrumentationConfig = InstrumentationConfig.ConfigProvider.get();
 
-      if (throwable instanceof HypertraceEvaluationException) {
-        httpResponse.setStatus(403);
-        // bytebuddy treats the reassignment of this variable to null as an instruction to suppress
-        // this exception, which is what we want
-        throwable = null;
-        return;
-      }
+      try {
+        // response context to capture body and clear the context
+        VirtualField<HttpServletResponse, SpanAndObjectPair> responseContextStore =
+            VirtualField.find(HttpServletResponse.class, SpanAndObjectPair.class);
+        VirtualField<ServletOutputStream, BoundedByteArrayOutputStream> outputStreamContextStore =
+            VirtualField.find(ServletOutputStream.class, BoundedByteArrayOutputStream.class);
+        VirtualField<PrintWriter, BoundedCharArrayWriter> writerContextStore =
+            VirtualField.find(PrintWriter.class, BoundedCharArrayWriter.class);
+
+        // request context to clear body buffer
+        VirtualField<HttpServletRequest, SpanAndObjectPair> requestContextStore =
+            VirtualField.find(HttpServletRequest.class, SpanAndObjectPair.class);
+        VirtualField<ServletInputStream, ByteBufferSpanPair> inputStreamContextStore =
+            VirtualField.find(ServletInputStream.class, ByteBufferSpanPair.class);
+        VirtualField<BufferedReader, CharBufferSpanPair> readerContextStore =
+            VirtualField.find(BufferedReader.class, CharBufferSpanPair.class);
+        VirtualField<HttpServletRequest, StringMapSpanPair> urlEncodedMapContextStore =
+            VirtualField.find(HttpServletRequest.class, StringMapSpanPair.class);
+
+        if (!request.isAsyncStarted()) {
+          if (instrumentationConfig.httpHeaders().response()) {
+            for (String headerName : httpResponse.getHeaderNames()) {
+              String headerValue = httpResponse.getHeader(headerName);
+              currentSpan.setAttribute(
+                  HypertraceSemanticAttributes.httpResponseHeader(headerName), headerValue);
+            }
+          }
 
-      // response context to capture body and clear the context
-      VirtualField<HttpServletResponse, SpanAndObjectPair> responseContextStore =
-          VirtualField.find(HttpServletResponse.class, SpanAndObjectPair.class);
-      VirtualField<ServletOutputStream, BoundedByteArrayOutputStream> outputStreamContextStore =
-          VirtualField.find(ServletOutputStream.class, BoundedByteArrayOutputStream.class);
-      VirtualField<PrintWriter, BoundedCharArrayWriter> writerContextStore =
-          VirtualField.find(PrintWriter.class, BoundedCharArrayWriter.class);
-
-      // request context to clear body buffer
-      VirtualField<HttpServletRequest, SpanAndObjectPair> requestContextStore =
-          VirtualField.find(HttpServletRequest.class, SpanAndObjectPair.class);
-      VirtualField<ServletInputStream, ByteBufferSpanPair> inputStreamContextStore =
-          VirtualField.find(ServletInputStream.class, ByteBufferSpanPair.class);
-      VirtualField<BufferedReader, CharBufferSpanPair> readerContextStore =
-          VirtualField.find(BufferedReader.class, CharBufferSpanPair.class);
-      VirtualField<HttpServletRequest, StringMapSpanPair> urlEncodedMapContextStore =
-          VirtualField.find(HttpServletRequest.class, StringMapSpanPair.class);
-
-      if (!request.isAsyncStarted()) {
-        if (instrumentationConfig.httpHeaders().response()) {
-          for (String headerName : httpResponse.getHeaderNames()) {
-            String headerValue = httpResponse.getHeader(headerName);
-            currentSpan.setAttribute(
-                HypertraceSemanticAttributes.httpResponseHeader(headerName), headerValue);
+          // capture response body
+          if (instrumentationConfig.httpBody().response()
+              && ContentTypeUtils.shouldCapture(httpResponse.getContentType())) {
+            Utils.captureResponseBody(
+                currentSpan,
+                httpResponse,
+                responseContextStore,
+                outputStreamContextStore,
+                writerContextStore);
           }
-        }
 
-        // capture response body
-        if (instrumentationConfig.httpBody().response()
-            && ContentTypeUtils.shouldCapture(httpResponse.getContentType())) {
-          Utils.captureResponseBody(
-              currentSpan,
-              httpResponse,
-              responseContextStore,
-              outputStreamContextStore,
-              writerContextStore);
+          // remove request body buffers from context stores, otherwise they might get reused
+          if (instrumentationConfig.httpBody().request()
+              && ContentTypeUtils.shouldCapture(httpRequest.getContentType())) {
+            Utils.resetRequestBodyBuffers(
+                httpRequest,
+                requestContextStore,
+                inputStreamContextStore,
+                readerContextStore,
+                urlEncodedMapContextStore);
+          }
         }
-
-        // remove request body buffers from context stores, otherwise they might get reused
-        if (instrumentationConfig.httpBody().request()
-            && ContentTypeUtils.shouldCapture(httpRequest.getContentType())) {
-          Utils.resetRequestBodyBuffers(
-              httpRequest,
-              requestContextStore,
-              inputStreamContextStore,
-              readerContextStore,
-              urlEncodedMapContextStore);
+      } finally {
+        if (throwable instanceof HypertraceEvaluationException) {
+          httpResponse.setStatus(403);
+          // bytebuddy treats the reassignment of this variable to null as an instruction to
+          // suppress
+          // this exception, which is what we want
+          throwable = null;
         }
       }
     }

instrumentation/servlet/servlet-3.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/servlet/v3_0/nowrapping/request/ServletInputStreamInstrumentation.java

@@ -115,7 +115,7 @@ public static void exit(
         if (read == -1) {
           bufferSpanPair.captureBody(HypertraceSemanticAttributes.HTTP_REQUEST_BODY);
         } else {
-          bufferSpanPair.buffer.write((byte) read);
+          bufferSpanPair.writeToBuffer((byte) read);
         }
       } catch (Throwable t) {
         if (t instanceof HypertraceEvaluationException) {
@@ -143,9 +143,11 @@ public static ByteBufferSpanPair enter(@Advice.This ServletInputStream thizz) {
 
     @Advice.OnMethodExit(onThrowable = Throwable.class)
     public static void exit(
+        @Advice.This ServletInputStream thizz,
         @Advice.Return int read,
         @Advice.Argument(0) byte b[],
-        @Advice.Enter ByteBufferSpanPair bufferSpanPair) {
+        @Advice.Enter ByteBufferSpanPair bufferSpanPair)
+        throws Throwable {
       try {
         if (bufferSpanPair == null) {
           return;
@@ -159,7 +161,10 @@ public static void exit(
         if (read == -1) {
           bufferSpanPair.captureBody(HypertraceSemanticAttributes.HTTP_REQUEST_BODY);
         } else {
-          bufferSpanPair.buffer.write(b, 0, read);
+          bufferSpanPair.writeToBuffer(b, 0, read);
+          if (thizz.available() == 0) {
+            bufferSpanPair.captureBody(HypertraceSemanticAttributes.HTTP_REQUEST_BODY);
+          }
         }
       } catch (Throwable t) {
         if (t instanceof HypertraceEvaluationException) {
@@ -187,11 +192,13 @@ public static ByteBufferSpanPair enter(@Advice.This ServletInputStream thizz) {
 
     @Advice.OnMethodExit(onThrowable = Throwable.class)
     public static void exit(
+        @Advice.This ServletInputStream thizz,
         @Advice.Return int read,
         @Advice.Argument(0) byte b[],
         @Advice.Argument(1) int off,
         @Advice.Argument(2) int len,
-        @Advice.Enter ByteBufferSpanPair bufferSpanPair) {
+        @Advice.Enter ByteBufferSpanPair bufferSpanPair)
+        throws Throwable {
       try {
 
         if (bufferSpanPair == null) {
@@ -206,7 +213,10 @@ public static void exit(
         if (read == -1) {
           bufferSpanPair.captureBody(HypertraceSemanticAttributes.HTTP_REQUEST_BODY);
         } else {
-          bufferSpanPair.buffer.write(b, off, read);
+          bufferSpanPair.writeToBuffer(b, off, read);
+          if (thizz.available() == 0) {
+            bufferSpanPair.captureBody(HypertraceSemanticAttributes.HTTP_REQUEST_BODY);
+          }
         }
       } catch (Throwable t) {
         if (t instanceof HypertraceEvaluationException) {
@@ -245,8 +255,7 @@ public static void exit(
         if (callDepth > 0) {
           return;
         }
-
-        bufferSpanPair.buffer.write(b);
+        bufferSpanPair.writeToBuffer(b);
         bufferSpanPair.captureBody(HypertraceSemanticAttributes.HTTP_REQUEST_BODY);
       } catch (Throwable t) {
         if (t instanceof HypertraceEvaluationException) {
@@ -274,11 +283,13 @@ public static ByteBufferSpanPair enter(@Advice.This ServletInputStream thizz) {
 
     @Advice.OnMethodExit(onThrowable = Throwable.class)
     public static void exit(
+        @Advice.This ServletInputStream thizz,
         @Advice.Return int read,
         @Advice.Argument(0) byte[] b,
         @Advice.Argument(1) int off,
         @Advice.Argument(2) int len,
-        @Advice.Enter ByteBufferSpanPair bufferSpanPair) {
+        @Advice.Enter ByteBufferSpanPair bufferSpanPair)
+        throws Throwable {
       try {
         if (bufferSpanPair == null) {
           return;
@@ -292,7 +303,10 @@ public static void exit(
         if (read == -1) {
           bufferSpanPair.captureBody(HypertraceSemanticAttributes.HTTP_REQUEST_BODY);
         } else {
-          bufferSpanPair.buffer.write(b, off, read);
+          bufferSpanPair.writeToBuffer(b, off, read);
+          if (thizz.available() == 0) {
+            bufferSpanPair.captureBody(HypertraceSemanticAttributes.HTTP_REQUEST_BODY);
+          }
         }
       } catch (Throwable t) {
         if (t instanceof HypertraceEvaluationException) {

instrumentation/servlet/servlet-rw/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/servlet/rw/reader/BufferedReaderInstrumentation.java

@@ -95,7 +95,7 @@ public static void exit(
       if (read == -1) {
         bufferSpanPair.captureBody(HypertraceSemanticAttributes.HTTP_REQUEST_BODY);
       } else {
-        bufferSpanPair.buffer.write(read);
+        bufferSpanPair.writeToBuffer((byte) read);
       }
     }
   }
@@ -129,7 +129,7 @@ public static void exit(
       if (read == -1) {
         bufferSpanPair.captureBody(HypertraceSemanticAttributes.HTTP_REQUEST_BODY);
       } else {
-        bufferSpanPair.buffer.write(c, 0, read);
+        bufferSpanPair.writeToBuffer(c, 0, read);
       }
     }
   }
@@ -165,7 +165,7 @@ public static void exit(
       if (read == -1) {
         bufferSpanPair.captureBody(HypertraceSemanticAttributes.HTTP_REQUEST_BODY);
       } else {
-        bufferSpanPair.buffer.write(c, off, read);
+        bufferSpanPair.writeToBuffer(c, off, read);
       }
     }
   }
@@ -198,7 +198,7 @@ public static void exit(
       if (line == null) {
         bufferSpanPair.captureBody(HypertraceSemanticAttributes.HTTP_REQUEST_BODY);
       } else {
-        bufferSpanPair.buffer.write(line);
+        bufferSpanPair.writeLine(line);
       }
     }
   }

javaagent-core/src/main/java/org/hypertrace/agent/core/instrumentation/buffer/ByteBufferSpanPair.java

@@ -18,6 +18,7 @@
 
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
+import java.io.IOException;
 import java.io.UnsupportedEncodingException;
 import java.util.Map;
 import java.util.Objects;
@@ -27,7 +28,7 @@
 public class ByteBufferSpanPair {
 
   public final Span span;
-  public final BoundedByteArrayOutputStream buffer;
+  private final BoundedByteArrayOutputStream buffer;
   private final Map<String, String> headers;
   private boolean bufferCaptured;
   private final TriFunction<Span, String, Map<String, String>, Boolean> filter;
@@ -62,4 +63,19 @@ public void captureBody(AttributeKey<String> attributeKey) {
       throw new HypertraceEvaluationException();
     }
   }
+
+  public void writeToBuffer(byte singleByte) {
+    bufferCaptured = false;
+    buffer.write(singleByte);
+  }
+
+  public void writeToBuffer(byte[] b, int offset, int len) {
+    bufferCaptured = false;
+    buffer.write(b, offset, len);
+  }
+
+  public void writeToBuffer(byte[] b) throws IOException {
+    bufferCaptured = false;
+    buffer.write(b);
+  }
 }

javaagent-core/src/main/java/org/hypertrace/agent/core/instrumentation/buffer/CharBufferSpanPair.java

@@ -18,6 +18,7 @@
 
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
+import java.io.IOException;
 import java.util.Map;
 import org.hypertrace.agent.core.TriFunction;
 import org.hypertrace.agent.core.instrumentation.HypertraceEvaluationException;
@@ -26,7 +27,7 @@ public class CharBufferSpanPair {
 
   public final Span span;
   public final Map<String, String> headers;
-  public final BoundedCharArrayWriter buffer;
+  private final BoundedCharArrayWriter buffer;
   private final TriFunction<Span, String, Map<String, String>, Boolean> filter;
 
   /**
@@ -59,4 +60,24 @@ public void captureBody(AttributeKey<String> attributeKey) {
       throw new HypertraceEvaluationException();
     }
   }
+
+  public void writeToBuffer(byte singleByte) {
+    bufferCaptured = false;
+    buffer.write(singleByte);
+  }
+
+  public void writeToBuffer(char[] c, int offset, int len) {
+    bufferCaptured = false;
+    buffer.write(c, offset, len);
+  }
+
+  public void writeToBuffer(char[] c) throws IOException {
+    bufferCaptured = false;
+    buffer.write(c);
+  }
+
+  public void writeLine(String line) throws IOException {
+    bufferCaptured = false;
+    buffer.write(line);
+  }
 }