Skip to content

Commit 883c505

Browse files
sezen-datadogsezen-datadog
committed
Xml util should only deal in maps
Signed-off-by: sezen.leblay <[email protected]>
1 parent 0a0cc9c commit 883c505

File tree

3 files changed

+131
-48
lines changed

3 files changed

+131
-48
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/XmlDomUtils.java

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -56,14 +56,7 @@ public static Object convertDocument(Document document, int maxRecursion) {
5656
return null;
5757
}
5858

59-
Element documentElement = document.getDocumentElement();
60-
if (documentElement == null) {
61-
return null;
62-
}
63-
64-
Object converted = convertW3cNode(documentElement, maxRecursion);
65-
// Wrap in a list for consistency with other XML processing patterns
66-
return converted != null ? Collections.singletonList(converted) : null;
59+
return convertW3cNode(document.getDocumentElement(), maxRecursion);
6760
}
6861

6962
/**
@@ -89,9 +82,7 @@ public static Object convertElement(Element element, int maxRecursion) {
8982
return null;
9083
}
9184

92-
Object converted = convertW3cNode(element, maxRecursion);
93-
// Wrap in a list for consistency with other XML processing patterns
94-
return converted != null ? Collections.singletonList(converted) : null;
85+
return convertW3cNode(element, maxRecursion);
9586
}
9687

9788
/**
@@ -270,8 +261,8 @@ public static Object processXmlForWaf(Object xmlObj, int maxRecursion) {
270261
}
271262

272263
if (xmlObj instanceof Node) {
273-
Object converted = convertW3cNode((Node) xmlObj, maxRecursion);
274-
return converted != null ? Collections.singletonList(converted) : null;
264+
// Return the converted node directly
265+
return convertW3cNode((Node) xmlObj, maxRecursion);
275266
}
276267

277268
// Handle XML strings by parsing them first

‎dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/instrumentation/XmlDomUtilsTest.java

Lines changed: 22 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -161,12 +161,11 @@ void testConvertDocument() throws Exception {
161161
Object result = XmlDomUtils.convertDocument(doc, MAX_RECURSION);
162162

163163
assertNotNull(result);
164-
assertTrue(result instanceof List);
164+
assertTrue(result instanceof Map);
165165

166166
@SuppressWarnings("unchecked")
167-
List<Object> resultList = (List<Object>) result;
168-
assertEquals(1, resultList.size());
169-
assertTrue(resultList.get(0) instanceof Map);
167+
Map<String, Object> resultMap = (Map<String, Object>) result;
168+
assertTrue(resultMap.containsKey("children"));
170169
}
171170

172171
@Test
@@ -184,12 +183,11 @@ void testConvertElement() throws Exception {
184183
Object result = XmlDomUtils.convertElement(element, MAX_RECURSION);
185184

186185
assertNotNull(result);
187-
assertTrue(result instanceof List);
186+
assertTrue(result instanceof Map);
188187

189188
@SuppressWarnings("unchecked")
190-
List<Object> resultList = (List<Object>) result;
191-
assertEquals(1, resultList.size());
192-
assertTrue(resultList.get(0) instanceof Map);
189+
Map<String, Object> resultMap = (Map<String, Object>) result;
190+
assertTrue(resultMap.containsKey("children"));
193191
}
194192

195193
@Test
@@ -206,7 +204,7 @@ void testProcessXmlForWaf_withDocument() throws Exception {
206204
Object result = XmlDomUtils.processXmlForWaf(doc, MAX_RECURSION);
207205

208206
assertNotNull(result);
209-
assertTrue(result instanceof List);
207+
assertTrue(result instanceof Map);
210208
}
211209

212210
@Test
@@ -218,7 +216,7 @@ void testProcessXmlForWaf_withElement() throws Exception {
218216
Object result = XmlDomUtils.processXmlForWaf(element, MAX_RECURSION);
219217

220218
assertNotNull(result);
221-
assertTrue(result instanceof List);
219+
assertTrue(result instanceof Map);
222220
}
223221

224222
@Test
@@ -228,12 +226,11 @@ void testProcessXmlForWaf_withXmlString() {
228226
Object result = XmlDomUtils.processXmlForWaf(xmlContent, MAX_RECURSION);
229227

230228
assertNotNull(result);
231-
assertTrue(result instanceof List);
229+
assertTrue(result instanceof Map);
232230

233231
@SuppressWarnings("unchecked")
234-
List<Object> resultList = (List<Object>) result;
235-
assertEquals(1, resultList.size());
236-
assertTrue(resultList.get(0) instanceof Map);
232+
Map<String, Object> resultMap = (Map<String, Object>) result;
233+
assertTrue(resultMap.containsKey("children"));
237234
}
238235

239236
@Test
@@ -258,19 +255,14 @@ void testParseXmlStringToWafFormat_validXml() throws Exception {
258255
Object result = XmlDomUtils.parseXmlStringToWafFormat(xmlContent, MAX_RECURSION);
259256

260257
assertNotNull(result);
261-
assertTrue(result instanceof List);
262-
263-
@SuppressWarnings("unchecked")
264-
List<Object> resultList = (List<Object>) result;
265-
assertEquals(1, resultList.size());
266-
assertTrue(resultList.get(0) instanceof Map);
258+
assertTrue(result instanceof Map);
267259

268260
@SuppressWarnings("unchecked")
269-
Map<String, Object> rootMap = (Map<String, Object>) resultList.get(0);
270-
assertTrue(rootMap.containsKey("children"));
261+
Map<String, Object> resultMap = (Map<String, Object>) result;
262+
assertTrue(resultMap.containsKey("children"));
271263

272264
@SuppressWarnings("unchecked")
273-
List<Object> children = (List<Object>) rootMap.get("children");
265+
List<Object> children = (List<Object>) resultMap.get("children");
274266
assertEquals(2, children.size()); // title and author elements
275267
}
276268

@@ -301,7 +293,7 @@ void testHandleXmlString_validXml() {
301293
Object result = XmlDomUtils.handleXmlString(xmlContent, MAX_RECURSION);
302294

303295
assertNotNull(result);
304-
assertTrue(result instanceof List);
296+
assertTrue(result instanceof Map);
305297
}
306298

307299
@Test
@@ -332,7 +324,7 @@ void testXmlWithNamespaces() throws Exception {
332324
Object result = XmlDomUtils.processXmlForWaf(xmlContent, MAX_RECURSION);
333325

334326
assertNotNull(result);
335-
assertTrue(result instanceof List);
327+
assertTrue(result instanceof Map);
336328
}
337329

338330
@Test
@@ -342,7 +334,7 @@ void testXmlWithCDATA() throws Exception {
342334
Object result = XmlDomUtils.processXmlForWaf(xmlContent, MAX_RECURSION);
343335

344336
assertNotNull(result);
345-
assertTrue(result instanceof List);
337+
assertTrue(result instanceof Map);
346338
}
347339

348340
@Test
@@ -371,19 +363,14 @@ void testComplexXmlStructure() throws Exception {
371363
Object result = XmlDomUtils.processXmlForWaf(xmlContent, MAX_RECURSION);
372364

373365
assertNotNull(result);
374-
assertTrue(result instanceof List);
375-
376-
@SuppressWarnings("unchecked")
377-
List<Object> resultList = (List<Object>) result;
378-
assertEquals(1, resultList.size());
379-
assertTrue(resultList.get(0) instanceof Map);
366+
assertTrue(result instanceof Map);
380367

381368
@SuppressWarnings("unchecked")
382-
Map<String, Object> rootMap = (Map<String, Object>) resultList.get(0);
383-
assertTrue(rootMap.containsKey("children"));
369+
Map<String, Object> resultMap = (Map<String, Object>) result;
370+
assertTrue(resultMap.containsKey("children"));
384371

385372
@SuppressWarnings("unchecked")
386-
List<Object> children = (List<Object>) rootMap.get("children");
373+
List<Object> children = (List<Object>) resultMap.get("children");
387374
// Should have 2 book elements (ignoring whitespace text nodes)
388375
long bookElements = children.stream().filter(child -> child instanceof Map).count();
389376
assertEquals(2, bookElements);

‎dd-java-agent/instrumentation/jersey-3-appsec/src/main/java/datadog/trace/instrumentation/jersey3/MessageBodyWriterInstrumentation.java

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
package datadog.trace.instrumentation.jersey3;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static datadog.trace.api.gateway.Events.EVENTS;
6+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
7+
8+
import com.google.auto.service.AutoService;
9+
import datadog.appsec.api.blocking.BlockingException;
10+
import datadog.trace.advice.ActiveRequestContext;
11+
import datadog.trace.advice.RequiresRequestContext;
12+
import datadog.trace.agent.tooling.Instrumenter;
13+
import datadog.trace.agent.tooling.InstrumenterModule;
14+
import datadog.trace.api.gateway.BlockResponseFunction;
15+
import datadog.trace.api.gateway.CallbackProvider;
16+
import datadog.trace.api.gateway.Flow;
17+
import datadog.trace.api.gateway.RequestContext;
18+
import datadog.trace.api.gateway.RequestContextSlot;
19+
import datadog.trace.bootstrap.instrumentation.XmlDomUtils;
20+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
21+
import jakarta.ws.rs.core.MediaType;
22+
import java.util.function.BiFunction;
23+
import net.bytebuddy.asm.Advice;
24+
import net.bytebuddy.description.type.TypeDescription;
25+
import net.bytebuddy.matcher.ElementMatcher;
26+
27+
@AutoService(InstrumenterModule.class)
28+
public class MessageBodyWriterInstrumentation extends InstrumenterModule.AppSec
29+
implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
30+
31+
public MessageBodyWriterInstrumentation() {
32+
super("jersey");
33+
}
34+
35+
@Override
36+
public String muzzleDirective() {
37+
return "common";
38+
}
39+
40+
@Override
41+
public String hierarchyMarkerType() {
42+
return "jakarta.ws.rs.ext.MessageBodyWriter";
43+
}
44+
45+
@Override
46+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
47+
return implementsInterface(named(hierarchyMarkerType()));
48+
}
49+
50+
@Override
51+
public void methodAdvice(MethodTransformer transformer) {
52+
transformer.applyAdvice(
53+
named("writeTo").and(takesArguments(7)), getClass().getName() + "$MessageBodyWriterAdvice");
54+
}
55+
56+
@RequiresRequestContext(RequestContextSlot.APPSEC)
57+
public static class MessageBodyWriterAdvice {
58+
@Advice.OnMethodEnter(suppress = Throwable.class)
59+
static void before(
60+
@Advice.Argument(0) Object entity,
61+
@Advice.Argument(4) MediaType mediaType,
62+
@ActiveRequestContext RequestContext reqCtx) {
63+
64+
// Handle both JSON and XML response bodies
65+
if (!MediaType.APPLICATION_JSON_TYPE.isCompatible(mediaType)
66+
&& !MediaType.APPLICATION_XML_TYPE.isCompatible(mediaType)
67+
&& !MediaType.TEXT_XML_TYPE.isCompatible(mediaType)) {
68+
return;
69+
}
70+
71+
CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
72+
BiFunction<RequestContext, Object, Flow<Void>> callback =
73+
cbp.getCallback(EVENTS.responseBody());
74+
if (callback == null) {
75+
return;
76+
}
77+
78+
// Process XML entities for WAF compatibility
79+
Object processedEntity = entity;
80+
if ((MediaType.APPLICATION_XML_TYPE.isCompatible(mediaType)
81+
|| MediaType.TEXT_XML_TYPE.isCompatible(mediaType))
82+
&& entity instanceof String) {
83+
Object xmlProcessed = XmlDomUtils.processXmlForWaf(entity);
84+
processedEntity = xmlProcessed != null ? xmlProcessed : entity;
85+
}
86+
87+
Flow<Void> flow = callback.apply(reqCtx, processedEntity);
88+
Flow.Action action = flow.getAction();
89+
if (action instanceof Flow.Action.RequestBlockingAction) {
90+
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
91+
if (blockResponseFunction == null) {
92+
return;
93+
}
94+
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
95+
blockResponseFunction.tryCommitBlockingResponse(
96+
reqCtx.getTraceSegment(),
97+
rba.getStatusCode(),
98+
rba.getBlockingContentType(),
99+
rba.getExtraHeaders());
100+
101+
throw new BlockingException("Blocked request (for MessageBodyWriter)");
102+
}
103+
}
104+
}
105+
}

0 commit comments

Comments
 (0)