Skip to content

Commit 3f811fd

Browse files
claponcetclaponcet
committed
forwarded headers parsing + downgrade log level
1 parent 240b5ab commit 3f811fd

File tree

2 files changed

+151
-50
lines changed

2 files changed

+151
-50
lines changed

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

Lines changed: 45 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -71,9 +71,12 @@ public static AgentSpanContext processRequestStart(Object event) {
7171

7272
try {
7373
LambdaEventData eventData = extractEventData((ByteArrayInputStream) event);
74+
if (eventData == LambdaEventData.EMPTY) {
75+
return null;
76+
}
7477
return processAppSecRequestData(eventData);
7578
} catch (Exception e) {
76-
log.error("Failed to process AppSec request data", e);
79+
log.debug("Failed to process AppSec request data", e);
7780
return null;
7881
}
7982
}
@@ -96,7 +99,7 @@ public static void processRequestEnd(AgentSpan span) {
9699
if (requestEndedCallback != null) {
97100
requestEndedCallback.apply(requestContext, span);
98101
} else {
99-
log.warn("requestEnded callback is null");
102+
log.debug("requestEnded callback is null");
100103
}
101104
}
102105
}
@@ -129,7 +132,7 @@ public static AgentSpanContext mergeContexts(
129132
return merged;
130133
}
131134

132-
log.warn(
135+
log.debug(
133136
"Cannot merge AppSec data: extension context is not a TagContext: {}",
134137
extensionContext.getClass());
135138
}
@@ -141,7 +144,7 @@ private static AgentSpanContext processAppSecRequestData(LambdaEventData eventDa
141144
Supplier<Flow<Object>> requestStartedCallback =
142145
tracer.getCallbackProvider(RequestContextSlot.APPSEC).getCallback(EVENTS.requestStarted());
143146
if (requestStartedCallback == null) {
144-
log.warn("requestStarted callback is null");
147+
log.debug("requestStarted callback is null");
145148
return null;
146149
}
147150

@@ -165,10 +168,10 @@ private static AgentSpanContext processAppSecRequestData(LambdaEventData eventDa
165168
if (methodUriCallback != null) {
166169
// Reconstruct full path with query string for AppSec analysis
167170
String fullPath = buildFullPath(eventData.path, eventData.queryParameters);
168-
LambdaURIDataAdapter uriAdapter = new LambdaURIDataAdapter(fullPath);
171+
LambdaURIDataAdapter uriAdapter = new LambdaURIDataAdapter(fullPath, eventData.headers);
169172
methodUriCallback.apply(requestContext, eventData.method, uriAdapter);
170173
} else {
171-
log.warn("requestMethodUriRaw callback is null");
174+
log.debug("requestMethodUriRaw callback is null");
172175
}
173176
}
174177

@@ -183,7 +186,7 @@ private static AgentSpanContext processAppSecRequestData(LambdaEventData eventDa
183186
headerCallback.accept(requestContext, header.getKey(), header.getValue());
184187
}
185188
} else {
186-
log.warn("requestHeader callback is null");
189+
log.debug("requestHeader callback is null");
187190
}
188191
}
189192

@@ -198,7 +201,7 @@ private static AgentSpanContext processAppSecRequestData(LambdaEventData eventDa
198201
Integer port = eventData.sourcePort != null ? eventData.sourcePort : 0;
199202
socketAddrCallback.apply(requestContext, eventData.sourceIp, port);
200203
} else {
201-
log.warn("requestClientSocketAddress callback is null");
204+
log.debug("requestClientSocketAddress callback is null");
202205
}
203206
}
204207

