Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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 @@ -118,7 +118,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