Skip to content

Commit fde6afa

Browse files
DDJavierSantosDDJavierSantos
committed
# This is a combination of 2 commits.
# This is the 1st commit message: Resteasy IAST instrumentation # This is the commit message #2: Addressing git comments
1 parent db10955 commit fde6afa

File tree

20 files changed

+925
-0
lines changed

20 files changed

+925
-0
lines changed

‎dd-java-agent/agent-iast/src/main/java/com/datadog/iast/source/WebModuleImpl.java

+101
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,15 @@
66
import com.datadog.iast.model.Source;
77
import com.datadog.iast.model.SourceType;
88
import com.datadog.iast.taint.Ranges;
9+
import com.datadog.iast.taint.TaintedObject;
910
import com.datadog.iast.taint.TaintedObjects;
1011
import datadog.trace.api.iast.source.WebModule;
1112
import java.io.BufferedReader;
1213
import java.io.InputStream;
14+
import java.util.Collection;
15+
import java.util.List;
16+
import java.util.Map;
17+
import java.util.Set;
1318
import javax.annotation.Nonnull;
1419
import javax.annotation.Nullable;
1520

@@ -131,4 +136,100 @@ public void onHeaderValue(@Nullable final String headerName, @Nullable final Str
131136
taintedObjects.taintInputString(
132137
headerValue, new Source(SourceType.REQUEST_HEADER_VALUE, headerName, headerValue));
133138
}
139+
140+
@Override
141+
public void onCookieValue(@Nullable final String cookieName, @Nullable final String cookieValue) {
142+
if (!canBeTainted(cookieName)) {
143+
return;
144+
}
145+
final IastRequestContext ctx = IastRequestContext.get();
146+
if (ctx == null) {
147+
return;
148+
}
149+
final TaintedObjects taintedObjects = ctx.getTaintedObjects();
150+
taintedObjects.taintInputString(
151+
cookieValue, new Source(SourceType.REQUEST_COOKIE_VALUE, cookieName, cookieValue));
152+
}
153+
154+
@Override
155+
public void onJaxGetQueryParameters(@Nonnull Object multiValuedMap) {
156+
final IastRequestContext ctx = IastRequestContext.get();
157+
if (ctx == null) {
158+
return;
159+
}
160+
final TaintedObjects taintedObjects = ctx.getTaintedObjects();
161+
taintedObjects.taint(multiValuedMap, Ranges.EMPTY);
162+
}
163+
164+
@Override
165+
public void onMultiValuedMapAccess(
166+
@Nonnull Object multiValuedMap, @Nonnull Object key, @Nonnull Object returnedValue) {
167+
final IastRequestContext ctx = IastRequestContext.get();
168+
if (ctx == null) {
169+
return;
170+
}
171+
final TaintedObjects taintedObjects = ctx.getTaintedObjects();
172+
TaintedObject taintedMultiValuedMap = taintedObjects.get(multiValuedMap);
173+
if (null != taintedMultiValuedMap && key instanceof String) {
174+
if (returnedValue instanceof String) {
175+
taintedObjects.taintInputString(
176+
(String) returnedValue,
177+
new Source(SourceType.REQUEST_PARAMETER_VALUE, (String) key, (String) returnedValue));
178+
} else if (returnedValue instanceof List) {
179+
for (Object o : ((List) returnedValue)) {
180+
if (o instanceof String) {
181+
taintedObjects.taintInputString(
182+
(String) o,
183+
new Source(SourceType.REQUEST_PARAMETER_VALUE, (String) key, (String) o));
184+
}
185+
}
186+
}
187+
}
188+
}
189+
190+
@Override
191+
public void taintSetIfInputIsTainted(@Nonnull Set toTaint, @Nullable Object input) {
192+
final IastRequestContext ctx = IastRequestContext.get();
193+
if (ctx == null) {
194+
return;
195+
}
196+
final TaintedObjects taintedObjects = ctx.getTaintedObjects();
197+
TaintedObject tainted = taintedObjects.get(input);
198+
if (tainted != null) {
199+
for (Object o : toTaint) {
200+
if (o instanceof String) {
201+
taintedObjects.taintInputString(
202+
(String) o, new Source(SourceType.REQUEST_PARAMETER_VALUE, null, (String) o));
203+
} else if (o instanceof Map.Entry) {
204+
Map.Entry entry = (Map.Entry) o;
205+
if (entry.getKey() instanceof String && entry.getValue() instanceof String) {
206+
taintedObjects.taintInputString(
207+
(String) entry.getValue(),
208+
new Source(
209+
SourceType.REQUEST_PARAMETER_VALUE,
210+
(String) entry.getKey(),
211+
(String) entry.getValue()));
212+
}
213+
}
214+
}
215+
}
216+
}
217+
218+
@Override
219+
public void taintCollectionIfInputIsTainted(@Nonnull Collection toTaint, @Nullable Object input) {
220+
final IastRequestContext ctx = IastRequestContext.get();
221+
if (ctx == null) {
222+
return;
223+
}
224+
final TaintedObjects taintedObjects = ctx.getTaintedObjects();
225+
TaintedObject tainted = taintedObjects.get(input);
226+
if (tainted != null) {
227+
for (Object o : toTaint) {
228+
if (o instanceof String) {
229+
taintedObjects.taintInputString(
230+
(String) o, new Source(SourceType.REQUEST_PARAMETER_VALUE, null, (String) o));
231+
}
232+
}
233+
}
234+
}
134235
}

