Add outbound blocking

.github/workflows/gradle-tests.yml

@@ -59,6 +59,7 @@ jobs:
       - name: Add local.aikido.io to /etc/hosts
         run: |
           echo "127.0.0.1 local.aikido.io" | sudo tee -a /etc/hosts
+          echo "127.0.0.1 app.local.aikido.io" | sudo tee -a /etc/hosts
 
       - name: Start databases
         working-directory: ./sample-apps/databases

agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/APIResponse.java

@@ -1,6 +1,7 @@
 package dev.aikido.agent_api.background.cloud.api;
 
 import dev.aikido.agent_api.background.Endpoint;
+import dev.aikido.agent_api.storage.service_configuration.Domain;
 
 import java.util.List;
 
@@ -11,6 +12,8 @@ public record APIResponse(
         List<Endpoint> endpoints,
         List<String> blockedUserIds,
         List<String> allowedIPAddresses,
+        boolean blockNewOutgoingRequests,
+        List<Domain> domains,
         boolean receivedAnyStats,
         boolean block
 ) {

agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/ReportingApiHTTP.java

@@ -152,7 +152,7 @@ private static APIResponse getUnsuccessfulAPIResponse(String error) {
         return new APIResponse(
                 false, // Success
                 error,
-                0, null, null, null, false, false // Unimportant values.
+                0, null, null, null, false, null, false, false // Unimportant values.
         );
     }
 }

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

@@ -1,11 +1,14 @@
 package dev.aikido.agent_api.collectors;
 
 import dev.aikido.agent_api.context.Context;
-import dev.aikido.agent_api.storage.Hostnames;
+import dev.aikido.agent_api.storage.HostnamesStore;
+import dev.aikido.agent_api.storage.PendingHostnamesStore;
+import dev.aikido.agent_api.storage.ServiceConfigStore;
 import dev.aikido.agent_api.storage.statistics.OperationKind;
 import dev.aikido.agent_api.storage.statistics.StatisticsStore;
 import dev.aikido.agent_api.vulnerabilities.Attack;
 import dev.aikido.agent_api.vulnerabilities.ssrf.SSRFDetector;
+import dev.aikido.agent_api.vulnerabilities.outbound_blocking.BlockedOutboundException;
 import dev.aikido.agent_api.vulnerabilities.ssrf.SSRFException;
 import dev.aikido.agent_api.helpers.logging.LogManager;
 import dev.aikido.agent_api.helpers.logging.Logger;
@@ -15,6 +18,7 @@
 import java.net.InetAddress;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
 
 import static dev.aikido.agent_api.helpers.ShouldBlockHelper.shouldBlock;
 import static dev.aikido.agent_api.storage.AttackQueue.attackDetected;
@@ -30,38 +34,49 @@ public static void report(String hostname, InetAddress[] 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();
+            }
+
             // Convert inetAddresses array to a List of IP strings :
             List<String> ipAddresses = new ArrayList<>();
             for (InetAddress inetAddress : inetAddresses) {
                 ipAddresses.add(inetAddress.getHostAddress());
             }
 
