Skip to content

Fix request context propagation for async executor tasks #305

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Mishenevd wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix request context propagation for async executor tasks #305

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions agent/src/main/java/dev/aikido/agent/Wrappers.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,11 @@
package dev.aikido.agent;

import dev.aikido.agent.wrappers.*;
import dev.aikido.agent.wrappers.executor.AbstractExecutorServiceWrapper;
import dev.aikido.agent.wrappers.executor.DelegatedExecutorServiceWrapper;
import dev.aikido.agent.wrappers.executor.ForkJoinPoolWrapper;
import dev.aikido.agent.wrappers.executor.ScheduledThreadPoolExecutorWrapper;
import dev.aikido.agent.wrappers.executor.ThreadPoolExecutorWrapper;
import dev.aikido.agent.wrappers.file.FileConstructorMultiArgumentWrapper;
import dev.aikido.agent.wrappers.file.FileConstructorSingleArgumentWrapper;
import dev.aikido.agent.wrappers.javalin.*;
Expand All @@ -9,6 +14,8 @@
import dev.aikido.agent.wrappers.spring.SpringWebfluxWrapper;
import dev.aikido.agent.wrappers.spring.SpringControllerWrapper;
import dev.aikido.agent.wrappers.spring.SpringMVCJakartaWrapper;
import dev.aikido.agent.wrappers.spring.SpringMVCJavaxWrapper;
import dev.aikido.agent.wrappers.spring.SpringWebfluxWrapper;