‎dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/source/WebModuleTest.groovy

+24
Original file line numberDiff line numberDiff line change
@@ -346,4 +346,28 @@ class WebModuleTest extends IastModuleImplTestBase {
346346
def to = ctx.getTaintedObjects().get(reader)
347347
assert to != null
348348
}
349+
350+
void 'test taintSetIfInputIsTainted'() {
351+
given:
352+
final span = Mock(AgentSpan)
353+
tracer.activeSpan() >> span
354+
final reqCtx = Mock(RequestContext)
355+
span.getRequestContext() >> reqCtx
356+
final ctx = new IastRequestContext()
357+
reqCtx.getData(RequestContextSlot.IAST) >> ctx
358+
def multiValuedMap = new Object()
359+
def retValue = new String()
360+
361+
when:
362+
module.taintSetIfInputIsTainted(retValue, multiValuedMap)
363+
364+
then:
365+
1 * tracer.activeSpan() >> span
366+
1 * span.getRequestContext() >> reqCtx
367+
1 * reqCtx.getData(RequestContextSlot.IAST) >> ctx
368+
0 * _
369+
370+
def to = ctx.getTaintedObjects().get(retValue)
371+
assert to != null
372+
}
349373
}

‎dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/CookieParamInjectorAdvice.java

+38
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package datadog.trace.instrumentation.resteasy;
2+
3+
import datadog.trace.api.iast.InstrumentationBridge;
4+
import datadog.trace.api.iast.source.WebModule;
5+
import java.util.Collection;
6+
import net.bytebuddy.asm.Advice;
7+
import org.jboss.resteasy.core.CookieParamInjector;
8+
9+
public class CookieParamInjectorAdvice {
10+
@Advice.OnMethodExit
11+
public static void onExit(
12+
@Advice.This CookieParamInjector self,
13+
@Advice.Return(readOnly = true) Object result,
14+
@Advice.FieldValue("paramName") String paramName) {
15+
if (result instanceof String || result instanceof Collection) {
16+
final WebModule module = InstrumentationBridge.WEB;
17+
18+
if (module != null) {
19+
try {
20+
if (result instanceof Collection) {
21+
Collection collection = (Collection) result;
22+
for (Object o : collection) {
23+
if (o instanceof String) {
24+
module.onCookieValue(paramName, (String) o);
25+
}
26+
}
27+
} else {
28+
module.onCookieValue(paramName, (String) result);
29+
}
30+
} catch (final Throwable e) {
31+
module.onUnexpectedException("CookieParamInjectorAdvice.onExit threw", e);
32+
}
33+
} else {
34+
System.out.println("Module is null");
35+
}
36+
}
37+
}
38+
}

‎dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/CookieParamInjectorInstrumentation.java

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
package datadog.trace.instrumentation.resteasy;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
5+
6+
import com.google.auto.service.AutoService;
7+
import datadog.trace.agent.tooling.Instrumenter;
8+
9+
@AutoService(Instrumenter.class)
10+
public class CookieParamInjectorInstrumentation extends Instrumenter.Iast
11+
implements Instrumenter.ForSingleType {
12+
13+
public CookieParamInjectorInstrumentation() {
14+
super("Cookie Param Value Injector for RestEasy");
15+
}
16+
17+
@Override
18+
public void adviceTransformations(AdviceTransformation transformation) {
19+
transformation.applyAdvice(
20+
named("inject").and(isPublic()), packageName + ".CookieParamInjectorAdvice");
21+
}
22+
23+
@Override
24+
public String instrumentedType() {
25+
return "org.jboss.resteasy.core.CookieParamInjector";
26+
}
27+
}

‎dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/HeaderParamInjectorAdvice.java

+38
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package datadog.trace.instrumentation.resteasy;
2+
3+
import datadog.trace.api.iast.InstrumentationBridge;
4+
import datadog.trace.api.iast.source.WebModule;
5+
import java.util.Collection;
6+
import net.bytebuddy.asm.Advice;
7+
import org.jboss.resteasy.core.HeaderParamInjector;
8+
9+
public class HeaderParamInjectorAdvice {
10+
@Advice.OnMethodExit
11+
public static void onExit(
12+
@Advice.This HeaderParamInjector self,
13+
@Advice.Return(readOnly = true) Object result,
14+
@Advice.FieldValue("paramName") String paramName) {
15+
if (result instanceof String || result instanceof Collection) {
16+
final WebModule module = InstrumentationBridge.WEB;
17+
18+
if (module != null) {
19+
try {
20+
if (result instanceof Collection) {
21+
Collection collection = (Collection) result;
22+
for (Object o : collection) {
23+
if (o instanceof String) {
24+
module.onHeaderValue(paramName, (String) o);
25+
}
26+
}
27+
} else {
28+
module.onHeaderValue(null, (String) result);
29+
}
30+
} catch (final Throwable e) {
31+
module.onUnexpectedException("PathParamInjectorAdvice.onExit threw", e);
32+
}
33+
} else {
34+
System.out.println("Module is null");
35+
}
36+
}
37+
}
38+
}

