Update dns record collector + test cases for outbound issues

Author: bitterpanda63

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

@@ -35,20 +35,24 @@ public static void report(String hostname, InetAddress[] inetAddresses) {
             // store stats
             StatisticsStore.registerCall("java.net.InetAddress.getAllByName", OperationKind.OUTGOING_HTTP_OP);
 
+            boolean bypassed = BypassedContextStore.isBypassed();
+
             // 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);
+            if (!bypassed) {
+                // Bypassed IPs are trusted — don't report their outbound hostnames in heartbeats.
+                if (!ports.isEmpty()) {
+                    for (int port : ports) {
+                        HostnamesStore.incrementHits(hostname, port);
+                    }
+                } else {
+                    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) && !BypassedContextStore.isBypassed()) {
+            if (ServiceConfigStore.shouldBlockOutgoingRequest(hostname) && !bypassed) {
                 logger.debug("Blocking DNS lookup for domain: %s", hostname);
                 throw BlockedOutboundException.get();
             }

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

@@ -2,8 +2,10 @@
 
 import dev.aikido.agent_api.storage.service_configuration.Domain;
 
+import java.net.IDN;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 
 public class OutboundDomains {
@@ -14,14 +16,14 @@ 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.domains.put(normalize(domain.hostname()), domain.mode());
             }
         }
         this.blockNewOutgoingRequests = blockNewOutgoingRequests;
     }
 
     public boolean shouldBlockOutgoingRequest(String hostname) {
-        String mode = this.domains.get(hostname);
+        String mode = this.domains.get(normalize(hostname));
 
         if (this.blockNewOutgoingRequests) {
             // Only allow outgoing requests if the mode is "allow"
@@ -32,4 +34,18 @@ public boolean shouldBlockOutgoingRequest(String hostname) {
         // Only block outgoing requests if the mode is "block"
         return "block".equals(mode);
     }
+
+    // Normalize to lowercased Unicode form so Punycode (xn--...) and Unicode
+    // variants of the same hostname compare equal.
+    private static String normalize(String hostname) {
+        if (hostname == null) {
+            return null;
+        }
+        String lower = hostname.toLowerCase(Locale.ROOT);
+        try {
+            return IDN.toUnicode(lower, IDN.ALLOW_UNASSIGNED);
+        } catch (IllegalArgumentException e) {
+            return lower;
+        }
+    }
 }

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

@@ -149,6 +149,73 @@ public void testBlockedDomainNotBlockedWhenIpBypassed() {
         );
     }
 
+    @Test
+    public void testUnicodeDomainBlockedForPunycodeRequest() {
+        // Blocklist stores Unicode; a Punycode request for the same domain must still be blocked.
+        ServiceConfigStore.updateFromAPIResponse(new APIResponse(
+            true, null, 0L, null, null, null,
+            false, List.of(new Domain("böse.example.com", "block")), true, true, List.of()
+        ));
+        assertThrows(BlockedOutboundException.class, () ->
+            DNSRecordCollector.report("xn--bse-sna.example.com", new InetAddress[]{inetAddress1})
+        );
+    }
+
+    @Test
+    public void testUnicodeDomainBlockedForUnicodeRequest() {
+        ServiceConfigStore.updateFromAPIResponse(new APIResponse(
+            true, null, 0L, null, null, null,
+            false, List.of(new Domain("münchen.example.com", "block")), true, true, List.of()
+        ));
+        assertThrows(BlockedOutboundException.class, () ->
+            DNSRecordCollector.report("münchen.example.com", new InetAddress[]{inetAddress1})
+        );
+    }
+
+    @Test
+    public void testAllowedUnicodeDomainNotBlockedForPunycodeRequest() {
+        ServiceConfigStore.updateFromAPIResponse(new APIResponse(
+            true, null, 0L, null, null, null,
+            true, List.of(new Domain("münchen-allowed.example.com", "allow")), true, true, List.of()
+        ));
+        assertDoesNotThrow(() ->
+            DNSRecordCollector.report("xn--mnchen-allowed-gsb.example.com", new InetAddress[]{inetAddress1})
+        );
+    }
+
+    @Test
+    public void testAllowedUnicodeDomainNotBlockedForUnicodeRequest() {
+        ServiceConfigStore.updateFromAPIResponse(new APIResponse(
+            true, null, 0L, null, null, null,
+            true, List.of(new Domain("münchen-allowed.example.com", "allow")), true, true, List.of()
+        ));
+        assertDoesNotThrow(() ->
+            DNSRecordCollector.report("münchen-allowed.example.com", new InetAddress[]{inetAddress1})
+        );
+    }
+
+    @Test
+    public void testBypassedIpDoesNotRecordHostname() {
+        BypassedContextStore.setBypassed(true);
+        DNSRecordCollector.report("domain1.example.com", new InetAddress[]{inetAddress1});
+        Hostnames.HostnameEntry[] entries = HostnamesStore.getHostnamesAsList();
+        assertEquals(0, entries.length);
+    }
+
+    @Test
+    public void testBlockedDomainStillRecordedWhenNotBypassed() {
+        ServiceConfigStore.updateFromAPIResponse(new APIResponse(
+            true, null, 0L, null, null, null,
+            false, List.of(new Domain("blocked.example.com", "block")), true, true, List.of()
+        ));
+        assertThrows(BlockedOutboundException.class, () ->
+            DNSRecordCollector.report("blocked.example.com", new InetAddress[]{inetAddress1})
+        );
+        Hostnames.HostnameEntry[] entries = HostnamesStore.getHostnamesAsList();
+        assertEquals(1, entries.length);
+        assertEquals("blocked.example.com", entries[0].getHostname());
+    }
+
     @Test
     public void testUnknownDomainBlockedWhenBlockNewOutgoingRequests() {
         ServiceConfigStore.updateFromAPIResponse(new APIResponse(