Skip to content

Commit

Permalink
Refactor propagation module to include two different APIs for strings…
Browse files Browse the repository at this point in the history 
… and objects (#6820)
  • Loading branch information
@manuel-alvarez-alvarez
manuel-alvarez-alvarez committed Apr 22, 2024
1 parent 6b60035 commit e9eb782
Show file tree
Hide file tree
Showing 25 changed files with 734 additions and 306 deletions.
643 changes: 472 additions & 171 deletions ...va-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java
View file

Large diffs are not rendered by default.

7 changes: 1 addition & 6 deletions ...-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/BaseCodecModuleTest.groovy
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ abstract class BaseCodecModuleTest extends IastModuleImplTestBase {
}
}

void '#method null or empty'() {
void '#method null'() {
when:
module.&"$method".call(args.toArray())

Expand All @@ -35,15 +35,10 @@ abstract class BaseCodecModuleTest extends IastModuleImplTestBase {
where:
method | args
'onUrlDecode' | ['test', 'utf-8', null]
'onUrlDecode' | ['test', 'utf-8', '']
'onStringGetBytes' | ['test', 'utf-8', null]
'onStringGetBytes' | ['test', 'utf-8', [] as byte[]]
'onStringFromBytes' | ['test'.bytes, 0, 2, 'utf-8', null]
'onStringFromBytes' | ['test'.bytes, 0, 2, 'utf-8', '']
'onBase64Encode' | ['test'.bytes, null]
'onBase64Encode' | ['test'.bytes, [] as byte[]]
'onBase64Decode' | ['test'.bytes, null]
'onBase64Decode' | ['test'.bytes, [] as byte[]]
}

void '#method no context'() {
Expand Down
11 changes: 8 additions & 3 deletions ...t/src/testFixtures/groovy/com/datadog/iast/test/IastRequestContextPreparationTrait.groovy
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,9 +69,14 @@ trait IastRequestContextPreparationTrait {
return tainted
}

private final static Logger LOGGER = LoggerFactory.getLogger("map tainted objects")
static {
((ch.qos.logback.classic.Logger) LOGGER).level = ch.qos.logback.classic.Level.DEBUG
private final static Logger LOGGER = withLogger("map tainted objects")

private static Logger withLogger(final String name) {
final logger = LoggerFactory.getLogger(name)
if (logger instanceof ch.qos.logback.classic.Logger) {
((ch.qos.logback.classic.Logger) logger).level = ch.qos.logback.classic.Level.DEBUG
}
return logger
}

private static void logTaint(Object o) {
Expand Down
15 changes: 7 additions & 8 deletions ...src/main/java/datadog/trace/instrumentation/akkahttp/iast/PathMatcherInstrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,24 +56,23 @@ static void onExit(
return;
}

scala.Tuple1 tuple = (scala.Tuple1) extractions;
Object value = tuple._1();

PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module == null) {
if (module == null || !(value instanceof String)) {
return;
}

scala.Tuple1 tuple = (scala.Tuple1) extractions;
Object value = tuple._1();
final String stringValue = (String) value;

final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST);

// in the test, 4 instances of PathMatcher$Match are created, all with the same value
if (module.isTainted(ctx, value)) {
if (module.isTainted(ctx, stringValue)) {
return;
}

if (value instanceof String) {
module.taint(ctx, value, SourceTypes.REQUEST_PATH_PARAMETER);
}
module.taint(ctx, stringValue, SourceTypes.REQUEST_PATH_PARAMETER);
}
}
}
4 changes: 2 additions & 2 deletions ...ava/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintSingleParameterFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,11 +50,11 @@ public Tuple1<Object> apply(Tuple1<Object> v1) {
while (iterator.hasNext()) {
Object o = iterator.next();
if (o instanceof String) {
mod.taint(ctx, o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
mod.taint(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
}
}
} else if (value instanceof String) {
mod.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
mod.taint(ctx, (String) value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
}

return v1;
Expand Down
4 changes: 2 additions & 2 deletions .../java/datadog/trace/instrumentation/akkahttp102/iast/helpers/TaintParametersFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,11 +41,11 @@ public Tuple1<T> apply(Tuple1<T> v1) {
while (iterator.hasNext()) {
Object o = iterator.next();
if (o instanceof String) {
mod.taint(ctx, o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
mod.taint(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
}
}
} else if (value instanceof String) {
mod.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
mod.taint(ctx, (String) value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
}

return v1;
Expand Down
6 changes: 5 additions & 1 deletion ...n/java-io/src/main/java/datadog/trace/instrumentation/java/lang/StringReaderCallSite.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,11 @@ public static StringReader afterInit(
@CallSite.Return @Nonnull final StringReader result) {
final PropagationModule propagationModule = InstrumentationBridge.PROPAGATION;
if (propagationModule != null) {
propagationModule.taintIfTainted(result, params[0]);
try {
propagationModule.taintIfTainted(result, params[0]);
} catch (Throwable e) {
propagationModule.onUnexpectedException("afterInit threw", e);
}
}
return result;
}
Expand Down
4 changes: 2 additions & 2 deletions ...ation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringCallSite.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ public static String aroundJoin(
final StringModule module = InstrumentationBridge.STRING;
if (module != null) {
try {
module.onStringJoin(result, delimiter, copy.toArray(new CharSequence[copy.size()]));
module.onStringJoin(result, delimiter, copy.toArray(new CharSequence[0]));
} catch (final Throwable e) {
module.onUnexpectedException("afterSubSequence threw", e);
}
Expand Down Expand Up @@ -412,7 +412,7 @@ public static String[] afterSplit(
public static String[] afterSplitWithLimit(
@CallSite.This @Nonnull final String self,
@CallSite.Argument(0) @Nonnull final String regex,
@CallSite.Argument(1) @Nonnull final int pos,
@CallSite.Argument(1) final int pos,
@CallSite.Return @Nonnull final String[] result) {
final StringModule module = InstrumentationBridge.STRING;
if (module != null) {
Expand Down
2 changes: 1 addition & 1 deletion ...java/datadog/trace/instrumentation/jersey/AbstractParamValueExtractorInstrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ public static void onExit(
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null) {
IastContext ctx = reqCtx.getData(RequestContextSlot.IAST);
module.taint(ctx, result, ThreadLocalSourceType.get(), parameterName);
module.taint(ctx, (String) result, ThreadLocalSourceType.get(), parameterName);
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion ...jersey/src/main/java/datadog/trace/instrumentation/jersey/AbstractStringReaderAdvice.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ public static void onExit(
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null) {
IastContext ctx = reqCtx.getData(RequestContextSlot.IAST);
module.taint(ctx, result, SourceTypes.REQUEST_PARAMETER_VALUE);
module.taint(ctx, (String) result, SourceTypes.REQUEST_PARAMETER_VALUE);
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion ...t/src/main/java/datadog/trace/instrumentation/josejwt/JSONObjectUtilsInstrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ public static void onEnter(
if (value instanceof String) {
// TODO: We could represent this source more accurately, perhaps tracking the original
// source, or using a special name.
module.taint(ctx, value, SourceTypes.REQUEST_HEADER_VALUE, name);
module.taint(ctx, (String) value, SourceTypes.REQUEST_HEADER_VALUE, name);
}
}
}
Expand Down
32 changes: 20 additions & 12 deletions ...n/org-json/src/main/java/datadog/trace/instrumentation/json/JSONArrayInstrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
import datadog.trace.api.iast.Propagation;
import datadog.trace.api.iast.propagation.PropagationModule;
import net.bytebuddy.asm.Advice;
import org.json.JSONArray;
import org.json.JSONObject;

@AutoService(InstrumenterModule.class)
public class JSONArrayInstrumentation extends InstrumenterModule.Iast
Expand Down Expand Up @@ -60,15 +62,18 @@ public static class GetAdvice {
@Advice.OnMethodExit(suppress = Throwable.class)
@Propagation
public static void afterMethod(@Advice.This Object self, @Advice.Return final Object result) {
if (result instanceof Integer
|| result instanceof Long
|| result instanceof Double
|| result instanceof Boolean) {
boolean isString = result instanceof String;
boolean isJson = !isString && (result instanceof JSONObject || result instanceof JSONArray);
if (!isString && !isJson) {
return;
}
final PropagationModule iastModule = InstrumentationBridge.PROPAGATION;
if (iastModule != null && result != null && result instanceof String) {
iastModule.taintIfTainted(result, self);
if (iastModule != null) {
if (isString) {
iastModule.taintIfTainted((String) result, self);
} else {
iastModule.taintIfTainted(result, self);
}
}
}
}
Expand All @@ -77,15 +82,18 @@ public static class OptAdvice {
@Advice.OnMethodExit(suppress = Throwable.class)
@Propagation
public static void afterMethod(@Advice.This Object self, @Advice.Return final Object result) {
if (result instanceof Integer
|| result instanceof Long
|| result instanceof Double
|| result instanceof Boolean) {
boolean isString = result instanceof String;
boolean isJson = !isString && (result instanceof JSONObject || result instanceof JSONArray);
if (!isString && !isJson) {
return;
}
final PropagationModule iastModule = InstrumentationBridge.PROPAGATION;
if (iastModule != null && result != null) {
iastModule.taintIfTainted(result, self);
if (iastModule != null) {
if (isString) {
iastModule.taintIfTainted((String) result, self);
} else {
iastModule.taintIfTainted(result, self);
}
}
}
}
Expand Down
32 changes: 20 additions & 12 deletions .../org-json/src/main/java/datadog/trace/instrumentation/json/JSONObjectInstrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
import datadog.trace.api.iast.Propagation;
import datadog.trace.api.iast.propagation.PropagationModule;
import net.bytebuddy.asm.Advice;
import org.json.JSONArray;
import org.json.JSONObject;

@AutoService(InstrumenterModule.class)
public class JSONObjectInstrumentation extends InstrumenterModule.Iast
Expand Down Expand Up @@ -62,15 +64,18 @@ public static class GetAdvice {
@Advice.OnMethodExit(suppress = Throwable.class)
@Propagation
public static void afterMethod(@Advice.This Object self, @Advice.Return final Object result) {
if (result instanceof Integer
|| result instanceof Long
|| result instanceof Double
|| result instanceof Boolean) {
boolean isString = result instanceof String;
boolean isJson = !isString && (result instanceof JSONObject || result instanceof JSONArray);
if (!isString && !isJson) {
return;
}
final PropagationModule iastModule = InstrumentationBridge.PROPAGATION;
if (iastModule != null && result != null) {
iastModule.taintIfTainted(result, self);
if (iastModule != null) {
if (isString) {
iastModule.taintIfTainted((String) result, self);
} else {
iastModule.taintIfTainted(result, self);
}
}
}
}
Expand All @@ -79,15 +84,18 @@ public static class OptAdvice {
@Advice.OnMethodExit(suppress = Throwable.class)
@Propagation
public static void afterMethod(@Advice.This Object self, @Advice.Return final Object result) {
if (result instanceof Integer
|| result instanceof Long
|| result instanceof Double
|| result instanceof Boolean) {
boolean isString = result instanceof String;
boolean isJson = !isString && (result instanceof JSONObject || result instanceof JSONArray);
if (!isString && !isJson) {
return;
}
final PropagationModule iastModule = InstrumentationBridge.PROPAGATION;
if (iastModule != null && result != null) {
iastModule.taintIfTainted(result, self);
if (iastModule != null) {
if (isString) {
iastModule.taintIfTainted((String) result, self);
} else {
iastModule.taintIfTainted(result, self);
}
}
}
}
Expand Down
17 changes: 8 additions & 9 deletions ...rc/main/java/datadog/trace/instrumentation/pekkohttp/iast/PathMatcherInstrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,24 +56,23 @@ static void onExit(
return;
}

scala.Tuple1 tuple = (scala.Tuple1) extractions;
Object value = tuple._1();

PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module == null) {
if (module == null || !(value instanceof String)) {
return;
}
final String stringValue = (String) value;

scala.Tuple1 tuple = (scala.Tuple1) extractions;
Object value = tuple._1();

IastContext ctx = reqCtx.getData(RequestContextSlot.IAST);
final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST);

// in the test, 4 instances of PathMatcher$Match are created, all with the same value
if (module.isTainted(ctx, value)) {
if (module.isTainted(ctx, stringValue)) {
return;
}

if (value instanceof String) {
module.taint(ctx, value, SourceTypes.REQUEST_PATH_PARAMETER);
}
module.taint(ctx, stringValue, SourceTypes.REQUEST_PATH_PARAMETER);
}
}
}
4 changes: 2 additions & 2 deletions ...in/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintParametersFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,11 +44,11 @@ public Tuple1<T> apply(Tuple1<T> v1) {
while (iterator.hasNext()) {
Object o = iterator.next();
if (o instanceof String) {
mod.taint(ctx, o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
mod.taint(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
}
}
} else if (value instanceof String) {
mod.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
mod.taint(ctx, (String) value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
}

return v1;
Expand Down
4 changes: 2 additions & 2 deletions ...va/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintSingleParameterFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,11 +51,11 @@ public Tuple1<Object> apply(Tuple1<Object> v1) {
while (iterator.hasNext()) {
Object o = iterator.next();
if (o instanceof String) {
mod.taint(ctx, o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
mod.taint(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
}
}
} else if (value instanceof String) {
mod.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
mod.taint(ctx, (String) value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
}

return v1;
Expand Down
4 changes: 2 additions & 2 deletions ...ppsec/src/main/java/datadog/trace/instrumentation/resteasy/CookieParamInjectorAdvice.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,11 +28,11 @@ public static void onExit(
Collection<?> collection = (Collection<?>) result;
for (Object o : collection) {
if (o instanceof String) {
module.taint(ctx, o, SourceTypes.REQUEST_COOKIE_VALUE, paramName);
module.taint(ctx, (String) o, SourceTypes.REQUEST_COOKIE_VALUE, paramName);
}
}
} else {
module.taint(ctx, result, SourceTypes.REQUEST_COOKIE_VALUE, paramName);
module.taint(ctx, (String) result, SourceTypes.REQUEST_COOKIE_VALUE, paramName);
}
}
}
Expand Down
4 changes: 2 additions & 2 deletions ...-appsec/src/main/java/datadog/trace/instrumentation/resteasy/FormParamInjectorAdvice.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,11 +28,11 @@ public static void onExit(
Collection<?> collection = (Collection<?>) result;
for (Object o : collection) {
if (o instanceof String) {
module.taint(ctx, o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
module.taint(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
}
}
} else {
module.taint(ctx, result, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
module.taint(ctx, (String) result, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
}
}
}
Expand Down
20 changes: 8 additions & 12 deletions ...ppsec/src/main/java/datadog/trace/instrumentation/resteasy/HeaderParamInjectorAdvice.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,20 +23,16 @@ public static void onExit(
if (result instanceof String || result instanceof Collection) {
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null) {
try {
IastContext ctx = reqCtx.getData(RequestContextSlot.IAST);
if (result instanceof Collection) {
Collection<?> collection = (Collection<?>) result;
for (Object o : collection) {
if (o instanceof String) {
module.taint(ctx, o, SourceTypes.REQUEST_HEADER_VALUE, paramName);
}
IastContext ctx = reqCtx.getData(RequestContextSlot.IAST);
if (result instanceof Collection) {
Collection<?> collection = (Collection<?>) result;
for (Object o : collection) {
if (o instanceof String) {
module.taint(ctx, (String) o, SourceTypes.REQUEST_HEADER_VALUE, paramName);
}
} else {
module.taint(ctx, result, SourceTypes.REQUEST_HEADER_VALUE, paramName);
}
} catch (final Throwable e) {
module.onUnexpectedException("HeaderParamInjectorAdvice.onExit threw", e);
} else {
module.taint(ctx, (String) result, SourceTypes.REQUEST_HEADER_VALUE, paramName);
}
}
}
Expand Down
Loading

0 comments on commit e9eb782

Please sign in to comment.