‎dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/HeaderParamInjectorInstrumentation.java

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
package datadog.trace.instrumentation.resteasy;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
5+
6+
import com.google.auto.service.AutoService;
7+
import datadog.trace.agent.tooling.Instrumenter;
8+
9+
@AutoService(Instrumenter.class)
10+
public class HeaderParamInjectorInstrumentation extends Instrumenter.Iast
11+
implements Instrumenter.ForSingleType {
12+
13+
public HeaderParamInjectorInstrumentation() {
14+
super("Header Value Injector for RestEasy");
15+
}
16+
17+
@Override
18+
public void adviceTransformations(AdviceTransformation transformation) {
19+
transformation.applyAdvice(
20+
named("inject").and(isPublic()), packageName + ".HeaderParamInjectorAdvice");
21+
}
22+
23+
@Override
24+
public String instrumentedType() {
25+
return "org.jboss.resteasy.core.HeaderParamInjector";
26+
}
27+
}

‎dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/MultivaluedMapInstrumentation.java

+58
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
package datadog.trace.instrumentation.resteasy;
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.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
6+
7+
import datadog.trace.agent.tooling.Instrumenter;
8+
import datadog.trace.api.iast.InstrumentationBridge;
9+
import datadog.trace.api.iast.source.WebModule;
10+
import net.bytebuddy.asm.Advice;
11+
import net.bytebuddy.description.type.TypeDescription;
12+
import net.bytebuddy.matcher.ElementMatcher;
13+
14+
public class MultivaluedMapInstrumentation extends Instrumenter.Iast
15+
implements Instrumenter.ForTypeHierarchy {
16+
17+
public MultivaluedMapInstrumentation() {
18+
super("MultivaluedMap");
19+
}
20+
21+
@Override
22+
public void adviceTransformations(AdviceTransformation transformation) {
23+
transformation.applyAdvice(
24+
namedOneOf("getFirst", "get"), getClass().getName() + "$MultiValuedMapAdviceGet");
25+
transformation.applyAdvice(
26+
namedOneOf("keySet", "values", "entrySet"),
27+
getClass().getName() + "$MultiValuedMapAdviceReturningMultipleValues");
28+
}
29+
30+
@Override
31+
public String hierarchyMarkerType() {
32+
return "jakarta.ws.rs.core.MultivaluedMap";
33+
}
34+
35+
@Override
36+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
37+
return implementsInterface(named(hierarchyMarkerType()));
38+
}
39+
40+
public static class MultiValuedMapAdviceGet {
41+
42+
@Advice.OnMethodExit
43+
public static void onExit(
44+
@Advice.This Object self, @Advice.Return(readOnly = true) Object result) {
45+
final WebModule module = InstrumentationBridge.WEB;
46+
47+
if (module != null) {
48+
try {
49+
module.onJaxGetQueryParameters(result);
50+
} catch (final Throwable e) {
51+
module.onUnexpectedException("UriInfoAdvice.onExit threw", e);
52+
}
53+
} else {
54+
System.out.println("Module is null");
55+
}
56+
}
57+
}
58+
}

‎dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/PathParamInjectorAdvice.java

+38
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package datadog.trace.instrumentation.resteasy;
2+
3+
import datadog.trace.api.iast.InstrumentationBridge;
4+
import datadog.trace.api.iast.source.WebModule;
5+
import java.util.Collection;
6+
import net.bytebuddy.asm.Advice;
7+
import org.jboss.resteasy.core.PathParamInjector;
8+
9+
public class PathParamInjectorAdvice {
10+
@Advice.OnMethodExit
11+
public static void onExit(
12+
@Advice.This PathParamInjector self,
13+
@Advice.Return(readOnly = true) Object result,
14+
@Advice.FieldValue("paramName") String paramName) {
15+
if (result instanceof String || result instanceof Collection) {
16+
final WebModule module = InstrumentationBridge.WEB;
17+
18+
if (module != null) {
19+
try {
20+
if (result instanceof Collection) {
21+
Collection collection = (Collection) result;
22+
for (Object o : collection) {
23+
if (o instanceof String) {
24+
module.onParameterValue(paramName, (String) o);
25+
}
26+
}
27+
} else {
28+
module.onParameterValue(paramName, (String) result);
29+
}
30+
} catch (final Throwable e) {
31+
module.onUnexpectedException("PathParamInjectorAdvice.onExit threw", e);
32+
}
33+
} else {
34+
System.out.println("Module is null");
35+
}
36+
}
37+
}
38+
}

0 commit comments

Comments
 (0)