Cleanup URLCollector & perform simplified check in DNSRecordCollector

also adds test cases for DNSRecordCollector

Author: bitterpanda63

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

@@ -7,6 +7,7 @@
 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;
@@ -31,6 +32,12 @@ public static void report(String hostname, InetAddress[] inetAddresses) {
             // store stats
             StatisticsStore.registerCall("java.net.InetAddress.getAllByName", OperationKind.OUTGOING_HTTP_OP);
 
+            // 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) {
@@ -77,7 +84,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

@@ -2,22 +2,15 @@
 
 import dev.aikido.agent_api.context.Context;
 import dev.aikido.agent_api.context.ContextObject;
-import dev.aikido.agent_api.context.User;
 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.ServiceConfigStore;
-import dev.aikido.agent_api.vulnerabilities.Attack;
-import dev.aikido.agent_api.vulnerabilities.Vulnerabilities;
 import dev.aikido.agent_api.vulnerabilities.outbound_blocking.BlockedOutboundException;
 
 import java.net.URL;
-import java.util.Map;
 
-import static dev.aikido.agent_api.helpers.ShouldBlockHelper.shouldBlock;
-import static dev.aikido.agent_api.helpers.StackTrace.getCurrentStackTrace;
 import static dev.aikido.agent_api.helpers.url.PortParser.getPortFromURL;
-import static dev.aikido.agent_api.storage.AttackQueue.attackDetected;
 
 public final class URLCollector {
     private static final Logger logger = LogManager.getLogger(URLCollector.class);
@@ -34,33 +27,7 @@ public static void report(URL url, String operation) {
             // 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.
 
-            // hostname blocking :
             String hostname = url.getHost();
-            if (ServiceConfigStore.shouldBlockOutgoingRequest(hostname)) {
-                ContextObject ctx = Context.get();
-
-                User currentUser = null;
-                if (ctx != null) {
-                    currentUser = ctx.getUser();
-                }
-
-                Attack attack = new Attack(
-                    operation,
-                    Vulnerabilities.SSRF,
-                    "",
-                    "",
-                    Map.of(),
-                    /* payload */ hostname,
-                    getCurrentStackTrace(),
-                    currentUser
-                );
-
-                attackDetected(attack, ctx);
-                if (shouldBlock()) {
-                    logger.debug("Blocking request to domain: %s", hostname);
-                    throw BlockedOutboundException.get();
-                }
-            };
 
             // Store (new) hostname hits
             HostnamesStore.incrementHits(hostname, port);

agent_api/src/test/java/collectors/DNSRecordCollectorTest.java

@@ -1,5 +1,6 @@
 package collectors;
 
+import dev.aikido.agent_api.background.cloud.api.APIResponse;
 import dev.aikido.agent_api.background.cloud.api.events.DetectedAttack;
 import dev.aikido.agent_api.collectors.DNSRecordCollector;
 import dev.aikido.agent_api.context.Context;
@@ -8,7 +9,9 @@
 import dev.aikido.agent_api.storage.Hostnames;
 import dev.aikido.agent_api.storage.HostnamesStore;
 import dev.aikido.agent_api.storage.ServiceConfigStore;
+import dev.aikido.agent_api.storage.service_configuration.Domain;
 import dev.aikido.agent_api.vulnerabilities.Attack;
+import dev.aikido.agent_api.vulnerabilities.outbound_blocking.BlockedOutboundException;
 import dev.aikido.agent_api.vulnerabilities.ssrf.SSRFException;
 import dev.aikido.agent_api.vulnerabilities.ssrf.StoredSSRFException;
 import org.junit.jupiter.api.*;
@@ -40,6 +43,10 @@ public void cleanup() {
         HostnamesStore.clear();
         Context.set(null);
         AttackQueue.clear();
+        // Reset domain config
+        ServiceConfigStore.updateFromAPIResponse(new APIResponse(
+            true, null, 0L, null, null, null, false, List.of(), true, false
+        ));
     }
 
     @Test
@@ -135,6 +142,39 @@ public void testHostnameSameWithContextAsAStoredSSRFAttack() {
         });
     }
 
+    @Test
+    public void testBlockedDomain() {
+        ServiceConfigStore.updateFromAPIResponse(new APIResponse(
+            true, null, 0L, null, null, null,
+            false, List.of(new Domain("blocked.example.com", "block")), true, true
+        ));
+        assertThrows(BlockedOutboundException.class, () ->
+            DNSRecordCollector.report("blocked.example.com", new InetAddress[]{inetAddress1})
+        );
+    }
+
+    @Test
+    public void testAllowedDomainNotBlocked() {
+        ServiceConfigStore.updateFromAPIResponse(new APIResponse(
+            true, null, 0L, null, null, null,
+            false, List.of(new Domain("allowed.example.com", "allow")), true, true
+        ));
+        assertDoesNotThrow(() ->
+            DNSRecordCollector.report("allowed.example.com", new InetAddress[]{inetAddress1})
+        );
+    }
+
+    @Test
+    public void testUnknownDomainBlockedWhenBlockNewOutgoingRequests() {
+        ServiceConfigStore.updateFromAPIResponse(new APIResponse(
+            true, null, 0L, null, null, null,
+            true, List.of(), true, true
+        ));
+        assertThrows(BlockedOutboundException.class, () ->
+            DNSRecordCollector.report("unknown.example.com", new InetAddress[]{inetAddress1})
+        );
+    }
+
     @Test
     public void testStoredSSRFWithNoContext() throws InterruptedException {
         ServiceConfigStore.updateBlocking(true);