Skip to content

Commit 47b551a

Browse files
committed
Fix and more tests
1 parent 102631f commit 47b551a

9 files changed

Lines changed: 157 additions & 0 deletions

File tree

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/groovy/AkkaHttpServerInstrumentationTest.groovy

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,22 @@ abstract class AkkaHttpServerInstrumentationTest extends HttpServerTest<AkkaHttp
199199
tags['response.header.content-type'] != null
200200
tags['response.header.content-length'] == SUCCESS.body.length() as String
201201
}
202+
203+
def 'blocking response sets http.response.headers.content-type span tag'() {
204+
setup:
205+
org.junit.jupiter.api.Assumptions.assumeTrue(testBlocking())
206+
207+
def request = request(SUCCESS, 'GET', null)
208+
.addHeader(IG_BLOCK_HEADER, 'auto')
209+
.build()
210+
client.newCall(request).execute()
211+
TEST_WRITER.waitForTraces(1)
212+
213+
expect:
214+
def rootSpan = TEST_WRITER.get(0).find { it.parentId == 0 }
215+
rootSpan != null
216+
rootSpan.tags['http.response.headers.content-type'] != null
217+
}
202218
}
203219

204220
abstract class AkkaHttpServerInstrumentationSyncTest extends AkkaHttpServerInstrumentationTest {

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,9 @@
44

55
import akka.http.javadsl.model.HttpHeader;
66
import akka.http.javadsl.model.headers.RawHeader;
7+
import akka.http.scaladsl.model.ContentType;
78
import akka.http.scaladsl.model.ContentTypes;
9+
import akka.http.scaladsl.model.HttpEntity;
810
import akka.http.scaladsl.model.HttpEntity$;
911
import akka.http.scaladsl.model.HttpRequest;
1012
import akka.http.scaladsl.model.HttpResponse;
@@ -37,6 +39,7 @@ public static HttpResponse handleFinishForWaf(final AgentSpan span, final HttpRe
3739
altResponse,
3840
altResponse.status().intValue(),
3941
AkkaHttpServerHeaders.responseGetter());
42+
writeBlockingResponseHeaderTags(span, altResponse);
4043
return altResponse;
4144
}
4245
}
@@ -64,6 +67,7 @@ public static HttpResponse maybeCreateBlockingResponse(AgentSpan span, HttpReque
6467
if (response != null) {
6568
DECORATE.callIGCallbackResponseAndHeaders(
6669
span, response, response.status().intValue(), AkkaHttpServerHeaders.responseGetter());
70+
writeBlockingResponseHeaderTags(span, response);
6771
}
6872
return response;
6973
}
@@ -110,4 +114,21 @@ public static HttpResponse maybeCreateBlockingResponse(
110114
}
111115
return HttpResponse.apply(code, headersList, entity, request.protocol());
112116
}
117+
118+
private static void writeBlockingResponseHeaderTags(AgentSpan span, HttpResponse response) {
119+
ResponseEntity entity = response.entity();
120+
if (entity instanceof HttpEntity.Strict) {
121+
HttpEntity.Strict strictEntity = (HttpEntity.Strict) entity;
122+
ContentType contentType = strictEntity.contentType();
123+
if (contentType != null) {
124+
span.getRequestContext()
125+
.getTraceSegment()
126+
.setTagTop("http.response.headers.content-type", contentType.value());
127+
}
128+
span.getRequestContext()
129+
.getTraceSegment()
130+
.setTagTop(
131+
"http.response.headers.content-length", Long.toString(strictEntity.contentLength()));
132+
}
133+
}
113134
}

dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/BlockingResponseHandler.java

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

3+
import static datadog.trace.instrumentation.netty38.server.NettyHttpServerDecorator.DECORATE;
34
import static org.jboss.netty.handler.codec.http.HttpHeaders.setContentLength;
45

56
import datadog.appsec.api.blocking.BlockingContentType;
67
import datadog.trace.api.gateway.Flow;
78
import datadog.trace.api.internal.TraceSegment;
89
import datadog.trace.bootstrap.blocking.BlockingActionHelper;
10+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
11+
import datadog.trace.instrumentation.netty38.ChannelTraceContext;
912
import java.util.Map;
1013
import java.util.NoSuchElementException;
1114
import org.jboss.netty.buffer.ChannelBuffer;
@@ -131,6 +134,22 @@ public void messageReceived(ChannelHandlerContext ctx, MessageEvent e) throws Ex
131134
response.setContent(buf);
132135
}
133136

