diff --git a/agent_api/src/main/java/dev/aikido/agent_api/collectors/DNSRecordCollector.java b/agent_api/src/main/java/dev/aikido/agent_api/collectors/DNSRecordCollector.java index ed8eb8e5..1e282eb4 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/collectors/DNSRecordCollector.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/collectors/DNSRecordCollector.java @@ -12,7 +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.helpers.net.IPValidator; import dev.aikido.agent_api.vulnerabilities.ssrf.StoredSSRFDetector; import dev.aikido.agent_api.vulnerabilities.ssrf.StoredSSRFException; @@ -42,8 +42,9 @@ public static void report(String hostname, InetAddress[] inetAddresses) { for (int port : ports) { HostnamesStore.incrementHits(hostname, port); } - } else if (!IsPrivateIP.isPrivateIp(hostname)) { - // Literal IPs reach this sink without a real DNS call, so skip private ones as noise. + } else if (!IPValidator.isIP(hostname)) { + // Skip literal IPs with no pending port: usually inbound IP parsing rather than an outbound request, + // and we can't tell them apart here. A real outbound one still hits the blocking/SSRF checks below. HostnamesStore.incrementHits(hostname, 0); } diff --git a/agent_api/src/test/java/collectors/DNSRecordCollectorTest.java b/agent_api/src/test/java/collectors/DNSRecordCollectorTest.java index 55558ad9..dabdcbf1 100644 --- a/agent_api/src/test/java/collectors/DNSRecordCollectorTest.java +++ b/agent_api/src/test/java/collectors/DNSRecordCollectorTest.java @@ -259,6 +259,42 @@ public void testNamedHostnameResolvingToPrivateIpWithNoPendingPortStillRecorded( assertEquals("internal-service.local", entries[0].getHostname()); } + @Test + public void testPublicIpLiteralWithNoPendingPortNotRecorded() { + Context.set(null); + DNSRecordCollector.report("8.8.8.8", new InetAddress[]{inetAddress1}); + assertEquals(0, HostnamesStore.getHostnamesAsList().length); + } + + @Test + public void testPublicIpv6LiteralWithNoPendingPortNotRecorded() { + Context.set(null); + DNSRecordCollector.report("2606:4700:4700::1111", new InetAddress[]{inetAddress1}); + assertEquals(0, HostnamesStore.getHostnamesAsList().length); + } + + @Test + public void testPublicIpLiteralWithPendingPortStillRecorded() { + PendingHostnamesStore.add("8.8.8.8", 443); + Context.set(mock(ContextObject.class)); + DNSRecordCollector.report("8.8.8.8", new InetAddress[]{inetAddress1}); + Hostnames.HostnameEntry[] entries = HostnamesStore.getHostnamesAsList(); + assertEquals(1, entries.length); + assertEquals(443, entries[0].getPort()); + } + + @Test + public void testPublicIpLiteralWithNoPendingPortNotRecordedButBlockedInLockdown() { + ServiceConfigStore.updateFromAPIResponse(new APIResponse( + true, null, 0L, null, null, null, true, List.of(), true, true, List.of() + )); + Context.set(null); + assertThrows(BlockedOutboundException.class, () -> + DNSRecordCollector.report("8.8.8.8", new InetAddress[]{inetAddress1}) + ); + assertEquals(0, HostnamesStore.getHostnamesAsList().length); + } + @Test public void testStoredSSRFWithNoContext() throws InterruptedException { ServiceConfigStore.updateBlocking(true);