Skip to content

Commit a45b404

Browse files
jandro996jandro996
committed
refactor(appsec): extract tryBlock() helper and add unit tests for GlassFishBlockingHelper
Extract the repeated brf/fallback blocking pattern from the GlassFish multipart advice into GlassFishBlockingHelper.tryBlock(). This eliminates the duplication between the filenames and content blocking paths, and fixes the ordering issue where blocked=true was assigned after effectivelyBlocked() — if effectivelyBlocked() threw, blocked would stay false and the content callback would fire after filenames already blocked. tryBlock() uses a two-phase try/catch: the outer block handles commit failures (returning false), while the inner block wraps effectivelyBlocked() so a thrown exception does not suppress the true return value when the response was already committed. Add GlassFishBlockingHelperTest covering all branches of commitBlocking() and tryBlock().
1 parent aaac59e commit a45b404

4 files changed

Lines changed: 255 additions & 18 deletions

File tree

‎dd-java-agent/instrumentation/tomcat/tomcat-appsec/tomcat-appsec-7.0/build.gradle‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ dependencies {
3636

3737
testImplementation group: 'org.apache.tomcat', name: 'tomcat-catalina', version: '7.0.4'
3838
testImplementation group: 'org.apache.tomcat', name: 'tomcat-coyote', version: '7.0.4'
39+
testImplementation group: 'javax.servlet', name: 'javax.servlet-api', version: '3.1.0'
40+
testImplementation libs.bundles.mockito
3941
}
4042

4143
// testing happens in tomcat-5.5 module

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@
22