-            // Fetch hostnames from Context (this is to get port number e.g.)
-            if (Context.get() != null && Context.get().getHostnames() != null) {
-                for (Hostnames.HostnameEntry hostnameEntry : Context.get().getHostnames().asArray()) {
-                    if (!hostnameEntry.getHostname().equals(hostname)) {
-                        continue;
-                    }
-                    logger.debug("Hostname: %s, Port: %s, IPs: %s", hostnameEntry.getHostname(), hostnameEntry.getPort(), ipAddresses);
-
-                    Attack attack = SSRFDetector.run(
-                        hostname, hostnameEntry.getPort(), ipAddresses, OPERATION_NAME
-                    );
-                    if (attack == null) {
-                        continue;
-                    }
-
-                    logger.debug("SSRF Attack detected due to: %s:%s", hostname, hostnameEntry.getPort());
-                    attackDetected(attack, Context.get());
-
-                    if (shouldBlock()) {
-                        logger.debug("Blocking SSRF attack...");
-                        throw SSRFException.get();
-                    }
-
-                    // We don't want to test for a stored SSRF attack.
-                    return;
+            // Run SSRF check for all ports found in the pending store (empty = no SSRF check)
+            for (int port : ports) {
+                logger.debug("Hostname: %s, Port: %s, IPs: %s", hostname, port, ipAddresses);
+
+                Attack attack = SSRFDetector.run(hostname, port, ipAddresses, OPERATION_NAME);
+                if (attack == null) {
+                    continue;
                 }
+
+                logger.debug("SSRF Attack detected due to: %s:%s", hostname, port);
+                attackDetected(attack, Context.get());
+
+                if (shouldBlock()) {
+                    logger.debug("Blocking SSRF attack...");
+                    throw SSRFException.get();
+                }
+
+                // We don't want to test for a stored SSRF attack.
+                return;
             }
 
             // We don't need the context object to check for stored ssrf, but we do want to run this after our other
@@ -76,7 +91,7 @@ public static void report(String hostname, InetAddress[] inetAddresses) {
                 }
             }
 
-        } catch (SSRFException | StoredSSRFException e) {
+        } catch (BlockedOutboundException | SSRFException | StoredSSRFException e) {
             throw e;
         } catch (Throwable e) {
             logger.trace(e);

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

@@ -1,10 +1,8 @@
 package dev.aikido.agent_api.collectors;
 
-import dev.aikido.agent_api.context.Context;
-import dev.aikido.agent_api.context.ContextObject;
-import dev.aikido.agent_api.storage.HostnamesStore;
 import dev.aikido.agent_api.helpers.logging.LogManager;
 import dev.aikido.agent_api.helpers.logging.Logger;
+import dev.aikido.agent_api.storage.PendingHostnamesStore;
 
 import java.net.URL;
 
@@ -15,25 +13,14 @@ public final class URLCollector {
 
     private URLCollector() {}
     public static void report(URL url) {
-        if(url != null) {
+        if (url != null) {
             if (!url.getProtocol().startsWith("http")) {
-                return; // Non-HTTP(S) URL 
+                return; // Non-HTTP(S) URL
             }
             logger.trace("Adding a new URL to the cache: %s", url);
-            int port = getPortFromURL(url);
-
-            // We store hostname and port in two places, HostnamesStore and Context. HostnamesStore is for reporting
-            // outbound domains. Context is to have a map of hostnames with used port numbers to detect SSRF attacks.
-
-            // Store (new) hostname hits
-            HostnamesStore.incrementHits(url.getHost(), port);
-
-            // Add to context :
-            ContextObject context = Context.get();
-            if (context != null) {
-                context.getHostnames().add(url.getHost(), port);
-                Context.set(context);
-            }
+            // Store hostname+port in the pending store so DNSRecordCollector can pick it
+            // up during the DNS lookup that follows, for SSRF detection and outbound hostnames
+            PendingHostnamesStore.add(url.getHost(), getPortFromURL(url));
         }
     }
 }

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

@@ -8,6 +8,7 @@
 import dev.aikido.agent_api.helpers.logging.LogManager;
 import dev.aikido.agent_api.helpers.logging.Logger;
 import dev.aikido.agent_api.storage.AttackQueue;
+import dev.aikido.agent_api.storage.PendingHostnamesStore;
 import dev.aikido.agent_api.storage.ServiceConfigStore;
 import dev.aikido.agent_api.storage.ServiceConfiguration;
 import dev.aikido.agent_api.storage.attack_wave_detector.AttackWaveDetectorStore;
@@ -36,7 +37,14 @@ private WebRequestCollector() {
      */
     public static Res report(ContextObject newContext) {
         ServiceConfiguration config = getConfig();
-        Context.reset(); // clear context
+
+        // clear context
+        Context.reset();
+
+        // Flush pending hostnames on every context change to prevent the store from
+        // growing unboundedly when a thread is reused across multiple requests.
+        PendingHostnamesStore.clear();
+
         if (config.isIpBypassed(newContext.getRemoteAddress())) {
             return null; // do not set context when the IP address is bypassed (zen = off)
         }

agent_api/src/main/java/dev/aikido/agent_api/context/Context.java

@@ -1,5 +1,7 @@
 package dev.aikido.agent_api.context;
 
+import dev.aikido.agent_api.storage.PendingHostnamesStore;
+
 public final class Context {
     private Context() {}
 

agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java

@@ -2,7 +2,6 @@
 
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
-import dev.aikido.agent_api.storage.Hostnames;
 import dev.aikido.agent_api.storage.RedirectNode;
 
 import java.util.*;
@@ -26,10 +25,6 @@ public class ContextObject {
     protected transient Map<String, Map<String, String>> cache = new HashMap<>();
     protected transient Optional<Boolean> forcedProtectionOff = Optional.empty();
 
-    // We store hostnames in the context object so we can match a given hostname (by DNS request)
-    // with its port number (which we know by instrumenting the URLs that get requested).
-    protected transient Hostnames hostnames = new Hostnames(1000); // max 1000 entries
-
     public boolean middlewareExecuted() {return executedMiddleware; }
     public void setExecutedMiddleware(boolean value) { executedMiddleware = value; }
 
@@ -97,7 +92,6 @@ public HashMap<String, List<String>> getCookies() {
         return cookies;
     }
     public Map<String, Map<String, String>> getCache() { return cache; }
-    public Hostnames getHostnames() { return hostnames; }
 
     public void setForcedProtectionOff(boolean forcedProtectionOff) {
         this.forcedProtectionOff = Optional.of(forcedProtectionOff);

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

@@ -0,0 +1,43 @@
+package dev.aikido.agent_api.storage;
+
+import java.util.*;
+
+/**
+ * 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.
+ */
+public final class PendingHostnamesStore {
+    private PendingHostnamesStore() {}
+
+    private static final ThreadLocal<Map<String, Set<Integer>>> store =
+            ThreadLocal.withInitial(LinkedHashMap::new);
+
+    public static void add(String hostname, int port) {
+        Map<String, Set<Integer>> map = store.get();
+        if (!map.containsKey(hostname)) {
+            map.put(hostname, new LinkedHashSet<>());
+        }
+        map.get(hostname).add(port);
+    }
+
+    public static Set<Integer> getAndRemove(String hostname) {
+        Set<Integer> ports = store.get().remove(hostname);
+        if (ports == null) {
+            return Collections.emptySet();
+        }
+        return ports;
+    }
+
+    public static Set<Integer> getPorts(String hostname) {
+        Set<Integer> ports = store.get().get(hostname);
+        if (ports == null) {
+            return Collections.emptySet();
+        }
+        return Collections.unmodifiableSet(ports);
+    }
+
+    public static void clear() {
+        store.get().clear();
+    }
+}

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

@@ -90,4 +90,13 @@ public static void setMiddlewareInstalled(boolean middlewareInstalled) {
             mutex.writeLock().unlock();
         }
     }
+
+    public static boolean shouldBlockOutgoingRequest(String hostname) {
+        mutex.readLock().lock();
+        try {
+            return config.shouldBlockOutgoingRequest(hostname);
+        } finally {
+            mutex.readLock().unlock();
+        }
+    }
 }

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

@@ -5,11 +5,10 @@
 import dev.aikido.agent_api.background.cloud.api.ReportingApi;
 import dev.aikido.agent_api.helpers.net.IPList;
 import dev.aikido.agent_api.storage.service_configuration.ParsedFirewallLists;
+import dev.aikido.agent_api.vulnerabilities.outbound_blocking.OutboundDomains;
 import dev.aikido.agent_api.storage.statistics.StatisticsStore;
 
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.List;
+import java.util.*;
 
 import static dev.aikido.agent_api.helpers.IPListBuilder.createIPList;
 import static dev.aikido.agent_api.vulnerabilities.ssrf.IsPrivateIP.isPrivateIp;
@@ -26,6 +25,7 @@ public class ServiceConfiguration {
     private IPList bypassedIPs = new IPList();
     private HashSet<String> blockedUserIDs = new HashSet<>();
     private List<Endpoint> endpoints = new ArrayList<>();
+    private OutboundDomains outboundDomains = new OutboundDomains();
 
     public ServiceConfiguration() {
         this.receivedAnyStats = true; // true by default, waiting for the startup event
@@ -46,6 +46,7 @@ public void updateConfig(APIResponse apiResponse) {
         if (apiResponse.endpoints() != null) {
             this.endpoints = apiResponse.endpoints();
         }
+        this.outboundDomains.update(apiResponse.domains(), apiResponse.blockNewOutgoingRequests());
         this.receivedAnyStats = apiResponse.receivedAnyStats();
     }
 
@@ -127,4 +128,8 @@ public boolean isBlockedUserAgent(String userAgent) {
 
     public record BlockedResult(boolean blocked, String description) {
     }
+
+    public boolean shouldBlockOutgoingRequest(String hostname) {
+        return this.outboundDomains.shouldBlockOutgoingRequest(hostname);
+    }
 }

agent_api/src/main/java/dev/aikido/agent_api/storage/service_configuration/Domain.java

@@ -0,0 +1,4 @@
+package dev.aikido.agent_api.storage.service_configuration;
+
+public record Domain(String hostname, String mode) {
+}

agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/outbound_blocking/BlockedOutboundException.java

@@ -0,0 +1,14 @@
+package dev.aikido.agent_api.vulnerabilities.outbound_blocking;
+
+import dev.aikido.agent_api.vulnerabilities.AikidoException;
+
+public class BlockedOutboundException extends AikidoException {
+    public BlockedOutboundException(String msg) {
+        super(msg);
+    }
+
+    public static BlockedOutboundException get() {
+        String defaultMsg = generateDefaultMessage("an outbound request");
+        return new BlockedOutboundException(defaultMsg);
+    }
+}

agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/outbound_blocking/OutboundDomains.java

@@ -0,0 +1,35 @@
+package dev.aikido.agent_api.vulnerabilities.outbound_blocking;
+
+import dev.aikido.agent_api.storage.service_configuration.Domain;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class OutboundDomains {
+    private Map<String, String> domains = new HashMap<>();
+    private boolean blockNewOutgoingRequests = false;
+
+    public void update(List<Domain> newDomains, boolean blockNewOutgoingRequests) {
+        if (newDomains != null) {
+            this.domains = new HashMap<>();
+            for (Domain domain : newDomains) {
+                this.domains.put(domain.hostname(), domain.mode());
+            }
+        }
+        this.blockNewOutgoingRequests = blockNewOutgoingRequests;
+    }
+
+    public boolean shouldBlockOutgoingRequest(String hostname) {
+        String mode = this.domains.get(hostname);
+
+        if (this.blockNewOutgoingRequests) {
+            // Only allow outgoing requests if the mode is "allow"
+            // null means unknown hostname, so they get blocked
+            return !"allow".equals(mode);
+        }
+
+        // Only block outgoing requests if the mode is "block"
+        return "block".equals(mode);
+    }
+}