import java.util.Arrays;
import java.util.List;
Expand All @@ -17,6 +24,13 @@ public final class Wrappers {
private Wrappers() {}
public static final List<Wrapper> WRAPPERS = Arrays.asList(
new PostgresWrapper(),

new DelegatedExecutorServiceWrapper(),
new ThreadPoolExecutorWrapper(),
new AbstractExecutorServiceWrapper(),
new ForkJoinPoolWrapper(),
new ScheduledThreadPoolExecutorWrapper(),

new SpringMVCJakartaWrapper(),
new SpringMVCJavaxWrapper(),
new SpringWebfluxWrapper(),
Expand Down
52 changes: 52 additions & 0 deletions agent/src/main/java/dev/aikido/agent/wrappers/executor/AbstractExecutorServiceWrapper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
package dev.aikido.agent.wrappers.executor;

import dev.aikido.agent.wrappers.Wrapper;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import java.util.concurrent.AbstractExecutorService;
import java.util.concurrent.Callable;

import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf;
import static net.bytebuddy.matcher.ElementMatchers.named;
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;

public class AbstractExecutorServiceWrapper implements Wrapper {
@Override
public String getName() {
return SubmitAdvice.class.getName();
}

@Override
public ElementMatcher getMatcher() {
return isMethod()
.and(named("submit"))
.and(
takesArguments(Runnable.class)
.or(takesArguments(Callable.class))
.or(takesArguments(Runnable.class, Object.class))
);
}

@Override
public ElementMatcher getTypeMatcher() {
return isSubTypeOf(AbstractExecutorService.class);
}

public static class SubmitAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void before(
@Advice.Argument(value = 0, readOnly = false, typing = DYNAMIC) Object task
) {
if (task instanceof Runnable) {
task = ExecutorContextPropagation.wrap((Runnable) task);
} else if (task instanceof Callable) {
task = ExecutorContextPropagation.wrap((Callable) task);
}
}
}
}
72 changes: 72 additions & 0 deletions agent/src/main/java/dev/aikido/agent/wrappers/executor/DelegatedExecutorServiceWrapper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
package dev.aikido.agent.wrappers.executor;

import dev.aikido.agent.wrappers.Wrapper;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import java.lang.reflect.Method;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.concurrent.Callable;

import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
import static net.bytebuddy.matcher.ElementMatchers.named;
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;

public class DelegatedExecutorServiceWrapper implements Wrapper {
@Override
public String getName() {
return DelegatedExecutorAdvice.class.getName();
}

@Override
public ElementMatcher getMatcher() {
return isMethod()
.and(named("execute").or(named("submit")))
.and(
takesArguments(Runnable.class)
.or(takesArguments(Callable.class))
.or(takesArguments(Runnable.class, Object.class))
);
}

@Override
public ElementMatcher getTypeMatcher() {
return nameStartsWith("java.util.concurrent.Executors$");
}

public static class DelegatedExecutorAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void before(
@Advice.Argument(value = 0, readOnly = false, typing = DYNAMIC) Object task
) throws Exception {
if (task == null) {
return;
}

// This advice is applied to JDK classes loaded by the bootstrap classloader.
// Load agent_api reflectively because bootstrap classes cannot directly reference agent classes.
String jarFilePath = System.getProperty("AIK_agent_api_jar");
if (jarFilePath == null || jarFilePath.isBlank()) {
return;
}

URLClassLoader classLoader = new URLClassLoader(new URL[] { new URL(jarFilePath) });

@aikido-pr-checks aikido-pr-checks Bot Jun 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dynamic loading of a JAR from System.getProperty("AIK_agent_api_jar") and reflective invocation (URLClassLoader + Class.getMethod + Method.invoke) hides call targets and control flow; this pattern can conceal malicious behavior. Avoid or clearly document and validate the loaded artifact.

Details

✨ AI Reasoning
​The change adds code that loads a JAR path from a system property and then constructs a URLClassLoader and uses reflection to load and invoke methods. Loading code dynamically from a path derived at runtime and invoking methods via reflection hides static call targets and control flow from readers and static analyzers, which can be used to conceal malicious behavior. This behavior was introduced in this PR where DelegatedExecutorServiceWrapper reflectively loads ContextPropagation and invokes its wrap methods using Method.invoke. The code paths accept an external system property (AIK_agent_api_jar) and dynamically load code from it, which is a typical pattern used to obscure execution and hide functionality.

🔧 How do I fix it?
Ensure code is transparent and not intentionally obfuscated. Avoid hiding functionality from code review. Focus on intent and deception, not specific patterns.

Reply @AikidoSec feedback: [FEEDBACK] to get better review comments in the future.
Reply @AikidoSec ignore: [REASON] to ignore this issue.
More info

Class<?> contextPropagationClass = classLoader.loadClass(
"dev.aikido.agent_api.context.ContextPropagation"
);

if (task instanceof Runnable) {
Method wrapRunnable = contextPropagationClass.getMethod("wrap", Runnable.class);
task = wrapRunnable.invoke(null, task);
} else if (task instanceof Callable) {
Method wrapCallable = contextPropagationClass.getMethod("wrap", Callable.class);
task = wrapCallable.invoke(null, task);
}
}
}
}
67 changes: 67 additions & 0 deletions agent/src/main/java/dev/aikido/agent/wrappers/executor/ExecutorContextPropagation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
package dev.aikido.agent.wrappers.executor;

import java.lang.reflect.Method;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.concurrent.Callable;

