Track and protect WebClient outbound requests, fix private-IP SSRF regression

Follow-up to the reverted #308/#310. Customer flood was InetAddress.getAllByName()
picking up Reactor Netty's own DNS-resolver bootstrap noise (0.0.0.0, ::,
/etc/resolv.conf nameservers) as "new outbound connections" on port 0, and
blocking them in lockdown mode. #310 fixed the flood with an early return in
DNSRecordCollector.report() that also skipped the SSRF check below it -
verified with a regression test that this let an attacker-supplied private-IP
literal (e.g. a webhook field pointing straight at 169.254.169.254) through
undetected.

Investigating further found the actual root cause is bigger: Spring's
WebClient was never instrumented at all, and Reactor Netty's default HTTP
client bypasses InetAddress.getAllByName() entirely (it uses its own async
DNS resolver). So even after wrapping WebClient to register pending ports,
DNSRecordCollector was never invoked for real WebClient targets - confirmed
empirically via trace logs against a live running app, with distinct markers
proving InetAddressWrapper never fires for WebClient/Reactor Netty traffic in
this configuration. WebClient had zero outbound-domain visibility and zero
SSRF protection, independent of the original bug.

- DNSRecordCollector: narrow the private-IP-literal gate to only skip
  recording + outbound blocking when there's no pending port (genuine infra
  noise). SSRF checks are unconditional again, fixing the bypass above.
- SpringWebClientWrapper: register pending host+port for every WebClient
  request by hooking ExchangeFunction.exchange(), the interface every
  WebClient call goes through, same pattern as the existing OkHttp/Apache/JDK
  HttpClient wrappers. Uses string-based ByteBuddy matchers
  (hasSuperType(named(...))) instead of .class literals, since spring-webflux
  is compileOnly and only present on the target app's classloader - a .class
  reference in the matcher crashes the agent at premain with
  NoClassDefFoundError.
- SocketChannelWrapper: hook java.nio.channels.SocketChannel.connect(), the
  JDK-level call every NIO-based client (including Reactor Netty) makes once
  it has a resolved address, regardless of which DNS resolver produced it.
  This is what actually closes the gap for WebClient, and it also catches
  literal IP targets that never go through any resolver at all. Not
  Netty-specific instrumentation - it's a generic JDK hook with no references
  to io.netty.* types.
