From 569142c5ec161e60c83b4dcbe6a67d8b82ee1321 Mon Sep 17 00:00:00 2001 From: "sezen.leblay" Date: Thu, 15 May 2025 14:25:15 +0200 Subject: [PATCH 1/3] HTTP response schema collection and data classification --- .../AbstractServletOutputStreamWrapper.java | 169 ++++++++++++ .../servlet/BufferedWriterWrapper.java | 244 ++++++++++++++++++ .../servlet2/ServletOutputStreamWrapper.java | 11 + .../ServletResponseBodyInstrumentation.java | 182 +++++++++++++ .../HttpServletGetOutputStreamAdvice.java | 73 ++++++ .../servlet3/HttpServletGetWriterAdvice.java | 70 +++++ .../Servlet31OutputStreamWrapper.java | 23 ++ .../Servlet31ResponseBodyInstrumentation.java | 78 ++++++ .../servlet3/Servlet3Decorator.java | 10 + .../datadog/trace/api/gateway/Events.java | 23 ++ 10 files changed, 883 insertions(+) create mode 100644 dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/AbstractServletOutputStreamWrapper.java create mode 100644 dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/BufferedWriterWrapper.java create mode 100644 dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/ServletOutputStreamWrapper.java create mode 100644 dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/ServletResponseBodyInstrumentation.java create mode 100644 dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetOutputStreamAdvice.java create mode 100644 dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetWriterAdvice.java create mode 100644 dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31OutputStreamWrapper.java create mode 100644 dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31ResponseBodyInstrumentation.java diff --git a/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/AbstractServletOutputStreamWrapper.java b/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/AbstractServletOutputStreamWrapper.java new file mode 100644 index 00000000000..49826af6edb --- /dev/null +++ b/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/AbstractServletOutputStreamWrapper.java @@ -0,0 +1,169 @@ +package datadog.trace.instrumentation.servlet; + +import datadog.trace.api.http.StoredByteBody; +import java.io.IOException; +import javax.servlet.ServletOutputStream; + +public abstract class AbstractServletOutputStreamWrapper extends ServletOutputStream { + protected final ServletOutputStream os; + private final StoredByteBody storedByteBody; + + public AbstractServletOutputStreamWrapper(ServletOutputStream os, StoredByteBody storedByteBody) { + this.os = os; + this.storedByteBody = storedByteBody; + } + + @Override + public void write(byte[] b) throws IOException { + storedByteBody.appendData(b, 0, b.length); + os.write(b); + } + + @Override + public void write(byte[] b, int off, int len) throws IOException { + storedByteBody.appendData(b, off, off + len); + os.write(b, off, len); + } + + @Override + public void write(int b) throws IOException { + storedByteBody.appendData(b); + os.write(b); + } + + @Override + public void flush() throws IOException { + os.flush(); + storedByteBody.maybeNotifyAndBlock(); + } + + @Override + public void close() throws IOException { + os.close(); + storedByteBody.maybeNotifyAndBlock(); + } + + @Override + public void print(String s) throws IOException { + if (s == null) { + s = "null"; + } + byte[] bytes = s.getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.print(s); + } + + @Override + public void print(boolean b) throws IOException { + String s = String.valueOf(b); + byte[] bytes = s.getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.print(b); + } + + @Override + public void print(char c) throws IOException { + String s = String.valueOf(c); + byte[] bytes = s.getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.print(c); + } + + @Override + public void print(int i) throws IOException { + String s = String.valueOf(i); + byte[] bytes = s.getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.print(i); + } + + @Override + public void print(long l) throws IOException { + String s = String.valueOf(l); + byte[] bytes = s.getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.print(l); + } + + @Override + public void print(float f) throws IOException { + String s = String.valueOf(f); + byte[] bytes = s.getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.print(f); + } + + @Override + public void print(double d) throws IOException { + String s = String.valueOf(d); + byte[] bytes = s.getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.print(d); + } + + @Override + public void println() throws IOException { + byte[] bytes = "\n".getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.println(); + } + + @Override + public void println(String s) throws IOException { + if (s == null) { + s = "null"; + } + String withNewline = s + "\n"; + byte[] bytes = withNewline.getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.println(s); + } + + @Override + public void println(boolean b) throws IOException { + String s = b + "\n"; + byte[] bytes = s.getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.println(b); + } + + @Override + public void println(char c) throws IOException { + String s = c + "\n"; + byte[] bytes = s.getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.println(c); + } + + @Override + public void println(int i) throws IOException { + String s = i + "\n"; + byte[] bytes = s.getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.println(i); + } + + @Override + public void println(long l) throws IOException { + String s = l + "\n"; + byte[] bytes = s.getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.println(l); + } + + @Override + public void println(float f) throws IOException { + String s = f + "\n"; + byte[] bytes = s.getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.println(f); + } + + @Override + public void println(double d) throws IOException { + String s = d + "\n"; + byte[] bytes = s.getBytes(); + storedByteBody.appendData(bytes, 0, bytes.length); + os.println(d); + } +} diff --git a/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/BufferedWriterWrapper.java b/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/BufferedWriterWrapper.java new file mode 100644 index 00000000000..d2163353a56 --- /dev/null +++ b/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/BufferedWriterWrapper.java @@ -0,0 +1,244 @@ +package datadog.trace.instrumentation.servlet; + +import datadog.trace.api.http.StoredCharBody; +import java.io.PrintWriter; +import java.util.Locale; + +public class BufferedWriterWrapper extends PrintWriter { + private final PrintWriter writer; + private final StoredCharBody storedCharBody; + + public BufferedWriterWrapper(PrintWriter writer, StoredCharBody storedCharBody) { + super(writer); + this.writer = writer; + this.storedCharBody = storedCharBody; + } + + @Override + public void write(int c) { + storedCharBody.appendData(c); + writer.write(c); + } + + @Override + public void write(char[] buf, int off, int len) { + storedCharBody.appendData(buf, off, off + len); + writer.write(buf, off, len); + } + + @Override + public void write(String s, int off, int len) { + storedCharBody.appendData(s.substring(off, off + len)); + writer.write(s, off, len); + } + + @Override + public void write(String s) { + storedCharBody.appendData(s); + writer.write(s); + } + + @Override + public void write(char[] buf) { + storedCharBody.appendData(buf, 0, buf.length); + writer.write(buf); + } + + @Override + public void print(boolean b) { + String s = String.valueOf(b); + storedCharBody.appendData(s); + writer.print(b); + } + + @Override + public void print(char c) { + storedCharBody.appendData(c); + writer.print(c); + } + + @Override + public void print(int i) { + String s = String.valueOf(i); + storedCharBody.appendData(s); + writer.print(i); + } + + @Override + public void print(long l) { + String s = String.valueOf(l); + storedCharBody.appendData(s); + writer.print(l); + } + + @Override + public void print(float f) { + String s = String.valueOf(f); + storedCharBody.appendData(s); + writer.print(f); + } + + @Override + public void print(double d) { + String s = String.valueOf(d); + storedCharBody.appendData(s); + writer.print(d); + } + + @Override + public void print(char[] s) { + storedCharBody.appendData(s, 0, s.length); + writer.print(s); + } + + @Override + public void print(String s) { + storedCharBody.appendData(s); + writer.print(s); + } + + @Override + public void print(Object obj) { + String s = String.valueOf(obj); + storedCharBody.appendData(s); + writer.print(obj); + } + + @Override + public void println() { + storedCharBody.appendData('\n'); + writer.println(); + } + + @Override + public void println(boolean x) { + String s = String.valueOf(x); + storedCharBody.appendData(s); + storedCharBody.appendData('\n'); + writer.println(x); + } + + @Override + public void println(char x) { + storedCharBody.appendData(x); + storedCharBody.appendData('\n'); + writer.println(x); + } + + @Override + public void println(int x) { + String s = String.valueOf(x); + storedCharBody.appendData(s); + storedCharBody.appendData('\n'); + writer.println(x); + } + + @Override + public void println(long x) { + String s = String.valueOf(x); + storedCharBody.appendData(s); + storedCharBody.appendData('\n'); + writer.println(x); + } + + @Override + public void println(float x) { + String s = String.valueOf(x); + storedCharBody.appendData(s); + storedCharBody.appendData('\n'); + writer.println(x); + } + + @Override + public void println(double x) { + String s = String.valueOf(x); + storedCharBody.appendData(s); + storedCharBody.appendData('\n'); + writer.println(x); + } + + @Override + public void println(char[] x) { + storedCharBody.appendData(x, 0, x.length); + storedCharBody.appendData('\n'); + writer.println(x); + } + + @Override + public void println(String x) { + storedCharBody.appendData(x); + storedCharBody.appendData('\n'); + writer.println(x); + } + + @Override + public void println(Object x) { + String s = String.valueOf(x); + storedCharBody.appendData(s); + storedCharBody.appendData('\n'); + writer.println(x); + } + + @Override + public PrintWriter printf(String format, Object... args) { + String s = String.format(format, args); + storedCharBody.appendData(s); + return writer.printf(format, args); + } + + @Override + public PrintWriter printf(Locale l, String format, Object... args) { + String s = String.format(l, format, args); + storedCharBody.appendData(s); + return writer.printf(l, format, args); + } + + @Override + public PrintWriter format(String format, Object... args) { + String s = String.format(format, args); + storedCharBody.appendData(s); + return writer.format(format, args); + } + + @Override + public PrintWriter format(Locale l, String format, Object... args) { + String s = String.format(l, format, args); + storedCharBody.appendData(s); + return writer.format(l, format, args); + } + + @Override + public PrintWriter append(CharSequence csq) { + storedCharBody.appendData(csq.toString()); + return writer.append(csq); + } + + @Override + public PrintWriter append(CharSequence csq, int start, int end) { + storedCharBody.appendData(csq.subSequence(start, end).toString()); + return writer.append(csq, start, end); + } + + @Override + public PrintWriter append(char c) { + storedCharBody.appendData(c); + return writer.append(c); + } + + @Override + public void flush() { + writer.flush(); + storedCharBody.maybeNotifyAndBlock(); + } + + @Override + public void close() { + writer.close(); + storedCharBody.maybeNotifyAndBlock(); + } + + @Override + public boolean checkError() { + return writer.checkError(); + } +} diff --git a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/ServletOutputStreamWrapper.java b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/ServletOutputStreamWrapper.java new file mode 100644 index 00000000000..ccfa3c4fcdd --- /dev/null +++ b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/ServletOutputStreamWrapper.java @@ -0,0 +1,11 @@ +package datadog.trace.instrumentation.servlet2; + +import datadog.trace.api.http.StoredByteBody; +import datadog.trace.instrumentation.servlet.AbstractServletOutputStreamWrapper; +import javax.servlet.ServletOutputStream; + +public class ServletOutputStreamWrapper extends AbstractServletOutputStreamWrapper { + public ServletOutputStreamWrapper(ServletOutputStream os, StoredByteBody storedByteBody) { + super(os, storedByteBody); + } +} diff --git a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/ServletResponseBodyInstrumentation.java b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/ServletResponseBodyInstrumentation.java new file mode 100644 index 00000000000..5014c12a7ef --- /dev/null +++ b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/ServletResponseBodyInstrumentation.java @@ -0,0 +1,182 @@ +package datadog.trace.instrumentation.servlet2; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed; +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass; +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedNoneOf; +import static datadog.trace.api.gateway.Events.EVENTS; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.not; +import static net.bytebuddy.matcher.ElementMatchers.returns; +import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; + +import com.google.auto.service.AutoService; +import datadog.trace.advice.ActiveRequestContext; +import datadog.trace.advice.RequiresRequestContext; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.gateway.CallbackProvider; +import datadog.trace.api.gateway.Flow; +import datadog.trace.api.gateway.RequestContext; +import datadog.trace.api.gateway.RequestContextSlot; +import datadog.trace.api.http.StoredBodySupplier; +import datadog.trace.api.http.StoredByteBody; +import datadog.trace.api.http.StoredCharBody; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.instrumentation.servlet.BufferedWriterWrapper; +import java.io.PrintWriter; +import java.nio.charset.Charset; +import java.util.function.BiFunction; +import javax.servlet.ServletOutputStream; +import javax.servlet.http.HttpServletResponse; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +/** + * Response bodies after servlet 3.0.x are covered by Servlet3ResponseBodyInstrumentation from the + * "request-3" module. Any changes to the behaviour here should also be reflected in "request-3". + */ +@AutoService(InstrumenterModule.class) +public class ServletResponseBodyInstrumentation extends InstrumenterModule.AppSec + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + public ServletResponseBodyInstrumentation() { + super("servlet-response-body"); + } + + @Override + public String muzzleDirective() { + return "servlet-2.x-and-3.0.x"; + } + + @Override + public ElementMatcher.Junction classLoaderMatcher() { + // Avoid matching response bodies after 3.0.x which have their own instrumentation + return not(hasClassNamed("javax.servlet.ReadListener")); + } + + @Override + public String hierarchyMarkerType() { + return "javax.servlet.http.HttpServletResponse"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named(hierarchyMarkerType())) + // ignore wrappers that ship with servlet-api + .and(namedNoneOf("javax.servlet.http.HttpServletResponseWrapper")) + .and(not(extendsClass(named("javax.servlet.http.HttpServletResponseWrapper")))); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("getOutputStream") + .and(takesNoArguments()) + .and(returns(named("javax.servlet.ServletOutputStream"))) + .and(isPublic()), + getClass().getName() + "$HttpServletGetOutputStreamAdvice"); + transformer.applyAdvice( + named("getWriter") + .and(takesNoArguments()) + .and(returns(named("java.io.PrintWriter"))) + .and(isPublic()), + getClass().getName() + "$HttpServletGetWriterAdvice"); + } + + @Override + public String[] helperClassNames() { + return new String[] { + "datadog.trace.instrumentation.servlet.BufferedWriterWrapper", + "datadog.trace.instrumentation.servlet.AbstractServletOutputStreamWrapper", + "datadog.trace.instrumentation.servlet2.ServletOutputStreamWrapper" + }; + } + + @SuppressWarnings("Duplicates") + @RequiresRequestContext(RequestContextSlot.APPSEC) + static class HttpServletGetOutputStreamAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + static void after( + @Advice.This final HttpServletResponse resp, + @Advice.Return(readOnly = false) ServletOutputStream os, + @ActiveRequestContext RequestContext reqCtx) { + if (os == null) { + return; + } + + if (os instanceof ServletOutputStreamWrapper) { + return; + } + + CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC); + BiFunction responseStartCb = + cbp.getCallback(EVENTS.responseBodyStart()); + BiFunction> responseEndedCb = + cbp.getCallback(EVENTS.responseBodyDone()); + if (responseStartCb == null || responseEndedCb == null) { + return; + } + + int lengthHint = 0; + + String encoding = resp.getCharacterEncoding(); + Charset charset = null; + try { + if (encoding != null) { + charset = Charset.forName(encoding); + } + } catch (IllegalArgumentException iae) { + // purposefully left blank + } + + StoredByteBody storedByteBody = + new StoredByteBody(reqCtx, responseStartCb, responseEndedCb, charset, lengthHint); + + os = new ServletOutputStreamWrapper(os, storedByteBody); + } + } + + @SuppressWarnings("Duplicates") + @RequiresRequestContext(RequestContextSlot.APPSEC) + static class HttpServletGetWriterAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + static void after( + @Advice.This final HttpServletResponse resp, + @Advice.Return(readOnly = false) PrintWriter writer) { + if (writer == null) { + return; + } + + AgentSpan agentSpan = activeSpan(); + if (agentSpan == null) { + return; + } + if (writer instanceof BufferedWriterWrapper) { + return; + } + RequestContext requestContext = agentSpan.getRequestContext(); + if (requestContext == null) { + return; + } + CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC); + BiFunction responseStartCb = + cbp.getCallback(EVENTS.responseBodyStart()); + BiFunction> responseEndedCb = + cbp.getCallback(EVENTS.responseBodyDone()); + if (responseStartCb == null || responseEndedCb == null) { + return; + } + + int lengthHint = 0; + + StoredCharBody storedCharBody = + new StoredCharBody(requestContext, responseStartCb, responseEndedCb, lengthHint); + + writer = new BufferedWriterWrapper(writer, storedCharBody); + } + } +} diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetOutputStreamAdvice.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetOutputStreamAdvice.java new file mode 100644 index 00000000000..938bab82896 --- /dev/null +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetOutputStreamAdvice.java @@ -0,0 +1,73 @@ +package datadog.trace.instrumentation.servlet3; + +import static datadog.trace.api.gateway.Events.EVENTS; + +import datadog.trace.advice.ActiveRequestContext; +import datadog.trace.advice.RequiresRequestContext; +import datadog.trace.api.gateway.CallbackProvider; +import datadog.trace.api.gateway.Flow; +import datadog.trace.api.gateway.RequestContext; +import datadog.trace.api.gateway.RequestContextSlot; +import datadog.trace.api.http.StoredBodySupplier; +import datadog.trace.api.http.StoredByteBody; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import java.nio.charset.Charset; +import java.util.function.BiFunction; +import javax.servlet.ServletOutputStream; +import javax.servlet.http.HttpServletResponse; +import net.bytebuddy.asm.Advice; + +@SuppressWarnings("Duplicates") +@RequiresRequestContext(RequestContextSlot.APPSEC) +class HttpServletGetOutputStreamAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + static void after( + @Advice.This final HttpServletResponse resp, + @Advice.Return(readOnly = false) ServletOutputStream os, + @ActiveRequestContext RequestContext reqCtx) { + if (os == null) { + return; + } + + String alreadyWrapped = resp.getHeader("datadog.wrapped_response_body"); + if (alreadyWrapped != null || os instanceof Servlet31OutputStreamWrapper) { + return; + } + + CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC); + BiFunction responseStartCb = + cbp.getCallback(EVENTS.responseBodyStart()); + BiFunction> responseEndedCb = + cbp.getCallback(EVENTS.responseBodyDone()); + if (responseStartCb == null || responseEndedCb == null) { + return; + } + + resp.setHeader("datadog.wrapped_response_body", Boolean.TRUE.toString()); + + int lengthHint = 0; + String lengthHeader = resp.getHeader("content-length"); + if (lengthHeader != null) { + try { + lengthHint = Integer.parseInt(lengthHeader); + } catch (NumberFormatException nfe) { + // purposefully left blank + } + } + + String encoding = resp.getCharacterEncoding(); + Charset charset = null; + try { + if (encoding != null) { + charset = Charset.forName(encoding); + } + } catch (IllegalArgumentException iae) { + // purposefully left blank + } + + StoredByteBody storedByteBody = + new StoredByteBody(reqCtx, responseStartCb, responseEndedCb, charset, lengthHint); + + os = new Servlet31OutputStreamWrapper(os, storedByteBody); + } +} diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetWriterAdvice.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetWriterAdvice.java new file mode 100644 index 00000000000..838bfd4aab0 --- /dev/null +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetWriterAdvice.java @@ -0,0 +1,70 @@ +package datadog.trace.instrumentation.servlet3; + +import static datadog.trace.api.gateway.Events.EVENTS; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; + +import datadog.trace.advice.RequiresRequestContext; +import datadog.trace.api.gateway.CallbackProvider; +import datadog.trace.api.gateway.Flow; +import datadog.trace.api.gateway.RequestContext; +import datadog.trace.api.gateway.RequestContextSlot; +import datadog.trace.api.http.StoredBodySupplier; +import datadog.trace.api.http.StoredCharBody; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.instrumentation.servlet.BufferedWriterWrapper; +import java.io.PrintWriter; +import java.util.function.BiFunction; +import javax.servlet.http.HttpServletResponse; +import net.bytebuddy.asm.Advice; + +@SuppressWarnings("Duplicates") +@RequiresRequestContext(RequestContextSlot.APPSEC) +class HttpServletGetWriterAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + static void after( + @Advice.This final HttpServletResponse resp, + @Advice.Return(readOnly = false) PrintWriter writer) { + if (writer == null) { + return; + } + + AgentSpan agentSpan = activeSpan(); + if (agentSpan == null) { + return; + } + String alreadyWrapped = resp.getHeader("datadog.wrapped_response_body"); + if (alreadyWrapped != null || writer instanceof BufferedWriterWrapper) { + return; + } + RequestContext requestContext = agentSpan.getRequestContext(); + if (requestContext == null) { + return; + } + CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC); + BiFunction responseStartCb = + cbp.getCallback(EVENTS.responseBodyStart()); + BiFunction> responseEndedCb = + cbp.getCallback(EVENTS.responseBodyDone()); + if (responseStartCb == null || responseEndedCb == null) { + return; + } + + resp.setHeader("datadog.wrapped_response_body", Boolean.TRUE.toString()); + + int lengthHint = 0; + String lengthHeader = resp.getHeader("content-length"); + if (lengthHeader != null) { + try { + lengthHint = Integer.parseInt(lengthHeader); + } catch (NumberFormatException nfe) { + // purposefully left blank + } + } + + StoredCharBody storedCharBody = + new StoredCharBody(requestContext, responseStartCb, responseEndedCb, lengthHint); + + writer = new BufferedWriterWrapper(writer, storedCharBody); + } +} diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31OutputStreamWrapper.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31OutputStreamWrapper.java new file mode 100644 index 00000000000..56939053a4b --- /dev/null +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31OutputStreamWrapper.java @@ -0,0 +1,23 @@ +package datadog.trace.instrumentation.servlet3; + +import datadog.trace.api.http.StoredByteBody; +import datadog.trace.instrumentation.servlet.AbstractServletOutputStreamWrapper; +import javax.servlet.ServletOutputStream; +import javax.servlet.WriteListener; + +/** Provides additional delegation for servlet 3.1 output streams */ +public class Servlet31OutputStreamWrapper extends AbstractServletOutputStreamWrapper { + public Servlet31OutputStreamWrapper(ServletOutputStream os, StoredByteBody storedByteBody) { + super(os, storedByteBody); + } + + @Override + public boolean isReady() { + return os.isReady(); + } + + @Override + public void setWriteListener(WriteListener writeListener) { + os.setWriteListener(writeListener); + } +} diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31ResponseBodyInstrumentation.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31ResponseBodyInstrumentation.java new file mode 100644 index 00000000000..07d4a25aecd --- /dev/null +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31ResponseBodyInstrumentation.java @@ -0,0 +1,78 @@ +package datadog.trace.instrumentation.servlet3; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed; +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass; +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedNoneOf; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.not; +import static net.bytebuddy.matcher.ElementMatchers.returns; +import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +/** + * Response bodies before servlet 3.1.x are covered by Servlet2ResponseBodyInstrumentation from the + * "request-2" module. Any changes to the behaviour here should also be reflected in "request-2". + */ +@AutoService(InstrumenterModule.class) +public class Servlet31ResponseBodyInstrumentation extends InstrumenterModule.AppSec + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + public Servlet31ResponseBodyInstrumentation() { + super("servlet-response-body"); + } + + @Override + public String muzzleDirective() { + return "servlet-3.1.x"; + } + + @Override + public ElementMatcher.Junction classLoaderMatcher() { + // Avoid matching response bodies before 3.1.x which have their own instrumentation + return hasClassNamed("javax.servlet.WriteListener"); + } + + @Override + public String hierarchyMarkerType() { + return "javax.servlet.http.HttpServletResponse"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named(hierarchyMarkerType())) + // ignore wrappers that ship with servlet-api + .and(namedNoneOf("javax.servlet.http.HttpServletResponseWrapper")) + .and(not(extendsClass(named("javax.servlet.http.HttpServletResponseWrapper")))); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("getOutputStream") + .and(takesNoArguments()) + .and(returns(named("javax.servlet.ServletOutputStream"))) + .and(isPublic()), + packageName + ".HttpServletGetOutputStreamAdvice"); + transformer.applyAdvice( + named("getWriter") + .and(takesNoArguments()) + .and(returns(named("java.io.PrintWriter"))) + .and(isPublic()), + packageName + ".HttpServletGetWriterAdvice"); + } + + @Override + public String[] helperClassNames() { + return new String[] { + "datadog.trace.instrumentation.servlet.BufferedWriterWrapper", + "datadog.trace.instrumentation.servlet.AbstractServletOutputStreamWrapper", + "datadog.trace.instrumentation.servlet3.Servlet31OutputStreamWrapper" + }; + } +} diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Decorator.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Decorator.java index f36907c76d9..c9458d3aeda 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Decorator.java +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Decorator.java @@ -7,6 +7,8 @@ import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter; import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; import datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator; +import java.io.IOException; +import java.io.OutputStream; import javax.servlet.ServletException; import javax.servlet.http.HttpServletRequest; import javax.servlet.http.HttpServletResponse; @@ -100,6 +102,14 @@ public AgentSpan onRequest( return super.onRequest(span, connection, request, context); } + protected OutputStream responseBody(HttpServletResponse response) { + try { + return response.getOutputStream(); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + @Override public AgentSpan onError(final AgentSpan span, final Throwable throwable) { if (throwable instanceof ServletException && throwable.getCause() != null) { diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/Events.java b/internal-api/src/main/java/datadog/trace/api/gateway/Events.java index d840cad01c3..9f982e692ff 100644 --- a/internal-api/src/main/java/datadog/trace/api/gateway/Events.java +++ b/internal-api/src/main/java/datadog/trace/api/gateway/Events.java @@ -312,6 +312,29 @@ public EventType>> shellCmd() { return (EventType>>) SHELL_CMD; } + static final int RESPONSE_BODY_START_ID = 26; + + @SuppressWarnings("rawtypes") + private static final EventType RESPONSE_BODY_START = + new ET<>("response.body.started", RESPONSE_BODY_START_ID); + /** The response body has started being read */ + @SuppressWarnings("unchecked") + public EventType> responseBodyStart() { + return (EventType>) RESPONSE_BODY_START; + } + + static final int RESPONSE_BODY_DONE_ID = 27; + + @SuppressWarnings("rawtypes") + private static final EventType RESPONSE_BODY_DONE = + new ET<>("response.body.done", RESPONSE_BODY_DONE_ID); + /** The response body is done being read */ + @SuppressWarnings("unchecked") + public EventType>> responseBodyDone() { + return (EventType>>) + RESPONSE_BODY_DONE; + } + static final int MAX_EVENTS = nextId.get(); private static final class ET extends EventType { From f8b5c97c52695b634a5d51ce003bf261fcc905e7 Mon Sep 17 00:00:00 2001 From: "sezen.leblay" Date: Wed, 4 Jun 2025 11:44:56 +0200 Subject: [PATCH 2/3] wip Signed-off-by: sezen.leblay --- ...et31ResponseBodyInstrumentationTest.groovy | 274 ++++++++++++++++++ 1 file changed, 274 insertions(+) create mode 100644 dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/Servlet31ResponseBodyInstrumentationTest.groovy diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/Servlet31ResponseBodyInstrumentationTest.groovy b/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/Servlet31ResponseBodyInstrumentationTest.groovy new file mode 100644 index 00000000000..4423fa03dc1 --- /dev/null +++ b/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/Servlet31ResponseBodyInstrumentationTest.groovy @@ -0,0 +1,274 @@ + +import datadog.trace.agent.test.base.HttpServerTest +import datadog.trace.api.gateway.CallbackProvider +import datadog.trace.api.gateway.Events +import datadog.trace.api.gateway.RequestContext +import datadog.trace.api.gateway.RequestContextSlot +import datadog.trace.api.http.StoredBodySupplier +import datadog.trace.bootstrap.instrumentation.api.AgentTracer +import datadog.trace.instrumentation.servlet.BufferedWriterWrapper +import datadog.trace.instrumentation.servlet3.Servlet31OutputStreamWrapper +import groovy.servlet.AbstractHttpServlet +import spock.lang.Shared + +import javax.servlet.ServletOutputStream +import javax.servlet.annotation.WebServlet +import javax.servlet.http.HttpServletRequest +import javax.servlet.http.HttpServletResponse +import java.util.function.BiFunction + +class Servlet31ResponseBodyInstrumentationTest extends TomcatServlet3Test { + + @Shared + CallbackProvider ig + + def setupSpec() { + // Set up AppSec callbacks to enable response body instrumentation + ig = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC) + Events events = Events.get() + + // Register minimal callbacks needed for response body instrumentation + ig.registerCallback(events.responseBodyStart(), { RequestContext ctx, StoredBodySupplier supplier -> + // Simple callback that just acknowledges the response body start + return null + } as BiFunction) + + ig.registerCallback(events.responseBodyDone(), { RequestContext ctx, StoredBodySupplier supplier -> + // Simple callback that just acknowledges the response body end + return datadog.trace.api.gateway.Flow.ResultFlow.empty() + } as BiFunction>) + } + + def cleanupSpec() { + if (ig != null) { + ig.reset() + } + } + + @Override + Class servlet() { + ResponseBodyTestServlet + } + + @Override + String getContext() { + return "response-body-test" + } + + def "test getOutputStream is wrapped with Servlet31OutputStreamWrapper"() { + setup: + def request = request(HttpServerTest.ServerEndpoint.SUCCESS, "GET", null).build() + + when: + def response = client.newCall(request).execute() + + then: + response.code() == 200 + response.body().string() == "output-stream-response" + + and: + // Verify that the instrumentation was applied by checking for the wrapper header + response.header("X-Stream-Wrapped") == "true" + } + + def "test getWriter is wrapped with BufferedWriterWrapper"() { + setup: + def request = request(HttpServerTest.ServerEndpoint.FORWARDED, "GET", null).build() + + when: + def response = client.newCall(request).execute() + + then: + response.code() == 200 + response.body().string() == "writer-response" + + and: + // Verify that the instrumentation was applied by checking for the wrapper header + response.header("X-Writer-Wrapped") == "true" + } + + def "test both getOutputStream and getWriter handling"() { + setup: + def request = request(HttpServerTest.ServerEndpoint.QUERY_PARAM, "GET", null).build() + + when: + def response = client.newCall(request).execute() + + then: + response.code() == 200 + response.body().string() == "both-response" + + and: + // Only one should be wrapped since they're mutually exclusive in servlet spec + (response.header("X-Stream-Wrapped") == "true") || (response.header("X-Writer-Wrapped") == "true") + } + + def "test response body instrumentation with content-length header"() { + setup: + def request = request(HttpServerTest.ServerEndpoint.CREATED, "GET", null).build() + + when: + def response = client.newCall(request).execute() + + then: + response.code() == 201 + response.body().string() == "content-length-response" + response.header("Content-Length") == "23" + response.header("X-Stream-Wrapped") == "true" + } + + def "test response body instrumentation with character encoding"() { + setup: + def request = request(HttpServerTest.ServerEndpoint.REDIRECT, "GET", null).build() + + when: + def response = client.newCall(request).execute() + + then: + response.code() == 302 + response.header("X-Writer-Wrapped") == "true" + response.header("Content-Type").contains("UTF-8") + } + + def "test response body instrumentation does not wrap when callbacks are not available"() { + setup: + // Temporarily reset callbacks to test the negative case + ig.reset() + def request = request(HttpServerTest.ServerEndpoint.SUCCESS, "GET", null).build() + + when: + def response = client.newCall(request).execute() + + then: + response.code() == 200 + response.body().string() == "output-stream-response" + + and: + // Verify that the instrumentation was NOT applied + response.header("X-Stream-Wrapped") == null + + cleanup: + // Restore callbacks for other tests + setupAppSecCallbacks() + } + + private void setupAppSecCallbacks() { + Events events = Events.get() + ig.registerCallback(events.responseBodyStart(), { RequestContext ctx, StoredBodySupplier supplier -> + return null + } as BiFunction) + + ig.registerCallback(events.responseBodyDone(), { RequestContext ctx, StoredBodySupplier supplier -> + return datadog.trace.api.gateway.Flow.ResultFlow.empty() + } as BiFunction>) + } + + @WebServlet + static class ResponseBodyTestServlet extends AbstractHttpServlet { + @Override + protected void service(HttpServletRequest req, HttpServletResponse resp) { + String path = req.getPathInfo() ?: req.getServletPath() + + resp.addHeader(HttpServerTest.IG_RESPONSE_HEADER, HttpServerTest.IG_RESPONSE_HEADER_VALUE) + + switch (path) { + case "/success": + testOutputStream(resp, "output-stream-response") + break + case "/forwarded": + testWriter(resp, "writer-response") + break + case "/query": + testBothOutputStreamAndWriter(resp, "both-response") + break + case "/created": + testOutputStreamWithContentLength(resp, "content-length-response") + break + case "/redirect": + testWriterWithEncoding(resp) + break + default: + resp.setStatus(404) + resp.getWriter().print("Not Found") + } + } + + private void testOutputStream(HttpServletResponse resp, String content) { + resp.setStatus(200) + resp.setContentType("text/plain") + + ServletOutputStream os = resp.getOutputStream() + + // Check if the output stream was wrapped + if (os instanceof Servlet31OutputStreamWrapper) { + resp.setHeader("X-Stream-Wrapped", "true") + } + + os.write(content.getBytes()) + } + + private void testWriter(HttpServletResponse resp, String content) { + resp.setStatus(200) + resp.setContentType("text/plain") + + PrintWriter writer = resp.getWriter() + + // Check if the writer was wrapped + if (writer instanceof BufferedWriterWrapper) { + resp.setHeader("X-Writer-Wrapped", "true") + } + + writer.print(content) + } + + private void testBothOutputStreamAndWriter(HttpServletResponse resp, String content) { + resp.setStatus(200) + resp.setContentType("text/plain") + + try { + // Try to get output stream first + ServletOutputStream os = resp.getOutputStream() + if (os instanceof Servlet31OutputStreamWrapper) { + resp.setHeader("X-Stream-Wrapped", "true") + } + os.write(content.getBytes()) + } catch (IllegalStateException e) { + // If output stream fails, try writer + PrintWriter writer = resp.getWriter() + if (writer instanceof BufferedWriterWrapper) { + resp.setHeader("X-Writer-Wrapped", "true") + } + writer.print(content) + } + } + + private void testOutputStreamWithContentLength(HttpServletResponse resp, String content) { + resp.setStatus(201) + resp.setContentType("text/plain") + resp.setContentLength(content.length()) + + ServletOutputStream os = resp.getOutputStream() + + if (os instanceof Servlet31OutputStreamWrapper) { + resp.setHeader("X-Stream-Wrapped", "true") + } + + os.write(content.getBytes()) + } + + private void testWriterWithEncoding(HttpServletResponse resp) { + resp.setStatus(302) + resp.setContentType("text/plain; charset=UTF-8") + resp.setCharacterEncoding("UTF-8") + resp.setHeader("Location", "/redirected") + + PrintWriter writer = resp.getWriter() + + if (writer instanceof BufferedWriterWrapper) { + resp.setHeader("X-Writer-Wrapped", "true") + } + + writer.print("Redirecting...") + } + } +} From cab7bfe297418cdc82ba7ea6d7d18032abc04591 Mon Sep 17 00:00:00 2001 From: "sezen.leblay" Date: Wed, 4 Jun 2025 13:51:38 +0200 Subject: [PATCH 3/3] some sort of test Signed-off-by: sezen.leblay --- ...et31ResponseBodyInstrumentationTest.groovy | 274 ------------------ ...et31ResponseBodyInstrumentationTest.groovy | 167 +++++++++++ 2 files changed, 167 insertions(+), 274 deletions(-) delete mode 100644 dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/Servlet31ResponseBodyInstrumentationTest.groovy create mode 100644 dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/datadog/trace/instrumentation/servlet3/Servlet31ResponseBodyInstrumentationTest.groovy diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/Servlet31ResponseBodyInstrumentationTest.groovy b/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/Servlet31ResponseBodyInstrumentationTest.groovy deleted file mode 100644 index 4423fa03dc1..00000000000 --- a/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/Servlet31ResponseBodyInstrumentationTest.groovy +++ /dev/null @@ -1,274 +0,0 @@ - -import datadog.trace.agent.test.base.HttpServerTest -import datadog.trace.api.gateway.CallbackProvider -import datadog.trace.api.gateway.Events -import datadog.trace.api.gateway.RequestContext -import datadog.trace.api.gateway.RequestContextSlot -import datadog.trace.api.http.StoredBodySupplier -import datadog.trace.bootstrap.instrumentation.api.AgentTracer -import datadog.trace.instrumentation.servlet.BufferedWriterWrapper -import datadog.trace.instrumentation.servlet3.Servlet31OutputStreamWrapper -import groovy.servlet.AbstractHttpServlet -import spock.lang.Shared - -import javax.servlet.ServletOutputStream -import javax.servlet.annotation.WebServlet -import javax.servlet.http.HttpServletRequest -import javax.servlet.http.HttpServletResponse -import java.util.function.BiFunction - -class Servlet31ResponseBodyInstrumentationTest extends TomcatServlet3Test { - - @Shared - CallbackProvider ig - - def setupSpec() { - // Set up AppSec callbacks to enable response body instrumentation - ig = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC) - Events events = Events.get() - - // Register minimal callbacks needed for response body instrumentation - ig.registerCallback(events.responseBodyStart(), { RequestContext ctx, StoredBodySupplier supplier -> - // Simple callback that just acknowledges the response body start - return null - } as BiFunction) - - ig.registerCallback(events.responseBodyDone(), { RequestContext ctx, StoredBodySupplier supplier -> - // Simple callback that just acknowledges the response body end - return datadog.trace.api.gateway.Flow.ResultFlow.empty() - } as BiFunction>) - } - - def cleanupSpec() { - if (ig != null) { - ig.reset() - } - } - - @Override - Class servlet() { - ResponseBodyTestServlet - } - - @Override - String getContext() { - return "response-body-test" - } - - def "test getOutputStream is wrapped with Servlet31OutputStreamWrapper"() { - setup: - def request = request(HttpServerTest.ServerEndpoint.SUCCESS, "GET", null).build() - - when: - def response = client.newCall(request).execute() - - then: - response.code() == 200 - response.body().string() == "output-stream-response" - - and: - // Verify that the instrumentation was applied by checking for the wrapper header - response.header("X-Stream-Wrapped") == "true" - } - - def "test getWriter is wrapped with BufferedWriterWrapper"() { - setup: - def request = request(HttpServerTest.ServerEndpoint.FORWARDED, "GET", null).build() - - when: - def response = client.newCall(request).execute() - - then: - response.code() == 200 - response.body().string() == "writer-response" - - and: - // Verify that the instrumentation was applied by checking for the wrapper header - response.header("X-Writer-Wrapped") == "true" - } - - def "test both getOutputStream and getWriter handling"() { - setup: - def request = request(HttpServerTest.ServerEndpoint.QUERY_PARAM, "GET", null).build() - - when: - def response = client.newCall(request).execute() - - then: - response.code() == 200 - response.body().string() == "both-response" - - and: - // Only one should be wrapped since they're mutually exclusive in servlet spec - (response.header("X-Stream-Wrapped") == "true") || (response.header("X-Writer-Wrapped") == "true") - } - - def "test response body instrumentation with content-length header"() { - setup: - def request = request(HttpServerTest.ServerEndpoint.CREATED, "GET", null).build() - - when: - def response = client.newCall(request).execute() - - then: - response.code() == 201 - response.body().string() == "content-length-response" - response.header("Content-Length") == "23" - response.header("X-Stream-Wrapped") == "true" - } - - def "test response body instrumentation with character encoding"() { - setup: - def request = request(HttpServerTest.ServerEndpoint.REDIRECT, "GET", null).build() - - when: - def response = client.newCall(request).execute() - - then: - response.code() == 302 - response.header("X-Writer-Wrapped") == "true" - response.header("Content-Type").contains("UTF-8") - } - - def "test response body instrumentation does not wrap when callbacks are not available"() { - setup: - // Temporarily reset callbacks to test the negative case - ig.reset() - def request = request(HttpServerTest.ServerEndpoint.SUCCESS, "GET", null).build() - - when: - def response = client.newCall(request).execute() - - then: - response.code() == 200 - response.body().string() == "output-stream-response" - - and: - // Verify that the instrumentation was NOT applied - response.header("X-Stream-Wrapped") == null - - cleanup: - // Restore callbacks for other tests - setupAppSecCallbacks() - } - - private void setupAppSecCallbacks() { - Events events = Events.get() - ig.registerCallback(events.responseBodyStart(), { RequestContext ctx, StoredBodySupplier supplier -> - return null - } as BiFunction) - - ig.registerCallback(events.responseBodyDone(), { RequestContext ctx, StoredBodySupplier supplier -> - return datadog.trace.api.gateway.Flow.ResultFlow.empty() - } as BiFunction>) - } - - @WebServlet - static class ResponseBodyTestServlet extends AbstractHttpServlet { - @Override - protected void service(HttpServletRequest req, HttpServletResponse resp) { - String path = req.getPathInfo() ?: req.getServletPath() - - resp.addHeader(HttpServerTest.IG_RESPONSE_HEADER, HttpServerTest.IG_RESPONSE_HEADER_VALUE) - - switch (path) { - case "/success": - testOutputStream(resp, "output-stream-response") - break - case "/forwarded": - testWriter(resp, "writer-response") - break - case "/query": - testBothOutputStreamAndWriter(resp, "both-response") - break - case "/created": - testOutputStreamWithContentLength(resp, "content-length-response") - break - case "/redirect": - testWriterWithEncoding(resp) - break - default: - resp.setStatus(404) - resp.getWriter().print("Not Found") - } - } - - private void testOutputStream(HttpServletResponse resp, String content) { - resp.setStatus(200) - resp.setContentType("text/plain") - - ServletOutputStream os = resp.getOutputStream() - - // Check if the output stream was wrapped - if (os instanceof Servlet31OutputStreamWrapper) { - resp.setHeader("X-Stream-Wrapped", "true") - } - - os.write(content.getBytes()) - } - - private void testWriter(HttpServletResponse resp, String content) { - resp.setStatus(200) - resp.setContentType("text/plain") - - PrintWriter writer = resp.getWriter() - - // Check if the writer was wrapped - if (writer instanceof BufferedWriterWrapper) { - resp.setHeader("X-Writer-Wrapped", "true") - } - - writer.print(content) - } - - private void testBothOutputStreamAndWriter(HttpServletResponse resp, String content) { - resp.setStatus(200) - resp.setContentType("text/plain") - - try { - // Try to get output stream first - ServletOutputStream os = resp.getOutputStream() - if (os instanceof Servlet31OutputStreamWrapper) { - resp.setHeader("X-Stream-Wrapped", "true") - } - os.write(content.getBytes()) - } catch (IllegalStateException e) { - // If output stream fails, try writer - PrintWriter writer = resp.getWriter() - if (writer instanceof BufferedWriterWrapper) { - resp.setHeader("X-Writer-Wrapped", "true") - } - writer.print(content) - } - } - - private void testOutputStreamWithContentLength(HttpServletResponse resp, String content) { - resp.setStatus(201) - resp.setContentType("text/plain") - resp.setContentLength(content.length()) - - ServletOutputStream os = resp.getOutputStream() - - if (os instanceof Servlet31OutputStreamWrapper) { - resp.setHeader("X-Stream-Wrapped", "true") - } - - os.write(content.getBytes()) - } - - private void testWriterWithEncoding(HttpServletResponse resp) { - resp.setStatus(302) - resp.setContentType("text/plain; charset=UTF-8") - resp.setCharacterEncoding("UTF-8") - resp.setHeader("Location", "/redirected") - - PrintWriter writer = resp.getWriter() - - if (writer instanceof BufferedWriterWrapper) { - resp.setHeader("X-Writer-Wrapped", "true") - } - - writer.print("Redirecting...") - } - } -} diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/datadog/trace/instrumentation/servlet3/Servlet31ResponseBodyInstrumentationTest.groovy b/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/datadog/trace/instrumentation/servlet3/Servlet31ResponseBodyInstrumentationTest.groovy new file mode 100644 index 00000000000..2abd6c6b4c0 --- /dev/null +++ b/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/datadog/trace/instrumentation/servlet3/Servlet31ResponseBodyInstrumentationTest.groovy @@ -0,0 +1,167 @@ +package datadog.trace.instrumentation.servlet3 + +import datadog.trace.agent.test.AgentTestRunner +import datadog.trace.api.gateway.CallbackProvider +import datadog.trace.api.gateway.Events +import datadog.trace.api.gateway.RequestContext +import datadog.trace.api.gateway.RequestContextSlot +import datadog.trace.api.http.StoredBodySupplier +import datadog.trace.api.http.StoredByteBody +import datadog.trace.bootstrap.instrumentation.api.AgentTracer +import spock.lang.Shared + +import javax.servlet.ServletOutputStream +import java.nio.charset.StandardCharsets +import java.util.concurrent.atomic.AtomicBoolean +import java.util.function.BiFunction + +class Servlet31ResponseBodyInstrumentationTest extends AgentTestRunner { + + @Shared + AtomicBoolean responseBodyStartCalled = new AtomicBoolean(false) + + @Shared + AtomicBoolean responseBodyDoneCalled = new AtomicBoolean(false) + + @Override + protected void configurePreAgent() { + super.configurePreAgent() + // Enable AppSec to activate response body instrumentation + injectSysConfig('dd.appsec.enabled', 'true') + injectSysConfig('dd.remote_config.enabled', 'false') + } + + def setupSpec() { + // Set up AppSec callbacks to enable response body instrumentation + CallbackProvider ig = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC) + Events events = Events.get() + + // Register callbacks that track when they're called + ig.registerCallback(events.responseBodyStart(), { RequestContext ctx, StoredBodySupplier supplier -> + responseBodyStartCalled.set(true) + return null + } as BiFunction) + + ig.registerCallback(events.responseBodyDone(), { RequestContext ctx, StoredBodySupplier supplier -> + responseBodyDoneCalled.set(true) + return datadog.trace.api.gateway.Flow.ResultFlow.empty() + } as BiFunction>) + } + + def setup() { + // Reset flags before each test + responseBodyStartCalled.set(false) + responseBodyDoneCalled.set(false) + } + + def "test Servlet31OutputStreamWrapper constructor and basic functionality"() { + given: + def mockStream = Mock(ServletOutputStream) + def mockRequestContext = Mock(RequestContext) + def storedByteBody = new StoredByteBody(mockRequestContext, + { ctx, supplier -> responseBodyStartCalled.set(true); return null } as BiFunction, + { ctx, supplier -> responseBodyDoneCalled.set(true); return datadog.trace.api.gateway.Flow.ResultFlow.empty() } as BiFunction, + StandardCharsets.UTF_8, 1024) + + when: + def wrapper = new Servlet31OutputStreamWrapper(mockStream, storedByteBody) + + then: + wrapper != null + wrapper instanceof Servlet31OutputStreamWrapper + } + + def "test maybeNotifyAndBlock is called on flush"() { + given: + def mockStream = Mock(ServletOutputStream) + def mockRequestContext = Mock(RequestContext) + def storedByteBody = new StoredByteBody(mockRequestContext, + { ctx, supplier -> responseBodyStartCalled.set(true); return null } as BiFunction, + { ctx, supplier -> responseBodyDoneCalled.set(true); return datadog.trace.api.gateway.Flow.ResultFlow.empty() } as BiFunction, + StandardCharsets.UTF_8, 1024) + def wrapper = new Servlet31OutputStreamWrapper(mockStream, storedByteBody) + + when: + wrapper.flush() + + then: + // The flush should call maybeNotifyAndBlock (line 35 in AbstractServletOutputStreamWrapper) + // This is verified by the wrapper not throwing an exception and completing normally + noExceptionThrown() + // Verify the underlying stream's flush was called + 1 * mockStream.flush() + } + + def "test maybeNotifyAndBlock is called on close"() { + given: + def mockStream = Mock(ServletOutputStream) + def mockRequestContext = Mock(RequestContext) + def storedByteBody = new StoredByteBody(mockRequestContext, + { ctx, supplier -> responseBodyStartCalled.set(true); return null } as BiFunction, + { ctx, supplier -> responseBodyDoneCalled.set(true); return datadog.trace.api.gateway.Flow.ResultFlow.empty() } as BiFunction, + StandardCharsets.UTF_8, 1024) + def wrapper = new Servlet31OutputStreamWrapper(mockStream, storedByteBody) + + when: + wrapper.close() + + then: + // The close should call maybeNotifyAndBlock (line 35 in AbstractServletOutputStreamWrapper) + // This is verified by the wrapper not throwing an exception and completing normally + noExceptionThrown() + // Verify the underlying stream's close was called + 1 * mockStream.close() + } + + def "test write operations trigger response body callbacks"() { + given: + def mockStream = Mock(ServletOutputStream) + def mockRequestContext = Mock(RequestContext) + def storedByteBody = new StoredByteBody(mockRequestContext, + { ctx, supplier -> responseBodyStartCalled.set(true); return null } as BiFunction, + { ctx, supplier -> responseBodyDoneCalled.set(true); return datadog.trace.api.gateway.Flow.ResultFlow.empty() } as BiFunction, + StandardCharsets.UTF_8, 1024) + def wrapper = new Servlet31OutputStreamWrapper(mockStream, storedByteBody) + + when: + wrapper.write("test data".getBytes()) + wrapper.close() + + then: + // Verify the underlying stream's write was called + 1 * mockStream.write(_ as byte[]) + 1 * mockStream.close() + noExceptionThrown() + } + + def "test servlet 3.1 specific methods are delegated"() { + given: + def mockStream = Mock(ServletOutputStream) + def mockRequestContext = Mock(RequestContext) + def storedByteBody = new StoredByteBody(mockRequestContext, + { ctx, supplier -> return null } as BiFunction, + { ctx, supplier -> return datadog.trace.api.gateway.Flow.ResultFlow.empty() } as BiFunction, + StandardCharsets.UTF_8, 1024) + def wrapper = new Servlet31OutputStreamWrapper(mockStream, storedByteBody) + + when: + def ready = wrapper.isReady() + + then: + ready == true + 1 * mockStream.isReady() >> true + } + + def "test instrumentation module configuration"() { + given: + def instrumentation = new Servlet31ResponseBodyInstrumentation() + + expect: + instrumentation.name() == "servlet-response-body" + instrumentation.muzzleDirective() == "servlet-3.1.x" + instrumentation.hierarchyMarkerType() == "javax.servlet.http.HttpServletResponse" + instrumentation.helperClassNames().contains("datadog.trace.instrumentation.servlet3.Servlet31OutputStreamWrapper") + instrumentation.helperClassNames().contains("datadog.trace.instrumentation.servlet.AbstractServletOutputStreamWrapper") + instrumentation.helperClassNames().contains("datadog.trace.instrumentation.servlet.BufferedWriterWrapper") + } +}