Skip to content

Commit 7b9feb1

Browse files
author
dmitrii
committed
Fix ports.isEmpty() reliability gap: make request correlation survive scheduler hops
PendingHostnamesStore was ThreadLocal, so a WebClient call's registered intent could get lost whenever intent and connect ran on different OS threads - either an app's own .publishOn() or Reactor Netty's internal event-loop dispatch. Confirmed live: an SSRF request to a private IP went out completely unblocked when preceded by a scheduler hop. Made PendingHostnamesStore a global bounded-LRU map instead (so it survives any thread hop), each entry now also carrying the ContextObject captured at registration time, bridged through Reactor's own Context (not ours) via a new reflection-only helper so it works without needing reactor-core as a hard runtime dependency. Also fixes a latent bug this surfaced: two wrappers located URLCollector.report via name-only reflection, which silently broke once a second overload existed.
1 parent b33ebce commit 7b9feb1

11 files changed

Lines changed: 259 additions & 59 deletions

File tree

‎agent/src/main/java/dev/aikido/agent/wrappers/HttpClientSendWrapper.java‎

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
import net.bytebuddy.matcher.ElementMatcher;
77
import net.bytebuddy.matcher.ElementMatchers;
88

9-
import java.lang.reflect.Method;
109
import java.net.HttpURLConnection;
1110
import java.net.MalformedURLException;
1211
import java.net.URL;
@@ -55,13 +54,9 @@ public static void before(
5554
// Load the class from the JAR
5655
Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.URLCollector");
5756

58-
// Run report with "argument"
59-
for (Method method2: clazz.getMethods()) {
60-
if(method2.getName().equals("report")) {
61-
method2.invoke(null, httpRequest.uri().toURL());
62-
break;
63-
}
64-
}
57+
// report(URL) is overloaded (also has a report(URL, ContextObject) variant), so it
58+
// must be looked up by exact signature - matching by name alone could pick either.
59+
clazz.getMethod("report", URL.class).invoke(null, httpRequest.uri().toURL());
6560
classLoader.close(); // Close the class loader
6661
}
6762
}

‎agent/src/main/java/dev/aikido/agent/wrappers/HttpURLConnectionWrapper.java‎

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
import net.bytebuddy.matcher.ElementMatcher;
66
import net.bytebuddy.matcher.ElementMatchers;
77

8-
import java.lang.reflect.Method;
98
import java.net.*;
109

1110
import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
@@ -50,13 +49,9 @@ public static void before(
5049
// Load the class from the JAR
5150
Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.URLCollector");
5251

53-
// Run report with "argument"
54-
for (Method method2: clazz.getMethods()) {
55-
if(method2.getName().equals("report")) {
56-
method2.invoke(null, url);
57-
break;
58-
}
59-
}
52+
// report(URL) is overloaded (also has a report(URL, ContextObject) variant), so it
53+
// must be looked up by exact signature - matching by name alone could pick either.
54+
clazz.getMethod("report", URL.class).invoke(null, url);
6055
classLoader.close(); // Close the class loader
6156
}
6257
}

