Skip to content

Commit 081ea99

Browse files
claponcetclaponcet
committed
add better support for query parameters
1 parent 52eea07 commit 081ea99

File tree

4 files changed

+115
-10
lines changed

4 files changed

+115
-10
lines changed

‎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: 2 additions & 0 deletions
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;
@@ -96,6 +97,7 @@ static AgentScope enter(
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);

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import datadog.trace.agent.test.naming.VersionedNamingTestBase
2+
import datadog.trace.api.DDSpanTypes
23
import java.nio.charset.StandardCharsets
34
import com.amazonaws.services.lambda.runtime.Context
45

@@ -30,6 +31,7 @@ abstract class LambdaHandlerInstrumentationTest extends VersionedNamingTestBase
3031
trace(1) {
3132
span {
3233
operationName operation()
34+
spanType DDSpanTypes.SERVERLESS
3335
errored false
3436
}
3537
}
@@ -51,6 +53,7 @@ abstract class LambdaHandlerInstrumentationTest extends VersionedNamingTestBase
5153
trace(1) {
5254
span {
5355
operationName operation()
56+
spanType DDSpanTypes.SERVERLESS
5457
errored true
5558
tags {
5659
tag "request_id", requestId

‎dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java‎

Lines changed: 109 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import java.nio.charset.StandardCharsets;
2727
import java.util.Base64;
2828
import java.util.Collections;
29+
import java.util.List;
2930
import java.util.Map;
3031
import java.util.function.BiFunction;
3132
import java.util.function.Function;
@@ -157,7 +158,9 @@ private static AgentSpanContext processAppSecRequestData(LambdaEventData eventDa
157158
datadog.trace.api.function.TriFunction<RequestContext, String, URIDataAdapter, Flow<Void>> methodUriCallback =
158159
tracer.getCallbackProvider(RequestContextSlot.APPSEC).getCallback(EVENTS.requestMethodUriRaw());
159160
if (methodUriCallback != null) {
160-
LambdaURIDataAdapter uriAdapter = new LambdaURIDataAdapter(eventData.path);
161+
// Reconstruct full path with query string for AppSec analysis
162+
String fullPath = buildFullPath(eventData.path, eventData.queryParameters);
163+
LambdaURIDataAdapter uriAdapter = new LambdaURIDataAdapter(fullPath);
161164
methodUriCallback.apply(requestContext, eventData.method, uriAdapter);
162165
} else {
163166
log.warn("requestMethodUriRaw callback is null");
@@ -235,7 +238,7 @@ private static LambdaEventData extractEventData(ByteArrayInputStream inputStream
235238
log.warn("Event size {} exceeds limit {} or is invalid, skipping AppSec processing",
236239
availableBytes, MAX_EVENT_SIZE);
237240
return new LambdaEventData(Collections.emptyMap(), null, null, null, null,
238-
LambdaTriggerType.UNKNOWN, Collections.emptyMap(), null);
241+
LambdaTriggerType.UNKNOWN, Collections.emptyMap(), Collections.emptyMap(), null);
239242
}
240243

241244
StringBuilder jsonBuilder = new StringBuilder(availableBytes);
@@ -259,7 +262,7 @@ private static LambdaEventData extractEventDataFromJson(String json) {
259262
log.debug("Event JSON parsed successfully");
260263

261264
if (event == null) {
262-
return new LambdaEventData(Collections.emptyMap(), null, null, null, null, LambdaTriggerType.UNKNOWN, Collections.emptyMap(), null);
265+
return new LambdaEventData(Collections.emptyMap(), null, null, null, null, LambdaTriggerType.UNKNOWN, Collections.emptyMap(), Collections.emptyMap(), null);
263266
}
264267

265268
// Detect trigger type
@@ -284,7 +287,7 @@ private static LambdaEventData extractEventDataFromJson(String json) {
284287
}
285288
} catch (Exception e) {
286289
log.error("Failed to parse event data from JSON", e);
287-
return new LambdaEventData(Collections.emptyMap(), null, null, null, null, LambdaTriggerType.UNKNOWN, Collections.emptyMap(), null);
290+
return new LambdaEventData(Collections.emptyMap(), null, null, null, null, LambdaTriggerType.UNKNOWN, Collections.emptyMap(), Collections.emptyMap(), null);
288291
}
289292
}
290293

@@ -339,6 +342,7 @@ static LambdaTriggerType detectTriggerType(Map<String, Object> event) {
339342
private static LambdaEventData extractApiGatewayV1Data(Map<String, Object> event) {
340343
Map<String, String> headers = extractHeaders(event.get("headers"));
341344
Map<String, String> pathParameters = extractPathParameters(event.get("pathParameters"));
345+
Map<String, List<String>> queryParameters = extractQueryParameters(event.get("queryStringParameters"));
342346
Object body = extractBody(event);
343347

344348
Map<?, ?> requestContext = (Map<?, ?>) event.get("requestContext");
@@ -352,7 +356,7 @@ private static LambdaEventData extractApiGatewayV1Data(Map<String, Object> event
352356
sourceIp = (String) identity.get("sourceIp");
353357
}
354358

355-
return new LambdaEventData(headers, method, path, sourceIp, null, LambdaTriggerType.API_GATEWAY_V1_REST, pathParameters, body);
359+
return new LambdaEventData(headers, method, path, sourceIp, null, LambdaTriggerType.API_GATEWAY_V1_REST, pathParameters, queryParameters, body);
356360
}
357361

358362
/**
@@ -361,6 +365,7 @@ private static LambdaEventData extractApiGatewayV1Data(Map<String, Object> event
361365
private static LambdaEventData extractApiGatewayV2HttpData(Map<String, Object> event, LambdaTriggerType triggerType) {
362366
Map<String, String> headers = extractHeadersWithCookies(event);
363367
Map<String, String> pathParameters = extractPathParameters(event.get("pathParameters"));
368+
Map<String, List<String>> queryParameters = extractQueryParameters(event.get("queryStringParameters"));
364369
Object body = extractBody(event);
365370

366371
Map<?, ?> requestContext = (Map<?, ?>) event.get("requestContext");
@@ -377,7 +382,7 @@ private static LambdaEventData extractApiGatewayV2HttpData(Map<String, Object> e
377382
sourcePort = ((Number) portObj).intValue();
378383
}
379384

380-
return new LambdaEventData(headers, method, path, sourceIp, sourcePort, triggerType, pathParameters, body);
385+
return new LambdaEventData(headers, method, path, sourceIp, sourcePort, triggerType, pathParameters, queryParameters, body);
381386
}
382387

383388
/**
@@ -386,6 +391,7 @@ private static LambdaEventData extractApiGatewayV2HttpData(Map<String, Object> e
386391
private static LambdaEventData extractApiGatewayV2WebSocketData(Map<String, Object> event) {
387392
Map<String, String> headers = extractHeadersWithCookies(event);
388393
Map<String, String> pathParameters = extractPathParameters(event.get("pathParameters"));
394+
Map<String, List<String>> queryParameters = extractQueryParameters(event.get("queryStringParameters"));
389395
Object body = extractBody(event);
390396

391397
Map<?, ?> requestContext = (Map<?, ?>) event.get("requestContext");
@@ -401,7 +407,7 @@ private static LambdaEventData extractApiGatewayV2WebSocketData(Map<String, Obje
401407
sourceIp = (String) identity.get("sourceIp");
402408
}
403409

404-
return new LambdaEventData(headers, method, path, sourceIp, null, LambdaTriggerType.API_GATEWAY_V2_WEBSOCKET, pathParameters, body);
410+
return new LambdaEventData(headers, method, path, sourceIp, null, LambdaTriggerType.API_GATEWAY_V2_WEBSOCKET, pathParameters, queryParameters, body);
405411
}
406412

407413
/**
@@ -437,13 +443,22 @@ private static LambdaEventData extractAlbData(Map<String, Object> event, LambdaT
437443
}
438444

439445
Map<String, String> pathParameters = extractPathParameters(event.get("pathParameters"));
446+
447+
// ALB can have both queryStringParameters and multiValueQueryStringParameters
448+
Map<String, List<String>> queryParameters;
449+
if (triggerType == LambdaTriggerType.ALB_MULTI_VALUE) {
450+
queryParameters = extractMultiValueQueryParameters(event.get("multiValueQueryStringParameters"));
451+
} else {
452+
queryParameters = extractQueryParameters(event.get("queryStringParameters"));
453+
}
454+
440455
Object body = extractBody(event);
441456

442457
String method = (String) event.get("httpMethod");
443458
String path = (String) event.get("path");
444459
String sourceIp = headers.get("x-forwarded-for");
445460

446-
return new LambdaEventData(headers, method, path, sourceIp, null, triggerType, pathParameters, body);
461+
return new LambdaEventData(headers, method, path, sourceIp, null, triggerType, pathParameters, queryParameters, body);
447462
}
448463

449464
/**
@@ -452,6 +467,7 @@ private static LambdaEventData extractAlbData(Map<String, Object> event, LambdaT
452467
private static LambdaEventData extractGenericData(Map<String, Object> event) {
453468
Map<String, String> headers = extractHeadersWithCookies(event);
454469
Map<String, String> pathParameters = extractPathParameters(event.get("pathParameters"));
470+
Map<String, List<String>> queryParameters = extractQueryParameters(event.get("queryStringParameters"));
455471
Object body = extractBody(event);
456472

457473
String method = null;
@@ -497,7 +513,7 @@ private static LambdaEventData extractGenericData(Map<String, Object> event) {
497513
}
498514
}
499515

500-
return new LambdaEventData(headers, method, path, sourceIp, null, LambdaTriggerType.UNKNOWN, pathParameters, body);
516+
return new LambdaEventData(headers, method, path, sourceIp, null, LambdaTriggerType.UNKNOWN, pathParameters, queryParameters, body);
501517
}
502518

503519
/**
@@ -540,6 +556,87 @@ private static Map<String, String> extractPathParameters(Object pathParamsObj) {
540556
return pathParams;
541557
}
542558

559+
/**
560+
* Helper method to extract query parameters from event.
561+
* Converts Map<String, String> to Map<String, List<String>> format expected by AppSec.
562+
*/
563+
private static Map<String, List<String>> extractQueryParameters(Object queryParamsObj) {
564+
Map<String, List<String>> result = new java.util.HashMap<>();
565+
if (queryParamsObj instanceof Map) {
566+
Map<?, ?> rawMap = (Map<?, ?>) queryParamsObj;
567+
for (Map.Entry<?, ?> entry : rawMap.entrySet()) {
568+
if (entry.getKey() != null && entry.getValue() != null) {
569+
String key = String.valueOf(entry.getKey());
570+
String value = String.valueOf(entry.getValue());
571+
result.put(key, java.util.Collections.singletonList(value));
572+
}
573+
}
574+
}
575+
log.debug("Extracted {} query parameters", result.size());
576+
return result;
577+
}
578+
579+
/**
580+
* Helper method to extract multi-value query parameters (used by ALB).
581+
* Handles Map<String, List<String>> format directly.
582+
*/
583+
private static Map<String, List<String>> extractMultiValueQueryParameters(Object queryParamsObj) {
584+
Map<String, List<String>> result = new java.util.HashMap<>();
585+
if (queryParamsObj instanceof Map) {
586+
Map<?, ?> rawMap = (Map<?, ?>) queryParamsObj;
587+
for (Map.Entry<?, ?> entry : rawMap.entrySet()) {
588+
if (entry.getKey() != null && entry.getValue() != null) {
589+
String key = String.valueOf(entry.getKey());
590+
if (entry.getValue() instanceof java.util.List) {
591+
java.util.List<?> values = (java.util.List<?>) entry.getValue();
592+
java.util.List<String> stringValues = new java.util.ArrayList<>();
593+
for (Object value : values) {
594+
if (value != null) {
595+
stringValues.add(String.valueOf(value));
596+
}
597+
}
598+
result.put(key, stringValues);
599+
} else {
600+
result.put(key, java.util.Collections.singletonList(String.valueOf(entry.getValue())));
601+
}
602+
}
603+
}
604+
}
605+
log.debug("Extracted {} multi-value query parameters", result.size());
606+
return result;
607+
}
608+
609+
/**
610+
* Helper method to build full path including query string.
611+
* Lambda events provide path and query parameters separately, so we need to reconstruct
612+
* the full URI for AppSec to parse.
613+
*/
614+
private static String buildFullPath(String path, Map<String, List<String>> queryParameters) {
615+
if (queryParameters == null || queryParameters.isEmpty()) {
616+
return path;
617+
}
618+
619+
StringBuilder fullPath = new StringBuilder(path);
620+
fullPath.append('?');
621+
622+
boolean first = true;
623+
for (Map.Entry<String, List<String>> entry : queryParameters.entrySet()) {
624+
String key = entry.getKey();
625+
for (String value : entry.getValue()) {
626+
if (!first) {
627+
fullPath.append('&');
628+
}
629+
first = false;
630+
fullPath.append(key);
631+
if (value != null) {
632+
fullPath.append('=').append(value);
633+
}
634+
}
635+
}
636+
637+
return fullPath.toString();
638+
}
639+
543640
/**
544641
* Helper method to extract and merge headers with cookies array from event.
545642
* API Gateway v2 provides a separate 'cookies' array that should be merged with headers.
@@ -689,16 +786,18 @@ static class LambdaEventData {
689786
final Integer sourcePort;
690787
final LambdaTriggerType triggerType;
691788
final Map<String, String> pathParameters;
789+
final Map<String, List<String>> queryParameters;
692790
final Object body;
693791

694-
LambdaEventData(Map<String, String> headers, String method, String path, String sourceIp, Integer sourcePort, LambdaTriggerType triggerType, Map<String, String> pathParameters, Object body) {
792+
LambdaEventData(Map<String, String> headers, String method, String path, String sourceIp, Integer sourcePort, LambdaTriggerType triggerType, Map<String, String> pathParameters, Map<String, List<String>> queryParameters, Object body) {
695793
this.headers = headers;
696794
this.method = method;
697795
this.path = path;
698796
this.sourceIp = sourceIp;
699797
this.sourcePort = sourcePort;
700798
this.triggerType = triggerType;
701799
this.pathParameters = pathParameters;
800+
this.queryParameters = queryParameters;
702801
this.body = body;
703802
}
704803
}

‎dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaAppSecHandlerTest.groovy‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package datadog.trace.lambda
22

33
import datadog.trace.api.Config
44
import datadog.trace.api.function.TriConsumer
5+
import datadog.trace.api.function.TriFunction
56
import datadog.trace.api.gateway.CallbackProvider
67
import datadog.trace.api.gateway.Flow
78
import datadog.trace.api.gateway.RequestContext

0 commit comments

Comments
 (0)