Skip to content

Commit f4163fa

Browse files
jandro996jandro996
authored
Extend support for SSRF in exploit prevention (#7376)
What Does This Do Add SSRF exploit prevention check to HttpClientDecorator Modify http client instrumentations that relay on CallDepthThreadLocalMap to avoid issues with blocking exception (only blocks the first time) -> Fix call depth counter for sqli blocking #7522 Add smoke tests for other libraries Motivation improve Exploit prevention for SSRF coverage
1 parent e5aa8fa commit f4163fa

File tree

14 files changed

+562
-87
lines changed

14 files changed

+562
-87
lines changed

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

Lines changed: 65 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,24 @@
11
package datadog.trace.bootstrap.instrumentation.decorator;
22

33
import static datadog.trace.api.cache.RadixTreeCache.UNSET_STATUS;
4+
import static datadog.trace.api.gateway.Events.EVENTS;
45
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
56
import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;
67

8+
import datadog.appsec.api.blocking.BlockingException;
79
import datadog.trace.api.Config;
810
import datadog.trace.api.DDTags;
911
import datadog.trace.api.InstrumenterConfig;
1012
import datadog.trace.api.ProductActivation;
13+
import datadog.trace.api.gateway.BlockResponseFunction;
14+
import datadog.trace.api.gateway.Flow;
15+
import datadog.trace.api.gateway.RequestContext;
16+
import datadog.trace.api.gateway.RequestContextSlot;
1117
import datadog.trace.api.iast.InstrumentationBridge;
1218
import datadog.trace.api.iast.sink.SsrfModule;
1319
import datadog.trace.api.naming.SpanNaming;
1420
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
21+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
1522
import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
1623
import datadog.trace.bootstrap.instrumentation.api.Tags;
1724
import datadog.trace.bootstrap.instrumentation.api.URIUtils;
@@ -21,6 +28,7 @@
2128
import java.util.BitSet;
2229
import java.util.LinkedHashMap;
2330
import java.util.Map;
31+
import java.util.function.BiFunction;
2432
import org.slf4j.Logger;
2533
import org.slf4j.LoggerFactory;
2634

@@ -42,6 +50,8 @@ public abstract class HttpClientDecorator<REQUEST, RESPONSE> extends UriBasedCli
4250

4351
private static final boolean CLIENT_TAG_HEADERS = Config.get().isHttpClientTagHeaders();
4452

53+
private static final boolean APPSEC_RASP_ENABLED = Config.get().isAppSecRaspEnabled();
54+
4555
protected abstract String method(REQUEST request);
4656

4757
protected abstract URI url(REQUEST request) throws URISyntaxException;
@@ -68,9 +78,20 @@ protected boolean shouldSetResourceName() {
6878

6979
public AgentSpan onRequest(final AgentSpan span, final REQUEST request) {
7080
if (request != null) {
81+
7182
String method = method(request);
7283
span.setTag(Tags.HTTP_METHOD, method);
7384

85+
if (CLIENT_TAG_HEADERS) {
86+
for (Map.Entry<String, String> headerTag :
87+
traceConfig(span).getRequestHeaderTags().entrySet()) {
88+
String headerValue = getRequestHeader(request, headerTag.getKey());
89+
if (null != headerValue) {
90+
span.setTag(headerTag.getValue(), headerValue);
91+
}
92+
}
93+
}
94+
7495
// Copy of HttpServerDecorator url handling
7596
try {
7697
final URI url = url(request);
@@ -86,24 +107,15 @@ public AgentSpan onRequest(final AgentSpan span, final REQUEST request) {
86107
if (shouldSetResourceName()) {
87108
HTTP_RESOURCE_DECORATOR.withClientPath(span, method, url.getPath());
88109
}
110+
// SSRF exploit prevention check
111+
onNetworkConnection(url.toString());
89112
} else if (shouldSetResourceName()) {
90113
span.setResourceName(DEFAULT_RESOURCE_NAME);
91114
}
92115
} catch (final Exception e) {
93116
log.debug("Error tagging url", e);
94117
}
95-
96118
ssrfIastCheck(request);
97-
98-
if (CLIENT_TAG_HEADERS) {
99-
for (Map.Entry<String, String> headerTag :
100-
traceConfig(span).getRequestHeaderTags().entrySet()) {
101-
String headerValue = getRequestHeader(request, headerTag.getKey());
102-
if (null != headerValue) {
103-
span.setTag(headerTag.getValue(), headerValue);
104-
}
105-
}
106-
}
107119
}
108120
return span;
109121
}
@@ -175,6 +187,48 @@ public long getResponseContentLength(final RESPONSE response) {
175187
return 0;
176188
}
177189

190+
private void onNetworkConnection(final String networkConnection) {
191+
if (!APPSEC_RASP_ENABLED) {
192+
return;
193+
}
194+
if (networkConnection == null) {
195+
return;
196+
}
197+
final BiFunction<RequestContext, String, Flow<Void>> networkConnectionCallback =
198+
AgentTracer.get()
199+
.getCallbackProvider(RequestContextSlot.APPSEC)
200+
.getCallback(EVENTS.networkConnection());
201+
202+
if (networkConnectionCallback == null) {
203+
return;
204+
}
205+
206+
final AgentSpan span = AgentTracer.get().activeSpan();
207+
if (span == null) {
208+
return;
209+
}
210+
211+
final RequestContext ctx = span.getRequestContext();
212+
if (ctx == null) {
213+
return;
214+
}
215+
216+
Flow<Void> flow = networkConnectionCallback.apply(ctx, networkConnection);
217+
Flow.Action action = flow.getAction();
218+
if (action instanceof Flow.Action.RequestBlockingAction) {
219+
BlockResponseFunction brf = ctx.getBlockResponseFunction();
220+
if (brf != null) {
221+
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
222+
brf.tryCommitBlockingResponse(
223+
ctx.getTraceSegment(),
224+
rba.getStatusCode(),
225+
rba.getBlockingContentType(),
226+
rba.getExtraHeaders());
227+
}
228+
throw new BlockingException("Blocked request (for SSRF attempt)");
229+
}
230+
}
231+
178232
/* This method must be overriden after making the proper propagations to the client before **/
179233
protected Object sourceUrl(REQUEST request) {
180234
return null;

‎dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/decorator/HttpClientDecoratorTest.groovy‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,12 @@
11
package datadog.trace.bootstrap.instrumentation.decorator
22

33
import datadog.trace.api.DDTags
4+
import datadog.trace.api.config.AppSecConfig
5+
import datadog.trace.api.gateway.CallbackProvider
6+
import static datadog.trace.api.gateway.Events.EVENTS
7+
import datadog.trace.api.gateway.RequestContext
8+
import datadog.trace.api.gateway.RequestContextSlot
9+
import datadog.trace.api.internal.TraceSegment
410
import datadog.trace.api.iast.InstrumentationBridge
511
import datadog.trace.api.iast.sink.SsrfModule
612
import datadog.trace.bootstrap.instrumentation.api.AgentSpan
@@ -12,11 +18,39 @@ import spock.lang.Shared
1218
import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN
1319
import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_CLIENT_TAG_QUERY_STRING
1420

21+
import java.util.function.BiFunction
22+
1523
class HttpClientDecoratorTest extends ClientDecoratorTest {
1624

1725
@Shared
1826
def testUrl = new URI("http://myhost:123/somepath")
1927

28+
@Shared
29+
protected static final ORIGINAL_TRACER = AgentTracer.get()
30+
31+
protected traceSegment
32+
protected reqCtx
33+
protected span2
34+
protected tracer
35+
36+
void setup() {
37+
traceSegment = Stub(TraceSegment)
38+
reqCtx = Stub(RequestContext) {
39+
getTraceSegment() >> traceSegment
40+
}
41+
span2 = Stub(AgentSpan) {
42+
getRequestContext() >> reqCtx
43+
}
44+
tracer = Stub(AgentTracer.TracerAPI) {
45+
activeSpan() >> span2
46+
}
47+
AgentTracer.forceRegister(tracer)
48+
}
49+
50+
void cleanup() {
51+
AgentTracer.forceRegister(ORIGINAL_TRACER)
52+
}
53+
2054
def span = Mock(AgentSpan)
2155

2256
def "test onRequest"() {
@@ -199,6 +233,26 @@ class HttpClientDecoratorTest extends ClientDecoratorTest {
199233
'false' | [method: "test-method", url: testUrl, path: '/somepath']
200234
}
201235

236+
void "test SSRF Exploit prevention onResponse"(){
237+
setup:
238+
injectSysConfig(AppSecConfig.APPSEC_ENABLED, 'true')
239+
injectSysConfig(AppSecConfig.APPSEC_RASP_ENABLED, 'true')
240+
241+
final callbackProvider = Mock(CallbackProvider)
242+
final listener = Mock(BiFunction)
243+
tracer.getCallbackProvider(RequestContextSlot.APPSEC) >> callbackProvider
244+
245+
final decorator = newDecorator()
246+
final req = [method: "GET", url: new URI("http://localhost:1234/somepath")]
247+
248+
when:
249+
decorator.onRequest(span2, req)
250+
251+
then:
252+
1 * callbackProvider.getCallback(EVENTS.networkConnection()) >> listener
253+
1 * listener.apply(reqCtx, _ as String)
254+
}
255+
202256
@Override
203257
def newDecorator(String serviceName = "test-service") {
204258
return new HttpClientDecorator<Map, Map>() {

‎dd-java-agent/instrumentation/apache-httpclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/ApacheHttpClientInstrumentation.java‎

Lines changed: 48 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
88

99
import com.google.auto.service.AutoService;
10+
import datadog.appsec.api.blocking.BlockingException;
1011
import datadog.trace.agent.tooling.Instrumenter;
1112
import datadog.trace.agent.tooling.InstrumenterModule;
1213
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
@@ -153,7 +154,13 @@ public void methodAdvice(MethodTransformer transformer) {
153154
public static class UriRequestAdvice {
154155
@Advice.OnMethodEnter(suppress = Throwable.class)
155156
public static AgentScope methodEnter(@Advice.Argument(0) final HttpUriRequest request) {
156-
return HelperMethods.doMethodEnter(request);
157+
try {
158+
return HelperMethods.doMethodEnter(request);
159+
} catch (BlockingException e) {
160+
HelperMethods.onBlockingRequest();
161+
// re-throw blocking exceptions
162+
throw e;
163+
}
157164
}
158165

159166
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
@@ -176,12 +183,20 @@ public static AgentScope methodEnter(
176183
typing = Assigner.Typing.DYNAMIC,
177184
readOnly = false)
178185
Object handler) {
179-
final AgentScope scope = HelperMethods.doMethodEnter(request);
180-
// Wrap the handler so we capture the status code
181-
if (null != scope && handler instanceof ResponseHandler) {
182-
handler = new WrappingStatusSettingResponseHandler(scope.span(), (ResponseHandler) handler);
186+
187+
try {
188+
final AgentScope scope = HelperMethods.doMethodEnter(request);
189+
// Wrap the handler so we capture the status code
190+
if (null != scope && handler instanceof ResponseHandler) {
191+
handler =
192+
new WrappingStatusSettingResponseHandler(scope.span(), (ResponseHandler) handler);
193+
}
194+
return scope;
195+
} catch (BlockingException e) {
196+
HelperMethods.onBlockingRequest();
197+
// re-throw blocking exceptions
198+
throw e;
183199
}
184-
return scope;
185200
}
186201

187202
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
@@ -197,10 +212,16 @@ public static class RequestAdvice {
197212
@Advice.OnMethodEnter(suppress = Throwable.class)
198213
public static AgentScope methodEnter(
199214
@Advice.Argument(0) final HttpHost host, @Advice.Argument(1) final HttpRequest request) {
200-
if (request instanceof HttpUriRequest) {
201-
return HelperMethods.doMethodEnter((HttpUriRequest) request);
202-
} else {
203-
return HelperMethods.doMethodEnter(host, request);
215+
try {
216+
if (request instanceof HttpUriRequest) {
217+
return HelperMethods.doMethodEnter((HttpUriRequest) request);
218+
} else {
219+
return HelperMethods.doMethodEnter(host, request);
220+
}
221+
} catch (BlockingException e) {
222+
HelperMethods.onBlockingRequest();
223+
// re-throw blocking exceptions
224+
throw e;
204225
}
205226
}
206227

@@ -225,17 +246,24 @@ public static AgentScope methodEnter(
225246
typing = Assigner.Typing.DYNAMIC,
226247
readOnly = false)
227248
Object handler) {
228-
final AgentScope scope;
229-
if (request instanceof HttpUriRequest) {
230-
scope = HelperMethods.doMethodEnter((HttpUriRequest) request);
231-
} else {
232-
scope = HelperMethods.doMethodEnter(host, request);
233-
}
234-
// Wrap the handler so we capture the status code
235-
if (null != scope && handler instanceof ResponseHandler) {
236-
handler = new WrappingStatusSettingResponseHandler(scope.span(), (ResponseHandler) handler);
249+
try {
250+
final AgentScope scope;
251+
if (request instanceof HttpUriRequest) {
252+
scope = HelperMethods.doMethodEnter((HttpUriRequest) request);
253+
} else {
254+
scope = HelperMethods.doMethodEnter(host, request);
255+
}
256+
// Wrap the handler so we capture the status code
257+
if (null != scope && handler instanceof ResponseHandler) {
258+
handler =
259+
new WrappingStatusSettingResponseHandler(scope.span(), (ResponseHandler) handler);
260+
}
261+
return scope;
262+
} catch (BlockingException e) {
263+
HelperMethods.onBlockingRequest();
264+
// re-throw blocking exceptions
265+
throw e;
237266
}
238-
return scope;
239267
}
240268

241269
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

‎dd-java-agent/instrumentation/apache-httpclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,4 +73,8 @@ public static void doMethodExit(
7373
CallDepthThreadLocalMap.reset(HttpClient.class);
7474
}
7575
}
76+
77+
public static void onBlockingRequest() {
78+
CallDepthThreadLocalMap.reset(HttpClient.class);
79+
}
7680
}

‎dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpClientInstrumentation.java‎

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
88

99
import com.google.auto.service.AutoService;
10+
import datadog.appsec.api.blocking.BlockingException;
1011
import datadog.trace.agent.tooling.Instrumenter;
1112
import datadog.trace.agent.tooling.InstrumenterModule;
1213
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
@@ -115,7 +116,13 @@ public void methodAdvice(MethodTransformer transformer) {
115116
public static class RequestAdvice {
116117
@Advice.OnMethodEnter(suppress = Throwable.class)
117118
public static AgentScope methodEnter(@Advice.Argument(0) final ClassicHttpRequest request) {
118-
return HelperMethods.doMethodEnter(request);
119+
try {
120+
return HelperMethods.doMethodEnter(request);
121+
} catch (BlockingException e) {
122+
HelperMethods.onBlockingRequest();
123+
// re-throw blocking exceptions
124+
throw e;
125+
}
119126
}
120127

121128
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
@@ -132,7 +139,13 @@ public static class HostRequestAdvice {
132139
public static AgentScope methodEnter(
133140
@Advice.Argument(0) final HttpHost host,
134141
@Advice.Argument(1) final ClassicHttpRequest request) {
135-
return HelperMethods.doMethodEnter(host, request);
142+
try {
143+
return HelperMethods.doMethodEnter(host, request);
144+
} catch (BlockingException e) {
145+
HelperMethods.onBlockingRequest();
146+
// re-throw blocking exceptions
147+
throw e;
148+
}
136149
}
137150

138151
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
@@ -157,14 +170,20 @@ public static AgentScope methodEnter(
157170
typing = Assigner.Typing.DYNAMIC,
158171
readOnly = false)
159172
Object handler) {
160-
final AgentScope scope = HelperMethods.doMethodEnter(host, request);
161-
// Wrap the handler so we capture the status code
162-
if (null != scope && handler instanceof HttpClientResponseHandler) {
163-
handler =
164-
new WrappingStatusSettingResponseHandler(
165-
scope.span(), (HttpClientResponseHandler) handler);
173+
try {
174+
final AgentScope scope = HelperMethods.doMethodEnter(host, request);
175+
// Wrap the handler so we capture the status code
176+
if (null != scope && handler instanceof HttpClientResponseHandler) {
177+
handler =
178+
new WrappingStatusSettingResponseHandler(
179+
scope.span(), (HttpClientResponseHandler) handler);
180+
}
181+
return scope;
182+
} catch (BlockingException e) {
183+
HelperMethods.onBlockingRequest();
184+
// re-throw blocking exceptions
185+
throw e;
166186
}
167-
return scope;
168187
}
169188

170189
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

0 commit comments

Comments
 (0)