‎agent/src/main/java/dev/aikido/agent/wrappers/spring/ReactorAikidoContext.java‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
package dev.aikido.agent.wrappers.spring;
2+
3+
import dev.aikido.agent_api.collectors.URLCollector;
4+
import dev.aikido.agent_api.context.ContextObject;
5+
6+
import java.lang.reflect.InvocationHandler;
7+
import java.lang.reflect.Method;
8+
import java.lang.reflect.Proxy;
9+
import java.net.URL;
10+
11+
/**
12+
* Carries the Aikido ContextObject through Reactor's own Context so it survives scheduler hops
13+
* (e.g. .publishOn()) between the incoming WebFlux request and any WebClient calls made while
14+
* handling it - unlike Context.get()'s ThreadLocal, which only sees the current OS thread.
15+
*
16+
* Everything here is Object-typed and goes through reflection, rooted at the classloader of a
17+
* live Mono instance passed in. reactor-core is compileOnly for this module: a *separate* class
18+
* (like this one, as opposed to an @Advice method whose parameter types ByteBuddy resolves
19+
* specially against the woven target's own classloader) that declares Mono, Context or
20+
* ContextView as a concrete parameter/return type throws NoClassDefFoundError at class
21+
* verification time on the agent's own classloader, which has no visibility into the target
22+
* application's classpath. Must be public: the woven target class (in a completely different
23+
* package) needs to call into it directly.
24+
*/
25+
public final class ReactorAikidoContext {
26+
private static final String KEY = "dev.aikido.agent.wrappers.spring.ReactorAikidoContextKey";
27+
28+
private ReactorAikidoContext() {}
29+
30+
// `mono` is a Mono<Void>, returned Object is that same Mono<Void> wrapped with .contextWrite().
31+
public static Object write(Object mono, ContextObject context) {
32+
try {
33+
ClassLoader cl = mono.getClass().getClassLoader();
34+
Class<?> contextClass = Class.forName("reactor.util.context.Context", false, cl);
35+
Class<?> contextViewClass = Class.forName("reactor.util.context.ContextView", false, cl);
36+
Object newContext = contextClass.getMethod("of", Object.class, Object.class)
37+
.invoke(null, KEY, context);
38+
Method contextWrite = mono.getClass().getMethod("contextWrite", contextViewClass);
39+
return contextWrite.invoke(mono, newContext);
40+
} catch (Throwable t) {
41+
return mono;
42+
}
43+
}
44+
45+
// `original` is a Mono<T>. Registers `url` once `original` is actually subscribed to, using
46+
// whatever ContextObject write() captured upstream in the same reactive chain (null if
47+
// none). Returns a Mono<T> equivalent to `original` (or `original` itself if anything here
48+
// fails - registration is best-effort, must never break the actual request).
49+
public static Object deferRegisterUrl(Object original, URL url) {
50+
try {
51+
ClassLoader cl = original.getClass().getClassLoader();
52+
Class<?> functionClass = Class.forName("java.util.function.Function", false, cl);
53+
Method deferContextual = original.getClass().getMethod("deferContextual", functionClass);
54+
InvocationHandler handler = new RegisterUrlHandler(original, url);
55+
Object proxy = Proxy.newProxyInstance(cl, new Class<?>[]{functionClass}, handler);
56+
return deferContextual.invoke(null, proxy);
57+
} catch (Throwable t) {
58+
return original;
59+
}
60+
}
61+
62+
// Not a lambda: constructed from advice code that ByteBuddy inlines into the *target*
63+
// class's bytecode, so a lambda here would become a private synthetic method that the
64+
// target class can't call back into (IllegalAccessError). A plain named class implementing
65+
// InvocationHandler - whose own methods only ever see java.lang.Object - avoids that.
66+
private static final class RegisterUrlHandler implements InvocationHandler {
67+
private final Object original;
68+
private final URL url;
69+
70+
RegisterUrlHandler(Object original, URL url) {
71+
this.original = original;
72+
this.url = url;
73+
}
74+
75+
@Override
76+
public Object invoke(Object proxy, Method method, Object[] args) {
77+
if (!"apply".equals(method.getName()) || args == null || args.length == 0) {
78+
return original;
79+
}
80+
Object ctxView = args[0];
81+
ContextObject context = null;
82+
try {
83+
Method getOrDefault = ctxView.getClass().getMethod("getOrDefault", Object.class, Object.class);
84+
context = (ContextObject) getOrDefault.invoke(ctxView, KEY, null);
85+
} catch (Throwable ignored) {
86+
}
87+
URLCollector.report(url, context);
88+
return original;
89+
}
90+
}
91+
}

‎agent/src/main/java/dev/aikido/agent/wrappers/spring/SpringWebClientWrapper.java‎

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,17 @@
11
package dev.aikido.agent.wrappers.spring;
22

33
import dev.aikido.agent.wrappers.Wrapper;
4-
import dev.aikido.agent_api.collectors.URLCollector;
54
import net.bytebuddy.asm.Advice;
65
import net.bytebuddy.description.method.MethodDescription;
76
import net.bytebuddy.description.type.TypeDescription;
87
import net.bytebuddy.matcher.ElementMatcher;
98
import net.bytebuddy.matcher.ElementMatchers;
109
import org.springframework.web.reactive.function.client.ClientRequest;
10+
import org.springframework.web.reactive.function.client.ClientResponse;
11+
import reactor.core.publisher.Mono;
1112

1213
import java.net.MalformedURLException;
14+
import java.net.URL;
1315