public final class ExecutorContextPropagation {
private static Method wrapRunnableMethod;
private static Method wrapCallableMethod;
private static boolean disabled;
Comment on lines +9 to +11

@aikido-pr-checks aikido-pr-checks Bot Jun 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Static fields wrapRunnableMethod/wrapCallableMethod/disabled are unsafely accessed from multiple threads; make accesses thread-safe (e.g., use synchronization or volatile) to avoid visibility races.

Show fix
Suggested change
private static Method wrapRunnableMethod;
private static Method wrapCallableMethod;
private static boolean disabled;
private static volatile Method wrapRunnableMethod;
private static volatile Method wrapCallableMethod;
private static volatile boolean disabled;
Details

✨ AI Reasoning
​ExecutorContextPropagation introduces static mutable fields (wrapRunnableMethod, wrapCallableMethod, disabled) that are accessed from multiple threads (executor/task-related code). Although init() is synchronized, other reads/writes (checks of disabled, setting disabled in catch, and reading wrapRunnableMethod/wrapCallableMethod) occur without proper synchronization or volatile, producing potential visibility and race conditions when used concurrently by worker threads. This can cause stale reads or lost updates of 'disabled' and method references across threads.

Reply @AikidoSec feedback: [FEEDBACK] to get better review comments in the future.
Reply @AikidoSec ignore: [REASON] to ignore this issue.
More info


private ExecutorContextPropagation() {}

public static Runnable wrap(Runnable task) {
if (task == null || disabled) {
return task;
}

try {
init();
if (wrapRunnableMethod == null) {
return task;
}
return (Runnable) wrapRunnableMethod.invoke(null, task);
} catch (Throwable ignored) {
disabled = true;
return task;
}
}

@SuppressWarnings("unchecked")
public static <T> Callable<T> wrap(Callable<T> task) {
if (task == null || disabled) {
return task;
}

try {
init();
if (wrapCallableMethod == null) {
return task;
}
return (Callable<T>) wrapCallableMethod.invoke(null, task);
} catch (Throwable ignored) {
disabled = true;
return task;
}
}

private static synchronized void init() throws Exception {
if (disabled || wrapRunnableMethod != null) {
return;
}

String jarFilePath = System.getProperty("AIK_agent_api_jar");
if (jarFilePath == null || jarFilePath.isBlank()) {
disabled = true;
return;
}

URLClassLoader classLoader = new URLClassLoader(new URL[] { new URL(jarFilePath) });

@aikido-pr-checks aikido-pr-checks Bot Jun 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reflective loading of ContextPropagation via URLClassLoader from System.getProperty("AIK_agent_api_jar") and resolving wrap methods obscures where behavior is implemented; this dynamic loading can hide malicious code. Prefer safe, explicit classloading or validate the artifact.

Details

✨ AI Reasoning
​ExecutorContextPropagation.init reads a jar path from System.getProperty("AIK_agent_api_jar"), creates a URLClassLoader, loads dev.aikido.agent_api.context.ContextPropagation, and obtains Method references for wrap. Centralizing dynamic class loading and reflective method resolution like this was added by the PR and makes runtime behavior dependent on an external, runtime-specified JAR, obscuring where wrap implementations come from and enabling hidden functionality.

🔧 How do I fix it?
Ensure code is transparent and not intentionally obfuscated. Avoid hiding functionality from code review. Focus on intent and deception, not specific patterns.

Reply @AikidoSec feedback: [FEEDBACK] to get better review comments in the future.
Reply @AikidoSec ignore: [REASON] to ignore this issue.
More info

Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.context.ContextPropagation");

wrapRunnableMethod = clazz.getMethod("wrap", Runnable.class);
wrapCallableMethod = clazz.getMethod("wrap", Callable.class);
}
}
52 changes: 52 additions & 0 deletions agent/src/main/java/dev/aikido/agent/wrappers/executor/ForkJoinPoolWrapper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
package dev.aikido.agent.wrappers.executor;

import dev.aikido.agent.wrappers.Wrapper;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import java.util.concurrent.Callable;
import java.util.concurrent.ForkJoinPool;

import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf;
import static net.bytebuddy.matcher.ElementMatchers.named;
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;

public class ForkJoinPoolWrapper implements Wrapper {
@Override
public String getName() {
return ForkJoinAdvice.class.getName();
}

@Override
public ElementMatcher getMatcher() {
return isMethod()
.and(named("execute").or(named("submit")))
.and(
takesArguments(Runnable.class)
.or(takesArguments(Callable.class))
.or(takesArguments(Runnable.class, Object.class))
);
}

@Override
public ElementMatcher getTypeMatcher() {
return isSubTypeOf(ForkJoinPool.class);
}

public static class ForkJoinAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void before(
@Advice.Argument(value = 0, readOnly = false, typing = DYNAMIC) Object task
) {
if (task instanceof Runnable) {
task = ExecutorContextPropagation.wrap((Runnable) task);
} else if (task instanceof Callable) {
task = ExecutorContextPropagation.wrap((Callable) task);
}
}
}
}
73 changes: 73 additions & 0 deletions .../src/main/java/dev/aikido/agent/wrappers/executor/ScheduledThreadPoolExecutorWrapper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
package dev.aikido.agent.wrappers.executor;