- DNSRecordCollector.reportConnect(): entry point for the new hook. Peeks the
  pending port instead of consuming it (report()'s getAndRemove), because a
  single request can trigger multiple connect() calls to the same hostname
  (e.g. the IPv4 then IPv6 address of a dual-stack host like localhost).
  Consuming on the first attempt let a blocked SSRF target succeed on the
  second attempt via the other address family - found live, fixed, covered by
  a regression test.
- PendingHostnamesStore: peeking instead of consuming means entries rely on
  WebRequestCollector's per-incoming-request clear() for cleanup, which never
  fires for WebClient calls made outside any incoming-request context (e.g. a
  @scheduled background task). Capped the store at 1000 entries per thread,
  evicting the least-recently-used one once exceeded - the same bounded-LRU
  pattern (LinkedHashMap with accessOrder=true + removeEldestEntry())
  already used by Hostnames.java for the same class of problem. Deliberately
  not a time-based TTL, to avoid a timing-dependent race reopening the
  dual-stack gap under load.
- RequestController (SpringWebfluxSampleApp): new /api/request endpoint used
  to validate all of the above against a real running app end to end.

Known limitation, not fixed here: Spring WebFlux has no request-body taint
tracking at all (SpringWebfluxContextObject never populates
ContextObject.body), so SSRF via JSON body can't be detected for WebFlux apps
regardless of this change - flagged separately, doesn't regress anything.

Author: dmitrii

agent/build.gradle

@@ -12,6 +12,7 @@ dependencies {
     compileOnly 'io.projectreactor.netty:reactor-netty-http:1.2.1' // For Spring Webflux
     compileOnly 'io.javalin:javalin:6.4.0'
     compileOnly 'org.springframework:spring-web:5.3.20'
+    compileOnly 'org.springframework:spring-webflux:5.3.20' // For Spring WebClient
 }
 
 shadowJar {

agent/src/main/java/dev/aikido/agent/Wrappers.java

@@ -9,6 +9,7 @@
 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.SpringWebClientWrapper;
 
 import java.util.Arrays;
 import java.util.List;
@@ -30,11 +31,13 @@ private Wrappers() {}
             // SSRF/HTTP wrappers
             new HttpURLConnectionWrapper(),
             new InetAddressWrapper(),
+            new SocketChannelWrapper(),
             new HttpClientWrapper(),
             new HttpConnectionRedirectWrapper(),
             new HttpClientSendWrapper(),
             new OkHttpWrapper(),
             new ApacheHttpClientWrapper(),
+            new SpringWebClientWrapper(),
 
             new PathWrapper(),
             new PathsWrapper(),

agent/src/main/java/dev/aikido/agent/wrappers/SocketChannelWrapper.java

@@ -0,0 +1,87 @@
+package dev.aikido.agent.wrappers;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.matcher.ElementMatchers;
+
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.net.InetAddress;
+import java.net.InetSocketAddress;
+import java.net.MalformedURLException;
+import java.net.SocketAddress;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.nio.channels.SocketChannel;
+
+public class SocketChannelWrapper implements Wrapper {
+    public String getName() {
+        // Wrap connect(SocketAddress) on SocketChannel. Clients that resolve hostnames with
+        // their own DNS resolver instead of InetAddress.getAllByName() (e.g. Reactor Netty's
+        // async resolver, used by default by Spring's WebClient) never trigger
+        // InetAddressWrapper, so this is the only point where we see the resolved address
+        // before the connection is made. Also catches literal IP targets, which never go
+        // through any resolver at all.
+        // https://docs.oracle.com/javase/8/docs/api/java/nio/channels/SocketChannel.html#connect-java.net.SocketAddress-
+        return SocketChannelAdvice.class.getName();
+    }
+    public ElementMatcher<? super MethodDescription> getMatcher() {
+        return ElementMatchers.named("connect");
+    }
+    @Override
+    public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+        return ElementMatchers.isSubTypeOf(SocketChannel.class);
+    }
+    public static class SocketChannelAdvice {
+        // Since we have to wrap a native Java Class stuff gets more complicated
+        // The classpath is not the same anymore, and we can't import our modules directly.
+        // To bypass this issue we load collectors from a .jar file
+        @Advice.OnMethodEnter
+        public static void before(
+                @Advice.Argument(0) SocketAddress remoteAddress
+        ) throws Throwable {
+            if (!(remoteAddress instanceof InetSocketAddress)) {
+                return;
+            }
+            InetSocketAddress inetSocketAddress = (InetSocketAddress) remoteAddress;
+            InetAddress resolvedAddress = inetSocketAddress.getAddress();
+            if (resolvedAddress == null) {
+                // Unresolved: nothing to report yet, connect() will throw on its own.
+                return;
+            }
+            String hostname = inetSocketAddress.getHostString();
+
+            String jarFilePath = System.getProperty("AIK_agent_api_jar");
+            URLClassLoader classLoader = null;
+            try {
+                URL[] urls = { new URL(jarFilePath) };
+                classLoader = new URLClassLoader(urls);
+            } catch (MalformedURLException ignored) {}
+            if (classLoader == null) {
+                return;
+            }
+
+            try {
+                // Load the class from the JAR
+                Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.DNSRecordCollector");
+
+                // Run reportConnect with "argument"
+                for (Method method2: clazz.getMethods()) {
+                    if(method2.getName().equals("reportConnect")) {
+                        method2.invoke(null, hostname, resolvedAddress);
+                        break;
+                    }
+                }
+                classLoader.close(); // Close the class loader
+            } catch (InvocationTargetException invocationTargetException) {
+                if(invocationTargetException.getCause().toString().startsWith("dev.aikido.agent_api.vulnerabilities")) {
+                    throw invocationTargetException.getCause();
+                }
+                // Ignore non-aikido throwables.
+            } catch(Throwable e) {
+                System.out.println("AIKIDO: " + e.getMessage());
+            }
+        }
+    }
+}

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

@@ -0,0 +1,48 @@
+package dev.aikido.agent.wrappers.spring;
+
+import dev.aikido.agent.wrappers.Wrapper;
+import dev.aikido.agent_api.collectors.URLCollector;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.matcher.ElementMatchers;
+import org.springframework.web.reactive.function.client.ClientRequest;
+
+import java.net.MalformedURLException;
+
+public class SpringWebClientWrapper implements Wrapper {
+    // Referenced by name (not by .class) in the matchers below: ExchangeFunction is only on
+    // the target application's classloader (spring-webflux is compileOnly here), not on the
+    // agent's own classloader, so a .class literal would throw NoClassDefFoundError at premain.
+    private static final String EXCHANGE_FUNCTION_CLASS_NAME =
+            "org.springframework.web.reactive.function.client.ExchangeFunction";
+
+    public String getName() {
+        // Wrap exchange(ClientRequest) on ExchangeFunction, the interface every WebClient
+        // request goes through before Reactor Netty resolves/connects.
+        // https://docs.spring.io/spring-framework/docs/5.3.20/javadoc-api/org/springframework/web/reactive/function/client/ExchangeFunction.html
+        return SpringWebClientAdvice.class.getName();
+    }
+    public ElementMatcher<? super MethodDescription> getMatcher() {
+        return ElementMatchers.isDeclaredBy(getTypeMatcher())
+                .and(ElementMatchers.named("exchange"));
+    }
+    @Override
+    public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+        return ElementMatchers.hasSuperType(ElementMatchers.named(EXCHANGE_FUNCTION_CLASS_NAME));
+    }
+    public static class SpringWebClientAdvice {
+        @Advice.OnMethodEnter(suppress = Throwable.class)
+        public static void before(
+                @Advice.Argument(0) ClientRequest request
+        ) throws MalformedURLException {
+            if (request == null || request.url() == null) {
+                return;
+            }
+            // Report the URL before the request is sent, so DNSRecordCollector can match the
+            // DNS lookup that follows to this outgoing request.
+            URLCollector.report(request.url().toURL());
+        }
+    }
+}

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

