Skip to content

Commit f6f2feb

Browse files
jandro996jandro996
committed
refactor: move GlassFish multipart processing into GlassFishBlockingHelper
Extract the part-iteration loop and IG callback dispatch from the advice into a new GlassFishBlockingHelper.processPartsAndBlock() static method. The advice now contains only the inlined reflection block (which must stay inlined due to Java module-system constraints) and a single call to the helper. Unit tests added for processPartsAndBlock(): form-field skip, empty filename, normal filename, content reading, maxFiles limit, getInputStream failure fallback, sequential callback ordering (filenames blocks → content not fired), content callback blocks → returns true, non-Part object skip.
1 parent a45b404 commit f6f2feb

3 files changed

Lines changed: 257 additions & 71 deletions

File tree

‎dd-java-agent/instrumentation/tomcat/tomcat-appsec/tomcat-appsec-7.0/src/main/java/datadog/trace/instrumentation/tomcat7/GlassFishBlockingHelper.java‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,17 @@
55
import datadog.trace.api.gateway.BlockResponseFunction;
66
import datadog.trace.api.gateway.Flow;
77
import datadog.trace.api.gateway.RequestContext;
8+
import datadog.trace.api.http.MultipartContentDecoder;
89
import datadog.trace.bootstrap.blocking.BlockingActionHelper;
10+
import java.io.InputStream;
11+
import java.util.ArrayList;
12+
import java.util.Collection;
13+
import java.util.List;
914
import java.util.Map;
15+
import java.util.function.BiFunction;
1016
import javax.servlet.http.HttpServletRequest;
1117
import javax.servlet.http.HttpServletResponse;
18+
import javax.servlet.http.Part;
1219

