diff --git a/agent_api/src/main/java/dev/aikido/agent_api/context/Context.java b/agent_api/src/main/java/dev/aikido/agent_api/context/Context.java index 0389cc90..0382f89b 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/context/Context.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/context/Context.java @@ -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(); + } + } } diff --git a/agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java index 33935838..42eeb099 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java @@ -6,7 +6,7 @@ import java.util.*; -public class ContextObject { +public class ContextObject implements Cloneable { protected String method; protected String source; protected String url; @@ -25,6 +25,46 @@ public class ContextObject { protected transient Map> cache = new HashMap<>(); protected transient Optional 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> copyCache(Map> cache) { + Map> copy = new HashMap<>(); + cache.forEach((key, strings) -> copy.put(key, new HashMap<>(strings))); + return copy; + } + + private static ArrayList copyRedirectChains(ArrayList nodes) { + if (nodes == null) { + return null; + } + ArrayList copy = new ArrayList<>(nodes.size()); + nodes.forEach(starter -> copy.add(starter.copyChain())); + return copy; + } + + private static HashMap> copyStringListMap(HashMap> map) { + if (map == null) { + return null; + } + HashMap> 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; } diff --git a/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingCallable.java b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingCallable.java new file mode 100644 index 00000000..f7f8b65b --- /dev/null +++ b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingCallable.java @@ -0,0 +1,24 @@ +package dev.aikido.agent_api.context; + +import java.util.concurrent.Callable; + +public final class ContextPropagatingCallable implements Callable { + private final Callable delegate; + private final ContextObject context; + + public ContextPropagatingCallable(Callable 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); + } + } +} diff --git a/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingRunnable.java b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingRunnable.java new file mode 100644 index 00000000..2c81e42e --- /dev/null +++ b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingRunnable.java @@ -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); + } + } +} diff --git a/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagation.java b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagation.java new file mode 100644 index 00000000..fb63a3c1 --- /dev/null +++ b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagation.java @@ -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; + } + + ContextObject context = Context.get(); + if (context == null) { + return task; + } + + return new ContextPropagatingRunnable(task, context.copyForPropagation()); + } + + public static Callable wrap(Callable task) { + if (task == null || task instanceof ContextPropagatingCallable) { + return task; + } + + ContextObject context = Context.get(); + if (context == null) { + return task; + } + + return new ContextPropagatingCallable<>(task, context.copyForPropagation()); + } +} diff --git a/agent_api/src/main/java/dev/aikido/agent_api/storage/RedirectNode.java b/agent_api/src/main/java/dev/aikido/agent_api/storage/RedirectNode.java index 690ee875..a7c372bb 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/storage/RedirectNode.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/storage/RedirectNode.java @@ -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()) { diff --git a/agent_api/src/test/java/context/ContextPropagationTest.java b/agent_api/src/test/java/context/ContextPropagationTest.java new file mode 100644 index 00000000..9623a102 --- /dev/null +++ b/agent_api/src/test/java/context/ContextPropagationTest.java @@ -0,0 +1,349 @@ +package context; + +import dev.aikido.agent_api.context.Context; +import dev.aikido.agent_api.context.ContextObject; +import dev.aikido.agent_api.context.ContextPropagatingCallable; +import dev.aikido.agent_api.context.ContextPropagatingRunnable; +import dev.aikido.agent_api.context.ContextPropagation; +import dev.aikido.agent_api.storage.RedirectNode; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.net.URL; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.Callable; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; + +class ContextPropagationTest { + @AfterEach + void tearDown() { + Context.reset(); + } + + @Test + void wrapRunnableReturnsNullForNullTask() { + Assertions.assertNull(ContextPropagation.wrap((Runnable) null)); + } + + @Test + void wrapCallableReturnsNullForNullTask() { + Assertions.assertNull(ContextPropagation.wrap((Callable) null)); + } + + @Test + void wrapRunnableReturnsOriginalTaskWhenNoContextIsSet() { + Runnable task = () -> {}; + + Runnable wrapped = ContextPropagation.wrap(task); + + Assertions.assertSame(task, wrapped); + } + + @Test + void wrapCallableReturnsOriginalTaskWhenNoContextIsSet() { + Callable task = () -> "ok"; + + Callable wrapped = ContextPropagation.wrap(task); + + Assertions.assertSame(task, wrapped); + } + + @Test + void wrapRunnableReturnsSameTaskWhenAlreadyWrapped() { + ContextObject contextObject = new ContextObject(); + Runnable task = new ContextPropagatingRunnable(() -> {}, contextObject); + + Runnable wrapped = ContextPropagation.wrap(task); + + Assertions.assertSame(task, wrapped); + } + + @Test + void wrapCallableReturnsSameTaskWhenAlreadyWrapped() { + ContextObject contextObject = new ContextObject(); + Callable task = new ContextPropagatingCallable<>(() -> "ok", contextObject); + + Callable wrapped = ContextPropagation.wrap(task); + + Assertions.assertSame(task, wrapped); + } + + @Test + void wrapRunnableCapturesCurrentContext() { + ContextObject requestContext = new ContextObject(); + requestContext.setRateLimitGroup("req"); + Context.set(requestContext); + + AtomicReference contextDuringRun = new AtomicReference<>(); + Runnable wrapped = ContextPropagation.wrap(() -> contextDuringRun.set(Context.get())); + + Context.reset(); + wrapped.run(); + + Assertions.assertEquals("req", contextDuringRun.get().getRateLimitGroup()); + Assertions.assertNull(Context.get(), "Expected worker context to be cleared after task execution"); + } + + @Test + void wrapCallableCapturesCurrentContext() throws Exception { + ContextObject requestContext = new ContextObject(); + requestContext.setRateLimitGroup("req"); + Context.set(requestContext); + + Callable wrapped = ContextPropagation.wrap(Context::get); + + Context.reset(); + ContextObject contextDuringCall = wrapped.call(); + + Assertions.assertEquals("req", contextDuringCall.getRateLimitGroup()); + Assertions.assertNull(Context.get(), "Expected worker context to be cleared after task execution"); + } + + @Test + void wrapPropagatesAnIsolatedSnapshotNotTheLiveContext() { + ContextObject requestContext = new ContextObject(); + Context.set(requestContext); + + AtomicReference contextDuringRun = new AtomicReference<>(); + Runnable wrapped = ContextPropagation.wrap(() -> contextDuringRun.set(Context.get())); + + Context.reset(); + wrapped.run(); + + ContextObject propagated = contextDuringRun.get(); + Assertions.assertNotSame(requestContext, propagated); + Assertions.assertNotSame(requestContext.getCache(), propagated.getCache()); + } + + @Test + void snapshotDeepCopiesRedirectChainNodes() throws Exception { + RedirectableContext requestContext = new RedirectableContext(); + RedirectNode starter = new RedirectNode(new URL("http://origin")); + new RedirectNode(starter, new URL("http://dest")); + requestContext.addRedirectNode(starter); + + ContextObject copy = requestContext.copyForPropagation(); + + RedirectNode originalStarter = requestContext.getRedirectStartNodes().get(0); + RedirectNode copiedStarter = copy.getRedirectStartNodes().get(0); + + Assertions.assertNotSame(originalStarter, copiedStarter); + Assertions.assertNotSame(originalStarter.getChild(), copiedStarter.getChild()); + Assertions.assertEquals("http://dest", copiedStarter.getChild().getUrl().toString()); + } + + @Test + void snapshotDeepCopiesCacheInnerMaps() { + ContextObject requestContext = new ContextObject(); + Map inner = new HashMap<>(); + inner.put("k", "v"); + requestContext.getCache().put("body", inner); + + ContextObject copy = requestContext.copyForPropagation(); + + Assertions.assertNotSame(requestContext.getCache().get("body"), copy.getCache().get("body")); + Assertions.assertEquals("v", copy.getCache().get("body").get("k")); + } + + @Test + void snapshotCacheMutationDoesNotLeakToOriginal() { + ContextObject requestContext = new ContextObject(); + Map inner = new HashMap<>(); + inner.put("k", "v"); + requestContext.getCache().put("body", inner); + + ContextObject copy = requestContext.copyForPropagation(); + copy.getCache().get("body").put("injected", "x"); + copy.getCache().put("query", new HashMap<>()); + + Assertions.assertFalse(requestContext.getCache().get("body").containsKey("injected")); + Assertions.assertFalse(requestContext.getCache().containsKey("query")); + } + + @Test + void snapshotRedirectChainMutationDoesNotLeakToOriginal() throws Exception { + RedirectableContext requestContext = new RedirectableContext(); + RedirectNode starter = new RedirectNode(new URL("http://origin")); + new RedirectNode(starter, new URL("http://dest")); + requestContext.addRedirectNode(starter); + + ContextObject copy = requestContext.copyForPropagation(); + RedirectNode copiedDest = copy.getRedirectStartNodes().get(0).getChild(); + new RedirectNode(copiedDest, new URL("http://extra")); + + RedirectNode originalDest = requestContext.getRedirectStartNodes().get(0).getChild(); + Assertions.assertNull(originalDest.getChild()); + } + + @Test + void copyChainPreservesMultiNodeChain() throws Exception { + RedirectNode a = new RedirectNode(new URL("http://a")); + RedirectNode b = new RedirectNode(a, new URL("http://b")); + RedirectNode c = new RedirectNode(b, new URL("http://c")); + + RedirectNode copy = a.copyChain(); + + Assertions.assertEquals("http://a", copy.getUrl().toString()); + Assertions.assertEquals("http://b", copy.getChild().getUrl().toString()); + Assertions.assertEquals("http://c", copy.getChild().getChild().getUrl().toString()); + Assertions.assertNotSame(b, copy.getChild()); + Assertions.assertNotSame(c, copy.getChild().getChild()); + } + + @Test + void snapshotHeaderListMutationDoesNotLeakToOriginal() { + HeaderfulContext requestContext = new HeaderfulContext(); + ArrayList values = new ArrayList<>(); + values.add("v"); + requestContext.getHeaders().put("h", values); + + ContextObject copy = requestContext.copyForPropagation(); + copy.getHeaders().get("h").add("injected"); + copy.getHeaders().put("h2", new ArrayList<>()); + + Assertions.assertFalse(requestContext.getHeaders().get("h").contains("injected")); + Assertions.assertFalse(requestContext.getHeaders().containsKey("h2")); + } + + private static class HeaderfulContext extends ContextObject { + HeaderfulContext() { + this.headers = new HashMap<>(); + } + } + + private static class RedirectableContext extends ContextObject { + RedirectableContext() { + this.redirectStartNodes = new ArrayList<>(); + } + } + + @Test + void contextPropagatingRunnableRestoresPreviousWorkerContext() { + ContextObject capturedContext = new ContextObject(); + ContextObject previousWorkerContext = new ContextObject(); + + ContextPropagatingRunnable task = new ContextPropagatingRunnable( + () -> Assertions.assertSame(capturedContext, Context.get()), + capturedContext + ); + + Context.set(previousWorkerContext); + task.run(); + + Assertions.assertSame(previousWorkerContext, Context.get()); + } + + @Test + void contextPropagatingCallableRestoresPreviousWorkerContext() throws Exception { + ContextObject capturedContext = new ContextObject(); + ContextObject previousWorkerContext = new ContextObject(); + + ContextPropagatingCallable task = new ContextPropagatingCallable<>( + Context::get, + capturedContext + ); + + Context.set(previousWorkerContext); + ContextObject contextDuringCall = task.call(); + + Assertions.assertSame(capturedContext, contextDuringCall); + Assertions.assertSame(previousWorkerContext, Context.get()); + } + + @Test + void contextPropagatingRunnableClearsContextWhenWorkerHadNoPreviousContext() { + ContextObject capturedContext = new ContextObject(); + + ContextPropagatingRunnable task = new ContextPropagatingRunnable( + () -> Assertions.assertSame(capturedContext, Context.get()), + capturedContext + ); + + Context.reset(); + task.run(); + + Assertions.assertNull(Context.get()); + } + + @Test + void contextPropagatingCallableClearsContextWhenWorkerHadNoPreviousContext() throws Exception { + ContextObject capturedContext = new ContextObject(); + + ContextPropagatingCallable task = new ContextPropagatingCallable<>( + Context::get, + capturedContext + ); + + Context.reset(); + ContextObject contextDuringCall = task.call(); + + Assertions.assertSame(capturedContext, contextDuringCall); + Assertions.assertNull(Context.get()); + } + + @Test + void contextPropagatingRunnableRestoresPreviousWorkerContextAfterException() { + ContextObject capturedContext = new ContextObject(); + ContextObject previousWorkerContext = new ContextObject(); + + ContextPropagatingRunnable task = new ContextPropagatingRunnable( + () -> { + throw new IllegalStateException("boom"); + }, + capturedContext + ); + + Context.set(previousWorkerContext); + + Assertions.assertThrows(IllegalStateException.class, task::run); + Assertions.assertSame(previousWorkerContext, Context.get()); + } + + @Test + void contextPropagatingCallableRestoresPreviousWorkerContextAfterException() { + ContextObject capturedContext = new ContextObject(); + ContextObject previousWorkerContext = new ContextObject(); + + ContextPropagatingCallable task = new ContextPropagatingCallable<>( + () -> { + throw new IllegalStateException("boom"); + }, + capturedContext + ); + + Context.set(previousWorkerContext); + + Assertions.assertThrows(IllegalStateException.class, task::call); + Assertions.assertSame(previousWorkerContext, Context.get()); + } + + @Test + void wrappedRunnableRunsDelegate() { + ContextObject requestContext = new ContextObject(); + Context.set(requestContext); + + AtomicBoolean delegateCalled = new AtomicBoolean(false); + Runnable wrapped = ContextPropagation.wrap(() -> delegateCalled.set(true)); + + Context.reset(); + wrapped.run(); + + Assertions.assertTrue(delegateCalled.get()); + } + + @Test + void wrappedCallableReturnsDelegateResult() throws Exception { + ContextObject requestContext = new ContextObject(); + Context.set(requestContext); + + Callable wrapped = ContextPropagation.wrap(() -> "ok"); + + Context.reset(); + + Assertions.assertEquals("ok", wrapped.call()); + } +}