Skip to content

Commit

Permalink
ENG-17140 Ensure requests can be filtered because of method body (#364)
Browse files Browse the repository at this point in the history
* ENG-17140  Ensure requests can be filtered because of method body

* Fix smoke tests
  • Loading branch information
bradAtTraceable authored Jun 9, 2022
1 parent d0c2e8c commit 3b74567
Show file tree
Hide file tree
Showing 5 changed files with 119 additions and 66 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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;
Expand All @@ -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) {
Expand Down Expand Up @@ -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) {
Expand All @@ -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) {
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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;
Expand All @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
}
Expand Down Expand Up @@ -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);
}
}
}
Expand Down Expand Up @@ -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);
}
}
}
Expand Down Expand Up @@ -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);
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;

/**
Expand Down Expand Up @@ -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);
}
}

0 comments on commit 3b74567

Please sign in to comment.