33
import datadog.appsec.api.blocking.BlockingContentType;
44
import datadog.trace.api.Config;
5+
import datadog.trace.api.gateway.BlockResponseFunction;
56
import datadog.trace.api.gateway.Flow;
7+
import datadog.trace.api.gateway.RequestContext;
68
import datadog.trace.bootstrap.blocking.BlockingActionHelper;
79
import java.util.Map;
810
import javax.servlet.http.HttpServletRequest;
@@ -13,6 +15,39 @@ public final class GlassFishBlockingHelper {
1315
public static final int MAX_FILE_CONTENT_COUNT = Config.get().getAppSecMaxFileContentCount();
1416
public static final int MAX_FILE_CONTENT_BYTES = Config.get().getAppSecMaxFileContentBytes();
1517

18+
/**
19+
* Attempts to commit a blocking response via the registered {@link BlockResponseFunction} or via
20+
* the Servlet API fallback, then marks the trace segment as effectively blocked.
21+
*
22+
* <p>Returns {@code true} if the response was committed (regardless of whether {@link
23+
* datadog.trace.api.internal.TraceSegment#effectivelyBlocked()} succeeded). Returns {@code false}
24+
* if no response could be committed.
25+
*/
26+
public static boolean tryBlock(
27+
RequestContext reqCtx,
28+
HttpServletRequest fallbackReq,
29+
HttpServletResponse fallbackResp,
30+
Flow.Action.RequestBlockingAction rba) {
31+
try {
32+
BlockResponseFunction brf = reqCtx.getBlockResponseFunction();
33+
if (brf != null) {
34+
brf.tryCommitBlockingResponse(reqCtx.getTraceSegment(), rba);
35+
} else if (!commitBlocking(fallbackReq, fallbackResp, rba)) {
36+
return false;
37+
}
38+
} catch (Exception ignored) {
39+
return false;
40+
}
41+
// Response was committed — mark as blocked on a best-effort basis.
42+
// effectivelyBlocked() can throw if the span is already finished; that must not suppress the
43+
// true return value since the response has already been sent to the client.
44+
try {
45+
reqCtx.getTraceSegment().effectivelyBlocked();
46+
} catch (Exception ignored) {
47+
}
48+
return true;
49+
}
50+
1651
public static boolean commitBlocking(
1752
HttpServletRequest request,
1853
HttpServletResponse response,

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

Lines changed: 7 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
import com.google.auto.service.AutoService;
99
import datadog.trace.agent.tooling.Instrumenter;
1010
import datadog.trace.agent.tooling.InstrumenterModule;
11-
import datadog.trace.api.gateway.BlockResponseFunction;
1211
import datadog.trace.api.gateway.CallbackProvider;
1312
import datadog.trace.api.gateway.Flow;
1413
import datadog.trace.api.gateway.RequestContext;
@@ -173,16 +172,9 @@ static void after(
173172
Flow<Void> flow = filenamesCb.apply(reqCtx, filenames);
174173
Flow.Action action = flow.getAction();
175174
if (action instanceof Flow.Action.RequestBlockingAction) {
176-
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
177-
BlockResponseFunction brf = reqCtx.getBlockResponseFunction();
178-
if (brf != null) {
179-
brf.tryCommitBlockingResponse(reqCtx.getTraceSegment(), rba);
175+
if (GlassFishBlockingHelper.tryBlock(
176+
reqCtx, fallbackReq, fallbackResp, (Flow.Action.RequestBlockingAction) action)) {
180177
parts = Collections.emptyList();
181-
reqCtx.getTraceSegment().effectivelyBlocked();
182-
blocked = true;
183-
} else if (GlassFishBlockingHelper.commitBlocking(fallbackReq, fallbackResp, rba)) {
184-
parts = Collections.emptyList();
185-
reqCtx.getTraceSegment().effectivelyBlocked();
186178
blocked = true;
187179
}
188180
}
@@ -192,15 +184,12 @@ static void after(
192184
Flow<Void> contentFlow = contentCb.apply(reqCtx, contents);
193185
Flow.Action contentAction = contentFlow.getAction();
194186
if (contentAction instanceof Flow.Action.RequestBlockingAction) {
195-
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) contentAction;
196-
BlockResponseFunction brf = reqCtx.getBlockResponseFunction();
197-
if (brf != null) {
198-
brf.tryCommitBlockingResponse(reqCtx.getTraceSegment(), rba);
199-
parts = Collections.emptyList();
200-
reqCtx.getTraceSegment().effectivelyBlocked();
201-
} else if (GlassFishBlockingHelper.commitBlocking(fallbackReq, fallbackResp, rba)) {
187+
if (GlassFishBlockingHelper.tryBlock(
188+
reqCtx,
189+
fallbackReq,
190+
fallbackResp,
191+
(Flow.Action.RequestBlockingAction) contentAction)) {
202192
parts = Collections.emptyList();
203-
reqCtx.getTraceSegment().effectivelyBlocked();
204193
}
205194
}
206195
}

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

Lines changed: 211 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,211 @@
1+
package datadog.trace.instrumentation.tomcat7;
2+
3+
import static org.junit.jupiter.api.Assertions.assertFalse;
4+
import static org.junit.jupiter.api.Assertions.assertTrue;
5+
import static org.mockito.ArgumentMatchers.any;
6+
import static org.mockito.ArgumentMatchers.contains;
7+
import static org.mockito.ArgumentMatchers.eq;
8+
import static org.mockito.Mockito.doThrow;
9+
import static org.mockito.Mockito.mock;
10+
import static org.mockito.Mockito.never;
11+
import static org.mockito.Mockito.verify;
12+
import static org.mockito.Mockito.when;
13+
14+
import datadog.appsec.api.blocking.BlockingContentType;
15+
import datadog.trace.api.gateway.BlockResponseFunction;
16+
import datadog.trace.api.gateway.Flow;
17+
import datadog.trace.api.gateway.RequestContext;
18+
import datadog.trace.api.internal.TraceSegment;
19+
import java.io.ByteArrayOutputStream;
20+
import java.io.IOException;
21+
import javax.servlet.ServletOutputStream;
22+
import javax.servlet.WriteListener;
23+
import javax.servlet.http.HttpServletRequest;
24+
import javax.servlet.http.HttpServletResponse;
25+
import org.junit.jupiter.api.Test;
26+
27+
class GlassFishBlockingHelperTest {
28+
29+
// ------- commitBlocking() -------
30+
31+
@Test
32+
void commitBlocking_nullResponse_returnsFalse() {
33+
assertFalse(GlassFishBlockingHelper.commitBlocking(null, null, rba(403)));
34+
}
35+
36+
@Test
37+
void commitBlocking_committedResponse_returnsFalse() {
38+
HttpServletResponse resp = mock(HttpServletResponse.class);
39+
when(resp.isCommitted()).thenReturn(true);
40+
assertFalse(GlassFishBlockingHelper.commitBlocking(null, resp, rba(403)));
41+
}
42+
43+
@Test
44+
void commitBlocking_blockingContentTypeNone_setsStatusWithoutBody() throws IOException {
45+
HttpServletResponse resp = mock(HttpServletResponse.class);
46+
when(resp.isCommitted()).thenReturn(false);
47+
48+
assertTrue(
49+
GlassFishBlockingHelper.commitBlocking(
50+
null, resp, new Flow.Action.RequestBlockingAction(403, BlockingContentType.NONE)));
51+
52+
verify(resp).setStatus(403);
53+
verify(resp).flushBuffer();
54+
verify(resp, never()).setHeader(eq("Content-Type"), any());
55+
verify(resp, never()).getOutputStream();
56+
}
57+
58+
@Test
59+
void commitBlocking_withJsonAccept_writesJsonBody() throws IOException {
60+
HttpServletRequest req = mock(HttpServletRequest.class);
61+
when(req.getHeader("Accept")).thenReturn("application/json");
62+
TestServletOutputStream out = new TestServletOutputStream();
63+
HttpServletResponse resp = mock(HttpServletResponse.class);
64+
when(resp.isCommitted()).thenReturn(false);
65+
when(resp.getOutputStream()).thenReturn(out);
66+
67+
assertTrue(GlassFishBlockingHelper.commitBlocking(req, resp, rba(403)));
68+
69+
verify(resp).setStatus(403);
70+
verify(resp).setHeader(eq("Content-Type"), contains("json"));
71+
verify(resp).setHeader(eq("Content-Length"), any());
72+
assertTrue(out.getBytes().length > 0);
73+
verify(resp).flushBuffer();
74+
}
75+
76+
@Test
77+
void commitBlocking_withHtmlAccept_writesHtmlBody() throws IOException {
78+
HttpServletRequest req = mock(HttpServletRequest.class);
79+
when(req.getHeader("Accept")).thenReturn("text/html");
80+
TestServletOutputStream out = new TestServletOutputStream();
81+
HttpServletResponse resp = mock(HttpServletResponse.class);
82+
when(resp.isCommitted()).thenReturn(false);
83+
when(resp.getOutputStream()).thenReturn(out);
84+
85+
assertTrue(GlassFishBlockingHelper.commitBlocking(req, resp, rba(403)));
86+
87+
verify(resp).setHeader(eq("Content-Type"), contains("html"));
88+
assertTrue(out.getBytes().length > 0);
89+
}
90+
91+
@Test
92+
void commitBlocking_nullRequest_defaultsToJsonBody() throws IOException {
93+
TestServletOutputStream out = new TestServletOutputStream();
94+
HttpServletResponse resp = mock(HttpServletResponse.class);
95+
when(resp.isCommitted()).thenReturn(false);
96+
when(resp.getOutputStream()).thenReturn(out);
97+
98+
assertTrue(GlassFishBlockingHelper.commitBlocking(null, resp, rba(403)));
99+
100+
verify(resp).setStatus(403);
101+
assertTrue(out.getBytes().length > 0);
102+
}
103+
104+
@Test
105+
void commitBlocking_ioException_returnsFalse() throws IOException {
106+
HttpServletResponse resp = mock(HttpServletResponse.class);
107+
when(resp.isCommitted()).thenReturn(false);
108+
when(resp.getOutputStream()).thenThrow(new IOException("stream error"));
109+
110+
assertFalse(GlassFishBlockingHelper.commitBlocking(null, resp, rba(403)));
111+
}
112+
113+
// ------- tryBlock() -------
114+
115+
@Test
116+
void tryBlock_withBrf_commitsViaFunctionAndReturnsTrue() throws Exception {
117+
TraceSegment segment = mock(TraceSegment.class);
118+
BlockResponseFunction brf = mock(BlockResponseFunction.class);
119+
RequestContext reqCtx = mockReqCtx(brf, segment);
120+
121+
Flow.Action.RequestBlockingAction action = rba(403);
122+
assertTrue(GlassFishBlockingHelper.tryBlock(reqCtx, null, null, action));
123+
124+
verify(brf).tryCommitBlockingResponse(segment, action);
125+
verify(segment).effectivelyBlocked();
126+
}
127+
128+
@Test
129+
void tryBlock_noBrf_fallbackSucceeds_returnsTrue() throws IOException {
130+
TraceSegment segment = mock(TraceSegment.class);
131+
RequestContext reqCtx = mockReqCtx(null, segment);
132+
TestServletOutputStream out = new TestServletOutputStream();
133+
HttpServletResponse resp = mock(HttpServletResponse.class);
134+
when(resp.isCommitted()).thenReturn(false);
135+
when(resp.getOutputStream()).thenReturn(out);
136+
137+
assertTrue(GlassFishBlockingHelper.tryBlock(reqCtx, null, resp, rba(403)));
138+
verify(segment).effectivelyBlocked();
139+
}
140+
141+
@Test
142+
void tryBlock_noBrf_nullFallbackResponse_returnsFalse() {
143+
RequestContext reqCtx = mock(RequestContext.class);
144+
when(reqCtx.getBlockResponseFunction()).thenReturn(null);
145+
146+
assertFalse(GlassFishBlockingHelper.tryBlock(reqCtx, null, null, rba(403)));
147+
verify(reqCtx, never()).getTraceSegment();
148+
}
149+
150+
@Test
151+
void tryBlock_brfThrows_returnsFalse() throws Exception {
152+
TraceSegment segment = mock(TraceSegment.class);
153+
BlockResponseFunction brf = mock(BlockResponseFunction.class);
154+
RequestContext reqCtx = mockReqCtx(brf, segment);
155+
doThrow(new RuntimeException("commit failed"))
156+
.when(brf)
157+
.tryCommitBlockingResponse(any(), any(Flow.Action.RequestBlockingAction.class));
158+
159+
assertFalse(GlassFishBlockingHelper.tryBlock(reqCtx, null, null, rba(403)));
160+
verify(segment, never()).effectivelyBlocked();
161+
}
162+
163+
@Test
164+
void tryBlock_effectivelyBlockedThrows_stillReturnsTrue() throws Exception {
165+
TraceSegment segment = mock(TraceSegment.class);
166+
BlockResponseFunction brf = mock(BlockResponseFunction.class);
167+
RequestContext reqCtx = mockReqCtx(brf, segment);
168+
doThrow(new RuntimeException("span already finished")).when(segment).effectivelyBlocked();
169+
170+
assertTrue(GlassFishBlockingHelper.tryBlock(reqCtx, null, null, rba(403)));
171+
}
172+
173+
// ------- Helpers -------
174+
175+
private static Flow.Action.RequestBlockingAction rba(int statusCode) {
176+
return new Flow.Action.RequestBlockingAction(statusCode, BlockingContentType.AUTO);
177+
}
178+
179+
private static RequestContext mockReqCtx(BlockResponseFunction brf, TraceSegment segment) {
180+
RequestContext reqCtx = mock(RequestContext.class);
181+
when(reqCtx.getBlockResponseFunction()).thenReturn(brf);
182+
when(reqCtx.getTraceSegment()).thenReturn(segment);
183+
return reqCtx;
184+
}
185+
186+
private static final class TestServletOutputStream extends ServletOutputStream {
187+
private final ByteArrayOutputStream buffer = new ByteArrayOutputStream();
188+
189+
@Override
190+
public boolean isReady() {
191+
return true;
192+
}
193+
194+
@Override
195+
public void setWriteListener(WriteListener listener) {}
196+
197+
@Override
198+
public void write(int b) throws IOException {
199+
buffer.write(b);
200+
}
201+
202+
@Override
203+
public void write(byte[] b, int off, int len) throws IOException {
204+
buffer.write(b, off, len);
205+
}
206+
207+
public byte[] getBytes() {
208+
return buffer.toByteArray();
209+
}
210+
}
211+
}

0 commit comments

Comments
 (0)