Skip to content

Commit cae8110

Browse files
claponcetdevflow.devflow-routing-intake
claponcet
and
devflow.devflow-routing-intake
authored
Add AppSec Java support for AWS Lambdas (#10570)
send lambda data to the WAF for analysis refactor + add appsec data to span unit tests add better support for query parameters apply spotless add condition to RC warning log to prevent logging multiple times when using Lambdas add appsec instrumentation tests fix test crash remove unused var forwarded headers parsing + downgrade log level formatting use substring rather than split to avoid regex matching address PR review: throttled warn log + simplify stream reading - Upgrade mergeContexts log from debug to throttled warn (RatelimitedLogger) so extension context type mismatches surface in production logs - Replace InputStreamReader/StringBuilder with direct byte[] read to reduce allocations and avoid implicit stream close Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Merge branch 'master' into clara.poncet/appsec-aws-lambda Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>
1 parent 6e28457 commit cae8110

File tree

8 files changed

+2596
-8
lines changed

8 files changed

+2596
-8
lines changed

‎dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecConfigServiceImpl.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -422,7 +422,7 @@ public void maybeSubscribeConfigPolling() {
422422
} else {
423423
subscribeConfigurationPoller();
424424
}
425-
} else {
425+
} else if (!tracerConfig.isAwsServerless()) {
426426
log.info("Remote config is disabled; AppSec will not be able to use it");
427427
}
428428
}

‎dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2424
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
2525
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
26+
import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
2627
import datadog.trace.config.inversion.ConfigHelper;
2728
import net.bytebuddy.asm.Advice;
2829
import net.bytebuddy.description.type.TypeDescription;
@@ -89,13 +90,14 @@ static AgentScope enter(
8990
return null;
9091
}
9192
String lambdaRequestId = awsContext.getAwsRequestId();
92-
AgentSpanContext lambdaContext = AgentTracer.get().notifyExtensionStart(in, lambdaRequestId);
93+
AgentSpanContext lambdaContext = AgentTracer.get().notifyLambdaStart(in, lambdaRequestId);
9394
final AgentSpan span;
9495
if (null == lambdaContext) {
9596
span = startSpan(INVOCATION_SPAN_NAME);
9697
} else {
9798
span = startSpan(INVOCATION_SPAN_NAME, lambdaContext);
9899
}
100+
span.setSpanType(InternalSpanTypes.SERVERLESS);
99101
span.setTag("request_id", lambdaRequestId);
100102

101103
final AgentScope scope = activateSpan(span);
@@ -123,6 +125,7 @@ static void exit(
123125
}
124126
String lambdaRequestId = awsContext.getAwsRequestId();
125127