137+
MaybeBlockResponseHandler maybeBlock =
138+
(MaybeBlockResponseHandler) ctx.getPipeline().get(MaybeBlockResponseHandler.class);
139+
if (maybeBlock != null) {
140+
ChannelTraceContext ctc = maybeBlock.getContextStore().get(ctx.getChannel());
141+
if (ctc != null) {
142+
AgentSpan span = ctc.getServerSpan();
143+
if (span != null) {
144+
DECORATE.callIGCallbackResponseAndHeaders(
145+
span, response, httpCode, ResponseExtractAdapter.GETTER);
146+
writeBlockingResponseHeaderTags(span, response);
147+
}
148+
ctc.setAnalyzedResponse(true);
149+
ctc.setBlockedResponse(true);
150+
}
151+
}
152+
134153
this.hasBlockedAlready = true;
135154
segment.effectivelyBlocked();
136155

@@ -147,4 +166,20 @@ public void messageReceived(ChannelHandlerContext ctx, MessageEvent e) throws Ex
147166
});
148167
Channels.write(ctxForDownstream, future, response);
149168
}
169+
170+
private static void writeBlockingResponseHeaderTags(
171+
AgentSpan span, DefaultHttpResponse response) {
172+
String contentType = response.headers().get("Content-type");
173+
if (contentType != null) {
174+
span.getRequestContext()
175+
.getTraceSegment()
176+
.setTagTop("http.response.headers.content-type", contentType);
177+
}
178+
String contentLength = response.headers().get("Content-Length");
179+
if (contentLength != null) {
180+
span.getRequestContext()
181+
.getTraceSegment()
182+
.setTagTop("http.response.headers.content-length", contentLength);
183+
}
184+
}
150185
}

dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/MaybeBlockResponseHandler.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,10 @@ public MaybeBlockResponseHandler(final ContextStore<Channel, ChannelTraceContext
3333
this.contextStore = contextStore;
3434
}
3535

36+
public ContextStore<Channel, ChannelTraceContext> getContextStore() {
37+
return contextStore;
38+
}
39+
3640
@Override
3741
public void writeRequested(ChannelHandlerContext ctx, MessageEvent msg) throws Exception {
3842
final ChannelTraceContext channelTraceContext =

dd-java-agent/instrumentation/netty/netty-3.8/src/test/groovy/datadog/trace/instrumentation/netty38/Netty38ServerTest.groovy

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ import org.jboss.netty.logging.InternalLogLevel
6363
import org.jboss.netty.logging.InternalLoggerFactory
6464
import org.jboss.netty.logging.Slf4JLoggerFactory
6565
import org.jboss.netty.util.CharsetUtil
66+
import org.junit.jupiter.api.Assumptions
6667
import spock.lang.Ignore
6768

6869
abstract class Netty38ServerTest extends HttpServerTest<ServerBootstrap> {
@@ -316,6 +317,22 @@ abstract class Netty38ServerTest extends HttpServerTest<ServerBootstrap> {
316317
boolean testBadUrl() {
317318
false
318319
}
320+
321+
def 'blocking response sets http.response.headers.content-type span tag'() {
322+
setup:
323+
Assumptions.assumeTrue(testBlocking())
324+
325+
def request = request(SUCCESS, 'GET', null)
326+
.addHeader(IG_BLOCK_HEADER, 'auto')
327+
.build()
328+
client.newCall(request).execute()
329+
TEST_WRITER.waitForTraces(1)
330+
331+
expect:
332+
def rootSpan = TEST_WRITER.get(0).find { it.parentId == 0 }
333+
rootSpan != null
334+
rootSpan.tags['http.response.headers.content-type'] != null
335+
}
319336
}
320337

321338
class Netty38ServerV0Test extends Netty38ServerTest implements TestingNettyHttpNamingConventions.ServerV0 {

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
125125
if (span != null) {
126126
DECORATE.callIGCallbackResponseAndHeaders(
127127
span, response, httpCode, ResponseExtractAdapter.GETTER);
128+
writeBlockingResponseHeaderTags(span, response.headers());
128129
}
129130
ctx.channel().attr(ANALYZED_RESPONSE_KEY).set(Boolean.TRUE);
130131
ctx.channel().attr(BLOCKED_RESPONSE_KEY).set(Boolean.TRUE);
@@ -159,6 +160,21 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
159160
});
160161
}
161162

