Skip to content

Context propagation primitives for async task execution #342

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 from
Open

Context propagation primitives for async task execution #342

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
7 changes: 7 additions & 0 deletions agent_api/src/main/java/dev/aikido/agent_api/context/Context.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,4 +15,11 @@ public static void set(ContextObject contextObject) {
public static void reset() {
threadLocalContext.remove();
}
public static void restore(ContextObject previous) {
if (previous != null) {
set(previous);
} else {
reset();
}
}
}
30 changes: 29 additions & 1 deletion agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

import java.util.*;

public class ContextObject {
public class ContextObject implements Cloneable {
protected String method;
protected String source;
protected String url;
Expand All @@ -25,6 +25,34 @@ public class ContextObject {
protected transient Map<String, Map<String, String>> cache = new HashMap<>();
protected transient Optional<Boolean> forcedProtectionOff = Optional.empty();

// Async tasks get their own copy so parallel workers sharing one request's
// context never race on its mutable working state (cache, redirect nodes).
public ContextObject copyForPropagation() {
try {
ContextObject copy = (ContextObject) super.clone();

@aikido-pr-checks aikido-pr-checks Bot Aug 17, 2026 •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

super.clone() shallow-copies mutable fields such as headers, query, cookies, params, body, and user into contexts used by parallel tasks. Deep-copy mutable state or make it immutable before propagation.

Show fix
Suggested change
ContextObject copy = (ContextObject) super.clone();
ContextObject copy = (ContextObject) super.clone();
if (this.headers != null) {
copy.headers = new HashMap<>();
for (Map.Entry<String, List<String>> entry : this.headers.entrySet()) {
copy.headers.put(entry.getKey(), new ArrayList<>(entry.getValue()));
}
}
if (this.query != null) {
copy.query = new HashMap<>();
for (Map.Entry<String, List<String>> entry : this.query.entrySet()) {
copy.query.put(entry.getKey(), new ArrayList<>(entry.getValue()));
}
}
if (this.cookies != null) {
copy.cookies = new HashMap<>();
for (Map.Entry<String, List<String>> entry : this.cookies.entrySet()) {
copy.cookies.put(entry.getKey(), new ArrayList<>(entry.getValue()));
}
}
Details

✨ AI Reasoning
​The propagation wrappers create a per-task context and therefore make the copied object's state available on worker threads. The clone operation is shallow for fields other than cache and redirectStartNodes, so headers, query, cookies, and potentially mutable params, body, and user objects remain shared with the original context. Existing getters expose several collection references directly, allowing concurrent task code to mutate the same underlying objects without synchronization. Deep-copy mutable fields or make them immutable before publishing the context to worker threads.

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

@Mishenevd Mishenevd Aug 17, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

headers/query/cookies/params/body/user are set once at request build time and only read afterwards, so sharing them across tasks is race-free. The only worker-mutated state (cache, redirect chain) is already deep-copied, and body/params are arbitrary Object types with no safe general deep-copy

copy.cache = copyCache(this.cache);
copy.redirectStartNodes = copyRedirectChains(this.redirectStartNodes);
return copy;
} catch (CloneNotSupportedException e) {
return this;
}
}

private static Map<String, Map<String, String>> copyCache(Map<String, Map<String, String>> cache) {
Map<String, Map<String, String>> copy = new HashMap<>();
cache.forEach((key, strings) -> copy.put(key, new HashMap<>(strings)));
return copy;
}

private static ArrayList<RedirectNode> copyRedirectChains(ArrayList<RedirectNode> nodes) {
if (nodes == null) {
return null;
}
ArrayList<RedirectNode> copy = new ArrayList<>(nodes.size());
nodes.forEach(starter -> copy.add(starter.copyChain()));
return copy;
}

public boolean middlewareExecuted() {return executedMiddleware; }
public void setExecutedMiddleware(boolean value) { executedMiddleware = value; }

Expand Down
24 changes: 24 additions & 0 deletions agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingCallable.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
package dev.aikido.agent_api.context;

import java.util.concurrent.Callable;

public final class ContextPropagatingCallable<T> implements Callable<T> {
private final Callable<T> delegate;
private final ContextObject context;

public ContextPropagatingCallable(Callable<T> delegate, ContextObject context) {
this.delegate = delegate;
this.context = context;
}

@Override
public T call() throws Exception {
ContextObject previous = Context.get();
try {
Context.set(context);
return delegate.call();
} finally {
Context.restore(previous);
}
}
}
22 changes: 22 additions & 0 deletions agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingRunnable.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
package dev.aikido.agent_api.context;

public final class ContextPropagatingRunnable implements Runnable {
private final Runnable delegate;
private final ContextObject context;

public ContextPropagatingRunnable(Runnable delegate, ContextObject context) {
this.delegate = delegate;
this.context = context;
}

@Override
public void run() {
ContextObject previous = Context.get();
try {
Context.set(context);
delegate.run();
} finally {
Context.restore(previous);
}
}
}
33 changes: 33 additions & 0 deletions agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
package dev.aikido.agent_api.context;

import java.util.concurrent.Callable;

public final class ContextPropagation {
private ContextPropagation() {}

public static Runnable wrap(Runnable task) {
if (task == null || task instanceof ContextPropagatingRunnable) {
return task;
Comment on lines +9 to +10

@aikido-pr-checks aikido-pr-checks Bot Aug 17, 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.

🟡 Medium - Re-wrapping a retained task keeps the first request's context

ContextPropagation.wrap(...) returns an existing ContextPropagatingRunnable/Callable unchanged instead of rebinding it to the current request. If an application keeps a task instance and calls wrap() before submitting it for later requests, the wrapper continues to install the original ContextObject, so later async work inherits stale route, user, IP, and forcedProtectionOff state. That can misattribute attack reports to the wrong request and, when the first request had protection forced off, skip vulnerability scanning for subsequent requests.

Show fix

Do not treat already-wrapped tasks as safe to reuse across requests. Either always create a fresh wrapper for each wrap() call, or unwrap and rebind the delegate when the current ContextObject differs from the one captured previously so each submission propagates the caller's current request context.

More info - Reply on this comment to give feedback or ignore the issue.

@Mishenevd Mishenevd Aug 17, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

the wrapper is created inside the agent at submit time and never handed back to the app, so a task is only "already wrapped" within the same submit (same request/context). An app can't retain and resubmit a wrapped task across requests, so the stale-context path is unreachable.

}

ContextObject context = Context.get();
if (context == null) {
return task;
}

return new ContextPropagatingRunnable(task, context.copyForPropagation());
}

public static <T> Callable<T> wrap(Callable<T> task) {
if (task == null || task instanceof ContextPropagatingCallable) {
return task;
}

ContextObject context = Context.get();
if (context == null) {
return task;
}

return new ContextPropagatingCallable<>(task, context.copyForPropagation());
}
}
11 changes: 11 additions & 0 deletions agent_api/src/main/java/dev/aikido/agent_api/storage/RedirectNode.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,17 @@ public void setChild(RedirectNode child) {
this.child = child;
}

public RedirectNode copyChain() {
RedirectNode copy = new RedirectNode(this.url);
RedirectNode source = this.child;
RedirectNode tail = copy;
while (source != null) {
tail = new RedirectNode(tail, source.url);
source = source.child;
}
return copy;
}

@Override
public boolean equals(Object obj) {
if (obj == null || getClass() != obj.getClass()) {
Expand Down
Loading
Loading