import dev.aikido.agent.wrappers.Wrapper;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import java.lang.reflect.Method;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.concurrent.Callable;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf;
import static net.bytebuddy.matcher.ElementMatchers.named;
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;

public class ScheduledThreadPoolExecutorWrapper implements Wrapper {
@Override
public String getName() {
return ScheduleAdvice.class.getName();
}

@Override
public ElementMatcher getMatcher() {
return isMethod()
.and(named("schedule"))
.and(
takesArguments(Runnable.class, long.class, TimeUnit.class)
.or(takesArguments(Callable.class, long.class, TimeUnit.class))
);
}

@Override
public ElementMatcher getTypeMatcher() {
return isSubTypeOf(ScheduledThreadPoolExecutor.class);
}

public static class ScheduleAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void before(
@Advice.Argument(value = 0, readOnly = false, typing = DYNAMIC) Object task
) throws Exception {
if (task == null) {
return;
}

// This advice is applied to JDK classes loaded by the bootstrap classloader.
// Load agent_api reflectively because bootstrap classes cannot directly reference agent classes.
String jarFilePath = System.getProperty("AIK_agent_api_jar");
if (jarFilePath == null || jarFilePath.isBlank()) {
return;
}

URLClassLoader classLoader = new URLClassLoader(new URL[] { new URL(jarFilePath) });

@aikido-pr-checks aikido-pr-checks Bot Jun 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dynamic loading of a JAR from System.getProperty("AIK_agent_api_jar") with URLClassLoader and reflective Method.invoke hides execution targets; this pattern can conceal malicious payloads. Validate and document the loaded artifact and minimize reflective use.

Details

✨ AI Reasoning
​The ScheduledThreadPoolExecutorWrapper also loads a JAR via a system property and uses a URLClassLoader to reflectively load ContextPropagation and invoke its wrap methods. This same dynamic reflective loading/hiding pattern was added in this PR and similarly obscures code execution and control flow, making it harder to audit and a vector for concealed behavior.

🔧 How do I fix it?
Ensure code is transparent and not intentionally obfuscated. Avoid hiding functionality from code review. Focus on intent and deception, not specific patterns.

Reply @AikidoSec feedback: [FEEDBACK] to get better review comments in the future.
Reply @AikidoSec ignore: [REASON] to ignore this issue.
More info

Class<?> contextPropagationClass = classLoader.loadClass(
"dev.aikido.agent_api.context.ContextPropagation"
);

if (task instanceof Runnable) {
Method wrapRunnable = contextPropagationClass.getMethod("wrap", Runnable.class);
task = wrapRunnable.invoke(null, task);
} else if (task instanceof Callable) {
Method wrapCallable = contextPropagationClass.getMethod("wrap", Callable.class);
task = wrapCallable.invoke(null, task);
}
}
}
}
42 changes: 42 additions & 0 deletions agent/src/main/java/dev/aikido/agent/wrappers/executor/ThreadPoolExecutorWrapper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
package dev.aikido.agent.wrappers.executor;

import dev.aikido.agent.wrappers.Wrapper;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import java.util.concurrent.ThreadPoolExecutor;

import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf;
import static net.bytebuddy.matcher.ElementMatchers.named;
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;

public class ThreadPoolExecutorWrapper implements Wrapper {
@Override
public String getName() {
return ExecuteAdvice.class.getName();
}

@Override
public ElementMatcher getMatcher() {
return isMethod()
.and(named("execute"))
.and(takesArguments(Runnable.class));
}

@Override
public ElementMatcher getTypeMatcher() {
return isSubTypeOf(ThreadPoolExecutor.class);
}

public static class ExecuteAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void before(
@Advice.Argument(value = 0, readOnly = false) Runnable task
) {
task = ExecutorContextPropagation.wrap(task);
}
}
}
Loading
Loading