128+
AgentTracer.get().notifyAppSecEnd(span);
126129
span.finish();
127130
AgentTracer.get().notifyExtensionEnd(span, result, null != throwable, lambdaRequestId);
128131
} finally {

‎dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/groovy/LambdaHandlerInstrumentationTest.groovy‎

Lines changed: 172 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,20 @@
1+
import static datadog.trace.api.gateway.Events.EVENTS
2+
13
import datadog.trace.agent.test.naming.VersionedNamingTestBase
2-
import java.nio.charset.StandardCharsets
4+
import datadog.trace.api.DDSpanTypes
5+
import datadog.trace.api.function.TriConsumer
6+
import datadog.trace.api.function.TriFunction
7+
import datadog.trace.api.gateway.Flow
8+
import datadog.trace.api.gateway.RequestContext
9+
import datadog.trace.api.gateway.RequestContextSlot
10+
import datadog.trace.bootstrap.ActiveSubsystems
11+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer
12+
import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter
313
import com.amazonaws.services.lambda.runtime.Context
14+
import java.nio.charset.StandardCharsets
15+
import java.util.function.BiFunction
16+
import java.util.function.Function
17+
import java.util.function.Supplier
418

519
abstract class LambdaHandlerInstrumentationTest extends VersionedNamingTestBase {
620
def requestId = "test-request-id"
@@ -16,6 +30,53 @@ abstract class LambdaHandlerInstrumentationTest extends VersionedNamingTestBase
1630
null
1731
}
1832

33+
def ig
34+
def appSecStarted = false
35+
def capturedMethod = null
36+
def capturedPath = null
37+
def capturedHeaders = [:]
38+
def capturedBody = null
39+
def appSecEnded = false
40+
41+
def setup() {
42+
ig = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC)
43+
ActiveSubsystems.APPSEC_ACTIVE = true
44+
appSecStarted = false
45+
capturedMethod = null
46+
capturedPath = null
47+
capturedHeaders = [:]
48+
capturedBody = null
49+
appSecEnded = false
50+
ig.registerCallback(EVENTS.requestStarted(), {
51+
appSecStarted = true
52+
new Flow.ResultFlow(new Object())
53+
} as Supplier)
54+
ig.registerCallback(EVENTS.requestMethodUriRaw(), { RequestContext ctx, String method, URIDataAdapter uri ->
55+
capturedMethod = method
56+
capturedPath = uri.path()
57+
Flow.ResultFlow.empty()
58+
} as TriFunction)
59+
ig.registerCallback(EVENTS.requestHeader(), { RequestContext ctx, String name, String value ->
60+
capturedHeaders[name] = value
61+
} as TriConsumer)
62+
ig.registerCallback(EVENTS.requestHeaderDone(), { RequestContext ctx ->
63+
Flow.ResultFlow.empty()
64+
} as Function)
65+
ig.registerCallback(EVENTS.requestBodyProcessed(), { RequestContext ctx, Object body ->
66+
capturedBody = body
67+
Flow.ResultFlow.empty()
68+
} as BiFunction)
69+
ig.registerCallback(EVENTS.requestEnded(), { RequestContext ctx, Object spanInfo ->
70+
appSecEnded = true
71+
Flow.ResultFlow.empty()
72+
} as BiFunction)
73+
}
74+
75+
def cleanup() {
76+
ig.reset()
77+
ActiveSubsystems.APPSEC_ACTIVE = false
78+
}
79+
1980
def "test lambda streaming handler"() {
2081
when:
2182
def input = new ByteArrayInputStream(StandardCharsets.UTF_8.encode("Hello").array())
@@ -30,6 +91,7 @@ abstract class LambdaHandlerInstrumentationTest extends VersionedNamingTestBase
3091
trace(1) {
3192
span {
3293
operationName operation()
94+
spanType DDSpanTypes.SERVERLESS
3395
errored false
3496
}
3597
}
@@ -51,6 +113,7 @@ abstract class LambdaHandlerInstrumentationTest extends VersionedNamingTestBase
51113
trace(1) {
52114
span {
53115
operationName operation()
116+
spanType DDSpanTypes.SERVERLESS
54117
errored true
55118
tags {
56119
tag "request_id", requestId
@@ -73,6 +136,114 @@ abstract class LambdaHandlerInstrumentationTest extends VersionedNamingTestBase
73136
}
74137
}
75138
}
139+
140+
def "appsec callbacks are invoked for API Gateway v1 event"() {
141+
given:
142+
def eventJson = """{
143+
"path": "/api/users/123",
144+
"headers": {"content-type": "application/json", "x-forwarded-for": "203.0.113.1"},
145+
"body": "{\\"key\\": \\"value\\"}",
146+
"requestContext": {
147+
"httpMethod": "GET",
148+
"requestId": "req-abc",
149+
"identity": {"sourceIp": "203.0.113.1"}
150+
}
151+
}"""
152+
153+
when:
154+
def input = new ByteArrayInputStream(eventJson.getBytes(StandardCharsets.UTF_8))
155+
def output = new ByteArrayOutputStream()
156+
def ctx = Stub(Context) { getAwsRequestId() >> requestId }
157+
new HandlerStreaming().handleRequest(input, output, ctx)
158+
159+
then:
160+
appSecStarted
161+
capturedMethod == "GET"
162+
capturedPath == "/api/users/123"
163+
capturedHeaders["content-type"] == "application/json"
164+
capturedBody instanceof Map
165+
appSecEnded
166+
assertTraces(1) {
167+
trace(1) {
168+
span {
169+
operationName operation()
170+
spanType DDSpanTypes.SERVERLESS
171+
errored false
172+
}
173+
}
174+
}
175+
}
176+
177+
def "appsec callbacks are invoked for API Gateway v2 HTTP event"() {
178+
given:
179+
def eventJson = """{
180+
"version": "2.0",
181+
"headers": {"content-type": "application/json", "accept": "application/json"},
182+
"cookies": ["session=abc123"],
183+
"body": "{\\"key\\": \\"value\\"}",
184+
"requestContext": {
185+
"http": {
186+
"method": "POST",
187+
"path": "/api/items",
188+
"sourceIp": "198.51.100.1"
189+
},
190+
"domainName": "api.example.com"
191+
}
192+
}"""
193+
194+
when:
195+
def input = new ByteArrayInputStream(eventJson.getBytes(StandardCharsets.UTF_8))
196+
def output = new ByteArrayOutputStream()
197+
def ctx = Stub(Context) { getAwsRequestId() >> requestId }
198+
new HandlerStreaming().handleRequest(input, output, ctx)
199+
200+
then:
201+
appSecStarted
202+
capturedMethod == "POST"
203+
capturedPath == "/api/items"
204+
capturedHeaders["content-type"] == "application/json"
205+
capturedHeaders["cookie"] == "session=abc123"
206+
capturedBody instanceof Map
207+
appSecEnded
208+
assertTraces(1) {
209+
trace(1) {
210+
span {
211+
operationName operation()
212+
spanType DDSpanTypes.SERVERLESS
213+
errored false
214+
}
215+
}
216+
}
217+
}
218+
219+
def "appsec callbacks are not invoked when appsec is disabled"() {
220+
given:
221+
ActiveSubsystems.APPSEC_ACTIVE = false
222+
223+
when:
224+
def eventJson = """{
225+
"path": "/api/test",
226+
"requestContext": {"httpMethod": "GET", "requestId": "req-xyz"}
227+
}"""
228+
def input = new ByteArrayInputStream(eventJson.getBytes(StandardCharsets.UTF_8))
229+
def output = new ByteArrayOutputStream()
230+
def ctx = Stub(Context) { getAwsRequestId() >> requestId }
231+
new HandlerStreaming().handleRequest(input, output, ctx)
232+
233+
then:
234+
!appSecStarted
235+
capturedMethod == null
236+
!appSecEnded
237+
assertTraces(1) {
238+
trace(1) {
239+
span {
240+
operationName operation()
241+
spanType DDSpanTypes.SERVERLESS
242+
errored false
243+
}
244+
}
245+
}
246+
}
76247
}
77248

