Skip to content
Merged
Original file line number Diff line number Diff line change
@@ -0,0 +1,167 @@
package datadog.trace.instrumentation.akkahttp.appsec;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertSame;

import datadog.appsec.api.blocking.BlockingContentType;
import datadog.appsec.api.blocking.BlockingException;
import datadog.trace.api.appsec.AppSecContext;
import datadog.trace.api.gateway.BlockResponseFunction;
import datadog.trace.api.gateway.Flow;
import datadog.trace.api.gateway.RequestContext;
import datadog.trace.api.gateway.RequestContextSlot;
import datadog.trace.api.internal.TraceSegment;
import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData;
import java.util.Map;
import java.util.function.Function;
import org.junit.jupiter.api.Test;

/**
* Covers the {@code UnmarshallerHelpers.tryBlock() -> AppSecContext.reportBlockFailure()} path.
* Hand-written test doubles are used because Mockito is only on this module's test runtime
* classpath, not its test compile classpath (see {@code gradle/java_deps.gradle}: {@code
* testRuntimeOnly libs.mokito.core}).
*/
class UnmarshallerHelpersBlockFailureTest {

private static final Flow.Action.RequestBlockingAction RBA =
new Flow.Action.RequestBlockingAction(403, BlockingContentType.AUTO);

@Test
void reportsBlockFailureWhenBlockingResponseCannotBeCommitted() {
CountingAppSecContext appSecCtx = new CountingAppSecContext();
TestRequestContext ctx =
new TestRequestContext(new TestBlockResponseFunction(false), appSecCtx);

BlockingException exception = UnmarshallerHelpers.tryBlock(ctx, RBA, "for test");

assertNull(exception);
assertEquals(1, appSecCtx.blockFailures);
assertSame(RBA, ctx.brf.lastAction);
assertSame(ctx.traceSegment, ctx.brf.lastSegment);
}

@Test
void doesNotReportBlockFailureWhenBlockingResponseIsCommitted() {
CountingAppSecContext appSecCtx = new CountingAppSecContext();
TestRequestContext ctx = new TestRequestContext(new TestBlockResponseFunction(true), appSecCtx);

BlockingException exception = UnmarshallerHelpers.tryBlock(ctx, RBA, "for test");

assertNotNull(exception);
assertEquals("Blocked request (for test)", exception.getMessage());
assertEquals(0, appSecCtx.blockFailures);
}

@Test
void doesNotReportOrThrowWhenNoBlockResponseFunctionIsRegistered() {
CountingAppSecContext appSecCtx = new CountingAppSecContext();
TestRequestContext ctx = new TestRequestContext(null, appSecCtx);

BlockingException exception = UnmarshallerHelpers.tryBlock(ctx, RBA, "for test");

assertNull(exception);
assertEquals(0, appSecCtx.blockFailures);
}

@Test
void doesNotThrowWhenAppSecSlotDoesNotHoldAnAppSecContext() {
TestRequestContext nullSlot =
new TestRequestContext(new TestBlockResponseFunction(false), null);
assertNull(UnmarshallerHelpers.tryBlock(nullSlot, RBA, "for test"));

TestRequestContext foreignSlot =
new TestRequestContext(new TestBlockResponseFunction(false), "not an AppSecContext");
assertNull(UnmarshallerHelpers.tryBlock(foreignSlot, RBA, "for test"));
}

private static final class CountingAppSecContext implements AppSecContext {
private int blockFailures;

@Override
public boolean isManuallyKept() {
return false;
}

@Override
public void reportBlockFailure() {
blockFailures++;
}
}

private static final class TestBlockResponseFunction implements BlockResponseFunction {
private final boolean committed;
private TraceSegment lastSegment;
private Flow.Action.RequestBlockingAction lastAction;

private TestBlockResponseFunction(boolean committed) {
this.committed = committed;
}

@Override
public boolean tryCommitBlockingResponse(
TraceSegment segment, Flow.Action.RequestBlockingAction rba) {
this.lastAction = rba;
return BlockResponseFunction.super.tryCommitBlockingResponse(segment, rba);
}

@Override
public boolean tryCommitBlockingResponse(
TraceSegment segment,
int statusCode,
BlockingContentType templateType,
Map<String, String> extraHeaders,
String securityResponseId) {
this.lastSegment = segment;
return committed;
}
}

private static final class TestRequestContext implements RequestContext {
private final TestBlockResponseFunction brf;
private final Object appSecData;
private final TraceSegment traceSegment = TraceSegment.NoOp.INSTANCE;

private TestRequestContext(TestBlockResponseFunction brf, Object appSecData) {
this.brf = brf;
this.appSecData = appSecData;
}

@SuppressWarnings("unchecked")
@Override
public <T> T getData(RequestContextSlot slot) {
return slot == RequestContextSlot.APPSEC ? (T) appSecData : null;
}

@Override
public TraceSegment getTraceSegment() {
return traceSegment;
}

@Override
public void setBlockResponseFunction(BlockResponseFunction blockResponseFunction) {}

@Override
public BlockResponseFunction getBlockResponseFunction() {
return brf;
}

@Override
public <T> T getOrCreateMetaStructTop(String key, Function<String, T> defaultValue) {
return null;
}

@Override
public void setClientIpAddressData(ClientIpAddressData clientIpAddressData) {}

@Override
public ClientIpAddressData getClientIpAddressData() {
return null;
}

@Override
public void close() {}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ public static HttpResponse handleFinishForWaf(final AgentSpan span, final HttpRe
if (action instanceof Flow.Action.RequestBlockingAction) {
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
if (brf instanceof AkkaBlockResponseFunction) {
brf.tryCommitBlockingResponse(requestContext.getTraceSegment(), rba);
brf.tryCommitBlockingResponse(requestContext, rba);
HttpResponse altResponse =
((AkkaBlockResponseFunction) brf).maybeCreateAlternativeResponse();
if (altResponse != null) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import datadog.trace.api.gateway.RequestContext;
import datadog.trace.api.gateway.RequestContextSlot;
import datadog.trace.api.http.MultipartContentDecoder;
import datadog.trace.api.internal.VisibleForTesting;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
import java.lang.reflect.Field;
Expand Down Expand Up @@ -607,13 +608,23 @@ private static void handleArbitraryPostData(Object o, String source) {
executeCallback(reqCtx, callback, o, source);
}

private static BlockingException tryBlock(
@VisibleForTesting
static BlockingException tryBlock(
RequestContext reqCtx, Flow.Action.RequestBlockingAction rba, String details) {
BlockResponseFunction brf = reqCtx.getBlockResponseFunction();
if (brf == null) {
return null;
}
boolean success = brf.tryCommitBlockingResponse(reqCtx.getTraceSegment(), rba);
// Conditional async-race gap (same class as netty-blocking.md §10/§11, but via
// Future.map/.recover/.thenApply on a Scala ExecutionContext instead of
// eventLoop().execute()): the block-failure report below is only guaranteed to run before
// GatewayBridge.onRequestEnded/end-of-request telemetry is emitted when the route's
// response Future causally depends (via flatMap) on this same unmarshalling Future - the
// idiomatic Akka HTTP usage. If the app decouples unmarshalling (used only for a side
// effect) from response production, or triggers toStrict() conversions independently of
// the main response chain, this report can arrive after end-of-request telemetry has
// already been emitted. This is not fixed here; see the KB entry for akka-http.
boolean success = brf.tryCommitBlockingResponse(reqCtx, rba);
if (!success) {
return null;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,10 @@

import datadog.appsec.api.blocking.BlockingContentType;
import datadog.context.Context;
import datadog.trace.api.appsec.AppSecContext;
import datadog.trace.api.gateway.Flow;
import datadog.trace.api.gateway.RequestContext;
import datadog.trace.api.gateway.RequestContextSlot;
import datadog.trace.bootstrap.blocking.BlockingActionHelper;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import java.io.OutputStream;
Expand Down Expand Up @@ -55,11 +58,20 @@ public static boolean block(
Map<String, String> extraHeaders,
String securityResponseId,
Context context) {
AgentSpan span = AgentSpan.fromContext(context);
if (GET_OUTPUT_STREAM == null) {
if (span != null) {
RequestContext reqCtx = span.getRequestContext();
if (reqCtx != null) {
Object rawAppSecCtx = reqCtx.getData(RequestContextSlot.APPSEC);
if (rawAppSecCtx instanceof AppSecContext) {
((AppSecContext) rawAppSecCtx).reportBlockFailure();
}
}
}
return false;
}

AgentSpan span = AgentSpan.fromContext(context);
try {
OutputStream os = (OutputStream) GET_OUTPUT_STREAM.invoke(response);
response.setStatus(BlockingActionHelper.getHttpCode(statusCode));
Expand All @@ -79,13 +91,34 @@ public static boolean block(
}
os.close();
response.finish();
} catch (Throwable e) {
log.info("Error committing blocking response", e);
if (span != null) {
// the response commit was attempted and failed; report it even though this method still
// returns true below (see known gap: the boolean contract can't signal this today)
RequestContext reqCtx = span.getRequestContext();
if (reqCtx != null) {
Object rawAppSecCtx = reqCtx.getData(RequestContextSlot.APPSEC);
if (rawAppSecCtx instanceof AppSecContext) {
((AppSecContext) rawAppSecCtx).reportBlockFailure();
Comment thread
jandro996 marked this conversation as resolved.
}
}
DECORATE.onError(span, e);
DECORATE.beforeFinish(context);
span.finish();
}
return true;
}

try {
if (span != null) {
span.getRequestContext().getTraceSegment().effectivelyBlocked();
}
SpanClosingListener.LISTENER.onAfterService(request);
} catch (Throwable e) {
log.info("Error committing blocking response", e);
// the response was already committed successfully; this is a finalization error, not a
// commit failure, so it must not be reported as a block failure
log.info("Error finalizing blocked request", e);
if (span != null) {
DECORATE.onError(span, e);
DECORATE.beforeFinish(context);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ static void after(
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
if (blockResponseFunction != null) {
blockResponseFunction.tryCommitBlockingResponse(reqCtx.getTraceSegment(), rba);
blockResponseFunction.tryCommitBlockingResponse(reqCtx, rba);
if (t == null) {
t = new BlockingException("Blocked request (for Parameters/processParameters)");
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ public static BlockingException fireFilesContentEvent(
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
BlockResponseFunction brf = reqCtx.getBlockResponseFunction();
if (brf != null) {
if (brf.tryCommitBlockingResponse(reqCtx.getTraceSegment(), rba)) {
if (brf.tryCommitBlockingResponse(reqCtx, rba)) {
reqCtx.getTraceSegment().effectivelyBlocked();
return new BlockingException("Blocked request (multipart file content)");
}
Expand Down Expand Up @@ -146,7 +146,7 @@ public static BlockingException fireFilenamesEvent(
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
BlockResponseFunction brf = reqCtx.getBlockResponseFunction();
if (brf != null) {
if (brf.tryCommitBlockingResponse(reqCtx.getTraceSegment(), rba)) {
if (brf.tryCommitBlockingResponse(reqCtx, rba)) {
reqCtx.getTraceSegment().effectivelyBlocked();
return new BlockingException("Blocked request (multipart file upload)");
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,7 @@ static void after(
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
if (blockResponseFunction != null) {
blockResponseFunction.tryCommitBlockingResponse(reqCtx.getTraceSegment(), rba);
blockResponseFunction.tryCommitBlockingResponse(reqCtx, rba);
if (t == null) {
t = new BlockingException("Blocked request (for Request/extractContentParameters)");
reqCtx.getTraceSegment().effectivelyBlocked();
Expand Down
Loading
Loading