Add tests for servlet-rw BufferedReader and PrintWriter instrumentations

instrumentation/servlet/servlet-rw/src/test/java/io/opentelemetry/javaagent/instrumentation/hypertrace/servlet/rw/reader/BufferedReaderContextAccessInstrumentationModule.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright The Hypertrace Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.opentelemetry.javaagent.instrumentation.hypertrace.servlet.rw.reader;
+
+import static net.bytebuddy.matcher.ElementMatchers.*;
+
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
+import io.opentelemetry.javaagent.tooling.InstrumentationModule;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.io.BufferedReader;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.hypertrace.agent.core.instrumentation.buffer.CharBufferSpanPair;
+
+// SPI explicitly added in META-INF/services/...
+public class BufferedReaderContextAccessInstrumentationModule extends InstrumentationModule {
+
+  public BufferedReaderContextAccessInstrumentationModule() {
+    super("test-buffered-reader");
+  }
+
+  @Override
+  protected Map<String, String> contextStore() {
+    return Collections.singletonMap("java.io.BufferedReader", CharBufferSpanPair.class.getName());
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return Collections.singletonList(new BufferedReaderTriggerInstrumentation());
+  }
+
+  class BufferedReaderTriggerInstrumentation implements TypeInstrumentation {
+
+    @Override
+    public ElementMatcher<? super TypeDescription> typeMatcher() {
+      return named("org.BufferedReaderPrintWriterContextAccess");
+    }
+
+    @Override
+    public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+      Map<ElementMatcher.Junction<MethodDescription>, String> matchers = new HashMap<>();
+      matchers.put(
+          named("addToBufferedReaderContext").and(takesArguments(2)).and(isPublic()),
+          BufferedReaderContextAccessInstrumentationModule.class.getName() + "$TestAdvice");
+      return matchers;
+    }
+  }
+
+  static class TestAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void enter(
+        @Advice.Argument(0) BufferedReader bufferedReader,
+        @Advice.Argument(1) CharBufferSpanPair metadata) {
+      ContextStore<BufferedReader, CharBufferSpanPair> contextStore =
+          InstrumentationContext.get(BufferedReader.class, CharBufferSpanPair.class);
+      contextStore.put(bufferedReader, metadata);
+    }
+  }
+}

instrumentation/servlet/servlet-rw/src/test/java/io/opentelemetry/javaagent/instrumentation/hypertrace/servlet/rw/reader/BufferedReaderInstrumentationTest.java

