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

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,8 @@
 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 dev.aikido.agent.wrappers.spring.SpringWebClientRedirectWrapper;
 
 import java.util.Arrays;
 import java.util.List;
@@ -30,11 +32,14 @@ private Wrappers() {}
             // SSRF/HTTP wrappers
             new HttpURLConnectionWrapper(),
             new InetAddressWrapper(),
+            new SocketChannelWrapper(),
             new HttpClientWrapper(),
             new HttpConnectionRedirectWrapper(),
             new HttpClientSendWrapper(),
             new OkHttpWrapper(),
             new ApacheHttpClientWrapper(),
+            new SpringWebClientWrapper(),
+            new SpringWebClientRedirectWrapper(),
 
             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/SpringWebClientRedirectWrapper.java

@@ -0,0 +1,66 @@
+package dev.aikido.agent.wrappers.spring;
+
+import dev.aikido.agent.wrappers.Wrapper;
+import dev.aikido.agent_api.collectors.RedirectCollector;
+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.Field;
+import java.lang.reflect.Method;
+import java.net.URL;
+
+public class SpringWebClientRedirectWrapper implements Wrapper {
+    // Package-private in Reactor Netty, referenced by name only. This is the internal method
+    // that runs once per redirect hop, for both WebClient and the Netty-backed RestClient -
+    // Spring's own request-adaptation layer (ExchangeFunction/ReactorClientHttpRequest) is
+    // only invoked once per top-level call and never sees redirect targets for bodiless (e.g.
+    // GET) requests, since Reactor Netty resends internally without going back through it.
+    // Mirrors HttpConnectionRedirectWrapper, which hooks the JDK's equally-internal
+    // followRedirect0 for the same reason.
+    private static final String HTTP_CLIENT_HANDLER_CLASS_NAME =
+            "reactor.netty.http.client.HttpClientConnect$HttpClientHandler";
+
+    public String getName() {
+        return RedirectAdvice.class.getName();
+    }
+    public ElementMatcher<? super MethodDescription> getMatcher() {
+        return ElementMatchers.isDeclaredBy(getTypeMatcher())
+                .and(ElementMatchers.named("redirect"))
+                .and(ElementMatchers.takesArguments(1));
+    }
+    @Override
+    public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+        return ElementMatchers.named(HTTP_CLIENT_HANDLER_CLASS_NAME);
+    }
+    public static class RedirectAdvice {
+        @Advice.OnMethodExit(suppress = Throwable.class)
+        public static void after(@Advice.This Object handler) throws Exception {
+            // fromURI/toURI are UriEndpoint (also package-private), both reassigned by
+            // redirect() before this advice runs: fromURI is the hostname that redirected,
+            // toURI is where it redirected to.
+            String origin = externalForm(handler, "fromURI");
+            String dest = externalForm(handler, "toURI");
+            if (origin == null || dest == null) {
+                return;
+            }
+            RedirectCollector.report(new URL(origin), new URL(dest));
+        }
+
+        // Must be public: after weaving, this is called as a real cross-class invocation from
+        // inside the target class's own bytecode, so a private method would raise IllegalAccessError.
+        public static String externalForm(Object handler, String fieldName) throws Exception {
+            Field field = handler.getClass().getDeclaredField(fieldName);
+            field.setAccessible(true);
+            Object uriEndpoint = field.get(handler);
+            if (uriEndpoint == null) {
+                return null;
+            }
+            Method toExternalForm = uriEndpoint.getClass().getDeclaredMethod("toExternalForm");
+            toExternalForm.setAccessible(true);
+            return (String) toExternalForm.invoke(uriEndpoint);
+        }
+    }
+}

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/build.gradle

@@ -39,6 +39,9 @@ dependencies {
     testImplementation 'org.springframework:spring-web:5.3.20'
     testImplementation 'org.springframework:spring-webmvc:5.3.20'
     testImplementation 'org.springframework:spring-test:5.3.20'
+    // Spring WebFlux for WebClient
+    testImplementation 'org.springframework:spring-webflux:5.3.20'
+    testImplementation 'io.projectreactor.netty:reactor-netty-http:1.2.1'
 
     testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.2'
 }

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,46 @@
 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) {
+        // InetAddress.getAllByName() resolves everything in one call, so it's safe to consume.
+        process(hostname, inetAddresses, PendingHostnamesStore.getAndRemove(hostname), INET_ADDRESS_OPERATION_NAME);
+    }
+
+    // For clients that resolve their own DNS (e.g. Reactor Netty, used by Spring's WebClient) or
+    // connect straight to an IP literal. A single request can trigger multiple connect() calls to
+    // the same hostname (IPv4 then IPv6), so unlike report(), this peeks the pending port instead
+    // of consuming it - consuming on the first attempt would let a later attempt bypass SSRF.
+    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);
+            StatisticsStore.registerCall(operationName, OperationKind.OUTGOING_HTTP_OP);
+
+            // No pending port + private IP literal = infrastructure noise (Netty resolver bootstrap
+            // etc), not a real request - skip recording/blocking. SSRF checks below still run regardless.
+            if (!ports.isEmpty() || !IsPrivateIP.isPrivateIp(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);
                 }
-            } 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();
+                // 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();
+                }
             }
 
             // Convert inetAddresses array to a List of IP strings :
@@ -62,7 +79,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 +98,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());
 

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();