@@ -210,7 +213,7 @@ private static AgentSpanContext processAppSecRequestData(LambdaEventData eventDa
210213
if (headerDoneCallback != null) {
211214
headerDoneCallback.apply(requestContext);
212215
} else {
213-
log.warn("requestHeaderDone callback is null");
216+
log.debug("requestHeaderDone callback is null");
214217
}
215218

216219
// Call requestPathParams
@@ -222,7 +225,7 @@ private static AgentSpanContext processAppSecRequestData(LambdaEventData eventDa
222225
if (pathParamsCallback != null) {
223226
pathParamsCallback.apply(requestContext, eventData.pathParameters);
224227
} else {
225-
log.warn("requestPathParams callback is null");
228+
log.debug("requestPathParams callback is null");
226229
}
227230
}
228231

@@ -235,7 +238,7 @@ private static AgentSpanContext processAppSecRequestData(LambdaEventData eventDa
235238
if (bodyCallback != null) {
236239
bodyCallback.apply(requestContext, eventData.body);
237240
} else {
238-
log.warn("requestBodyProcessed callback is null");
241+
log.debug("requestBodyProcessed callback is null");
239242
}
240243
}
241244
}
@@ -249,20 +252,11 @@ private static LambdaEventData extractEventData(ByteArrayInputStream inputStream
249252
int availableBytes = inputStream.available();
250253

251254
if (availableBytes <= 0 || availableBytes > MAX_EVENT_SIZE) {
252-
log.warn(
255+
log.debug(
253256
"Event size {} exceeds limit {} or is invalid, skipping AppSec processing",
254257
availableBytes,
255258
MAX_EVENT_SIZE);
256-
return new LambdaEventData(
257-
Collections.emptyMap(),
258-
null,
259-
null,
260-
null,
261-
null,
262-
LambdaTriggerType.UNKNOWN,
263-
Collections.emptyMap(),
264-
Collections.emptyMap(),
265-
null);
259+
return LambdaEventData.EMPTY;
266260
}
267261

268262
StringBuilder jsonBuilder = new StringBuilder(availableBytes);
@@ -286,16 +280,7 @@ private static LambdaEventData extractEventDataFromJson(String json) {
286280
log.debug("Event JSON parsed successfully");
287281

288282
if (event == null) {
289-
return new LambdaEventData(
290-
Collections.emptyMap(),
291-
null,
292-
null,
293-
null,
294-
null,
295-
LambdaTriggerType.UNKNOWN,
296-
Collections.emptyMap(),
297-
Collections.emptyMap(),
298-
null);
283+
return LambdaEventData.EMPTY;
299284
}
300285

301286
// Detect trigger type
@@ -319,17 +304,8 @@ private static LambdaEventData extractEventDataFromJson(String json) {
319304
return extractGenericData(event);
320305
}
321306
} catch (Exception e) {
322-
log.error("Failed to parse event data from JSON", e);
323-
return new LambdaEventData(
324-
Collections.emptyMap(),
325-
null,
326-
null,
327-
null,
328-
null,
329-
LambdaTriggerType.UNKNOWN,
330-
Collections.emptyMap(),
331-
Collections.emptyMap(),
332-
null);
307+
log.debug("Failed to parse event data from JSON", e);
308+
return LambdaEventData.EMPTY;
333309
}
334310
}
335311

@@ -524,7 +500,8 @@ private static LambdaEventData extractAlbData(
524500

525501
String method = (String) event.get("httpMethod");
526502
String path = (String) event.get("path");
527-
String sourceIp = headers.get("x-forwarded-for");
503+
String xff = headers.get("x-forwarded-for");
504+
String sourceIp = xff != null ? xff.split(",")[0].trim() : null;
528505

529506
return new LambdaEventData(
530507
headers, method, path, sourceIp, null, triggerType, pathParameters, queryParameters, body);
@@ -779,8 +756,7 @@ private static Object parseBodyAsJson(String body) {
779756
}
780757

781758
try {
782-
Object parsed = OBJECT_ADAPTER.fromJson(body);
783-
return parsed;
759+
return OBJECT_ADAPTER.fromJson(body);
784760
} catch (Exception e) {
785761
return null;
786762
}
@@ -853,6 +829,11 @@ static class LambdaEventData {
853829
final Map<String, List<String>> queryParameters;
854830
final Object body;
855831

832+
static final LambdaEventData EMPTY = new LambdaEventData(
833+
Collections.emptyMap(), null, null, null, null,
834+
LambdaTriggerType.UNKNOWN,
835+
Collections.emptyMap(), Collections.emptyMap(), null);
836+
856837
LambdaEventData(
857838
Map<String, String> headers,
858839
String method,
@@ -879,8 +860,10 @@ static class LambdaEventData {
879860
private static class LambdaURIDataAdapter extends URIDataAdapterBase {
880861
private final String path;
881862
private final String query;
863+
private final String scheme;
864+
private final int port;
882865

883-
LambdaURIDataAdapter(String pathWithQuery) {
866+
LambdaURIDataAdapter(String pathWithQuery, Map<String, String> headers) {
884867
if (pathWithQuery != null) {
885868
int queryIndex = pathWithQuery.indexOf('?');
886869
if (queryIndex != -1) {
@@ -894,11 +877,24 @@ private static class LambdaURIDataAdapter extends URIDataAdapterBase {
894877
this.path = "/";
895878
this.query = null;
896879
}
880+
881+
String forwardedProto = headers != null ? headers.get("x-forwarded-proto") : null;
882+
this.scheme = (forwardedProto != null && !forwardedProto.isEmpty()) ? forwardedProto : "https";
883+
884+
String forwardedPort = headers != null ? headers.get("x-forwarded-port") : null;
885+
int parsedPort = -1;
886+
if (forwardedPort != null && !forwardedPort.isEmpty()) {
887+
try {
888+
parsedPort = Integer.parseInt(forwardedPort.trim());
889+
} catch (NumberFormatException ignored) {
890+
}
891+
}
892+
this.port = parsedPort > 0 ? parsedPort : 443;
897893
}
898894

899895
@Override
900896
public String scheme() {
901-
return "https";
897+
return scheme;
902898
}
903899

904900
@Override
@@ -908,7 +904,7 @@ public String host() {
908904

909905
@Override
910906
public int port() {
911-
return 443;
907+
return port;
912908
}
913909

914910
@Override

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

Lines changed: 106 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -730,6 +730,111 @@ class LambdaAppSecHandlerTest extends DDCoreSpecification {
730730
capturedQuery == "id=123&filter=active"
731731
}
732732

733+
def "extracts scheme and port from X-Forwarded headers"() {
734+
given:
735+
def eventJson = '''
736+
{
737+
"path": "/api/test",
738+
"headers": {
739+
"x-forwarded-proto": "http",
740+
"x-forwarded-port": "8080"
741+
},
742+
"requestContext": {
743+
"httpMethod": "GET",
744+
"requestId": "req-123"
745+
}
746+
}
747+
'''
748+
def event = createInputStream(eventJson)
749+
750+
def capturedScheme = null
751+
def capturedPort = null
752+
753+
setupMockCallbacks(
754+
onMethodUri: { method, uri ->
755+
capturedScheme = uri.scheme()
756+
capturedPort = uri.port()
757+
}
758+
)
759+
760+
when:
761+
def result = LambdaAppSecHandler.processRequestStart(event)
762+
763+
then:
764+
result != null
765+
capturedScheme == "http"
766+
capturedPort == 8080
767+
}
768+
769+
def "falls back to https/443 when X-Forwarded headers are absent"() {
770+
given:
771+
def eventJson = '''
772+
{
773+
"path": "/api/test",
774+
"headers": {},
775+
"requestContext": {
776+
"httpMethod": "GET",
777+
"requestId": "req-123"
778+
}
779+
}
780+
'''
781+
def event = createInputStream(eventJson)
782+
783+
def capturedScheme = null
784+
def capturedPort = null
785+
786+
setupMockCallbacks(
787+
onMethodUri: { method, uri ->
788+
capturedScheme = uri.scheme()
789+
capturedPort = uri.port()
790+
}
791+
)
792+
793+
when:
794+
def result = LambdaAppSecHandler.processRequestStart(event)
795+
796+
then:
797+
result != null
798+
capturedScheme == "https"
799+
capturedPort == 443
800+
}
801+
802+
def "handles invalid X-Forwarded-Port gracefully"() {
803+
given:
804+
def eventJson = '''
805+
{
806+
"path": "/api/test",
807+
"headers": {
808+
"x-forwarded-proto": "https",
809+
"x-forwarded-port": "not-a-number"
810+
},
811+
"requestContext": {
812+
"httpMethod": "GET",
813+
"requestId": "req-123"
814+
}
815+
}
816+
'''
817+
def event = createInputStream(eventJson)
818+
819+
def capturedScheme = null
820+
def capturedPort = null
821+
822+
setupMockCallbacks(
823+
onMethodUri: { method, uri ->
824+
capturedScheme = uri.scheme()
825+
capturedPort = uri.port()
826+
}
827+
)
828+
829+
when:
830+
def result = LambdaAppSecHandler.processRequestStart(event)
831+
832+
then:
833+
result != null
834+
capturedScheme == "https"
835+
capturedPort == 443
836+
}
837+
733838
def "handles invalid base64 body gracefully"() {
734839
given:
735840
def eventJson = '''
@@ -1271,7 +1376,7 @@ class LambdaAppSecHandlerTest extends DDCoreSpecification {
12711376
get() >> new Flow.ResultFlow<>(mockAppSecContext)
12721377
}
12731378

1274-
def mockMethodUriCallback = callbacks.onMethodUri ? Mock(datadog.trace.api.function.TriFunction) {
1379+
def mockMethodUriCallback = callbacks.onMethodUri ? Mock(TriFunction) {
12751380
apply(_ as RequestContext, _ as String, _ as URIDataAdapter) >> {
12761381
RequestContext ctx, String method, URIDataAdapter uri ->
12771382
callbacks.onMethodUri(method, uri)

0 commit comments

Comments
 (0)