1416
public class SpringWebClientWrapper implements Wrapper {
1517
// Referenced by name (not by .class) in the matchers below: ExchangeFunction is only on
@@ -33,16 +35,21 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
3335
return ElementMatchers.hasSuperType(ElementMatchers.named(EXCHANGE_FUNCTION_CLASS_NAME));
3436
}
3537
public static class SpringWebClientAdvice {
36-
@Advice.OnMethodEnter(suppress = Throwable.class)
37-
public static void before(
38-
@Advice.Argument(0) ClientRequest request
38+
// Registration happens in onExit, wrapped around the returned Mono via
39+
// deferContextual(), rather than eagerly in onEnter. That way it runs at subscribe
40+
// time, reading back whatever ContextObject SpringWebfluxWrapper wrote into Reactor's
41+
// Context (see ReactorAikidoContext) - reliable regardless of scheduler hops between
42+
// the incoming request and this WebClient call, unlike Context.get()'s ThreadLocal.
43+
@Advice.OnMethodExit(suppress = Throwable.class)
44+
public static void after(
45+
@Advice.Argument(0) ClientRequest request,
46+
@Advice.Return(readOnly = false) Mono<ClientResponse> returnValue
3947
) throws MalformedURLException {
40-
if (request == null || request.url() == null) {
48+
if (request == null || request.url() == null || returnValue == null) {
4149
return;
4250
}
43-
// Report the URL before the request is sent, so DNSRecordCollector can match the
44-
// DNS lookup that follows to this outgoing request.
45-
URLCollector.report(request.url().toURL());
51+
URL url = request.url().toURL();
52+
returnValue = (Mono<ClientResponse>) ReactorAikidoContext.deferRegisterUrl(returnValue, url);
4653
}
4754
}
4855
}

‎agent/src/main/java/dev/aikido/agent/wrappers/spring/SpringWebfluxWrapper.java‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,12 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
5555
public record SkipOnWrapper(Mono<Void> newReturnValue) {
5656
}
5757

58+
// Non-skip path result: carries the ContextObject alongside the response so onExit() can
59+
// write it into Reactor's own Context (see ReactorAikidoContext), letting it survive
60+
// scheduler hops before any WebClient call made while handling this request.
61+
public record EnterResult(ServerHttpResponse res, ContextObject context) {
62+
}
63+
5864
public static class SpringWebfluxAdvice {
5965
@Advice.OnMethodEnter(skipOn = SkipOnWrapper.class, suppress = Throwable.class)
6066
public static Object onEnter(
@@ -94,7 +100,7 @@ public static Object onEnter(
94100
return new SkipOnWrapper(res.writeWith(Mono.just(dataBuffer)));
95101
}
96102

97-
return res; // Return to analyze status code in OnMethodExit.
103+
return new EnterResult(res, context); // Return to analyze status code in OnMethodExit.
98104
}
99105

100106
/** onExit()
@@ -105,17 +111,20 @@ public static void onExit(
105111
@Advice.Enter Object enterResult,
106112
@Advice.Return(readOnly = false) Mono<Void> returnValue
107113
) {
108-
// enterResult can be two things : Either the SkipOnWrapper or the ServerHttpResponse
109-
// ServerHttpResponse -> Extract status code.
114+
// enterResult can be two things : Either the SkipOnWrapper or the EnterResult
115+
// EnterResult -> Extract status code, write the context into Reactor's Context.
110116
// SkipOnWrapper -> we blocked a request (e.g. IP Blocking), and are returning the value below
111117
if (enterResult instanceof SkipOnWrapper wrapper && wrapper.newReturnValue() != null) {
112118
returnValue = wrapper.newReturnValue();
113-
} else if (enterResult instanceof ServerHttpResponse res) {
119+
} else if (enterResult instanceof EnterResult er) {
114120
// Report status code of response :
115-
Integer statusCode = res.getRawStatusCode();
121+
Integer statusCode = er.res() != null ? er.res().getRawStatusCode() : null;
116122
if (statusCode != null) {
117123
WebResponseCollector.report(statusCode);
118124
}
125+
if (returnValue != null) {
126+
returnValue = (Mono<Void>) ReactorAikidoContext.write(returnValue, er.context());
127+
}
119128
}
120129
}
121130
}

‎agent_api/src/main/java/dev/aikido/agent_api/collectors/DNSRecordCollector.java‎

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package dev.aikido.agent_api.collectors;
22

33
import dev.aikido.agent_api.context.Context;
4+
import dev.aikido.agent_api.context.ContextObject;
45
import dev.aikido.agent_api.storage.HostnamesStore;
56
import dev.aikido.agent_api.storage.PendingHostnamesStore;
67
import dev.aikido.agent_api.storage.ServiceConfigStore;
@@ -32,15 +33,40 @@ private DNSRecordCollector() {}
3233

3334
public static void report(String hostname, InetAddress[] inetAddresses) {
3435
// InetAddress.getAllByName() resolves everything in one call, so it's safe to consume.
35-
process(hostname, inetAddresses, PendingHostnamesStore.getAndRemove(hostname), INET_ADDRESS_OPERATION_NAME);
36+
withCapturedContext(hostname, () ->
37+
process(hostname, inetAddresses, PendingHostnamesStore.getAndRemove(hostname), INET_ADDRESS_OPERATION_NAME));
3638
}
3739

3840
// For clients that resolve their own DNS (e.g. Reactor Netty, used by Spring's WebClient) or
3941
// connect straight to an IP literal. A single request can trigger multiple connect() calls to
4042
// the same hostname (IPv4 then IPv6), so unlike report(), this peeks the pending port instead
4143
// of consuming it - consuming on the first attempt would let a later attempt bypass SSRF.
4244
public static void reportConnect(String hostname, InetAddress resolvedAddress) {
43-
process(hostname, new InetAddress[]{resolvedAddress}, PendingHostnamesStore.getPorts(hostname), SOCKET_CHANNEL_OPERATION_NAME);
45+
withCapturedContext(hostname, () ->
46+
process(hostname, new InetAddress[]{resolvedAddress}, PendingHostnamesStore.getPorts(hostname), SOCKET_CHANNEL_OPERATION_NAME));
47+
}
48+
49+
// Restores the ContextObject captured when this hostname's pending entry was registered
50+
// (PendingHostnamesStore is global, not thread-local) so SSRFDetector's Context.get() sees
51+
// the request that actually triggered the outbound call, even if we're running on a
52+
// different thread than the one that registered it.
53+
private static void withCapturedContext(String hostname, Runnable action) {
54+
ContextObject capturedContext = PendingHostnamesStore.getContext(hostname);
55+
if (capturedContext == null) {
56+
action.run();
57+
return;
58+
}
59+
ContextObject previous = Context.get();
60+
Context.set(capturedContext);
61+
try {
62+
action.run();
63+
} finally {
64+
if (previous != null) {
65+
Context.set(previous);
66+
} else {
67+
Context.reset();
68+
}
69+
}
4470
}
4571

4672
private static void process(String hostname, InetAddress[] inetAddresses, Set<Integer> ports, String operationName) {

‎agent_api/src/main/java/dev/aikido/agent_api/collectors/URLCollector.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package dev.aikido.agent_api.collectors;
22

3+
import dev.aikido.agent_api.context.Context;
4+
import dev.aikido.agent_api.context.ContextObject;
35
import dev.aikido.agent_api.helpers.logging.LogManager;
46
import dev.aikido.agent_api.helpers.logging.Logger;
57
import dev.aikido.agent_api.storage.PendingHostnamesStore;
@@ -13,14 +15,21 @@ public final class URLCollector {
1315

1416
private URLCollector() {}
1517
public static void report(URL url) {
18+
report(url, Context.get());
19+
}
20+
21+
// Used where the caller already resolved the correct context itself instead of relying on
22+
// Context.get() (e.g. Spring WebClient reading it back from Reactor's own Context, which
23+
// survives scheduler hops that break Context.get()'s ThreadLocal).
24+
public static void report(URL url, ContextObject context) {
1625
if (url != null) {
1726
if (!url.getProtocol().startsWith("http")) {
1827
return; // Non-HTTP(S) URL
1928
}
2029
logger.trace("Adding a new URL to the cache: %s", url);
2130
// Store hostname+port in the pending store so DNSRecordCollector can pick it
2231
// up during the DNS lookup that follows, for SSRF detection and outbound hostnames
23-
PendingHostnamesStore.add(url.getHost(), getPortFromURL(url));
32+
PendingHostnamesStore.add(url.getHost(), getPortFromURL(url), context);
2433
}
2534
}
2635
}

‎agent_api/src/main/java/dev/aikido/agent_api/collectors/WebRequestCollector.java‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
import dev.aikido.agent_api.helpers.logging.LogManager;
99
import dev.aikido.agent_api.helpers.logging.Logger;
1010
import dev.aikido.agent_api.storage.AttackQueue;
11-
import dev.aikido.agent_api.storage.PendingHostnamesStore;
1211
import dev.aikido.agent_api.storage.ServiceConfigStore;
1312
import dev.aikido.agent_api.storage.ServiceConfiguration;
1413
import dev.aikido.agent_api.storage.attack_wave_detector.AttackWaveDetectorStore;
@@ -41,10 +40,6 @@ public static Res report(ContextObject newContext) {
4140
// clear context
4241
Context.reset();
4342

44-
// Flush pending hostnames on every context change to prevent the store from
45-
// growing unboundedly when a thread is reused across multiple requests.
46-
PendingHostnamesStore.clear();
47-
4843
if (config.isIpBypassed(newContext.getRemoteAddress())) {
4944
return null; // do not set context when the IP address is bypassed (zen = off)
5045
}

0 commit comments

Comments
 (0)