@@ -0,0 +1,133 @@
+/*
+ * Copyright The Hypertrace Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.opentelemetry.javaagent.instrumentation.hypertrace.servlet.rw.reader;
+
+import io.opentelemetry.api.trace.Span;
+import java.io.BufferedReader;
+import java.io.CharArrayReader;
+import java.io.IOException;
+import org.BufferedReaderPrintWriterContextAccess;
+import org.hypertrace.agent.core.instrumentation.buffer.*;
+import org.hypertrace.agent.testing.AbstractInstrumenterTest;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class BufferedReaderInstrumentationTest extends AbstractInstrumenterTest {
+
+  private static final String TEST_SPAN_NAME = "foo";
+  private static final String BODY = "boobar";
+
+  @Test
+  public void read() throws IOException {
+    Span span = TEST_TRACER.spanBuilder(TEST_SPAN_NAME).startSpan();
+
+    BufferedReader bufferedReader = new BufferedReader(new CharArrayReader(BODY.toCharArray()));
+
+    BoundedCharArrayWriter buffer = BoundedBuffersFactory.createWriter();
+    CharBufferSpanPair bufferSpanPair = new CharBufferSpanPair(span, buffer);
+
+    BufferedReaderPrintWriterContextAccess.addToBufferedReaderContext(
+        bufferedReader, bufferSpanPair);
+
+    while (bufferedReader.read() != -1) {}
+    Assertions.assertEquals(BODY, buffer.toString());
+  }
+
+  @Test
+  public void read_callDepth_isCleared() throws IOException {
+    Span span = TEST_TRACER.spanBuilder(TEST_SPAN_NAME).startSpan();
+
+    BufferedReader bufferedReader = new BufferedReader(new CharArrayReader(BODY.toCharArray()));
+    bufferedReader.read();
+
+    BoundedCharArrayWriter buffer = BoundedBuffersFactory.createWriter();
+    CharBufferSpanPair bufferSpanPair = new CharBufferSpanPair(span, buffer);
+
+    BufferedReaderPrintWriterContextAccess.addToBufferedReaderContext(
+        bufferedReader, bufferSpanPair);
+
+    while (bufferedReader.read() != -1) {}
+    Assertions.assertEquals(BODY.substring(1), buffer.toString());
+  }
+
+  @Test
+  public void read_char_arr() throws IOException {
+    Span span = TEST_TRACER.spanBuilder(TEST_SPAN_NAME).startSpan();
+
+    BufferedReader bufferedReader = new BufferedReader(new CharArrayReader(BODY.toCharArray()));
+
+    BoundedCharArrayWriter buffer = BoundedBuffersFactory.createWriter();
+    CharBufferSpanPair bufferSpanPair = new CharBufferSpanPair(span, buffer);
+
+    BufferedReaderPrintWriterContextAccess.addToBufferedReaderContext(
+        bufferedReader, bufferSpanPair);
+
+    while (bufferedReader.read(new char[BODY.length()]) != -1) {}
+    Assertions.assertEquals(BODY, buffer.toString());
+  }
+
+  @Test
+  public void read_callDepth_char_arr() throws IOException {
+    Span span = TEST_TRACER.spanBuilder(TEST_SPAN_NAME).startSpan();
+
+    BufferedReader bufferedReader = new BufferedReader(new CharArrayReader(BODY.toCharArray()));
+    bufferedReader.read(new char[2]);
+
+    BoundedCharArrayWriter buffer = BoundedBuffersFactory.createWriter();
+    CharBufferSpanPair bufferSpanPair = new CharBufferSpanPair(span, buffer);
+
+    BufferedReaderPrintWriterContextAccess.addToBufferedReaderContext(
+        bufferedReader, bufferSpanPair);
+
+    while (bufferedReader.read(new char[BODY.length()]) != -1) {}
+    Assertions.assertEquals(BODY.substring(2), buffer.toString());
+  }
+
+  @Test
+  public void read_char_arr_offset() throws IOException {
+    Span span = TEST_TRACER.spanBuilder(TEST_SPAN_NAME).startSpan();
+
+    BufferedReader bufferedReader = new BufferedReader(new CharArrayReader(BODY.toCharArray()));
+
+    BoundedCharArrayWriter buffer = BoundedBuffersFactory.createWriter();
+    CharBufferSpanPair bufferSpanPair = new CharBufferSpanPair(span, buffer);
+
+    BufferedReaderPrintWriterContextAccess.addToBufferedReaderContext(
+        bufferedReader, bufferSpanPair);
+
+    bufferedReader.read(new char[BODY.length()], 0, 2);
+    bufferedReader.read(new char[BODY.length()], 2, BODY.length() - 2);
+    Assertions.assertEquals(BODY, buffer.toString());
+  }
+
+  @Test
+  public void readLine() throws IOException {
+    Span span = TEST_TRACER.spanBuilder(TEST_SPAN_NAME).startSpan();
+
+    BufferedReader bufferedReader =
+        new BufferedReader(new CharArrayReader((BODY + "\n").toCharArray()));
+
+    BoundedCharArrayWriter buffer = BoundedBuffersFactory.createWriter();
+    CharBufferSpanPair bufferSpanPair = new CharBufferSpanPair(span, buffer);
+
+    BufferedReaderPrintWriterContextAccess.addToBufferedReaderContext(
+        bufferedReader, bufferSpanPair);
+
+    bufferedReader.readLine();
+    Assertions.assertEquals(BODY, buffer.toString());
+  }
+}

instrumentation/servlet/servlet-rw/src/test/java/io/opentelemetry/javaagent/instrumentation/hypertrace/servlet/rw/writer/PrintWriterContextAccessInstrumentationModule.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright The Hypertrace Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.opentelemetry.javaagent.instrumentation.hypertrace.servlet.rw.writer;
+
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
+import io.opentelemetry.javaagent.tooling.InstrumentationModule;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.io.PrintWriter;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedCharArrayWriter;
+
+// SPI explicitly added in META-INF/services/...
+public class PrintWriterContextAccessInstrumentationModule extends InstrumentationModule {
+
+  public PrintWriterContextAccessInstrumentationModule() {
+    super("test-print-writer");
+  }
+
+  @Override
+  protected Map<String, String> contextStore() {
+    return Collections.singletonMap("java.io.PrintWriter", BoundedCharArrayWriter.class.getName());
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return Collections.singletonList(new PrintWriterTriggerInstrumentation());
+  }
+
+  class PrintWriterTriggerInstrumentation implements TypeInstrumentation {
+
+    @Override
+    public ElementMatcher<? super TypeDescription> typeMatcher() {
+      return named("org.BufferedReaderPrintWriterContextAccess");
+    }
+
+    @Override
+    public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+      Map<ElementMatcher.Junction<MethodDescription>, String> matchers = new HashMap<>();
+      matchers.put(
+          named("addToPrintWriterContext").and(takesArguments(2)).and(isPublic()),
+          PrintWriterContextAccessInstrumentationModule.class.getName() + "$TestAdvice");
+      return matchers;
+    }
+  }
+
+  static class TestAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void enter(
+        @Advice.Argument(0) PrintWriter printWriter,
+        @Advice.Argument(1) BoundedCharArrayWriter metadata) {
+      ContextStore<PrintWriter, BoundedCharArrayWriter> contextStore =
+          InstrumentationContext.get(PrintWriter.class, BoundedCharArrayWriter.class);
+      contextStore.put(printWriter, metadata);
+    }
+  }
+}