Skip to content

Commit 281e32f

Browse files
jandro996jandro996
committed
Fix missing http.response.headers.content-type on blocking responses for Netty
1 parent edc4559 commit 281e32f

8 files changed

Lines changed: 132 additions & 6 deletions

File tree

‎dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -529,6 +529,11 @@ public boolean isFinishedResponseHeaders() {
529529
return finishedResponseHeaders;
530530
}
531531

532+
void clearResponseHeadersForBlocking() {
533+
responseHeaders.clear();
534+
finishedResponseHeaders = false;
535+
}
536+
532537
Map<String, List<String>> getResponseHeaders() {
533538
return responseHeaders;
534539
}

‎dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -657,7 +657,11 @@ private Flow<Void> onResponseHeaderDone(RequestContext ctx_) {
657657
return NoopFlow.INSTANCE;
658658
}
659659
ctx.finishResponseHeaders();
660-
return maybePublishResponseData(ctx);
660+
Flow<Void> flow = maybePublishResponseData(ctx);
661+
if (flow.getAction() instanceof Flow.Action.RequestBlockingAction) {
662+
ctx.clearResponseHeadersForBlocking();
663+
}
664+
return flow;
661665
}
662666

663667
private void onResponseHeader(RequestContext ctx_, String name, String value) {

‎dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/AppSecRequestContextSpecification.groovy‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,37 @@ class AppSecRequestContextSpecification extends DDSpecification {
8181
thrown(IllegalStateException)
8282
}
8383

84+
void 'clearResponseHeadersForBlocking clears response headers and resets finished flag'() {
85+
given:
86+
ctx.addResponseHeader('content-type', 'text/html')
87+
ctx.finishResponseHeaders()
88+
89+
expect:
90+
!ctx.responseHeaders.isEmpty()
91+
ctx.isFinishedResponseHeaders()
92+
93+
when:
94+
ctx.clearResponseHeadersForBlocking()
95+
96+
then:
97+
ctx.responseHeaders.isEmpty()
98+
!ctx.isFinishedResponseHeaders()
99+
}
100+
101+
void 'after clearResponseHeadersForBlocking new response headers can be added'() {
102+
given:
103+
ctx.addResponseHeader('content-type', 'text/html')
104+
ctx.finishResponseHeaders()
105+
ctx.clearResponseHeadersForBlocking()
106+
107+
when:
108+
ctx.addResponseHeader('content-type', 'application/json')
109+
110+
then:
111+
ctx.responseHeaders == ['content-type': ['application/json']]
112+
notThrown(IllegalStateException)
113+
}
114+
84115
void 'setting uri a second time is ignored, first value wins'() {
85116
when:
86117
ctx.rawURI = '/a'

‎dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import datadog.trace.api.appsec.MediaType
1818
import datadog.trace.api.config.GeneralConfig
1919
import datadog.trace.api.function.TriConsumer
2020
import datadog.trace.api.function.TriFunction
21+
import datadog.appsec.api.blocking.BlockingContentType
2122
import datadog.trace.api.gateway.BlockResponseFunction
2223
import datadog.trace.api.gateway.Flow
2324
import datadog.trace.api.gateway.IGSpanInfo
@@ -225,6 +226,49 @@ class GatewayBridgeSpecification extends DDSpecification {
225226
1 * traceSegment.setTagTop('actor.ip', '8.8.8.8')
226227
}
227228

229+
void 'request_end writes response headers even when no appsec events'() {
230+
AppSecRequestContext mockAppSecCtx = Mock(AppSecRequestContext)
231+
mockAppSecCtx.requestHeaders >> [:]
232+
mockAppSecCtx.responseHeaders >> ['content-type': ['text/plain']]
233+
RequestContext mockCtx = Stub(RequestContext) {
234+
getData(RequestContextSlot.APPSEC) >> mockAppSecCtx
235+
getTraceSegment() >> traceSegment
236+
}
237+
IGSpanInfo spanInfo = Mock(AgentSpan)
238+
239+
when:
240+
def flow = requestEndedCB.apply(mockCtx, spanInfo)
241+
242+
then:
243+
1 * spanInfo.getTags() >> TagMap.fromMap([:])
244+
1 * mockAppSecCtx.transferCollectedEvents() >> []
245+
1 * mockAppSecCtx.close()
246+
1 * traceSegment.setTagTop("_dd.appsec.enabled", 1)
247+
1 * traceSegment.setTagTop("_dd.runtime_family", "jvm")
248+
1 * traceSegment.setTagTop('http.response.headers.content-type', 'text/plain')
249+
1 * wafMetricCollector.wafRequest(_, _, _, _, _, _, _)
250+
flow.result == null
251+
flow.action == Flow.Action.Noop.INSTANCE
252+
}
253+
254+
void 'response_header_done clears response headers for blocking when WAF blocks'() {
255+
given:
256+
def blockingFlow = Stub(Flow) {
257+
getAction() >> new Flow.Action.RequestBlockingAction(403, BlockingContentType.AUTO)
258+
}
259+
eventDispatcher.getDataSubscribers(_) >> nonEmptyDsInfo
260+
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >> blockingFlow
261+
262+
when:
263+
respHeaderCB.accept(ctx, 'content-type', 'text/html')
264+
responseStartedCB.apply(ctx, 403)
265+
respHeadersDoneCB.apply(ctx)
266+
267+
then:
268+
ctx.data.responseHeaders.isEmpty()
269+
!ctx.data.finishedResponseHeaders
270+
}
271+
228272
void 'bridge can collect headers'() {
229273
when:
230274
reqHeaderCB.accept(ctx, 'header1', 'value 1.1')

‎dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
import akka.stream.Materializer;
99
import datadog.context.Context;
1010
import datadog.context.ContextScope;
11-
import datadog.trace.api.gateway.Flow;
1211
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1312
import datadog.trace.instrumentation.akkahttp.appsec.BlockingResponseHelper;
1413
import scala.Function1;
@@ -35,10 +34,9 @@ public Future<HttpResponse> apply(final HttpRequest request) {
3534
Future<HttpResponse> futureResponse;
3635

3736
// handle blocking in the beginning of the request
38-
Flow.Action.RequestBlockingAction rba;
39-
if ((rba = span.getRequestBlockingAction()) != null) {
37+
if (span.getRequestBlockingAction() != null) {
4038
request.discardEntityBytes(materializer);
41-
HttpResponse response = BlockingResponseHelper.maybeCreateBlockingResponse(rba, request);
39+
HttpResponse response = BlockingResponseHelper.maybeCreateBlockingResponse(span, request);
4240
span.getRequestContext().getTraceSegment().effectivelyBlocked();
4341
DatadogWrapperHelper.finishSpan(context, response);
4442
return FastFuture$.MODULE$.<HttpResponse>successful().apply(response);

‎dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/BlockingResponseHelper.java‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,11 @@ public static HttpResponse handleFinishForWaf(final AgentSpan span, final HttpRe
3232
HttpResponse altResponse = ((AkkaBlockResponseFunction) brf).maybeCreateAlternativeResponse();
3333
if (altResponse != null) {
3434
// we already blocked during the request
35+
DECORATE.callIGCallbackResponseAndHeaders(
36+
span,
37+
altResponse,
38+
altResponse.status().intValue(),
39+
AkkaHttpServerHeaders.responseGetter());
3540
return altResponse;
3641
}
3742
}
@@ -55,7 +60,12 @@ public static HttpResponse handleFinishForWaf(final AgentSpan span, final HttpRe
5560
}
5661

5762
public static HttpResponse maybeCreateBlockingResponse(AgentSpan span, HttpRequest request) {
58-
return maybeCreateBlockingResponse(span.getRequestBlockingAction(), request);
63+
HttpResponse response = maybeCreateBlockingResponse(span.getRequestBlockingAction(), request);
64+
if (response != null) {
65+
DECORATE.callIGCallbackResponseAndHeaders(
66+
span, response, response.status().intValue(), AkkaHttpServerHeaders.responseGetter());
67+
}
68+
return response;
5969
}
6070

6171
public static HttpResponse maybeCreateBlockingResponse(

‎dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/BlockingResponseHandler.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,18 @@
11
package datadog.trace.instrumentation.netty40.server;
22

3+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
4+
import static datadog.trace.instrumentation.netty40.AttributeKeys.ANALYZED_RESPONSE_KEY;
5+
import static datadog.trace.instrumentation.netty40.AttributeKeys.BLOCKED_RESPONSE_KEY;
6+
import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY;
7+
import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.DECORATE;
38
import static io.netty.handler.codec.http.HttpHeaders.setContentLength;
49

510
import datadog.appsec.api.blocking.BlockingContentType;
11+
import datadog.context.Context;
612
import datadog.trace.api.gateway.Flow;
713
import datadog.trace.api.internal.TraceSegment;
814
import datadog.trace.bootstrap.blocking.BlockingActionHelper;
15+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
916
import io.netty.channel.ChannelHandler;
1017
import io.netty.channel.ChannelHandlerContext;
1118
import io.netty.channel.ChannelInboundHandlerAdapter;
@@ -112,6 +119,16 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
112119
}
113120

114121
this.hasBlockedAlready = true;
122+
123+
Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get();
124+
AgentSpan span = spanFromContext(storedContext);
125+
if (span != null) {
126+
DECORATE.callIGCallbackResponseAndHeaders(
127+
span, response, httpCode, ResponseExtractAdapter.GETTER);
128+
}
129+
ctx.channel().attr(ANALYZED_RESPONSE_KEY).set(Boolean.TRUE);
130+
ctx.channel().attr(BLOCKED_RESPONSE_KEY).set(Boolean.TRUE);
131+
115132
ReferenceCountUtil.release(msg);
116133

117134
// write starts in the handler before the one associated with ctx

‎dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/BlockingResponseHandler.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,17 @@
11
package datadog.trace.instrumentation.netty41.server;
22

3+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
4+
import static datadog.trace.instrumentation.netty41.AttributeKeys.ANALYZED_RESPONSE_KEY;
5+
import static datadog.trace.instrumentation.netty41.AttributeKeys.BLOCKED_RESPONSE_KEY;
6+
import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY;
7+
import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE;
8+
39
import datadog.appsec.api.blocking.BlockingContentType;
10+
import datadog.context.Context;
411
import datadog.trace.api.gateway.Flow;
512
import datadog.trace.api.internal.TraceSegment;
613
import datadog.trace.bootstrap.blocking.BlockingActionHelper;
14+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
715
import io.netty.channel.ChannelHandler;
816
import io.netty.channel.ChannelHandlerContext;
917
import io.netty.channel.ChannelInboundHandlerAdapter;
@@ -112,6 +120,15 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
112120

113121
this.hasBlockedAlready = true;
114122

123+
Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get();
124+
AgentSpan span = spanFromContext(storedContext);
125+
if (span != null) {
126+
DECORATE.callIGCallbackResponseAndHeaders(
127+
span, response, httpCode, ResponseExtractAdapter.GETTER);
128+
}
129+
ctx.channel().attr(ANALYZED_RESPONSE_KEY).set(Boolean.TRUE);
130+
ctx.channel().attr(BLOCKED_RESPONSE_KEY).set(Boolean.TRUE);
131+
115132
ReferenceCountUtil.release(msg);
116133

117134
// write starts in the handler before the one associated with ctx

0 commit comments

Comments
 (0)