diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java index c67a316b659..e8302a7623c 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java @@ -6,6 +6,7 @@ import com.squareup.moshi.Moshi; import datadog.logging.RatelimitedLogger; import datadog.trace.api.Config; +import datadog.trace.api.DDTags; import datadog.trace.api.ProductTraceSource; import datadog.trace.api.appsec.AppSecContext; import datadog.trace.api.function.TriConsumer; @@ -25,6 +26,7 @@ import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter; import datadog.trace.bootstrap.instrumentation.api.URIDataAdapterBase; +import datadog.trace.bootstrap.instrumentation.api.URIUtils; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; @@ -34,6 +36,7 @@ import java.util.Base64; import java.util.Collections; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Locale; import java.util.Map; @@ -60,10 +63,37 @@ public class LambdaAppSecHandler { private static final int MAX_EVENT_SIZE = Config.get().getAppSecBodyParsingSizeLimit(); + // Serverless span metric marking an event AppSec cannot analyze as HTTP. Mutually exclusive with + // _dd.appsec.enabled. Package-private so tests can assert against the same name. + static final String TAG_UNSUPPORTED_EVENT_TYPE = "_dd.appsec.unsupported_event_type"; + + // HTTP header names (lowercased, matching the normalised header maps). + private static final String HEADER_CONTENT_TYPE = "content-type"; + private static final String HEADER_COOKIE = "cookie"; + private static final String HEADER_USER_AGENT = "user-agent"; + private static final String HEADER_HOST = "host"; + private static final String HEADER_X_FORWARDED_PROTO = "x-forwarded-proto"; + private static final String HEADER_X_FORWARDED_PORT = "x-forwarded-port"; + private static final String HEADER_X_FORWARDED_FOR = "x-forwarded-for"; + + private static final String SCHEME_HTTPS = "https"; + + // Content types used for body-parsing dispatch. JSON/JAVASCRIPT are matched as substrings to + // tolerate charset parameters (e.g. "application/json; charset=utf-8"). + private static final String CONTENT_TYPE_FORM_URLENCODED = "application/x-www-form-urlencoded"; + private static final String CONTENT_TYPE_APPLICATION_JSON = "application/json"; + private static final String CONTENT_TYPE_TEXT_PLAIN = "text/plain"; + private static final String CONTENT_TYPE_JSON_TOKEN = "json"; + private static final String CONTENT_TYPE_JAVASCRIPT_TOKEN = "javascript"; + // Carries the detected trigger type from processRequestStart to processResponseData within the // same Lambda invocation. Cleared in processRequestEnd. private static final ThreadLocal CURRENT_TRIGGER_TYPE = new ThreadLocal<>(); + // Carries the extracted event data from processRequestStart to processRequestEnd, where it is + // used to set HTTP span tags once the span exists. Cleared in processRequestEnd. + private static final ThreadLocal CURRENT_EVENT_DATA = new ThreadLocal<>(); + /** * Process AppSec request data at the start of a Lambda invocation. Extract event data and invokes * all relevant AppSec gateway callbacks. @@ -79,11 +109,16 @@ public static AgentSpanContext processRequestStart(Object event) { } CURRENT_TRIGGER_TYPE.set(LambdaTriggerType.UNKNOWN); + CURRENT_EVENT_DATA.remove(); if (!(event instanceof ByteArrayInputStream)) { log.debug( "Event is not a ByteArrayInputStream, type: {}", event != null ? event.getClass().getName() : "null"); + // A non-stream event carries no raw HTTP payload AppSec can analyze. + // Record EMPTY so processRequestEnd marks the span with + // _dd.appsec.unsupported_event_type. + CURRENT_EVENT_DATA.set(LambdaEventData.EMPTY); return null; } @@ -93,6 +128,10 @@ public static AgentSpanContext processRequestStart(Object event) { return null; } CURRENT_TRIGGER_TYPE.set(eventData.triggerType); + CURRENT_EVENT_DATA.set(eventData); + if (!isSupportedHttpEvent(eventData)) { + return null; + } return processAppSecRequestData(eventData); } catch (Exception e) { log.debug("Failed to process AppSec request data", e); @@ -106,7 +145,9 @@ public static AgentSpanContext processRequestStart(Object event) { * @param span the current span */ public static void processRequestEnd(AgentSpan span) { + LambdaEventData eventData = CURRENT_EVENT_DATA.get(); CURRENT_TRIGGER_TYPE.remove(); + CURRENT_EVENT_DATA.remove(); if (!ActiveSubsystems.APPSEC_ACTIVE || span == null) { return; @@ -135,6 +176,93 @@ public static void processRequestEnd(AgentSpan span) { traceSeg.setTagTop(Tags.ASM_KEEP, true); traceSeg.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM); } + + if (eventData != null) { + if (!isSupportedHttpEvent(eventData)) { + span.setMetric(TAG_UNSUPPORTED_EVENT_TYPE, 1); + } else { + applyHttpSpanTags(span, eventData); + } + } + } + } + + /** + * Returns true if the event carries enough HTTP-like data (a known trigger type, or a best-effort + * method/path extracted via generic extraction) for the static HTTP security rules to run + * against. + */ + private static boolean isSupportedHttpEvent(LambdaEventData eventData) { + return eventData.triggerType != LambdaTriggerType.UNKNOWN + || (eventData.method != null && eventData.path != null); + } + + /** Sets HTTP span tags (http.url, http.route, http.useragent) derived from the Lambda event. */ + private static void applyHttpSpanTags(AgentSpan span, LambdaEventData eventData) { + if (eventData.method != null && !eventData.method.isEmpty()) { + span.setTag(Tags.HTTP_METHOD, eventData.method); + } + + String userAgent = eventData.headers != null ? eventData.headers.get(HEADER_USER_AGENT) : null; + if (userAgent != null && !userAgent.isEmpty()) { + span.setTag(Tags.HTTP_USER_AGENT, userAgent); + } + + if (eventData.path != null && !eventData.path.isEmpty()) { + String host = eventData.headers != null ? eventData.headers.get(HEADER_HOST) : null; + String scheme = firstForwardedValue(eventData.headers, HEADER_X_FORWARDED_PROTO); + if (scheme == null || scheme.isEmpty()) { + scheme = SCHEME_HTTPS; + } + // buildURL omits the port when it is <= 0 or the scheme's default (80/443), so a missing or + // default X-Forwarded-Port yields a clean URL while a non-standard port is preserved. + int port = forwardedPort(eventData.headers); + String url = + (host != null && !host.isEmpty()) + ? URIUtils.buildURL(scheme, host, port, eventData.path) + : eventData.path; + span.setTag(Tags.HTTP_URL, url); + + if (Config.get().isHttpServerTagQueryString()) { + String query = buildQueryString(eventData.queryParameters); + if (query != null && !query.isEmpty()) { + span.setTag(DDTags.HTTP_QUERY, query); + } + } + } + + if (eventData.route != null && !eventData.route.isEmpty()) { + span.setTag(Tags.HTTP_ROUTE, eventData.route); + } + } + + /** + * Returns the first value of a potentially comma-separated forwarded header, trimmed (e.g. {@code + * X-Forwarded-Proto: "https, http"} behind multiple proxies yields {@code "https"}). Returns null + * if the header is absent. + */ + private static String firstForwardedValue(Map headers, String name) { + String value = headers != null ? headers.get(name) : null; + if (value == null) { + return null; + } + int commaIdx = value.indexOf(','); + return (commaIdx >= 0 ? value.substring(0, commaIdx) : value).trim(); + } + + /** + * Parses the first {@code X-Forwarded-Port} value as an int, or -1 if the header is absent, empty, + * or not a number. + */ + private static int forwardedPort(Map headers) { + String forwardedPort = firstForwardedValue(headers, HEADER_X_FORWARDED_PORT); + if (forwardedPort == null || forwardedPort.isEmpty()) { + return -1; + } + try { + return Integer.parseInt(forwardedPort); + } catch (NumberFormatException ignored) { + return -1; } } @@ -171,27 +299,27 @@ public static void processResponseData(AgentSpan span, Object result) { return; } - if (responseData == null || responseData.statusCode == 0) { - // No statusCode means this is not an API-GW formatted response, or JSON parsing failed. - if (responseData == null || (responseData.headers.isEmpty() && responseData.body == null)) { - // Parse failed or response has no API-GW structure (plain JSON body). - // Treat the full response as the body - Object fallbackBody; - String fallbackContentType; - try { - fallbackBody = OBJECT_ADAPTER.fromJson(json); - fallbackContentType = "application/json"; - } catch (Exception e) { - fallbackBody = json; - fallbackContentType = "text/plain"; - } - Map fallbackHeaders = - Collections.singletonMap("content-type", fallbackContentType); - responseData = new LambdaResponseData(0, fallbackHeaders, fallbackBody); + // No statusCode means this is not an API-GW formatted response, or JSON parsing failed. + // If there is also no other API-GW structure (headers/body), treat the full response as + // the body. Otherwise (responseData has explicit headers/body fields) keep them and just + // skip responseStarted below (statusCode remains 0, so the responseStarted guard will not + // fire). + if (responseData == null + || (responseData.statusCode == 0 + && responseData.headers.isEmpty() + && responseData.body == null)) { + Object fallbackBody; + String fallbackContentType; + try { + fallbackBody = OBJECT_ADAPTER.fromJson(json); + fallbackContentType = CONTENT_TYPE_APPLICATION_JSON; + } catch (Exception e) { + fallbackBody = json; + fallbackContentType = CONTENT_TYPE_TEXT_PLAIN; } - // else: responseData has explicit headers/body fields — keep them, just skip - // responseStarted - // (statusCode remains 0, so the responseStarted guard below will not fire). + Map fallbackHeaders = + Collections.singletonMap(HEADER_CONTENT_TYPE, fallbackContentType); + responseData = new LambdaResponseData(0, fallbackHeaders, fallbackBody); } RequestContext requestContext = span.getRequestContext(); @@ -261,11 +389,7 @@ static LambdaResponseData extractResponseData(String json) { } // Extract headers — keys are lowercased to normalise casing across API GW / ALB variants - Map headers = new HashMap<>(); - Map rawHeaders = extractStringMap(response.get("headers")); - for (Map.Entry entry : rawHeaders.entrySet()) { - headers.put(entry.getKey().toLowerCase(Locale.ROOT), entry.getValue()); - } + Map headers = extractLowercasedStringMap(response.get("headers")); // Merge multiValueHeaders if present (API GW v1 / ALB), also lowercasing keys Object multiValueHeadersObj = response.get("multiValueHeaders"); @@ -300,20 +424,8 @@ static LambdaResponseData extractResponseData(String json) { } if (bodyString != null) { - String contentType = headers.get("content-type"); - - // If JSON content-type or unknown, attempt JSON parsing - // Normalise casing: media type tokens are case-insensitive per RFC 7231 - String contentTypeLower = - contentType == null ? null : contentType.toLowerCase(Locale.ROOT); - if (contentTypeLower == null - || contentTypeLower.contains("json") - || contentTypeLower.contains("javascript")) { - Object parsed = parseBodyAsJson(bodyString); - body = parsed != null ? parsed : bodyString; - } else { - body = bodyString; - } + String contentType = headers.get(HEADER_CONTENT_TYPE); + body = parseBodyByContentType(bodyString, contentType); } } @@ -547,20 +659,14 @@ static LambdaTriggerType detectTriggerType(Map event) { return LambdaTriggerType.API_GATEWAY_V2_WEBSOCKET; } - // Check for API Gateway v2 format + // Check for API Gateway v2 / Lambda Function URL format Object httpObj = requestContext.get("http"); if (httpObj instanceof Map) { Object domainNameObj = requestContext.get("domainName"); - if (domainNameObj instanceof String) { - String domainName = (String) domainNameObj; - if (domainName.contains("lambda-url")) { - return LambdaTriggerType.LAMBDA_URL; - } else { - return LambdaTriggerType.API_GATEWAY_V2_HTTP; - } - } else { + if (domainNameObj instanceof String && ((String) domainNameObj).contains("lambda-url")) { return LambdaTriggerType.LAMBDA_URL; } + return LambdaTriggerType.API_GATEWAY_V2_HTTP; } // Check for API Gateway v1 REST API @@ -577,7 +683,7 @@ private static LambdaEventData extractApiGatewayV1Data(Map event Map pathParameters = extractPathParameters(event.get("pathParameters")); Map> queryParameters = extractQueryParameters(event.get("queryStringParameters")); - Object body = extractBody(event); + Object body = extractBody(event, headers); Map requestContext = (Map) event.get("requestContext"); String method = (String) requestContext.get("httpMethod"); @@ -590,6 +696,8 @@ private static LambdaEventData extractApiGatewayV1Data(Map event sourceIp = (String) identity.get("sourceIp"); } + String route = (String) event.get("resource"); + return new LambdaEventData( headers, method, @@ -599,7 +707,8 @@ private static LambdaEventData extractApiGatewayV1Data(Map event LambdaTriggerType.API_GATEWAY_V1_REST, pathParameters, queryParameters, - body); + body, + route); } /** Extracts data from API Gateway v2 (HTTP API) or Lambda URL event */ @@ -609,7 +718,7 @@ private static LambdaEventData extractApiGatewayV2HttpData( Map pathParameters = extractPathParameters(event.get("pathParameters")); Map> queryParameters = extractQueryParameters(event.get("queryStringParameters")); - Object body = extractBody(event); + Object body = extractBody(event, headers); Map requestContext = (Map) event.get("requestContext"); Map http = (Map) requestContext.get("http"); @@ -625,6 +734,19 @@ private static LambdaEventData extractApiGatewayV2HttpData( sourcePort = ((Number) portObj).intValue(); } + // routeKey carries the method-prefixed route template for API Gateway v2 HTTP APIs (e.g. + // "GET /pets/{petId}"). Lambda Function URLs report "$default", which is not a real route. + // Strip the method prefix so http.route is a bare path template, consistent with the v1 REST + // extraction above and with HttpResourceDecorator's convention elsewhere in the tracer. + String route = null; + if (triggerType == LambdaTriggerType.API_GATEWAY_V2_HTTP) { + String routeKey = (String) requestContext.get("routeKey"); + if (routeKey != null && !"$default".equals(routeKey)) { + int spaceIdx = routeKey.indexOf(' '); + route = spaceIdx >= 0 ? routeKey.substring(spaceIdx + 1) : routeKey; + } + } + return new LambdaEventData( headers, method, @@ -634,7 +756,8 @@ private static LambdaEventData extractApiGatewayV2HttpData( triggerType, pathParameters, queryParameters, - body); + body, + route); } /** Extracts data from API Gateway v2 WebSocket event */ @@ -643,7 +766,7 @@ private static LambdaEventData extractApiGatewayV2WebSocketData(Map pathParameters = extractPathParameters(event.get("pathParameters")); Map> queryParameters = extractQueryParameters(event.get("queryStringParameters")); - Object body = extractBody(event); + Object body = extractBody(event, headers); Map requestContext = (Map) event.get("requestContext"); @@ -667,7 +790,8 @@ private static LambdaEventData extractApiGatewayV2WebSocketData(Map rawHeaders = (Map) multiValueHeadersObj; for (Map.Entry entry : rawHeaders.entrySet()) { if (entry.getKey() != null && entry.getValue() != null) { - String key = String.valueOf(entry.getKey()); + String key = String.valueOf(entry.getKey()).toLowerCase(Locale.ROOT); if (entry.getValue() instanceof List) { List values = (List) entry.getValue(); // Join multiple values with comma @@ -711,19 +835,24 @@ private static LambdaEventData extractAlbData( queryParameters = extractQueryParameters(event.get("queryStringParameters")); } - Object body = extractBody(event); + Object body = extractBody(event, headers); String method = (String) event.get("httpMethod"); String path = (String) event.get("path"); - String xff = headers.get("x-forwarded-for"); - String sourceIp = null; - if (xff != null) { - int commaIdx = xff.indexOf(','); - sourceIp = (commaIdx >= 0 ? xff.substring(0, commaIdx) : xff).trim(); - } + // x-forwarded-for may carry a comma-separated proxy chain; take the first (client) hop. + String sourceIp = firstForwardedValue(headers, HEADER_X_FORWARDED_FOR); return new LambdaEventData( - headers, method, path, sourceIp, null, triggerType, pathParameters, queryParameters, body); + headers, + method, + path, + sourceIp, + null, + triggerType, + pathParameters, + queryParameters, + body, + null); } /** Generic data extraction for unknown trigger types (fallback) */ @@ -732,7 +861,7 @@ private static LambdaEventData extractGenericData(Map event) { Map pathParameters = extractPathParameters(event.get("pathParameters")); Map> queryParameters = extractQueryParameters(event.get("queryStringParameters")); - Object body = extractBody(event); + Object body = extractBody(event, headers); String method = null; String path = null; @@ -786,7 +915,8 @@ private static LambdaEventData extractGenericData(Map event) { LambdaTriggerType.UNKNOWN, pathParameters, queryParameters, - body); + body, + null); } /** @@ -808,12 +938,20 @@ private static Map extractStringMap(Object mapObj) { return result; } - /** Helper method to extract headers from event */ + private static Map extractLowercasedStringMap(Object mapObj) { + Map rawMap = extractStringMap(mapObj); + Map result = new HashMap<>(); + for (Map.Entry entry : rawMap.entrySet()) { + result.put(entry.getKey().toLowerCase(Locale.ROOT), entry.getValue()); + } + return result; + } + private static Map extractHeaders(Object headersObj) { - Map headers = extractStringMap(headersObj); + Map headers = extractLowercasedStringMap(headersObj); log.debug("Extracted {} headers", headers.size()); - if (headers.containsKey("cookie")) { - log.debug("Cookie header found with value length: {}", headers.get("cookie").length()); + if (headers.containsKey(HEADER_COOKIE)) { + log.debug("Cookie header found with value length: {}", headers.get(HEADER_COOKIE).length()); } return headers; } @@ -880,39 +1018,48 @@ private static Map> extractMultiValueQueryParameters(Object * parameters separately, so we need to reconstruct the full URI for AppSec to parse. */ private static String buildFullPath(String path, Map> queryParameters) { - if (queryParameters == null || queryParameters.isEmpty()) { + String query = buildQueryString(queryParameters); + if (query == null || query.isEmpty()) { return path; } + return path + "?" + query; + } - StringBuilder fullPath = new StringBuilder(path); - fullPath.append('?'); + /** + * Builds a URL-encoded query string (without the leading {@code ?}) from the parsed query + * parameters, or null if there are none. Keys and values are percent-encoded so that special + * characters (e.g. {@code &} inside a value) are not mistaken for query-string delimiters. + */ + private static String buildQueryString(Map> queryParameters) { + if (queryParameters == null || queryParameters.isEmpty()) { + return null; + } + StringBuilder query = new StringBuilder(); boolean first = true; for (Map.Entry> entry : queryParameters.entrySet()) { String key = entry.getKey(); for (String value : entry.getValue()) { if (!first) { - fullPath.append('&'); + query.append('&'); } first = false; try { - // URL-encode key and value so that special characters (e.g. '&' inside a value) are not - // mistaken for query string delimiters when AppSec parses the raw query string. - fullPath.append(URLEncoder.encode(key, "UTF-8")); + query.append(URLEncoder.encode(key, "UTF-8")); if (value != null) { - fullPath.append('=').append(URLEncoder.encode(value, "UTF-8")); + query.append('=').append(URLEncoder.encode(value, "UTF-8")); } } catch (java.io.UnsupportedEncodingException e) { // UTF-8 is always available; fall back to unencoded - fullPath.append(key); + query.append(key); if (value != null) { - fullPath.append('=').append(value); + query.append('=').append(value); } } } } - return fullPath.toString(); + return query.toString(); } /** @@ -932,11 +1079,11 @@ private static Map extractHeadersWithCookies(Map cookiesList.stream().map(String::valueOf).collect(Collectors.joining("; ")); // Merge with existing cookie header if present - String existingCookie = headers.get("cookie"); + String existingCookie = headers.get(HEADER_COOKIE); if (existingCookie != null && !existingCookie.isEmpty()) { - headers.put("cookie", existingCookie + "; " + cookieValue); + headers.put(HEADER_COOKIE, existingCookie + "; " + cookieValue); } else { - headers.put("cookie", cookieValue); + headers.put(HEADER_COOKIE, cookieValue); } } } @@ -944,8 +1091,11 @@ private static Map extractHeadersWithCookies(Map return headers; } - /** Helper method to extract and parse body from event */ - private static Object extractBody(Map event) { + /** + * Helper method to extract and parse body from event. Dispatches on the request's Content-Type + * header (see {@link #parseBodyByContentType}). + */ + private static Object extractBody(Map event, Map headers) { Object bodyObj = event.get("body"); if (bodyObj == null) { return null; @@ -964,15 +1114,49 @@ private static Object extractBody(Map event) { } } - // Try to parse as JSON - Object parsedBody = parseBodyAsJson(bodyString); - if (parsedBody != null) { - log.debug("Body parsed as JSON successfully"); - return parsedBody; + String contentType = headers != null ? headers.get(HEADER_CONTENT_TYPE) : null; + return parseBodyByContentType(bodyString, contentType); + } + + /** + * Parses a raw body string according to its Content-Type, dispatching strictly on the declared + * type rather than guessing. + * + *
    + *
  • {@code application/x-www-form-urlencoded} → structured map. + *
  • A JSON content-type (contains {@code json} or {@code javascript}) → parsed as JSON, or + * dropped (null) if it fails to parse — a body that is malformed for its declared type is + * not analyzed. + *
  • A missing content-type → best-effort JSON, falling back to the raw string so the body + * stays scannable by string-based WAF rules. + *
  • Any other content-type (including {@code multipart/form-data}) → the raw string. + * Multipart bodies are not structurally parsed; the raw payload still stays scannable by + * string-based WAF rules. + *
+ */ + private static Object parseBodyByContentType(String bodyString, String contentType) { + String contentTypeLower = contentType == null ? null : contentType.toLowerCase(Locale.ROOT); + + if (contentTypeLower != null && contentTypeLower.startsWith(CONTENT_TYPE_FORM_URLENCODED)) { + return parseUrlEncodedBody(bodyString); } - // If not JSON, return the raw string - log.debug("Body is not JSON, returning raw string"); + if (contentTypeLower != null + && (contentTypeLower.contains(CONTENT_TYPE_JSON_TOKEN) + || contentTypeLower.contains(CONTENT_TYPE_JAVASCRIPT_TOKEN))) { + // Explicit JSON content-type: parse as JSON. A body that fails to parse is malformed for its + // declared type, so drop it (null) rather than forwarding a raw string + return parseBodyAsJson(bodyString); + } + + if (contentTypeLower == null) { + // No declared content-type: best-effort JSON parse, falling back to the raw string so the + // body stays scannable by string-based WAF rules. + Object parsed = parseBodyAsJson(bodyString); + return parsed != null ? parsed : bodyString; + } + + // Any other (non-JSON) content-type: keep the raw string; do not guess JSON. return bodyString; } @@ -989,6 +1173,38 @@ private static Object parseBodyAsJson(String body) { } } + /** + * Parses an {@code application/x-www-form-urlencoded} body into a map of decoded keys to their + * (possibly repeated) decoded values, e.g. {@code a=1&a=2&b=3} becomes {@code {a: [1, 2], b: + * [3]}}. + */ + private static Map> parseUrlEncodedBody(String body) { + Map> result = new LinkedHashMap<>(); + if (body == null || body.isEmpty()) { + return result; + } + + int start = 0; + int len = body.length(); + while (start <= len) { + int ampIdx = body.indexOf('&', start); + String pair = ampIdx >= 0 ? body.substring(start, ampIdx) : body.substring(start); + if (!pair.isEmpty()) { + int eqIdx = pair.indexOf('='); + String rawKey = eqIdx >= 0 ? pair.substring(0, eqIdx) : pair; + String rawValue = eqIdx >= 0 ? pair.substring(eqIdx + 1) : ""; + String key = URIUtils.decode(rawKey, true); + String value = URIUtils.decode(rawValue, true); + result.computeIfAbsent(key, k -> new ArrayList<>()).add(value); + } + if (ampIdx < 0) { + break; + } + start = ampIdx + 1; + } + return result; + } + /** Sets the current trigger type thread-local. Package-private for use in tests only. */ static void setCurrentTriggerType(LambdaTriggerType type) { if (type == null) { @@ -1078,6 +1294,7 @@ static class LambdaEventData { final Map pathParameters; final Map> queryParameters; final Object body; + final String route; static final LambdaEventData EMPTY = new LambdaEventData( @@ -1089,6 +1306,7 @@ static class LambdaEventData { LambdaTriggerType.UNKNOWN, Collections.emptyMap(), Collections.emptyMap(), + null, null); LambdaEventData( @@ -1100,7 +1318,8 @@ static class LambdaEventData { LambdaTriggerType triggerType, Map pathParameters, Map> queryParameters, - Object body) { + Object body, + String route) { this.headers = headers; this.method = method; this.path = path; @@ -1110,6 +1329,7 @@ static class LambdaEventData { this.pathParameters = pathParameters; this.queryParameters = queryParameters; this.body = body; + this.route = route; } } @@ -1148,18 +1368,11 @@ private static class LambdaURIDataAdapter extends URIDataAdapterBase { this.query = null; } - String forwardedProto = headers != null ? headers.get("x-forwarded-proto") : null; + String forwardedProto = firstForwardedValue(headers, HEADER_X_FORWARDED_PROTO); this.scheme = - (forwardedProto != null && !forwardedProto.isEmpty()) ? forwardedProto : "https"; + (forwardedProto != null && !forwardedProto.isEmpty()) ? forwardedProto : SCHEME_HTTPS; - String forwardedPort = headers != null ? headers.get("x-forwarded-port") : null; - int parsedPort = -1; - if (forwardedPort != null && !forwardedPort.isEmpty()) { - try { - parsedPort = Integer.parseInt(forwardedPort.trim()); - } catch (NumberFormatException ignored) { - } - } + int parsedPort = forwardedPort(headers); this.port = parsedPort > 0 ? parsedPort : 443; } diff --git a/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java b/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java index 75dfb9a6928..0288c09295e 100644 --- a/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java @@ -11,6 +11,7 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -19,6 +20,7 @@ import static org.mockito.Mockito.when; import datadog.trace.api.Config; +import datadog.trace.api.DDTags; import datadog.trace.api.ProductTraceSource; import datadog.trace.api.appsec.AppSecContext; import datadog.trace.api.function.TriConsumer; @@ -46,6 +48,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Base64; +import java.util.Collections; import java.util.HashMap; import java.util.LinkedHashMap; import java.util.List; @@ -55,11 +58,15 @@ import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; +import java.util.stream.Stream; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; class LambdaAppSecHandlerTest extends DDCoreJavaSpecification { @@ -146,135 +153,93 @@ void streamCanBeReadMultipleTimesAfterProcessing() throws IOException { // Trigger Type Detection Tests // ============================================================================ - @Test - void detectsApiGatewayV1RestTriggerType() { - Map event = - mapOf("requestContext", mapOf("httpMethod", "GET", "requestId", "abc123")); - - LambdaAppSecHandler.LambdaTriggerType triggerType = - LambdaAppSecHandler.detectTriggerType(event); - - assertEquals(LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST, triggerType); - } - - @Test - void detectsApiGatewayV2HttpTriggerType() { - Map event = - mapOf( - "requestContext", + @ParameterizedTest(name = "{0}") + @MethodSource("triggerTypeDetectionCases") + void detectsTriggerType( + String description, + Map event, + LambdaAppSecHandler.LambdaTriggerType expected) { + assertEquals(expected, LambdaAppSecHandler.detectTriggerType(event)); + } + + static Stream triggerTypeDetectionCases() { + return Stream.of( + Arguments.of( + "API Gateway v1 REST", + mapOf("requestContext", mapOf("httpMethod", "GET", "requestId", "abc123")), + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST), + Arguments.of( + "API Gateway v2 HTTP", mapOf( - "http", mapOf("method", "POST", "path", "/api"), "domainName", "api.example.com")); - - LambdaAppSecHandler.LambdaTriggerType triggerType = - LambdaAppSecHandler.detectTriggerType(event); - - assertEquals(LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V2_HTTP, triggerType); - } - - @Test - void detectsLambdaFunctionUrlTriggerType() { - Map event = - mapOf( - "requestContext", + "requestContext", + mapOf( + "http", + mapOf("method", "POST", "path", "/api"), + "domainName", + "api.example.com")), + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V2_HTTP), + Arguments.of( + "Lambda Function URL", mapOf( - "http", - mapOf("method", "GET", "path", "/"), - "domainName", - "xyz123.lambda-url.us-east-1.on.aws")); - - LambdaAppSecHandler.LambdaTriggerType triggerType = - LambdaAppSecHandler.detectTriggerType(event); - - assertEquals(LambdaAppSecHandler.LambdaTriggerType.LAMBDA_URL, triggerType); - } - - @Test - void detectsAlbTriggerTypeWithoutMultiValueHeaders() { - Map event = - mapOf( - "httpMethod", - "GET", - "path", - "/", - "requestContext", - mapOf("elb", mapOf("targetGroupArn", "arn:aws:..."))); - - LambdaAppSecHandler.LambdaTriggerType triggerType = - LambdaAppSecHandler.detectTriggerType(event); - - assertEquals(LambdaAppSecHandler.LambdaTriggerType.ALB, triggerType); - } - - @Test - void detectsAlbTriggerTypeWithMultiValueHeaders() { - Map event = - mapOf( - "httpMethod", - "GET", - "path", - "/", - "multiValueHeaders", - mapOf("accept", Arrays.asList("text/html", "application/json")), - "requestContext", - mapOf("elb", mapOf("targetGroupArn", "arn:aws:..."))); - - LambdaAppSecHandler.LambdaTriggerType triggerType = - LambdaAppSecHandler.detectTriggerType(event); - - assertEquals(LambdaAppSecHandler.LambdaTriggerType.ALB_MULTI_VALUE, triggerType); - } - - @Test - void detectsWebSocketTriggerTypeWithRouteKey() { - Map event = - mapOf("requestContext", mapOf("connectionId", "conn-123", "routeKey", "$connect")); - - LambdaAppSecHandler.LambdaTriggerType triggerType = - LambdaAppSecHandler.detectTriggerType(event); - - assertEquals(LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V2_WEBSOCKET, triggerType); - } - - @Test - void detectsWebSocketTriggerTypeWithEventType() { - Map event = - mapOf("requestContext", mapOf("connectionId", "conn-456", "eventType", "CONNECT")); - - LambdaAppSecHandler.LambdaTriggerType triggerType = - LambdaAppSecHandler.detectTriggerType(event); - - assertEquals(LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V2_WEBSOCKET, triggerType); - } - - @Test - void detectsUnknownTriggerTypeForUnrecognizedEvents() { - Map event = mapOf("someUnknownField", "value"); - - LambdaAppSecHandler.LambdaTriggerType triggerType = - LambdaAppSecHandler.detectTriggerType(event); - - assertEquals(LambdaAppSecHandler.LambdaTriggerType.UNKNOWN, triggerType); - } - - @Test - void detectsUnknownTriggerTypeForEmptyRequestContext() { - Map event = mapOf("requestContext", mapOf()); - - LambdaAppSecHandler.LambdaTriggerType triggerType = - LambdaAppSecHandler.detectTriggerType(event); - - assertEquals(LambdaAppSecHandler.LambdaTriggerType.UNKNOWN, triggerType); - } - - @Test - void detectsLambdaUrlWhenHttpPresentButNoDomainName() { - Map event = - mapOf("requestContext", mapOf("http", mapOf("method", "GET", "path", "/ambiguous"))); - - LambdaAppSecHandler.LambdaTriggerType triggerType = - LambdaAppSecHandler.detectTriggerType(event); - - assertEquals(LambdaAppSecHandler.LambdaTriggerType.LAMBDA_URL, triggerType); + "requestContext", + mapOf( + "http", + mapOf("method", "GET", "path", "/"), + "domainName", + "xyz123.lambda-url.us-east-1.on.aws")), + LambdaAppSecHandler.LambdaTriggerType.LAMBDA_URL), + Arguments.of( + "ALB without multi-value headers", + mapOf( + "httpMethod", + "GET", + "path", + "/", + "requestContext", + mapOf("elb", mapOf("targetGroupArn", "arn:aws:..."))), + LambdaAppSecHandler.LambdaTriggerType.ALB), + Arguments.of( + "ALB with multi-value headers", + mapOf( + "httpMethod", + "GET", + "path", + "/", + "multiValueHeaders", + mapOf("accept", Arrays.asList("text/html", "application/json")), + "requestContext", + mapOf("elb", mapOf("targetGroupArn", "arn:aws:..."))), + LambdaAppSecHandler.LambdaTriggerType.ALB_MULTI_VALUE), + Arguments.of( + "WebSocket via routeKey", + mapOf("requestContext", mapOf("connectionId", "conn-123", "routeKey", "$connect")), + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V2_WEBSOCKET), + Arguments.of( + "WebSocket via eventType", + mapOf("requestContext", mapOf("connectionId", "conn-456", "eventType", "CONNECT")), + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V2_WEBSOCKET), + Arguments.of( + "Unknown for unrecognized event", + mapOf("someUnknownField", "value"), + LambdaAppSecHandler.LambdaTriggerType.UNKNOWN), + Arguments.of( + "Unknown for empty requestContext", + mapOf("requestContext", mapOf()), + LambdaAppSecHandler.LambdaTriggerType.UNKNOWN), + // No domainName => no positive evidence of a Lambda Function URL (matching the Rust + // extension and Python's datadog-lambda layer), so this falls through to + // API_GATEWAY_V2_HTTP rather than defaulting to LAMBDA_URL. + Arguments.of( + "v2 HTTP when http present but no domainName", + mapOf("requestContext", mapOf("http", mapOf("method", "GET", "path", "/ambiguous"))), + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V2_HTTP), + // A non-string domainName is likewise not positive evidence of a Function URL. + Arguments.of( + "v2 HTTP when domainName is not a string", + mapOf( + "requestContext", + mapOf("http", mapOf("method", "GET", "path", "/ambiguous"), "domainName", 12345)), + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V2_HTTP)); } // ============================================================================ @@ -285,18 +250,11 @@ void detectsLambdaUrlWhenHttpPresentButNoDomainName() { @SuppressWarnings("unchecked") void extractsApiGatewayV1RestDataCorrectly() { String eventJson = - "{" - + "\"path\": \"/api/users/123\"," - + "\"httpMethod\": \"POST\"," - + "\"headers\": {\"Content-Type\": \"application/json\", \"Authorization\": \"Bearer token123\"}," - + "\"pathParameters\": {\"userId\": \"123\"}," - + "\"body\": \"{\\\"name\\\": \\\"John\\\"}\"," - + "\"requestContext\": {" - + " \"httpMethod\": \"POST\"," - + " \"requestId\": \"req-123\"," - + " \"identity\": {\"sourceIp\": \"192.168.1.100\"}" - + "}" - + "}"; + "{\"path\": \"/api/users/123\",\"httpMethod\": \"POST\",\"headers\": {\"Content-Type\":" + + " \"application/json\", \"Authorization\": \"Bearer token123\"},\"pathParameters\":" + + " {\"userId\": \"123\"},\"body\": \"{\\\"name\\\":" + + " \\\"John\\\"}\",\"requestContext\": { \"httpMethod\": \"POST\", \"requestId\":" + + " \"req-123\", \"identity\": {\"sourceIp\": \"192.168.1.100\"}}}"; ByteArrayInputStream event = createInputStream(eventJson); String[] capturedMethod = {null}; @@ -329,8 +287,8 @@ void extractsApiGatewayV1RestDataCorrectly() { assertInstanceOf(TagContext.class, result); assertEquals("POST", capturedMethod[0]); assertEquals("/api/users/123", capturedPath[0]); - assertEquals("application/json", capturedHeaders.get("Content-Type")); - assertEquals("Bearer token123", capturedHeaders.get("Authorization")); + assertEquals("application/json", capturedHeaders.get("content-type")); + assertEquals("Bearer token123", capturedHeaders.get("authorization")); assertEquals("192.168.1.100", capturedSourceIp[0]); assertEquals(0, capturedSourcePort[0]); assertNotNull(capturedPathParams[0]); @@ -342,17 +300,12 @@ void extractsApiGatewayV1RestDataCorrectly() { @Test void extractsApiGatewayV2HttpDataCorrectly() { String eventJson = - "{" - + "\"version\": \"2.0\"," - + "\"headers\": {\"content-type\": \"application/json\", \"x-custom-header\": \"custom-value\"}," - + "\"cookies\": [\"session=abc123\", \"user=john\"]," - + "\"pathParameters\": {\"id\": \"456\"}," - + "\"body\": \"test body\"," - + "\"requestContext\": {" - + " \"http\": {\"method\": \"PUT\", \"path\": \"/api/items/456\", \"sourceIp\": \"10.0.0.50\", \"sourcePort\": 54321}," - + " \"domainName\": \"api.example.com\"" - + "}" - + "}"; + "{\"version\": \"2.0\",\"headers\": {\"content-type\": \"application/json\"," + + " \"x-custom-header\": \"custom-value\"},\"cookies\": [\"session=abc123\"," + + " \"user=john\"],\"pathParameters\": {\"id\": \"456\"},\"body\": \"test" + + " body\",\"requestContext\": { \"http\": {\"method\": \"PUT\", \"path\":" + + " \"/api/items/456\", \"sourceIp\": \"10.0.0.50\", \"sourcePort\": 54321}, " + + " \"domainName\": \"api.example.com\"}}"; ByteArrayInputStream event = createInputStream(eventJson); String[] capturedMethod = {null}; @@ -394,14 +347,10 @@ void extractsApiGatewayV2HttpDataCorrectly() { @Test void extractsLambdaFunctionUrlDataCorrectly() { String eventJson = - "{" - + "\"version\": \"2.0\"," - + "\"headers\": {\"host\": \"xyz.lambda-url.us-east-1.on.aws\"}," - + "\"requestContext\": {" - + " \"http\": {\"method\": \"GET\", \"path\": \"/function/path\", \"sourceIp\": \"1.2.3.4\"}," - + " \"domainName\": \"xyz.lambda-url.us-east-1.on.aws\"" - + "}" - + "}"; + "{\"version\": \"2.0\",\"headers\": {\"host\":" + + " \"xyz.lambda-url.us-east-1.on.aws\"},\"requestContext\": { \"http\": {\"method\":" + + " \"GET\", \"path\": \"/function/path\", \"sourceIp\": \"1.2.3.4\"}, \"domainName\":" + + " \"xyz.lambda-url.us-east-1.on.aws\"}}"; ByteArrayInputStream event = createInputStream(eventJson); String[] capturedMethod = {null}; @@ -425,14 +374,11 @@ void extractsLambdaFunctionUrlDataCorrectly() { @Test void extractsAlbDataCorrectly() { String eventJson = - "{" - + "\"path\": \"/alb/test\"," - + "\"httpMethod\": \"DELETE\"," - + "\"headers\": {\"x-forwarded-for\": \"203.0.113.42\", \"user-agent\": \"curl/7.64.1\"}," - + "\"requestContext\": {" - + " \"elb\": {\"targetGroupArn\": \"arn:aws:elasticloadbalancing:us-east-1:123456789012:targetgroup/tg/50dc6c495c0c9188\"}" - + "}" - + "}"; + "{\"path\": \"/alb/test\",\"httpMethod\": \"DELETE\",\"headers\": {\"x-forwarded-for\":" + + " \"203.0.113.42\", \"user-agent\": \"curl/7.64.1\"},\"requestContext\": { \"elb\":" + + " {\"targetGroupArn\":" + + " \"arn:aws:elasticloadbalancing:us-east-1:123456789012:targetgroup/tg/50dc6c495c0c9188\"}" + + "}}"; ByteArrayInputStream event = createInputStream(eventJson); String[] capturedMethod = {null}; @@ -459,12 +405,9 @@ void extractsAlbDataCorrectly() { @Test void extractsAlbMultiValueHeadersCorrectly() { String eventJson = - "{" - + "\"path\": \"/test\"," - + "\"httpMethod\": \"GET\"," - + "\"multiValueHeaders\": {\"accept\": [\"text/html\", \"application/json\"], \"x-custom\": [\"value1\", \"value2\"]}," - + "\"requestContext\": {\"elb\": {\"targetGroupArn\": \"arn:aws:...\"}}" - + "}"; + "{\"path\": \"/test\",\"httpMethod\": \"GET\",\"multiValueHeaders\": {\"accept\":" + + " [\"text/html\", \"application/json\"], \"x-custom\": [\"value1\"," + + " \"value2\"]},\"requestContext\": {\"elb\": {\"targetGroupArn\": \"arn:aws:...\"}}}"; ByteArrayInputStream event = createInputStream(eventJson); Map capturedHeaders = new HashMap<>(); @@ -481,13 +424,10 @@ void extractsAlbMultiValueHeadersCorrectly() { @Test void albMultiValueQueryParamsHandlesListValues() { String eventJson = - "{" - + "\"path\": \"/test\"," - + "\"httpMethod\": \"GET\"," - + "\"multiValueHeaders\": {\"accept\": [\"text/html\"]}," - + "\"multiValueQueryStringParameters\": {\"foo\": [\"bar\", \"baz\"], \"x\": [null, \"val\"]}," - + "\"requestContext\": {\"elb\": {\"targetGroupArn\": \"arn:aws:...\"}}" - + "}"; + "{\"path\": \"/test\",\"httpMethod\": \"GET\",\"multiValueHeaders\": {\"accept\":" + + " [\"text/html\"]},\"multiValueQueryStringParameters\": {\"foo\": [\"bar\", \"baz\"]," + + " \"x\": [null, \"val\"]},\"requestContext\": {\"elb\": {\"targetGroupArn\":" + + " \"arn:aws:...\"}}}"; ByteArrayInputStream event = createInputStream(eventJson); String[] capturedQuery = {null}; @@ -510,12 +450,9 @@ void albMultiValueQueryParamsHandlesListValues() { @Test void albMultiValueHeadersHandlesNonListValue() { String eventJson = - "{" - + "\"path\": \"/test\"," - + "\"httpMethod\": \"GET\"," - + "\"multiValueHeaders\": {\"content-type\": \"text/plain\", \"accept\": [\"application/json\"]}," - + "\"requestContext\": {\"elb\": {\"targetGroupArn\": \"arn:aws:...\"}}" - + "}"; + "{\"path\": \"/test\",\"httpMethod\": \"GET\",\"multiValueHeaders\": {\"content-type\":" + + " \"text/plain\", \"accept\": [\"application/json\"]},\"requestContext\": {\"elb\":" + + " {\"targetGroupArn\": \"arn:aws:...\"}}}"; ByteArrayInputStream event = createInputStream(eventJson); Map capturedHeaders = new HashMap<>(); @@ -779,11 +716,9 @@ void buildFullPathEncodesSpecialCharactersInQueryParams() { // Keys/values containing '&', '=', and spaces must be percent-encoded so they are not // misinterpreted as query string delimiters when AppSec parses the raw query string. String eventJson = - "{" - + "\"path\": \"/search\"," - + "\"queryStringParameters\": {\"q\": \"hello world\", \"filter\": \"a&b\", \"eq\": \"x=y\"}," - + "\"requestContext\": {\"httpMethod\": \"GET\", \"requestId\": \"req-special\"}" - + "}"; + "{\"path\": \"/search\",\"queryStringParameters\": {\"q\": \"hello world\", \"filter\":" + + " \"a&b\", \"eq\": \"x=y\"},\"requestContext\": {\"httpMethod\": \"GET\"," + + " \"requestId\": \"req-special\"}}"; ByteArrayInputStream event = createInputStream(eventJson); String[] capturedQuery = {null}; @@ -868,10 +803,40 @@ void fallsBackToHttps443WhenXForwardedHeadersAreAbsent() { @Test void handlesInvalidXForwardedPortGracefully() { + String eventJson = + "{\"path\": \"/api/test\",\"headers\": {\"x-forwarded-proto\": \"https\"," + + " \"x-forwarded-port\": \"not-a-number\"},\"requestContext\": {\"httpMethod\":" + + " \"GET\", \"requestId\": \"req-123\"}}"; + ByteArrayInputStream event = createInputStream(eventJson); + + String[] capturedScheme = {null}; + int[] capturedPort = {-1}; + + setupMockCallbacks( + new Callbacks() + .onMethodUri( + (method, uri) -> { + capturedScheme[0] = uri.scheme(); + capturedPort[0] = uri.port(); + })); + + AgentSpanContext result = LambdaAppSecHandler.processRequestStart(event); + + assertNotNull(result); + assertEquals("https", capturedScheme[0]); + assertEquals(443, capturedPort[0]); + } + + @Test + void extractsFirstSchemeAndPortFromCommaSeparatedXForwardedHeaders() { + // Behind a chain of proxies these headers carry a comma-separated list; the first (client-most) + // hop must be used. Previously x-forwarded-port took the raw value, so a comma-separated value + // threw in Integer.parseInt and silently fell back to 443. String eventJson = "{" + "\"path\": \"/api/test\"," - + "\"headers\": {\"x-forwarded-proto\": \"https\", \"x-forwarded-port\": \"not-a-number\"}," + + "\"headers\": {\"x-forwarded-proto\": \"http, https\"," + + " \"x-forwarded-port\": \"8080, 443\"}," + "\"requestContext\": {\"httpMethod\": \"GET\", \"requestId\": \"req-123\"}" + "}"; ByteArrayInputStream event = createInputStream(eventJson); @@ -890,8 +855,8 @@ void handlesInvalidXForwardedPortGracefully() { AgentSpanContext result = LambdaAppSecHandler.processRequestStart(event); assertNotNull(result); - assertEquals("https", capturedScheme[0]); - assertEquals(443, capturedPort[0]); + assertEquals("http", capturedScheme[0]); + assertEquals(8080, capturedPort[0]); } @Test @@ -958,6 +923,114 @@ void handlesBodyWithSpecialCharacters() { assertEquals("Hello 世界 🌍", ((Map) capturedBody[0]).get("text")); } + // ============================================================================ + // Body Content-Type Parsing Tests (URL-encoded / multipart) + // ============================================================================ + + @Test + void parsesUrlEncodedBodyIntoMultiValueMap() { + String eventJson = + "{" + + "\"body\": \"a=1&a=2&b=hello+world&c\"," + + "\"headers\": {\"Content-Type\": \"application/x-www-form-urlencoded\"}," + + "\"requestContext\": {\"httpMethod\": \"POST\"}" + + "}"; + ByteArrayInputStream event = createInputStream(eventJson); + + Object[] capturedBody = {null}; + setupMockCallbacks(new Callbacks().onBody(body -> capturedBody[0] = body)); + + AgentSpanContext result = LambdaAppSecHandler.processRequestStart(event); + + assertNotNull(result); + assertInstanceOf(Map.class, capturedBody[0]); + @SuppressWarnings("unchecked") + Map> body = (Map>) capturedBody[0]; + assertEquals(Arrays.asList("1", "2"), body.get("a")); + assertEquals(Collections.singletonList("hello world"), body.get("b")); + assertEquals(Collections.singletonList(""), body.get("c")); + } + + @Test + void multipartBodyIsKeptAsRawString() { + // multipart/form-data is not structurally parsed; the raw payload is forwarded so string-based + // WAF rules can still scan it. (Structured multipart schema extraction is not a milestone-1 + // system-test gap for Java Lambda.) + String multipartBody = + "--BOUNDARY\r\n" + + "Content-Disposition: form-data; name=\"field1\"\r\n" + + "\r\n" + + "value1\r\n" + + "--BOUNDARY--\r\n"; + String eventJson = + "{" + + "\"body\": \"" + + escapeJson(multipartBody) + + "\"," + + "\"headers\": {\"Content-Type\": \"multipart/form-data; boundary=BOUNDARY\"}," + + "\"requestContext\": {\"httpMethod\": \"POST\"}" + + "}"; + ByteArrayInputStream event = createInputStream(eventJson); + + Object[] capturedBody = {null}; + setupMockCallbacks(new Callbacks().onBody(body -> capturedBody[0] = body)); + + AgentSpanContext result = LambdaAppSecHandler.processRequestStart(event); + + assertNotNull(result); + assertEquals(multipartBody, capturedBody[0]); + } + + @Test + void keepsTextPlainRequestBodyWithJsonContentAsRawString() { + // A text/plain request body that happens to be valid JSON must NOT be JSON-parsed; it is kept + // as the raw string so it stays scannable (matches the extension; more protective than + // datadog-lambda-python which drops it). + String eventJson = + "{" + + "\"body\": \"{\\\"user\\\": \\\"admin\\\"}\"," + + "\"headers\": {\"Content-Type\": \"text/plain\"}," + + "\"requestContext\": {\"httpMethod\": \"POST\"}" + + "}"; + ByteArrayInputStream event = createInputStream(eventJson); + + Object[] capturedBody = {null}; + setupMockCallbacks(new Callbacks().onBody(body -> capturedBody[0] = body)); + + AgentSpanContext result = LambdaAppSecHandler.processRequestStart(event); + + assertNotNull(result); + assertEquals("{\"user\": \"admin\"}", capturedBody[0]); + } + + @Test + void dropsMalformedRequestBodyUnderJsonContentType() { + // Explicit application/json content-type with an invalid JSON body: dropped, so + // requestBodyProcessed does not fire (matches the extension). + String eventJson = + "{" + + "\"body\": \"not json {\"," + + "\"headers\": {\"Content-Type\": \"application/json\"}," + + "\"requestContext\": {\"httpMethod\": \"POST\"}" + + "}"; + ByteArrayInputStream event = createInputStream(eventJson); + + String[] capturedBody = {"NOT_CALLED"}; + setupMockCallbacks(new Callbacks().onBody(body -> capturedBody[0] = String.valueOf(body))); + + AgentSpanContext result = LambdaAppSecHandler.processRequestStart(event); + + assertNotNull(result); + assertEquals("NOT_CALLED", capturedBody[0]); + } + + private static String escapeJson(String raw) { + return raw.replace("\\", "\\\\") + .replace("\"", "\\\"") + .replace("\r", "\\r") + .replace("\n", "\\n"); + } + // ============================================================================ // Generic Data Extraction Tests // ============================================================================ @@ -1001,11 +1074,8 @@ void extractsDataFromUnknownTriggerTypeUsingGenericExtraction() { @Test void extractsDataFromUnknownTriggerWithHttpInRequestContext() { String eventJson = - "{" - + "\"requestContext\": {" - + " \"http\": {\"method\": \"OPTIONS\", \"path\": \"/options/path\", \"sourceIp\": \"198.51.100.50\"}" - + "}" - + "}"; + "{\"requestContext\": { \"http\": {\"method\": \"OPTIONS\", \"path\": \"/options/path\"," + + " \"sourceIp\": \"198.51.100.50\"}}}"; ByteArrayInputStream event = createInputStream(eventJson); String[] capturedMethod = {null}; @@ -1201,6 +1271,261 @@ void processRequestEndSetsAsmKeepTagWhenAppSecContextIsManuallyKept() { verify(mockTraceSegment).setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM); } + @Test + void processRequestStartReturnsNullAndSkipsWafForNonHttpEvent() { + String eventJson = "{\"Records\": [{\"body\": \"hello\"}]}"; + ByteArrayInputStream event = createInputStream(eventJson); + setupMockCallbacks(new Callbacks()); + + AgentSpanContext result = LambdaAppSecHandler.processRequestStart(event); + + assertNull(result); + } + + @Test + void processRequestEndSetsUnsupportedEventTypeMetricForNonHttpEvent() { + String eventJson = "{\"Records\": [{\"body\": \"hello\"}]}"; + ByteArrayInputStream event = createInputStream(eventJson); + setupMockCallbacks(new Callbacks()); + LambdaAppSecHandler.processRequestStart(event); + + AgentSpan span = setupSpanForRequestEnd(); + LambdaAppSecHandler.processRequestEnd(span); + + verify(span).setMetric(LambdaAppSecHandler.TAG_UNSUPPORTED_EVENT_TYPE, 1); + verify(span, never()).setTag(eq(Tags.HTTP_URL), any(String.class)); + verify(span, never()).setTag(eq(Tags.HTTP_ROUTE), any(String.class)); + verify(span, never()).setTag(eq(Tags.HTTP_USER_AGENT), any(String.class)); + } + + @Test + void processRequestEndSetsUnsupportedEventTypeMetricForNonByteArrayInputStreamEvent() { + // A typed POJO handler event (not a ByteArrayInputStream) carries no raw HTTP payload AppSec + // can analyze, so processRequestStart returns null without starting a WAF context but the span + // must still be marked unsupported. + setupMockCallbacks(new Callbacks()); + assertNull(LambdaAppSecHandler.processRequestStart("not a stream")); + + AgentSpan span = setupSpanForRequestEnd(); + LambdaAppSecHandler.processRequestEnd(span); + + verify(span).setMetric(LambdaAppSecHandler.TAG_UNSUPPORTED_EVENT_TYPE, 1); + verify(span, never()).setTag(eq(Tags.HTTP_URL), any(String.class)); + verify(span, never()).setTag(eq(Tags.HTTP_ROUTE), any(String.class)); + verify(span, never()).setTag(eq(Tags.HTTP_USER_AGENT), any(String.class)); + } + + @Test + void processRequestEndSetsHttpSpanTagsForApiGatewayV1RestEvent() { + String eventJson = + "{" + + "\"resource\": \"/pets/{petId}\"," + + "\"path\": \"/pets/123\"," + + "\"httpMethod\": \"GET\"," + + "\"headers\": {\"Host\": \"api.example.com\", \"User-Agent\": \"curl/8.0\"}," + + "\"requestContext\": {" + + " \"httpMethod\": \"GET\"," + + " \"requestId\": \"req-1\"," + + " \"identity\": {\"sourceIp\": \"1.2.3.4\"}" + + "}" + + "}"; + ByteArrayInputStream event = createInputStream(eventJson); + setupMockCallbacks(new Callbacks()); + LambdaAppSecHandler.processRequestStart(event); + + AgentSpan span = setupSpanForRequestEnd(); + LambdaAppSecHandler.processRequestEnd(span); + + verify(span, never()).setMetric(eq(LambdaAppSecHandler.TAG_UNSUPPORTED_EVENT_TYPE), anyInt()); + verify(span).setTag(Tags.HTTP_USER_AGENT, "curl/8.0"); + verify(span).setTag(Tags.HTTP_URL, "https://api.example.com/pets/123"); + verify(span).setTag(Tags.HTTP_ROUTE, "/pets/{petId}"); + } + + @Test + void processRequestEndOmitsHttpRouteForAlbEvent() { + String eventJson = + "{\"httpMethod\": \"GET\",\"path\": \"/health\",\"headers\": {\"host\":" + + " \"my-alb.example.com\", \"user-agent\":" + + " \"ALB-HealthChecker/2.0\"},\"requestContext\": {\"elb\": {\"targetGroupArn\":" + + " \"arn:aws:elasticloadbalancing:x\"}}}"; + ByteArrayInputStream event = createInputStream(eventJson); + setupMockCallbacks(new Callbacks()); + LambdaAppSecHandler.processRequestStart(event); + + AgentSpan span = setupSpanForRequestEnd(); + LambdaAppSecHandler.processRequestEnd(span); + + verify(span).setTag(Tags.HTTP_URL, "https://my-alb.example.com/health"); + verify(span).setTag(Tags.HTTP_USER_AGENT, "ALB-HealthChecker/2.0"); + verify(span, never()).setTag(eq(Tags.HTTP_ROUTE), any(String.class)); + } + + @Test + void processRequestEndOmitsHttpRouteForApiGatewayV2DefaultRoute() { + String eventJson = + "{" + + "\"headers\": {\"host\": \"xyz.execute-api.us-east-1.amazonaws.com\"}," + + "\"requestContext\": {" + + " \"routeKey\": \"$default\"," + + " \"domainName\": \"xyz.execute-api.us-east-1.amazonaws.com\"," + + " \"http\": {\"method\": \"GET\", \"path\": \"/anything\"}" + + "}" + + "}"; + ByteArrayInputStream event = createInputStream(eventJson); + setupMockCallbacks(new Callbacks()); + LambdaAppSecHandler.processRequestStart(event); + + AgentSpan span = setupSpanForRequestEnd(); + LambdaAppSecHandler.processRequestEnd(span); + + verify(span, never()).setTag(eq(Tags.HTTP_ROUTE), any(String.class)); + } + + @Test + void processRequestEndSetsHttpRouteForApiGatewayV2NonDefaultRoute() { + String eventJson = + "{" + + "\"headers\": {\"host\": \"xyz.execute-api.us-east-1.amazonaws.com\"}," + + "\"requestContext\": {" + + " \"routeKey\": \"GET /pets/{petId}\"," + + " \"domainName\": \"xyz.execute-api.us-east-1.amazonaws.com\"," + + " \"http\": {\"method\": \"GET\", \"path\": \"/pets/123\"}" + + "}" + + "}"; + ByteArrayInputStream event = createInputStream(eventJson); + setupMockCallbacks(new Callbacks()); + LambdaAppSecHandler.processRequestStart(event); + + AgentSpan span = setupSpanForRequestEnd(); + LambdaAppSecHandler.processRequestEnd(span); + + verify(span).setTag(Tags.HTTP_ROUTE, "/pets/{petId}"); + } + + @Test + void processRequestEndSetsHttpMethodForHttpEvent() { + String eventJson = + "{" + + "\"resource\": \"/pets/{petId}\"," + + "\"path\": \"/pets/123\"," + + "\"httpMethod\": \"POST\"," + + "\"headers\": {\"Host\": \"api.example.com\"}," + + "\"requestContext\": {\"httpMethod\": \"POST\", \"requestId\": \"req-1\"}" + + "}"; + ByteArrayInputStream event = createInputStream(eventJson); + setupMockCallbacks(new Callbacks()); + LambdaAppSecHandler.processRequestStart(event); + + AgentSpan span = setupSpanForRequestEnd(); + LambdaAppSecHandler.processRequestEnd(span); + + verify(span).setTag(Tags.HTTP_METHOD, "POST"); + } + + @Test + void processRequestEndTagsQueryStringSeparatelyAndKeepsItOutOfUrl() { + // http.url must not carry the raw query string; the query is tagged as http.query.string so + // QueryObfuscator can redact secrets before re-appending it to http.url. + String eventJson = + "{" + + "\"resource\": \"/login\"," + + "\"path\": \"/login\"," + + "\"httpMethod\": \"GET\"," + + "\"headers\": {\"Host\": \"api.example.com\"}," + + "\"queryStringParameters\": {\"password\": \"hunter2\"}," + + "\"requestContext\": {\"httpMethod\": \"GET\", \"requestId\": \"req-1\"}" + + "}"; + ByteArrayInputStream event = createInputStream(eventJson); + setupMockCallbacks(new Callbacks()); + LambdaAppSecHandler.processRequestStart(event); + + AgentSpan span = setupSpanForRequestEnd(); + LambdaAppSecHandler.processRequestEnd(span); + + verify(span).setTag(Tags.HTTP_URL, "https://api.example.com/login"); + verify(span).setTag(DDTags.HTTP_QUERY, "password=hunter2"); + } + + @Test + void processRequestEndUsesFirstForwardedProtoValueForUrlScheme() { + // A double-proxy hop can produce a comma-separated X-Forwarded-Proto; only the first value is + // a valid scheme. + String eventJson = + "{\"resource\": \"/pets/{petId}\",\"path\": \"/pets/123\",\"httpMethod\":" + + " \"GET\",\"headers\": {\"Host\": \"api.example.com\", \"X-Forwarded-Proto\": \"http," + + " https\"},\"requestContext\": {\"httpMethod\": \"GET\", \"requestId\": \"req-1\"}}"; + ByteArrayInputStream event = createInputStream(eventJson); + setupMockCallbacks(new Callbacks()); + LambdaAppSecHandler.processRequestStart(event); + + AgentSpan span = setupSpanForRequestEnd(); + LambdaAppSecHandler.processRequestEnd(span); + + verify(span).setTag(Tags.HTTP_URL, "http://api.example.com/pets/123"); + } + + @Test + void processRequestEndIncludesNonDefaultForwardedPortInUrl() { + // A non-standard X-Forwarded-Port must be preserved in http.url (only the first hop is used). + String eventJson = + "{\"resource\": \"/pets/{petId}\",\"path\": \"/pets/123\",\"httpMethod\":" + + " \"GET\",\"headers\": {\"Host\": \"api.example.com\", \"X-Forwarded-Port\": \"8443," + + " 443\"},\"requestContext\": {\"httpMethod\": \"GET\", \"requestId\": \"req-1\"}}"; + ByteArrayInputStream event = createInputStream(eventJson); + setupMockCallbacks(new Callbacks()); + LambdaAppSecHandler.processRequestStart(event); + + AgentSpan span = setupSpanForRequestEnd(); + LambdaAppSecHandler.processRequestEnd(span); + + verify(span).setTag(Tags.HTTP_URL, "https://api.example.com:8443/pets/123"); + } + + @Test + void processRequestEndOmitsDefaultForwardedPortFromUrl() { + // The scheme's default port (443 for https) must not appear in http.url. + String eventJson = + "{\"resource\": \"/pets/{petId}\",\"path\": \"/pets/123\",\"httpMethod\":" + + " \"GET\",\"headers\": {\"Host\": \"api.example.com\", \"X-Forwarded-Port\":" + + " \"443\"},\"requestContext\": {\"httpMethod\": \"GET\", \"requestId\": \"req-1\"}}"; + ByteArrayInputStream event = createInputStream(eventJson); + setupMockCallbacks(new Callbacks()); + LambdaAppSecHandler.processRequestStart(event); + + AgentSpan span = setupSpanForRequestEnd(); + LambdaAppSecHandler.processRequestEnd(span); + + verify(span).setTag(Tags.HTTP_URL, "https://api.example.com/pets/123"); + } + + @SuppressWarnings("unchecked") + private AgentSpan setupSpanForRequestEnd() { + AppSecContext mockAppSecContext = mock(AppSecContext.class); + when(mockAppSecContext.isManuallyKept()).thenReturn(false); + TraceSegment mockTraceSegment = mock(TraceSegment.class); + RequestContext mockRequestContext = mock(RequestContext.class); + when(mockRequestContext.getData(RequestContextSlot.APPSEC)).thenReturn(mockAppSecContext); + when(mockRequestContext.getTraceSegment()).thenReturn(mockTraceSegment); + + AgentSpan span = mock(AgentSpan.class); + when(span.getRequestContext()).thenReturn(mockRequestContext); + + BiFunction> requestEndedCallback = + mock(BiFunction.class); + when(requestEndedCallback.apply(any(), any())).thenReturn(new Flow.ResultFlow<>(null)); + + CallbackProvider mockCallbackProvider = mock(CallbackProvider.class); + when(mockCallbackProvider.getCallback(EVENTS.requestEnded())).thenReturn(requestEndedCallback); + + AgentTracer.TracerAPI mockTracer = mock(AgentTracer.TracerAPI.class); + when(mockTracer.getCallbackProvider(RequestContextSlot.APPSEC)) + .thenReturn(mockCallbackProvider); + AgentTracer.forceRegister(mockTracer); + + return span; + } + // ============================================================================ // mergeContexts Tests // ============================================================================ @@ -1341,7 +1666,8 @@ void processRequestStartHandlesNullHeaderDoneCallbackGracefully() { @SuppressWarnings("unchecked") void processRequestStartHandlesNullPathParamsCallbackGracefully() { String eventJson = - "{\"path\": \"/test\", \"pathParameters\": {\"id\": \"42\"}, \"requestContext\": {\"httpMethod\": \"GET\"}}"; + "{\"path\": \"/test\", \"pathParameters\": {\"id\": \"42\"}, \"requestContext\":" + + " {\"httpMethod\": \"GET\"}}"; ByteArrayInputStream event = createInputStream(eventJson); Object mockAppSecContext = new Object(); @@ -1373,12 +1699,6 @@ void processRequestStartHandlesNullPathParamsCallbackGracefully() { assertInstanceOf(TagContext.class, result); } - @Test - void processRequestStartHandlesExceptionDuringJsonParsing() { - ByteArrayInputStream event = createInputStream("{this is not valid JSON at all"); - assertNull(LambdaAppSecHandler.processRequestStart(event)); - } - @Test void processRequestStartHandlesExceptionDuringStreamReading() { ByteArrayInputStream mockStream = @@ -1510,12 +1830,14 @@ void processResponseDataDoesNothingForNullResult() { @Test void processResponseDataDoesNothingWhenSpanHasNoRequestContext() { + LambdaAppSecHandler.setCurrentTriggerType( + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); AgentSpan span = mock(AgentSpan.class); when(span.getRequestContext()).thenReturn(null); ByteArrayOutputStream result = createOutputStream("{\"statusCode\": 200}"); setupMockResponseCallbacks(null, null, null, null); LambdaAppSecHandler.processResponseData(span, result); - // no exception expected + // reaches the requestContext-null guard and returns without firing callbacks } @Test @@ -1589,7 +1911,8 @@ void processResponseDataKeepsParsedHeadersAndBodyWhenStatusCodeIsZero() { LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); ByteArrayOutputStream result = createOutputStream( - "{\"statusCode\": 0, \"headers\": {\"content-type\": \"text/plain\"}, \"body\": \"hello\"}"); + "{\"statusCode\": 0, \"headers\": {\"content-type\": \"text/plain\"}, \"body\":" + + " \"hello\"}"); Integer[] capturedStatus = {null}; Map capturedHeaders = new HashMap<>(); boolean[] headerDoneCalled = {false}; @@ -1704,18 +2027,10 @@ void processResponseDataExtractsStatusCodeAsIntegerFromDouble() { assertEquals(404, capturedStatus[0]); } - @Test - void processResponseDataHandlesMissingStatusCode() { - ByteArrayOutputStream result = createOutputStream("{\"body\": \"ok\"}"); - Integer[] capturedStatus = {null}; - AgentSpan span = - setupMockResponseCallbacks(status -> capturedStatus[0] = status, null, null, null); - LambdaAppSecHandler.processResponseData(span, result); - assertNull(capturedStatus[0]); - } - @Test void processResponseDataHandlesNonNumericStatusCode() { + LambdaAppSecHandler.setCurrentTriggerType( + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); ByteArrayOutputStream result = createOutputStream("{\"statusCode\": \"bad\", \"body\": \"ok\"}"); Integer[] capturedStatus = {null}; @@ -1732,7 +2047,8 @@ void processResponseDataForwardsAllResponseHeaders() { LambdaAppSecHandler.setCurrentTriggerType( LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); String json = - "{\"statusCode\": 200, \"headers\": {\"content-type\": \"application/json\", \"x-custom\": \"val\", \"content-length\": \"42\", \"set-cookie\": \"a=1\"}}"; + "{\"statusCode\": 200, \"headers\": {\"content-type\": \"application/json\", \"x-custom\":" + + " \"val\", \"content-length\": \"42\", \"set-cookie\": \"a=1\"}}"; ByteArrayOutputStream result = createOutputStream(json); Map capturedHeaders = new HashMap<>(); AgentSpan span = setupMockResponseCallbacks(null, capturedHeaders::put, null, null); @@ -1749,7 +2065,8 @@ void processResponseDataLowercasesHeaderKeys() { LambdaAppSecHandler.setCurrentTriggerType( LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); String json = - "{\"statusCode\": 200, \"headers\": {\"Content-Type\": \"text/html\", \"CONTENT-LENGTH\": \"10\"}}"; + "{\"statusCode\": 200, \"headers\": {\"Content-Type\": \"text/html\", \"CONTENT-LENGTH\":" + + " \"10\"}}"; ByteArrayOutputStream result = createOutputStream(json); Map capturedHeaders = new HashMap<>(); AgentSpan span = setupMockResponseCallbacks(null, capturedHeaders::put, null, null); @@ -1763,7 +2080,8 @@ void processResponseDataMergesMultiValueHeadersWithSingleValueHeaders() { LambdaAppSecHandler.setCurrentTriggerType( LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); String json = - "{\"statusCode\": 200, \"headers\": {\"content-type\": \"text/html\"}, \"multiValueHeaders\": {\"content-encoding\": [\"gzip\", \"br\"]}}"; + "{\"statusCode\": 200, \"headers\": {\"content-type\": \"text/html\"}," + + " \"multiValueHeaders\": {\"content-encoding\": [\"gzip\", \"br\"]}}"; ByteArrayOutputStream result = createOutputStream(json); Map capturedHeaders = new HashMap<>(); AgentSpan span = setupMockResponseCallbacks(null, capturedHeaders::put, null, null); @@ -1795,7 +2113,8 @@ void processResponseDataParsesJsonBody() { LambdaAppSecHandler.setCurrentTriggerType( LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); String json = - "{\"statusCode\": 200, \"headers\": {\"content-type\": \"application/json\"}, \"body\": \"{\\\"key\\\": \\\"value\\\"}\"}"; + "{\"statusCode\": 200, \"headers\": {\"content-type\": \"application/json\"}, \"body\":" + + " \"{\\\"key\\\": \\\"value\\\"}\"}"; ByteArrayOutputStream result = createOutputStream(json); Object[] capturedBody = {null}; AgentSpan span = setupMockResponseCallbacks(null, null, null, body -> capturedBody[0] = body); @@ -1809,7 +2128,8 @@ void processResponseDataHandlesNonJsonBodyAsRawString() { LambdaAppSecHandler.setCurrentTriggerType( LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); String json = - "{\"statusCode\": 200, \"headers\": {\"content-type\": \"text/plain\"}, \"body\": \"plain text\"}"; + "{\"statusCode\": 200, \"headers\": {\"content-type\": \"text/plain\"}, \"body\": \"plain" + + " text\"}"; ByteArrayOutputStream result = createOutputStream(json); Object[] capturedBody = {null}; AgentSpan span = setupMockResponseCallbacks(null, null, null, body -> capturedBody[0] = body); @@ -1837,6 +2157,8 @@ void processResponseDataHandlesBase64EncodedBody() { @Test void processResponseDataHandlesNullBody() { + LambdaAppSecHandler.setCurrentTriggerType( + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); ByteArrayOutputStream result = createOutputStream("{\"statusCode\": 200, \"body\": null}"); String[] capturedBody = {"NOT_CALLED"}; AgentSpan span = @@ -1848,6 +2170,8 @@ void processResponseDataHandlesNullBody() { @Test void processResponseDataHandlesMissingBodyField() { + LambdaAppSecHandler.setCurrentTriggerType( + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); ByteArrayOutputStream result = createOutputStream("{\"statusCode\": 200}"); String[] capturedBody = {"NOT_CALLED"}; AgentSpan span = @@ -1883,6 +2207,41 @@ void processResponseDataFallsBackToRawStringWhenJsonParseFails() { assertEquals("not json {", capturedBody[0]); } + @Test + void processResponseDataDropsMalformedBodyUnderJsonContentType() { + LambdaAppSecHandler.setCurrentTriggerType( + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); + // Explicit application/json content-type but the body is not valid JSON: it is malformed for + // its declared type, so responseBody must not fire (dropped, matching the extension) rather + // than being forwarded as a raw string. + ByteArrayOutputStream result = + createOutputStream( + "{\"statusCode\": 200, \"headers\": {\"content-type\": \"application/json\"}," + + " \"body\": \"not json {\"}"); + String[] capturedBody = {"NOT_CALLED"}; + AgentSpan span = + setupMockResponseCallbacks( + null, null, null, body -> capturedBody[0] = String.valueOf(body)); + LambdaAppSecHandler.processResponseData(span, result); + assertEquals("NOT_CALLED", capturedBody[0]); + } + + @Test + void processResponseDataKeepsTextPlainJsonBodyAsRawString() { + LambdaAppSecHandler.setCurrentTriggerType( + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); + // text/plain body that happens to be valid JSON must NOT be JSON-parsed; it is kept as the raw + // string so string-based WAF rules can still scan it (matches the extension). + ByteArrayOutputStream result = + createOutputStream( + "{\"statusCode\": 200, \"headers\": {\"content-type\": \"text/plain\"}," + + " \"body\": \"{\\\"a\\\": 1}\"}"); + Object[] capturedBody = {null}; + AgentSpan span = setupMockResponseCallbacks(null, null, null, body -> capturedBody[0] = body); + LambdaAppSecHandler.processResponseData(span, result); + assertEquals("{\"a\": 1}", capturedBody[0]); + } + // --- Event ordering --- @Test @@ -1890,7 +2249,8 @@ void processResponseDataFiresEventsInCorrectOrder() { LambdaAppSecHandler.setCurrentTriggerType( LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); String json = - "{\"statusCode\": 200, \"headers\": {\"content-type\": \"application/json\"}, \"body\": \"{\\\"k\\\": \\\"v\\\"}\"}"; + "{\"statusCode\": 200, \"headers\": {\"content-type\": \"application/json\"}, \"body\":" + + " \"{\\\"k\\\": \\\"v\\\"}\"}"; ByteArrayOutputStream result = createOutputStream(json); List order = new ArrayList<>(); @@ -1913,6 +2273,8 @@ void processResponseDataFiresEventsInCorrectOrder() { @Test void processResponseDataHandlesInvalidBase64ResponseBodyGracefully() { + LambdaAppSecHandler.setCurrentTriggerType( + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); String json = "{\"statusCode\": 200, \"body\": \"not-valid-base64!!!\", \"isBase64Encoded\": true}"; ByteArrayOutputStream result = createOutputStream(json); @@ -1930,7 +2292,8 @@ void processResponseDataParsesBodyAsJsonForJavascriptContentType() { LambdaAppSecHandler.setCurrentTriggerType( LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); String json = - "{\"statusCode\": 200, \"headers\": {\"content-type\": \"application/javascript\"}, \"body\": \"{\\\"key\\\": \\\"val\\\"}\"}"; + "{\"statusCode\": 200, \"headers\": {\"content-type\": \"application/javascript\"}," + + " \"body\": \"{\\\"key\\\": \\\"val\\\"}\"}"; ByteArrayOutputStream result = createOutputStream(json); Object[] capturedBody = {null}; AgentSpan span = setupMockResponseCallbacks(null, null, null, body -> capturedBody[0] = body); @@ -1939,12 +2302,57 @@ void processResponseDataParsesBodyAsJsonForJavascriptContentType() { assertEquals("val", ((Map) capturedBody[0]).get("key")); } + @Test + void processResponseDataParsesUrlEncodedBody() { + LambdaAppSecHandler.setCurrentTriggerType( + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); + String json = + "{\"statusCode\": 200, \"headers\": {\"content-type\":" + + " \"application/x-www-form-urlencoded\"}, \"body\": \"a=1&b=hello+world\"}"; + ByteArrayOutputStream result = createOutputStream(json); + Object[] capturedBody = {null}; + AgentSpan span = setupMockResponseCallbacks(null, null, null, body -> capturedBody[0] = body); + LambdaAppSecHandler.processResponseData(span, result); + assertInstanceOf(Map.class, capturedBody[0]); + @SuppressWarnings("unchecked") + Map> body = (Map>) capturedBody[0]; + assertEquals(Collections.singletonList("1"), body.get("a")); + assertEquals(Collections.singletonList("hello world"), body.get("b")); + } + + @Test + void processResponseDataMultipartBodyIsKeptAsRawString() { + // multipart/form-data response bodies are not structurally parsed; the raw payload is forwarded + // so string-based WAF rules can still scan it. + LambdaAppSecHandler.setCurrentTriggerType( + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); + String multipartBody = + "--BOUNDARY\r\n" + + "Content-Disposition: form-data; name=\"field1\"\r\n" + + "\r\n" + + "value1\r\n" + + "--BOUNDARY--\r\n"; + String json = + "{\"statusCode\": 200, " + + "\"headers\": {\"content-type\": \"multipart/form-data; boundary=BOUNDARY\"}, " + + "\"body\": \"" + + escapeJson(multipartBody) + + "\"}"; + ByteArrayOutputStream result = createOutputStream(json); + Object[] capturedBody = {null}; + AgentSpan span = setupMockResponseCallbacks(null, null, null, body -> capturedBody[0] = body); + LambdaAppSecHandler.processResponseData(span, result); + assertEquals(multipartBody, capturedBody[0]); + } + @Test void processResponseDataSkipsMultiValueHeadersEntryWithNonListValue() { LambdaAppSecHandler.setCurrentTriggerType( LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); String json = - "{\"statusCode\": 200, \"headers\": {\"content-type\": \"text/html\"}, \"multiValueHeaders\": {\"x-scalar\": \"not-a-list\", \"x-valid\": [\"v1\", \"v2\"]}}"; + "{\"statusCode\": 200, \"headers\": {\"content-type\": \"text/html\"}," + + " \"multiValueHeaders\": {\"x-scalar\": \"not-a-list\", \"x-valid\": [\"v1\"," + + " \"v2\"]}}"; ByteArrayOutputStream result = createOutputStream(json); Map capturedHeaders = new HashMap<>(); AgentSpan span = setupMockResponseCallbacks(null, capturedHeaders::put, null, null); @@ -1959,7 +2367,9 @@ void processResponseDataMultiValueHeadersOverrideSingleValueHeaders() { LambdaAppSecHandler.setCurrentTriggerType( LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); String json = - "{\"statusCode\": 200, \"headers\": {\"content-type\": \"text/html\"}, \"multiValueHeaders\": {\"content-type\": [\"application/json\", \"charset=utf-8\"]}}"; + "{\"statusCode\": 200, \"headers\": {\"content-type\": \"text/html\"}," + + " \"multiValueHeaders\": {\"content-type\": [\"application/json\"," + + " \"charset=utf-8\"]}}"; ByteArrayOutputStream result = createOutputStream(json); Map capturedHeaders = new HashMap<>(); AgentSpan span = setupMockResponseCallbacks(null, capturedHeaders::put, null, null); @@ -1970,13 +2380,21 @@ void processResponseDataMultiValueHeadersOverrideSingleValueHeaders() { // --- Error handling --- @Test - void processResponseDataHandlesMalformedJsonResponse() { + void processResponseDataFallsBackToTextPlainForMalformedJsonResponse() { + // Not valid API-GW-formatted JSON and not valid JSON at all, so extractResponseData returns + // null and the OBJECT_ADAPTER.fromJson retry in the fallback path also throws — the raw + // string must be surfaced as the response body with a text/plain content-type. + LambdaAppSecHandler.setCurrentTriggerType( + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); ByteArrayOutputStream result = createOutputStream("{not valid json"); - Integer[] capturedStatus = {null}; + Object[] capturedBody = {null}; + Map capturedHeaders = new HashMap<>(); AgentSpan span = - setupMockResponseCallbacks(status -> capturedStatus[0] = status, null, null, null); + setupMockResponseCallbacks( + null, capturedHeaders::put, null, body -> capturedBody[0] = body); LambdaAppSecHandler.processResponseData(span, result); - assertNull(capturedStatus[0]); + assertEquals("{not valid json", capturedBody[0]); + assertEquals("text/plain", capturedHeaders.get("content-type")); } @Test @@ -1993,6 +2411,8 @@ void processResponseDataHandlesEmptyStringResponse() { @Test void processResponseDataHandlesNullResponseHeaderDoneCallbackGracefully() { + LambdaAppSecHandler.setCurrentTriggerType( + LambdaAppSecHandler.LambdaTriggerType.API_GATEWAY_V1_REST); String json = "{\"statusCode\": 200, \"headers\": {\"content-type\": \"text/plain\"}, \"body\": \"ok\"}"; ByteArrayOutputStream result = createOutputStream(json);