@@ -12,6 +12,7 @@
 import dev.aikido.agent_api.vulnerabilities.ssrf.SSRFException;
 import dev.aikido.agent_api.helpers.logging.LogManager;
 import dev.aikido.agent_api.helpers.logging.Logger;
+import dev.aikido.agent_api.vulnerabilities.ssrf.IsPrivateIP;
 import dev.aikido.agent_api.vulnerabilities.ssrf.StoredSSRFDetector;
 import dev.aikido.agent_api.vulnerabilities.ssrf.StoredSSRFException;
 
@@ -26,30 +27,44 @@
 public final class DNSRecordCollector {
     private DNSRecordCollector() {}
     private static final Logger logger = LogManager.getLogger(DNSRecordCollector.class);
-    private static final String OPERATION_NAME = "java.net.InetAddress.getAllByName";
+    private static final String INET_ADDRESS_OPERATION_NAME = "java.net.InetAddress.getAllByName";
+    private static final String SOCKET_CHANNEL_OPERATION_NAME = "java.nio.channels.SocketChannel.connect";
+
     public static void report(String hostname, InetAddress[] inetAddresses) {
+        // Consume pending ports recorded by URLCollector for this hostname. InetAddress.getAllByName()
+        // resolves every address for the hostname in one call, so removing them here ensures each
+        // (hostname, port) pair is counted exactly once per DNS lookup.
+        process(hostname, inetAddresses, PendingHostnamesStore.getAndRemove(hostname), INET_ADDRESS_OPERATION_NAME);
+    }
+
+    // Used for clients that resolve their own DNS (bypassing InetAddress.getAllByName entirely,
+    // e.g. Reactor Netty's async resolver used by default by Spring's WebClient) or that connect
+    // straight to an IP literal without any resolution step. The resolved address is only known
+    // at the point of the actual connection, one address at a time.
+    //
+    // A single outbound request can trigger multiple connect() calls to the same hostname (e.g.
+    // trying the IPv4 then the IPv6 address of a dual-stack host). Unlike report(), this peeks the
+    // pending port instead of consuming it, so every connect attempt for that hostname is still
+    // checked; consuming on the first attempt would let a later attempt to a different resolved
+    // address for the same hostname bypass SSRF/blocking entirely. The entry is cleared per
+    // incoming request by WebRequestCollector instead of being consumed here.
+    public static void reportConnect(String hostname, InetAddress resolvedAddress) {
+        process(hostname, new InetAddress[]{resolvedAddress}, PendingHostnamesStore.getPorts(hostname), SOCKET_CHANNEL_OPERATION_NAME);
+    }
+
+    private static void process(String hostname, InetAddress[] inetAddresses, Set<Integer> ports, String operationName) {
         try {
             logger.trace("DNSRecordCollector called with %s & inet addresses: %s", hostname, List.of(inetAddresses));
 
             // store stats
-            StatisticsStore.registerCall("java.net.InetAddress.getAllByName", OperationKind.OUTGOING_HTTP_OP);
-
-            // Consume pending ports recorded by URLCollector for this hostname.
-            // Removing them here ensures each (hostname, port) pair is counted exactly once.
-            Set<Integer> ports = PendingHostnamesStore.getAndRemove(hostname);
-            if (!ports.isEmpty()) {
-                for (int port : ports) {
-                    HostnamesStore.incrementHits(hostname, port);
-                }
-            } else {
-                // We still need to report a hit to the hostname for outbound domain blocking
-                HostnamesStore.incrementHits(hostname, 0);
-            }
-
-            // Block if the hostname is in the blocked domains list
-            if (ServiceConfigStore.shouldBlockOutgoingRequest(hostname)) {
-                logger.debug("Blocking DNS lookup for domain: %s", hostname);
-                throw BlockedOutboundException.get();
+            StatisticsStore.registerCall(operationName, OperationKind.OUTGOING_HTTP_OP);
+
+            // No pending port + a private IP literal = infrastructure noise (Netty's resolver
+            // bootstrap resolving nameserver/bind addresses, a client connecting directly by
+            // IP, etc), not a real outbound request. SSRF checks below still run regardless.
+            boolean isInfrastructureNoise = ports.isEmpty() && IsPrivateIP.isPrivateIp(hostname);
+            if (!isInfrastructureNoise) {
+                recordHitAndEnforceBlocklist(hostname, ports);
             }
 
             // Convert inetAddresses array to a List of IP strings :
@@ -62,7 +77,7 @@ public static void report(String hostname, InetAddress[] inetAddresses) {
             for (int port : ports) {
                 logger.debug("Hostname: %s, Port: %s, IPs: %s", hostname, port, ipAddresses);
 
-                Attack attack = SSRFDetector.run(hostname, port, ipAddresses, OPERATION_NAME);
+                Attack attack = SSRFDetector.run(hostname, port, ipAddresses, operationName);
                 if (attack == null) {
                     continue;
                 }
@@ -81,7 +96,7 @@ public static void report(String hostname, InetAddress[] inetAddresses) {
 
             // We don't need the context object to check for stored ssrf, but we do want to run this after our other
             // SSRF checks, making sure if it's a normal ssrf attack it gets reported like that.
-            Attack storedSsrfAttack = new StoredSSRFDetector().run(hostname, ipAddresses, OPERATION_NAME);
+            Attack storedSsrfAttack = new StoredSSRFDetector().run(hostname, ipAddresses, operationName);
             if (storedSsrfAttack != null) {
                 attackDetected(storedSsrfAttack, Context.get());
 
@@ -97,4 +112,21 @@ public static void report(String hostname, InetAddress[] inetAddresses) {
             logger.trace(e);
         }
     }
+
+    // Throws BlockedOutboundException, propagated by process()'s catch clause above.
+    private static void recordHitAndEnforceBlocklist(String hostname, Set<Integer> ports) {
+        if (ports.isEmpty()) {
+            // No pending port: still record a hit so lockdown mode can still block unknown domains.
+            HostnamesStore.incrementHits(hostname, 0);
+        } else {
+            for (int port : ports) {
+                HostnamesStore.incrementHits(hostname, port);
+            }
+        }
+
+        if (ServiceConfigStore.shouldBlockOutgoingRequest(hostname)) {
+            logger.debug("Blocking DNS lookup for domain: %s", hostname);
+            throw BlockedOutboundException.get();
+        }
+    }
 }

agent_api/src/main/java/dev/aikido/agent_api/storage/PendingHostnamesStore.java

@@ -4,14 +4,30 @@
 
 /**
  * Thread-local bridge between URLCollector and DNSRecordCollector.
- * URLCollector records hostname+port here; DNSRecordCollector reads and removes the entry
- * so each (hostname, port) pair is processed exactly once per DNS lookup.
+ * URLCollector records hostname+port here; DNSRecordCollector.report() (fed by
+ * InetAddress.getAllByName(), which resolves everything in one call) reads and removes the
+ * entry so each (hostname, port) pair is processed exactly once per DNS lookup.
+ * DNSRecordCollector.reportConnect() (fed by SocketChannel.connect(), which fires once per
+ * connect attempt) instead peeks the entry, since a single outbound request can trigger
+ * multiple connect attempts to the same hostname (e.g. IPv4 then IPv6 for a dual-stack host).
+ *
+ * Entries are normally cleared per incoming request by WebRequestCollector, but a peeked
+ * entry added outside any incoming-request context (e.g. a WebClient call from a @Scheduled
+ * task) would never be cleared that way. Capped at MAX_ENTRIES per thread, evicting the least
+ * recently used entry once exceeded, same bounded-LRU pattern as Hostnames.
  */
 public final class PendingHostnamesStore {
     private PendingHostnamesStore() {}
 
+    private static final int MAX_ENTRIES = 1000;
+
     private static final ThreadLocal<Map<String, Set<Integer>>> store =
-            ThreadLocal.withInitial(LinkedHashMap::new);
+            ThreadLocal.withInitial(() -> new LinkedHashMap<>(16, 0.75f, true) {
+                @Override
+                protected boolean removeEldestEntry(Map.Entry<String, Set<Integer>> eldest) {
+                    return size() > MAX_ENTRIES;
+                }
+            });
 
     public static void add(String hostname, int port) {
         Map<String, Set<Integer>> map = store.get();