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 d33c165c..203e1445 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,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; @@ -27,6 +28,7 @@ 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"; + public static void report(String hostname, InetAddress[] inetAddresses) { try { logger.trace("DNSRecordCollector called with %s & inet addresses: %s", hostname, List.of(inetAddresses)); @@ -37,12 +39,9 @@ public static void report(String hostname, InetAddress[] inetAddresses) { // Consume pending ports recorded by URLCollector for this hostname. // Removing them here ensures each (hostname, port) pair is counted exactly once. Set 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 + // URLCollector already recorded known-port hits. This is the fallback for hostnames + // it never saw (JDBC, raw sockets, ...) - skip only private IP literals, infra noise. + if (ports.isEmpty() && !IsPrivateIP.isPrivateIp(hostname)) { HostnamesStore.incrementHits(hostname, 0); } diff --git a/agent_api/src/main/java/dev/aikido/agent_api/collectors/URLCollector.java b/agent_api/src/main/java/dev/aikido/agent_api/collectors/URLCollector.java index e1c1244b..c85c4ab2 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/collectors/URLCollector.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/collectors/URLCollector.java @@ -2,6 +2,7 @@ import dev.aikido.agent_api.helpers.logging.LogManager; import dev.aikido.agent_api.helpers.logging.Logger; +import dev.aikido.agent_api.storage.HostnamesStore; import dev.aikido.agent_api.storage.PendingHostnamesStore; import java.net.URL; @@ -18,9 +19,14 @@ public static void report(URL url) { return; // Non-HTTP(S) URL } logger.trace("Adding a new URL to the cache: %s", url); + int port = getPortFromURL(url); // 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)); + // up during the DNS lookup that follows, for SSRF detection + PendingHostnamesStore.add(url.getHost(), port); + // Record the hit here, at the real call site, where the port is always known - + // instead of waiting for the DNS lookup, which can't distinguish this call from + // unrelated infra noise resolving the same hostname. + HostnamesStore.incrementHits(url.getHost(), port); } } } diff --git a/agent_api/src/test/java/collectors/DNSRecordCollectorTest.java b/agent_api/src/test/java/collectors/DNSRecordCollectorTest.java index c7cdd4b3..d589a439 100644 --- a/agent_api/src/test/java/collectors/DNSRecordCollectorTest.java +++ b/agent_api/src/test/java/collectors/DNSRecordCollectorTest.java @@ -155,31 +155,6 @@ public void testHostnamesStorePort0WhenNoPendingEntry() { assertEquals(0, entries[0].getPort()); } - @Test - public void testHostnamesStoreUsesPortFromPendingStore() { - PendingHostnamesStore.add("dev.aikido", 8080); - Context.set(mock(ContextObject.class)); - - DNSRecordCollector.report("dev.aikido", new InetAddress[]{inetAddress1}); - Hostnames.HostnameEntry[] entries = HostnamesStore.getHostnamesAsList(); - assertEquals(1, entries.length); - assertEquals("dev.aikido", entries[0].getHostname()); - assertEquals(8080, entries[0].getPort()); - } - - @Test - public void testHostnamesStoreIncrementedForAllPendingPorts() { - PendingHostnamesStore.add("dev.aikido", 80); - PendingHostnamesStore.add("dev.aikido", 443); - Context.set(mock(ContextObject.class)); - - DNSRecordCollector.report("dev.aikido", new InetAddress[]{inetAddress1}); - Hostnames.HostnameEntry[] entries = HostnamesStore.getHostnamesAsList(); - assertEquals(2, entries.length); - assertEquals(80, entries[0].getPort()); - assertEquals(443, entries[1].getPort()); - } - @Test public void testPendingEntryRemovedAfterDNSLookup() { PendingHostnamesStore.add("dev.aikido", 8080); @@ -223,4 +198,60 @@ public void testStoredSSRFWithNoContext() throws InterruptedException { DNSRecordCollector.report("metadata.google.internal", new InetAddress[]{imdsAddress1, inetAddress2}); }); } + + @Test + public void testPrivateIpLiteralWithNoPendingPortNotRecorded() { + Context.set(null); + DNSRecordCollector.report("10.20.11.143", new InetAddress[]{inetAddress2}); + assertEquals(0, HostnamesStore.getHostnamesAsList().length); + } + + @Test + public void testPrivateIpLiteralWithNoPendingPortNotRecordedButBlockedInLockdown() { + // Outbound domain blocking is never gated by the noise check - only the dashboard hit is. + 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("10.20.11.143", new InetAddress[]{inetAddress2}) + ); + assertEquals(0, HostnamesStore.getHostnamesAsList().length); + } + + @Test + public void testPrivateIpLiteralWithPendingPortStillRecordedAndBlockedInLockdown() { + ServiceConfigStore.updateFromAPIResponse(new APIResponse( + true, null, 0L, null, null, null, true, List.of(), true, true, List.of() + )); + PendingHostnamesStore.add("10.20.11.143", 443); + Context.set(mock(ContextObject.class)); + + assertThrows(BlockedOutboundException.class, () -> + DNSRecordCollector.report("10.20.11.143", new InetAddress[]{inetAddress2}) + ); + } + + @Test + public void testSsrfStillDetectedForPrivateIpLiteralWithPendingPort() { + ServiceConfigStore.updateBlocking(true); + PendingHostnamesStore.add("169.254.169.254", 80); + Context.set(new EmptySampleContextObject("http://169.254.169.254:80/latest/meta-data/")); + + Exception exception = assertThrows(SSRFException.class, () -> + DNSRecordCollector.report("169.254.169.254", new InetAddress[]{imdsAddress1}) + ); + assertEquals("Aikido Zen has blocked a server-side request forgery", exception.getMessage()); + } + + @Test + public void testNamedHostnameResolvingToPrivateIpWithNoPendingPortStillRecorded() { + // e.g. an RDS/JDBC hostname resolving to a VPC-private IP, no port because it + // never went through an instrumented HTTP client - still real, wanted visibility. + Context.set(null); + DNSRecordCollector.report("internal-service.local", new InetAddress[]{inetAddress2}); + Hostnames.HostnameEntry[] entries = HostnamesStore.getHostnamesAsList(); + assertEquals(1, entries.length); + assertEquals("internal-service.local", entries[0].getHostname()); + } } diff --git a/agent_api/src/test/java/collectors/URLCollectorTest.java b/agent_api/src/test/java/collectors/URLCollectorTest.java index 140065a3..aa9c8fe7 100644 --- a/agent_api/src/test/java/collectors/URLCollectorTest.java +++ b/agent_api/src/test/java/collectors/URLCollectorTest.java @@ -2,6 +2,7 @@ import dev.aikido.agent_api.collectors.URLCollector; 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 org.junit.jupiter.api.AfterAll; @@ -93,11 +94,13 @@ public void testWithNullContext() throws IOException { } @Test - public void testOnlyPendingStore() throws IOException { + public void testRecordsHitAtTheCallSite() throws IOException { setContextAndLifecycle(""); URLCollector.report(new URL("https://aikido.dev")); - // HostnamesStore is only written by DNSRecordCollector, not URLCollector - assertEquals(0, HostnamesStore.getHostnamesAsList().length); + Hostnames.HostnameEntry[] entries = HostnamesStore.getHostnamesAsList(); + assertEquals(1, entries.length); + assertEquals("aikido.dev", entries[0].getHostname()); + assertEquals(443, entries[0].getPort()); Set ports = PendingHostnamesStore.getPorts("aikido.dev"); assertEquals(1, ports.size()); assertTrue(ports.contains(443));