163+
private static void writeBlockingResponseHeaderTags(AgentSpan span, HttpHeaders headers) {
164+
String contentType = headers.get("Content-type");
165+
if (contentType != null) {
166+
span.getRequestContext()
167+
.getTraceSegment()
168+
.setTagTop("http.response.headers.content-type", contentType);
169+
}
170+
String contentLength = headers.get("Content-Length");
171+
if (contentLength != null) {
172+
span.getRequestContext()
173+
.getTraceSegment()
174+
.setTagTop("http.response.headers.content-length", contentLength);
175+
}
176+
}
177+
162178
@ChannelHandler.Sharable
163179
public static class IgnoreAllWritesHandler extends ChannelOutboundHandlerAdapter {
164180
public static final IgnoreAllWritesHandler INSTANCE = new IgnoreAllWritesHandler();

dd-java-agent/instrumentation/netty/netty-4.0/src/test/groovy/Netty40ServerTest.groovy

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,22 @@ abstract class Netty40ServerTest extends HttpServerTest<EventLoopGroup> {
253253
boolean testBadUrl() {
254254
false
255255
}
256+
257+
def 'blocking response sets http.response.headers.content-type span tag'() {
258+
setup:
259+
org.junit.jupiter.api.Assumptions.assumeTrue(testBlocking())
260+
261+
def request = request(SUCCESS, 'GET', null)
262+
.addHeader(IG_BLOCK_HEADER, 'auto')
263+
.build()
264+
client.newCall(request).execute()
265+
TEST_WRITER.waitForTraces(1)
266+
267+
expect:
268+
def rootSpan = TEST_WRITER.get(0).find { it.parentId == 0 }
269+
rootSpan != null
270+
rootSpan.tags['http.response.headers.content-type'] != null
271+
}
256272
}
257273

258274
class Netty40ServerV0Test extends Netty40ServerTest implements TestingNettyHttpNamingConventions.ServerV0 {

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
125125
if (span != null) {
126126
DECORATE.callIGCallbackResponseAndHeaders(
127127
span, response, httpCode, ResponseExtractAdapter.GETTER);
128+
writeBlockingResponseHeaderTags(span, response.headers());
128129
}
129130
ctx.channel().attr(ANALYZED_RESPONSE_KEY).set(Boolean.TRUE);
130131
ctx.channel().attr(BLOCKED_RESPONSE_KEY).set(Boolean.TRUE);
@@ -159,6 +160,21 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
159160
});
160161
}
161162

163+
private static void writeBlockingResponseHeaderTags(AgentSpan span, HttpHeaders headers) {
164+
String contentType = headers.get("Content-type");
165+
if (contentType != null) {
166+
span.getRequestContext()
167+
.getTraceSegment()
168+
.setTagTop("http.response.headers.content-type", contentType);
169+
}
170+
String contentLength = headers.get("Content-Length");
171+
if (contentLength != null) {
172+
span.getRequestContext()
173+
.getTraceSegment()
174+
.setTagTop("http.response.headers.content-length", contentLength);
175+
}
176+
}
177+
162178
@ChannelHandler.Sharable
163179
public static class IgnoreAllWritesHandler extends ChannelOutboundHandlerAdapter {
164180
public static final IgnoreAllWritesHandler INSTANCE = new IgnoreAllWritesHandler();

dd-java-agent/instrumentation/netty/netty-4.1/src/test/groovy/Netty41ServerTest.groovy

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -285,6 +285,22 @@ abstract class Netty41ServerTest extends HttpServerTest<EventLoopGroup> {
285285
boolean testBadUrl() {
286286
false
287287
}
288+
289+
def 'blocking response sets http.response.headers.content-type span tag'() {
290+
setup:
291+
org.junit.jupiter.api.Assumptions.assumeTrue(testBlocking())
292+
293+
def request = request(SUCCESS, 'GET', null)
294+
.addHeader(IG_BLOCK_HEADER, 'auto')
295+
.build()
296+
client.newCall(request).execute()
297+
TEST_WRITER.waitForTraces(1)
298+
299+
expect:
300+
def rootSpan = TEST_WRITER.get(0).find { it.parentId == 0 }
301+
rootSpan != null
302+
rootSpan.tags['http.response.headers.content-type'] != null
303+
}
288304
}
289305

290306
class Netty41ServerV0Test extends Netty41ServerTest implements TestingNettyHttpNamingConventions.ServerV0 {

0 commit comments

Comments
 (0)