1320
public final class GlassFishBlockingHelper {
1421

@@ -48,6 +55,79 @@ public static boolean tryBlock(
4855
return true;
4956
}
5057

58+
/**
59+
* Collects filenames and file contents from the given multipart parts, fires the AppSec IG
60+
* callbacks, and commits a blocking response if the WAF requests one.
61+
*
62+
* <p>Returns {@code true} if a blocking response was committed (the caller should replace the
63+
* parts collection with an empty list to prevent further processing).
64+
*/
65+
public static boolean processPartsAndBlock(
66+
Collection<?> parts,
67+
RequestContext reqCtx,
68+
HttpServletRequest fallbackReq,
69+
HttpServletResponse fallbackResp,
70+
BiFunction<RequestContext, List<String>, Flow<Void>> filenamesCb,
71+
BiFunction<RequestContext, List<String>, Flow<Void>> contentCb) {
72+
List<String> filenames = null;
73+
List<String> contents = null;
74+
for (Object partObj : parts) {
75+
try {
76+
if (!(partObj instanceof Part)) {
77+
continue;
78+
}
79+
Part part = (Part) partObj;
80+
String filename = part.getSubmittedFileName();
81+
if (filename == null) {
82+
continue;
83+
}
84+
if (filenamesCb != null && !filename.isEmpty()) {
85+
if (filenames == null) {
86+
filenames = new ArrayList<>();
87+
}
88+
filenames.add(filename);
89+
}
90+
if (contentCb != null) {
91+
if (contents == null) {
92+
contents = new ArrayList<>();
93+
}
94+
if (contents.size() < MAX_FILE_CONTENT_COUNT) {
95+
try (InputStream is = part.getInputStream()) {
96+
contents.add(
97+
MultipartContentDecoder.readInputStream(
98+
is, MAX_FILE_CONTENT_BYTES, part.getContentType()));
99+
} catch (Exception ignored) {
100+
contents.add("");
101+
}
102+
}
103+
}
104+
} catch (Exception ignored) {
105+
}
106+
}
107+
108+
if (filenames != null && !filenames.isEmpty()) {
109+
Flow<Void> flow = filenamesCb.apply(reqCtx, filenames);
110+
Flow.Action action = flow.getAction();
111+
if (action instanceof Flow.Action.RequestBlockingAction) {
112+
if (tryBlock(
113+
reqCtx, fallbackReq, fallbackResp, (Flow.Action.RequestBlockingAction) action)) {
114+
return true;
115+
}
116+
}
117+
}
118+
119+
if (contents != null && !contents.isEmpty()) {
120+
Flow<Void> contentFlow = contentCb.apply(reqCtx, contents);
121+
Flow.Action contentAction = contentFlow.getAction();
122+
if (contentAction instanceof Flow.Action.RequestBlockingAction) {
123+
return tryBlock(
124+
reqCtx, fallbackReq, fallbackResp, (Flow.Action.RequestBlockingAction) contentAction);
125+
}
126+
}
127+
128+
return false;
129+
}
130+
51131
public static boolean commitBlocking(
52132
HttpServletRequest request,
53133
HttpServletResponse response,

‎dd-java-agent/instrumentation/tomcat/tomcat-appsec/tomcat-appsec-7.0/src/main/java/datadog/trace/instrumentation/tomcat7/GlassFishMultipartInstrumentation.java‎

Lines changed: 3 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -12,20 +12,16 @@
1212
import datadog.trace.api.gateway.Flow;
1313
import datadog.trace.api.gateway.RequestContext;
1414
import datadog.trace.api.gateway.RequestContextSlot;
15-
import datadog.trace.api.http.MultipartContentDecoder;
1615
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1716
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
18-
import java.io.InputStream;
1917
import java.lang.reflect.Field;
2018
import java.lang.reflect.Method;
21-
import java.util.ArrayList;
2219
import java.util.Collection;
2320
import java.util.Collections;
2421
import java.util.List;
2522
import java.util.function.BiFunction;
2623
import javax.servlet.http.HttpServletRequest;
2724
import javax.servlet.http.HttpServletResponse;
28-
import javax.servlet.http.Part;
2925
import net.bytebuddy.asm.Advice;
3026

3127
/**
@@ -125,73 +121,9 @@ static void after(
125121
} catch (Exception ignored) {
126122
}
127123

128-
int maxFiles = GlassFishBlockingHelper.MAX_FILE_CONTENT_COUNT;
129-
int maxBytes = GlassFishBlockingHelper.MAX_FILE_CONTENT_BYTES;
130-
131-
List<String> filenames = null;
132-
List<String> contents = null;
133-
134-
for (Object partObj : parts) {
135-
try {
136-
if (!(partObj instanceof Part)) {
137-
continue;
138-
}
139-
Part part = (Part) partObj;
140-
String filename = part.getSubmittedFileName();
141-
// null means no filename parameter → form field, skip
142-
// empty string means filename="" was sent → file upload without a name
143-
if (filename == null) {
144-
continue;
145-
}
146-
if (filenamesCb != null && !filename.isEmpty()) {
147-
if (filenames == null) {
148-
filenames = new ArrayList<>();
149-
}
150-
filenames.add(filename);
151-
}
152-
if (contentCb != null) {
153-
if (contents == null) {
154-
contents = new ArrayList<>();
155-
}
156-
if (contents.size() < maxFiles) {
157-
try (InputStream is = part.getInputStream()) {
158-
contents.add(
159-
MultipartContentDecoder.readInputStream(is, maxBytes, part.getContentType()));
160-
} catch (Exception ignored) {
161-
contents.add("");
162-
}
163-
}
164-
}
165-
} catch (Exception ignored) {
166-
}
167-
}
168-
169-
boolean blocked = false;
170-
171-
if (filenames != null && !filenames.isEmpty() && filenamesCb != null) {
172-
Flow<Void> flow = filenamesCb.apply(reqCtx, filenames);
173-
Flow.Action action = flow.getAction();
174-
if (action instanceof Flow.Action.RequestBlockingAction) {
175-
if (GlassFishBlockingHelper.tryBlock(
176-
reqCtx, fallbackReq, fallbackResp, (Flow.Action.RequestBlockingAction) action)) {
177-
parts = Collections.emptyList();
178-
blocked = true;
179-
}
180-
}
181-
}
182-
183-
if (!blocked && contents != null && !contents.isEmpty() && contentCb != null) {
184-
Flow<Void> contentFlow = contentCb.apply(reqCtx, contents);
185-
Flow.Action contentAction = contentFlow.getAction();
186-
if (contentAction instanceof Flow.Action.RequestBlockingAction) {
187-
if (GlassFishBlockingHelper.tryBlock(
188-
reqCtx,
189-
fallbackReq,
190-
fallbackResp,
191-
(Flow.Action.RequestBlockingAction) contentAction)) {
192-
parts = Collections.emptyList();
193-
}
194-
}
124+
if (GlassFishBlockingHelper.processPartsAndBlock(
125+
parts, reqCtx, fallbackReq, fallbackResp, filenamesCb, contentCb)) {
126+
parts = Collections.emptyList();
195127
}
196128
}
197129
}

‎dd-java-agent/instrumentation/tomcat/tomcat-appsec/tomcat-appsec-7.0/src/test/java/datadog/trace/instrumentation/tomcat7/GlassFishBlockingHelperTest.java‎

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,15 @@
1818
import datadog.trace.api.internal.TraceSegment;
1919
import java.io.ByteArrayOutputStream;
2020
import java.io.IOException;
21+
import java.util.Arrays;
22+
import java.util.Collections;
23+
import java.util.List;
24+
import java.util.function.BiFunction;
2125
import javax.servlet.ServletOutputStream;
2226
import javax.servlet.WriteListener;
2327
import javax.servlet.http.HttpServletRequest;
2428
import javax.servlet.http.HttpServletResponse;
29+
import javax.servlet.http.Part;
2530
import org.junit.jupiter.api.Test;
2631

2732
class GlassFishBlockingHelperTest {
@@ -170,6 +175,146 @@ void tryBlock_effectivelyBlockedThrows_stillReturnsTrue() throws Exception {
170175
assertTrue(GlassFishBlockingHelper.tryBlock(reqCtx, null, null, rba(403)));
171176
}
172177

178+
// ------- processPartsAndBlock() -------
179+
180+
@Test
181+
void processPartsAndBlock_formField_skipped() throws Exception {
182+
Part formField = mockPart(null, "text/plain", new byte[0]);
183+
RequestContext reqCtx = mockReqCtx(null, mock(TraceSegment.class));
184+
BiFunction<RequestContext, List<String>, Flow<Void>> filenamesCb = mockPassThroughCb();
185+
186+
assertFalse(
187+
GlassFishBlockingHelper.processPartsAndBlock(
188+
Collections.singletonList(formField), reqCtx, null, null, filenamesCb, null));
189+
verify(formField).getSubmittedFileName();
190+
verify(formField, never()).getInputStream();
191+
}
192+
193+
@Test
194+
void processPartsAndBlock_emptyFilename_notAddedToFilenames_butContentRead() throws Exception {
195+
byte[] content = "data".getBytes();
196+
Part filePart = mockPart("", "application/octet-stream", content);
197+
RequestContext reqCtx = mockReqCtx(null, mock(TraceSegment.class));
198+
BiFunction<RequestContext, List<String>, Flow<Void>> filenamesCb = mockPassThroughCb();
199+
BiFunction<RequestContext, List<String>, Flow<Void>> contentCb = mockPassThroughCb();
200+
201+
assertFalse(
202+
GlassFishBlockingHelper.processPartsAndBlock(
203+
Collections.singletonList(filePart), reqCtx, null, null, filenamesCb, contentCb));
204+
205+
verify(filePart).getInputStream();
206+
verify(filenamesCb, never()).apply(any(), any());
207+
verify(contentCb).apply(eq(reqCtx), any());
208+
}
209+
210+
@Test
211+
void processPartsAndBlock_normalFilename_reportedViaFilenamesCb() throws Exception {
212+
Part filePart = mockPart("file.txt", "text/plain", "hello".getBytes());
213+
RequestContext reqCtx = mockReqCtx(null, mock(TraceSegment.class));
214+
BiFunction<RequestContext, List<String>, Flow<Void>> filenamesCb = mockPassThroughCb();
215+
216+
assertFalse(
217+
GlassFishBlockingHelper.processPartsAndBlock(
218+
Collections.singletonList(filePart), reqCtx, null, null, filenamesCb, null));
219+
220+
verify(filenamesCb).apply(eq(reqCtx), eq(Collections.singletonList("file.txt")));
221+
}
222+
223+
@Test
224+
void processPartsAndBlock_contentRead_reportedViaContentCb() throws Exception {
225+
Part filePart = mockPart("file.bin", "application/octet-stream", new byte[] {1, 2, 3});
226+
RequestContext reqCtx = mockReqCtx(null, mock(TraceSegment.class));
227+
BiFunction<RequestContext, List<String>, Flow<Void>> contentCb = mockPassThroughCb();
228+
229+
assertFalse(
230+
GlassFishBlockingHelper.processPartsAndBlock(
231+
Collections.singletonList(filePart), reqCtx, null, null, null, contentCb));
232+
233+
verify(contentCb).apply(eq(reqCtx), any());
234+
}
235+
236+
@Test
237+
void processPartsAndBlock_maxFilesLimit_enforced() throws Exception {
238+
int limit = GlassFishBlockingHelper.MAX_FILE_CONTENT_COUNT;
239+
Part[] tooMany = new Part[limit + 1];
240+
for (int i = 0; i <= limit; i++) {
241+
tooMany[i] = mockPart("f" + i + ".bin", "application/octet-stream", new byte[0]);
242+
}
243+
RequestContext reqCtx = mockReqCtx(null, mock(TraceSegment.class));
244+
BiFunction<RequestContext, List<String>, Flow<Void>> contentCb = mockPassThroughCb();
245+
246+
assertFalse(
247+
GlassFishBlockingHelper.processPartsAndBlock(
248+
Arrays.asList(tooMany), reqCtx, null, null, null, contentCb));
249+
250+
verify(contentCb).apply(eq(reqCtx), any(List.class));
251+
verify(tooMany[limit], never()).getInputStream();
252+
}
253+
254+
@Test
255+
@SuppressWarnings("unchecked")
256+
void processPartsAndBlock_getInputStreamThrows_emptyStringFallback() throws Exception {
257+
Part filePart = mock(Part.class);
258+
when(filePart.getSubmittedFileName()).thenReturn("bad.bin");
259+
when(filePart.getInputStream()).thenThrow(new IOException("disk error"));
260+
RequestContext reqCtx = mockReqCtx(null, mock(TraceSegment.class));
261+
BiFunction<RequestContext, List<String>, Flow<Void>> contentCb = mockPassThroughCb();
262+
263+
assertFalse(
264+
GlassFishBlockingHelper.processPartsAndBlock(
265+
Collections.singletonList(filePart), reqCtx, null, null, null, contentCb));
266+
267+
verify(contentCb).apply(eq(reqCtx), eq(Collections.singletonList("")));
268+
}
269+
270+
@Test
271+
void processPartsAndBlock_filenamesCbBlocks_contentCbNotFired() throws Exception {
272+
Part filePart = mockPart("evil.exe", "application/octet-stream", "content".getBytes());
273+
TraceSegment segment = mock(TraceSegment.class);
274+
RequestContext reqCtx = mockReqCtx(null, segment);
275+
TestServletOutputStream out = new TestServletOutputStream();
276+
HttpServletResponse resp = mock(HttpServletResponse.class);
277+
when(resp.isCommitted()).thenReturn(false);
278+
when(resp.getOutputStream()).thenReturn(out);
279+
BiFunction<RequestContext, List<String>, Flow<Void>> filenamesCb = mockBlockingCb(403);
280+
BiFunction<RequestContext, List<String>, Flow<Void>> contentCb = mockPassThroughCb();
281+
282+
assertTrue(
283+
GlassFishBlockingHelper.processPartsAndBlock(
284+
Collections.singletonList(filePart), reqCtx, null, resp, filenamesCb, contentCb));
285+
286+
verify(contentCb, never()).apply(any(), any());
287+
}
288+
289+
@Test
290+
void processPartsAndBlock_contentCbBlocks_returnsTrue() throws Exception {
291+
Part filePart = mockPart("upload.bin", "application/octet-stream", "payload".getBytes());
292+
TraceSegment segment = mock(TraceSegment.class);
293+
RequestContext reqCtx = mockReqCtx(null, segment);
294+
TestServletOutputStream out = new TestServletOutputStream();
295+
HttpServletResponse resp = mock(HttpServletResponse.class);
296+
when(resp.isCommitted()).thenReturn(false);
297+
when(resp.getOutputStream()).thenReturn(out);
298+
BiFunction<RequestContext, List<String>, Flow<Void>> filenamesCb = mockPassThroughCb();
299+
BiFunction<RequestContext, List<String>, Flow<Void>> contentCb = mockBlockingCb(403);
300+
301+
assertTrue(
302+
GlassFishBlockingHelper.processPartsAndBlock(
303+
Collections.singletonList(filePart), reqCtx, null, resp, filenamesCb, contentCb));
304+
}
305+
306+
@Test
307+
void processPartsAndBlock_nonPartObject_skipped() {
308+
RequestContext reqCtx = mockReqCtx(null, mock(TraceSegment.class));
309+
BiFunction<RequestContext, List<String>, Flow<Void>> filenamesCb = mockPassThroughCb();
310+
311+
assertFalse(
312+
GlassFishBlockingHelper.processPartsAndBlock(
313+
Collections.singletonList("not-a-part"), reqCtx, null, null, filenamesCb, null));
314+
315+
verify(filenamesCb, never()).apply(any(), any());
316+
}
317+
173318
// ------- Helpers -------
174319

175320
private static Flow.Action.RequestBlockingAction rba(int statusCode) {
@@ -183,6 +328,35 @@ private static RequestContext mockReqCtx(BlockResponseFunction brf, TraceSegment
183328
return reqCtx;
184329
}
185330

331+
private static Part mockPart(String submittedFilename, String contentType, byte[] content)
332+
throws Exception {
333+
Part part = mock(Part.class);
334+
when(part.getSubmittedFileName()).thenReturn(submittedFilename);
335+
when(part.getContentType()).thenReturn(contentType);
336+
when(part.getInputStream()).thenAnswer(ignored -> new java.io.ByteArrayInputStream(content));
337+
return part;
338+
}
339+
340+
@SuppressWarnings("unchecked")
341+
private static BiFunction<RequestContext, List<String>, Flow<Void>> mockPassThroughCb() {
342+
BiFunction<RequestContext, List<String>, Flow<Void>> cb = mock(BiFunction.class);
343+
Flow<Void> flow = mock(Flow.class);
344+
when(flow.getAction()).thenReturn(Flow.Action.Noop.INSTANCE);
345+
when(cb.apply(any(), any())).thenReturn(flow);
346+
return cb;
347+
}
348+
349+
@SuppressWarnings("unchecked")
350+
private static BiFunction<RequestContext, List<String>, Flow<Void>> mockBlockingCb(
351+
int statusCode) {
352+
BiFunction<RequestContext, List<String>, Flow<Void>> cb = mock(BiFunction.class);
353+
Flow<Void> flow = mock(Flow.class);
354+
when(flow.getAction())
355+
.thenReturn(new Flow.Action.RequestBlockingAction(statusCode, BlockingContentType.AUTO));
356+
when(cb.apply(any(), any())).thenReturn(flow);
357+
return cb;
358+
}
359+
186360
private static final class TestServletOutputStream extends ServletOutputStream {
187361
private final ByteArrayOutputStream buffer = new ByteArrayOutputStream();
188362

0 commit comments

Comments
 (0)