Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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";

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OPERATION_NAME constant was added but the same string literal is still used in registerCall. Use OPERATION_NAME in registerCall to avoid duplicated literals and simplify maintenance.

Details

✨ AI Reasoning
​A constant OPERATION_NAME was added and later used in SSRFDetector and StoredSSRFDetector calls, but the call to StatisticsStore.registerCall still uses the literal string "java.net.InetAddress.getAllByName". This duplicates the same string in two places and was introduced by the current diff (the added OPERATION_NAME constant at the top of the file). Using the constant in the registerCall call would be a simpler, less error-prone, and consistent approach. This change is purely readability/maintenance (no behavior change) and is fixable within this PR.

🔧 How do I fix it?
Rewrite the snippet in the simpler, behavior-equivalent form: return a boolean expression directly instead of if cond return true else return false, avoid using lists when they are guaranteed to contain one element, etc.

Reply @AikidoSec feedback: [FEEDBACK] to get better review comments in the future.
Reply @AikidoSec ignore: [REASON] to ignore this issue.
More info

public static void report(String hostname, InetAddress[] inetAddresses) {
try {
logger.trace("DNSRecordCollector called with %s & inet addresses: %s", hostname, List.of(inetAddresses));
Expand All @@ -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<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
// 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)) {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

DNSRecordCollector.report's behavior now conditionally skips recording hostname hits for private IPs when ports are absent, increasing mixed responsibilities and making the method's purpose less clear; consider renaming or extracting the fallback hit-recording into a well-named function.

Details

✨ AI Reasoning
​The DNSRecordCollector.report method's role became less obvious: it now conditionally records a fallback HostnamesStore hit only when no pending ports exist and the hostname is not a private IP. The method already handled multiple concerns (port consumption, blocking, SSRF detection); the added conditional logic increases cognitive load and makes 'report' a vague name for these responsibilities.

🔧 How do I fix it?
Use descriptive verb-noun function names, add docstrings explaining the function's purpose, or provide meaningful return type hints.

Reply @AikidoSec feedback: [FEEDBACK] to get better review comments in the future.
Reply @AikidoSec ignore: [REASON] to ignore this issue.
More info

HostnamesStore.incrementHits(hostname, 0);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

URLCollector.report now both adds to PendingHostnamesStore and records HostnamesStore.incrementHits, mixing responsibilities and adding an unexpected side-effect; consider renaming or splitting into clearer methods.

Details

✨ AI Reasoning
​A public method named report in a collector class now both records a pending hostname/port and also increments the HostnamesStore hit counter. The method's name doesn't make these distinct side-effects obvious, and the change increased its responsibilities (state mutation of two different stores). This makes the function's purpose less self-evident to callers and maintainers.

🔧 How do I fix it?
Use descriptive verb-noun function names, add docstrings explaining the function's purpose, or provide meaningful return type hints.

Reply @AikidoSec feedback: [FEEDBACK] to get better review comments in the future.
Reply @AikidoSec ignore: [REASON] to ignore this issue.
More info

}
}
}
81 changes: 56 additions & 25 deletions agent_api/src/test/java/collectors/DNSRecordCollectorTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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());
}
}
9 changes: 6 additions & 3 deletions agent_api/src/test/java/collectors/URLCollectorTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<Integer> ports = PendingHostnamesStore.getPorts("aikido.dev");
assertEquals(1, ports.size());
assertTrue(ports.contains(443));
Expand Down
Loading