78249

‎dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@
105105
import datadog.trace.core.taginterceptor.RuleFlags;
106106
import datadog.trace.core.taginterceptor.TagInterceptor;
107107
import datadog.trace.core.traceinterceptor.LatencyTraceInterceptor;
108+
import datadog.trace.lambda.LambdaAppSecHandler;
108109
import datadog.trace.lambda.LambdaHandler;
109110
import datadog.trace.util.AgentTaskScheduler;
110111
import java.io.IOException;
@@ -1195,8 +1196,15 @@ public void closeActive() {
11951196
}
11961197

11971198
@Override
1198-
public AgentSpanContext notifyExtensionStart(Object event, String lambdaRequestId) {
1199-
return LambdaHandler.notifyStartInvocation(event, lambdaRequestId);
1199+
public AgentSpanContext notifyLambdaStart(Object event, String lambdaRequestId) {
1200+
// Get context from AppSec
1201+
AgentSpanContext appSecContext = LambdaAppSecHandler.processRequestStart(event);
1202+
1203+
// Get context from extension
1204+
AgentSpanContext extensionContext = LambdaHandler.notifyStartInvocation(event, lambdaRequestId);
1205+
1206+
// Merge contexts
1207+
return LambdaAppSecHandler.mergeContexts(extensionContext, appSecContext);
12001208
}
12011209

12021210
@Override
@@ -1205,6 +1213,11 @@ public void notifyExtensionEnd(
12051213
LambdaHandler.notifyEndInvocation(span, result, isError, lambdaRequestId);
12061214
}
12071215

1216+
@Override
1217+
public void notifyAppSecEnd(AgentSpan span) {
1218+
LambdaAppSecHandler.processRequestEnd(span);
1219+
}
1220+
12081221
@Override
12091222
public AgentDataStreamsMonitoring getDataStreamsMonitoring() {
12101223
return dataStreamsMonitoring;

0 commit comments

Comments
 (0)