Skip to content

Commit

Permalink
Use inner class inside IastContext to prevent issues in instrumentations
Browse files Browse the repository at this point in the history
  • Loading branch information
manuel-alvarez-alvarez committed Oct 22, 2023
1 parent b74e1b3 commit 995f5c4
Show file tree
Hide file tree
Showing 46 changed files with 94 additions and 85 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ public class GrpcRequestMessageHandler implements BiFunction<RequestContext, Obj
public Flow<Void> apply(final RequestContext ctx, final Object o) {
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null && o != null) {
final IastContext iastCtx = IastContext.get(ctx);
final IastContext iastCtx = IastContext.Provider.get(ctx);
module.taintDeeply(
iastCtx, o, SourceTypes.GRPC_BODY, GrpcRequestMessageHandler::isProtobufArtifact);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -98,17 +98,17 @@ public IastMetricCollector getMetricCollector() {

@Nullable
public static IastRequestContext get() {
return asRequestContext(IastContext.get());
return asRequestContext(IastContext.Provider.get());
}

@Nullable
public static IastRequestContext get(final AgentSpan span) {
return asRequestContext(IastContext.get(span));
return asRequestContext(IastContext.Provider.get(span));
}

@Nullable
public static IastRequestContext get(final RequestContext reqCtx) {
return asRequestContext(IastContext.get(reqCtx));
return asRequestContext(IastContext.Provider.get(reqCtx));
}

private static IastRequestContext asRequestContext(final IastContext ctx) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@ static void after(
return;
}

final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Iterator<HttpCookiePair> iterator = cookiePairs.iterator();
while (iterator.hasNext()) {
HttpCookiePair pair = iterator.next();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ static void onExit(
return;
}

final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Iterator<HttpHeader> iterator = headers.iterator();
while (iterator.hasNext()) {
HttpHeader h = iterator.next();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ static void after(@Advice.This /*Uri*/ Object uri, @Advice.Return Uri.Query ret)
return;
}

final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Iterator<Tuple2<String, String>> iterator = ret.iterator();
while (iterator.hasNext()) {
Tuple2<String, String> pair = iterator.next();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ public Tuple1<Map<String, String>> apply(Tuple1<Map<String, String>> v1) {
return v1;
}

final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Iterator<Tuple2<String, String>> iterator = m.iterator();
while (iterator.hasNext()) {
Tuple2<String, String> e = iterator.next();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ public Tuple1<Map<String, List<String>>> apply(Tuple1<Map<String, List<String>>>
return v1;
}

final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Iterator<Tuple2<String, List<String>>> entriesIterator = m.iterator();
while (entriesIterator.hasNext()) {
Tuple2<String, List<String>> e = entriesIterator.next();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ public Tuple1<Seq<Tuple2<String, String>>> apply(Tuple1<Seq<Tuple2<String, Strin
return v1;
}

final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Iterator<Tuple2<String, String>> iterator = seq.iterator();
Set<String> seenKeys = Collections.newSetFromMap(new IdentityHashMap<>());
while (iterator.hasNext()) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ public Tuple1<Object> apply(Tuple1<Object> v1) {
}

if (value instanceof Iterable) {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Iterator<?> iterator = ((Iterable<?>) value).iterator();
while (iterator.hasNext()) {
Object o = iterator.next();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ public Tuple1<T> apply(Tuple1<T> v1) {
}

if (value instanceof Iterable) {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Iterator<?> iterator = ((Iterable<?>) value).iterator();
while (iterator.hasNext()) {
Object o = iterator.next();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ public static void onExit(@Advice.Return Map<String, List<String>> result) {
if (prop == null || result == null || result.isEmpty()) {
return;
}
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (Map.Entry<String, List<String>> entry : result.entrySet()) {
final String name = entry.getKey();
prop.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ public static class InstrumenterAdviceGetHeaders {
public static void onExit(@Advice.Return Map<String, List<String>> headers) {
final PropagationModule prop = InstrumentationBridge.PROPAGATION;
if (prop != null && headers != null && !headers.isEmpty()) {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (Map.Entry<String, List<String>> entry : headers.entrySet()) {
final String name = entry.getKey();
prop.taint(ctx, name, SourceTypes.REQUEST_HEADER_NAME);
Expand All @@ -64,7 +64,7 @@ public static class InstrumenterAdviceGetRequestCookies {
public static void onExit(@Advice.Return Map<String, Object> cookies) {
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null && cookies != null && !cookies.isEmpty()) {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (Map.Entry<String, Object> entry : cookies.entrySet()) {
final String name = entry.getKey();
module.taint(ctx, name, SourceTypes.REQUEST_COOKIE_NAME);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ static void after(
if (prop == null || cookiePairs == null || cookiePairs.isEmpty()) {
return;
}
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
if (!prop.isTainted(ctx, cookie)) {
return;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ static void onExit(
return;
}

final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Iterator<HttpHeader> iterator = headers.iterator();
while (iterator.hasNext()) {
HttpHeader h = iterator.next();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ static void after(@Advice.This /*Uri*/ Object uri, @Advice.Return Uri.Query ret)
return;
}

final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
if (!prop.isTainted(ctx, uri)) {
return;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ public Tuple1<Map<String, String>> apply(Tuple1<Map<String, String>> v1) {
return v1;
}

final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Iterator<Tuple2<String, String>> iterator = m.iterator();
while (iterator.hasNext()) {
Tuple2<String, String> e = iterator.next();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ public Tuple1<Map<String, List<String>>> apply(Tuple1<Map<String, List<String>>>
return v1;
}

final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Iterator<Tuple2<String, List<String>>> entriesIterator = m.iterator();
while (entriesIterator.hasNext()) {
Tuple2<String, List<String>> e = entriesIterator.next();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ public Tuple1<T> apply(Tuple1<T> v1) {
}

if (value instanceof Iterable) {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Iterator<?> iterator = ((Iterable<?>) value).iterator();
while (iterator.hasNext()) {
Object o = iterator.next();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ public Tuple1<Seq<Tuple2<String, String>>> apply(Tuple1<Seq<Tuple2<String, Strin
return v1;
}

final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Iterator<Tuple2<String, String>> iterator = seq.iterator();
Set<String> seenKeys = Collections.newSetFromMap(new IdentityHashMap<>());
while (iterator.hasNext()) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ public Tuple1<Object> apply(Tuple1<Object> v1) {
}

if (value instanceof Iterable) {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Iterator<?> iterator = ((Iterable<?>) value).iterator();
while (iterator.hasNext()) {
Object o = iterator.next();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ public static void onExit(
if (module != null) {
if (result instanceof Collection) {
Collection<?> collection = (Collection<?>) result;
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (Object o : collection) {
if (o instanceof String) {
module.taint(ctx, (String) o, SourceTypes.REQUEST_COOKIE_VALUE, paramName);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ public static void onExit(
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null) {
if (result instanceof Collection) {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
Collection<?> collection = (Collection<?>) result;
for (Object o : collection) {
if (o instanceof String) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ public static void onExit(
try {
if (result instanceof Collection) {
Collection<?> collection = (Collection<?>) result;
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (Object o : collection) {
if (o instanceof String) {
module.taint(ctx, (String) o, SourceTypes.REQUEST_HEADER_VALUE, paramName);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ public static void onExit(
if (module != null) {
if (result instanceof Collection) {
Collection<?> collection = (Collection<?>) result;
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (Object o : collection) {
if (o instanceof String) {
module.taint(ctx, o, SourceTypes.REQUEST_PATH_PARAMETER, paramName);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ public static void onExit(
if (module != null) {
if (result instanceof Collection) {
Collection<?> collection = (Collection<?>) result;
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (Object o : collection) {
if (o instanceof String) {
module.taint(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ public static void onExit(
}
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null) {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final String value : headerValues) {
module.taint(ctx, value, SourceTypes.REQUEST_MULTIPART_PARAMETER, headerName);
}
Expand All @@ -108,7 +108,7 @@ public static void onExit(@Advice.Return final Collection<String> headerNames) {
}
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null) {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final String name : headerNames) {
module.taint(ctx, name, SourceTypes.REQUEST_MULTIPART_PARAMETER);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ public static java.util.Map<java.lang.String, java.lang.String[]> afterGetParame
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null && map != null && !map.isEmpty()) {
try {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final Map.Entry<String, String[]> entry : map.entrySet()) {
final String name = entry.getKey();
module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ public static java.util.Map<java.lang.String, java.lang.String[]> afterGetParame
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null && map != null && !map.isEmpty()) {
try {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final Map.Entry<String, String[]> entry : map.entrySet()) {
final String name = entry.getKey();
module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ public static Enumeration<String> afterGetParameterNames(
parameterNames.add(paramName);
}
try {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final String name : parameterNames) {
module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);
}
Expand Down Expand Up @@ -89,7 +89,7 @@ public static String[] afterGetParameterValues(
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null) {
try {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final String value : parameterValues) {
module.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
}
Expand All @@ -110,7 +110,7 @@ public static java.util.Map<java.lang.String, java.lang.String[]> afterGetParame
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null && map != null && !map.isEmpty()) {
try {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final Map.Entry<String, String[]> entry : map.entrySet()) {
final String name = entry.getKey();
module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ public static void onExit(
}
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null) {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final String value : headerValues) {
module.taint(ctx, value, SourceTypes.REQUEST_MULTIPART_PARAMETER, headerName);
}
Expand All @@ -104,7 +104,7 @@ public static void onExit(@Advice.Return final Collection<String> headerNames) {
}
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null) {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final String name : headerNames) {
module.taint(ctx, name, SourceTypes.REQUEST_MULTIPART_PARAMETER);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ public static Enumeration<String> afterGetParameterNames(
parameterNames.add(paramName);
}
try {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final String name : parameterNames) {
module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);
}
Expand Down Expand Up @@ -87,7 +87,7 @@ public static String[] afterGetParameterValues(
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null) {
try {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final String value : parameterValues) {
module.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
}
Expand All @@ -107,7 +107,7 @@ public static java.util.Map<java.lang.String, java.lang.String[]> afterGetParame
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null && map != null && !map.isEmpty()) {
try {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final Map.Entry<String, String[]> entry : map.entrySet()) {
final String name = entry.getKey();
module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,7 @@ public static Enumeration<String> afterGetHeaders(
headerValues.add(headerValue);
}
try {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final String value : headerValues) {
module.taint(ctx, value, SourceTypes.REQUEST_HEADER_VALUE, headerName);
}
Expand Down Expand Up @@ -102,7 +102,7 @@ public static Enumeration<String> afterGetHeaderNames(
headerNames.add(headerName);
}
try {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final String name : headerNames) {
module.taint(ctx, name, SourceTypes.REQUEST_HEADER_NAME);
}
Expand All @@ -126,7 +126,7 @@ public static Cookie[] afterGetCookies(
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null) {
try {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final Cookie cookie : cookies) {
module.taint(ctx, cookie, SourceTypes.REQUEST_COOKIE_VALUE);
}
Expand Down Expand Up @@ -193,7 +193,7 @@ public static Enumeration<String> afterGetParameterNames(
parameterNames.add(paramName);
}
try {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final String name : parameterNames) {
module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);
}
Expand Down Expand Up @@ -221,7 +221,7 @@ public static String[] afterGetParameterValues(
final PropagationModule module = InstrumentationBridge.PROPAGATION;
if (module != null) {
try {
final IastContext ctx = IastContext.get();
final IastContext ctx = IastContext.Provider.get();
for (final String value : parameterValues) {
module.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
}
Expand Down
Loading

0 comments on commit 995f5c4

Please sign in to comment.