Skip to content

Commit 4d929d6

Browse files
claponcetclaude
claponcet
and
claude
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> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 6e28457 commit 4d929d6

8 files changed

Lines changed: 2596 additions & 8 deletions

File tree

‎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)