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();
}
}
}
42 changes: 41 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,46 @@ 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();
copy.cache = copyCache(this.cache);
copy.redirectStartNodes = copyRedirectChains(this.redirectStartNodes);
copy.headers = copyStringListMap(this.headers);
copy.query = copyStringListMap(this.query);
copy.cookies = copyStringListMap(this.cookies);
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;
}

private static HashMap<String, List<String>> copyStringListMap(HashMap<String, List<String>> map) {
if (map == null) {
return null;
}
HashMap<String, List<String>> copy = new HashMap<>();
map.forEach((key, values) -> copy.put(key, new ArrayList<>(values)));
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