Skip to content

Commit fa7a2b9

Browse files
Merge pull request #110 from AikidoSec/AIK-4330
AIK-4330 Java fix SSRF not working when IPC fails
2 parents 8ad3e9c + 5064496 commit fa7a2b9

File tree

11 files changed

+124
-103
lines changed

11 files changed

+124
-103
lines changed

.github/workflows/end2end.yml

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -36,14 +36,15 @@ jobs:
3636
strategy:
3737
matrix:
3838
app:
39-
- { name: SpringBootPostgres, test_file: end2end/spring_boot_postgres.py }
40-
- { name: SpringBootMySQL, test_file: end2end/spring_boot_mysql.py }
41-
- { name: SpringBootMSSQL, test_file: end2end/spring_boot_mssql.py }
42-
- { name: SpringWebfluxSampleApp, test_file: end2end/spring_webflux_postgres.py }
43-
- { name: SpringMVCPostgresKotlin, test_file: end2end/spring_mvc_postgres_kotlin.py }
44-
- { name: SpringMVCPostgresGroovy, test_file: end2end/spring_mvc_postgres_groovy.py }
45-
- { name: JavalinPostgres, test_file: end2end/javalin_postgres.py }
46-
- { name: JavalinMySQLKotlin, test_file: end2end/javalin_mysql_kotlin.py }
39+
- { name: SpringBootPostgres, test_file: end2end/spring_boot_postgres.py, with_ipc: true }
40+
- { name: SpringBootPostgres, test_file: end2end/spring_boot_postgres.py, with_ipc: false }
41+
- { name: SpringBootMySQL, test_file: end2end/spring_boot_mysql.py, with_ipc: true }
42+
- { name: SpringBootMSSQL, test_file: end2end/spring_boot_mssql.py, with_ipc: true }
43+
- { name: SpringWebfluxSampleApp, test_file: end2end/spring_webflux_postgres.py, with_ipc: true }
44+
- { name: SpringMVCPostgresKotlin, test_file: end2end/spring_mvc_postgres_kotlin.py, with_ipc: true }
45+
- { name: SpringMVCPostgresGroovy, test_file: end2end/spring_mvc_postgres_groovy.py, with_ipc: true }
46+
- { name: JavalinPostgres, test_file: end2end/javalin_postgres.py, with_ipc: true }
47+
- { name: JavalinMySQLKotlin, test_file: end2end/javalin_mysql_kotlin.py, with_ipc: true }
4748
java-version: [17, 18, 19, 20, 21]
4849
steps:
4950
- name: Download build artifacts
@@ -63,6 +64,7 @@ jobs:
6364
docker build -t mock_core .
6465
docker run --name mock_core -d -p 5000:5000 mock_core
6566
- name: Setup IPC Folder
67+
if: matrix.app.with_ipc
6668
run: mkdir /opt/aikido && chmod a+rw /opt/aikido
6769
- name: Start databases
6870
working-directory: ./sample-apps/databases

agent/src/main/java/dev/aikido/agent/wrappers/InetAddressWrapper.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ public static void after(
4949

5050
try {
5151
// Load the class from the JAR
52-
Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.HostnameCollector");
52+
Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.DNSRecordCollector");
5353

5454
// Run report with "argument"
5555
for (Method method2: clazz.getMethods()) {

agent_api/src/main/java/dev/aikido/agent_api/collectors/HostnameCollector.java renamed to agent_api/src/main/java/dev/aikido/agent_api/collectors/DNSRecordCollector.java

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
import dev.aikido.agent_api.background.utilities.ThreadIPCClient;
55
import dev.aikido.agent_api.context.Context;
66
import dev.aikido.agent_api.storage.Hostnames;
7-
import dev.aikido.agent_api.thread_cache.ThreadCache;
87
import dev.aikido.agent_api.vulnerabilities.Attack;
98
import dev.aikido.agent_api.vulnerabilities.ssrf.SSRFDetector;
109
import dev.aikido.agent_api.vulnerabilities.ssrf.SSRFException;
@@ -18,25 +17,25 @@
1817
import static dev.aikido.agent_api.background.utilities.ThreadIPCClientFactory.getDefaultThreadIPCClient;
1918
import static dev.aikido.agent_api.helpers.ShouldBlockHelper.shouldBlock;
2019

21-
public final class HostnameCollector {
22-
private HostnameCollector() {}
23-
private static final Logger logger = LogManager.getLogger(HostnameCollector.class);
20+
public final class DNSRecordCollector {
21+
private DNSRecordCollector() {}
22+
private static final Logger logger = LogManager.getLogger(DNSRecordCollector.class);
2423
public static void report(String hostname, InetAddress[] inetAddresses) {
2524
try {
26-
logger.trace("HostnameCollector called with %s & inet addresses: %s", hostname, List.of(inetAddresses));
27-
25+
logger.trace("DNSRecordCollector called with %s & inet addresses: %s", hostname, List.of(inetAddresses));
2826

2927
// Convert inetAddresses array to a List of IP strings :
3028
List<String> ipAddresses = new ArrayList<>();
3129
for (InetAddress inetAddress : inetAddresses) {
3230
ipAddresses.add(inetAddress.getHostAddress());
3331
}
34-
// Currently using hostnames from thread cache, might not be as accurate as using Context-dependant hostnames.
35-
if (ThreadCache.get() == null || ThreadCache.get().getHostnames() == null) {
36-
logger.trace("Thread cache is empty, returning.");
32+
33+
// Fetch hostnames from Context (this is to get port number e.g.)
34+
if (Context.get() == null || Context.get().getHostnames() == null) {
35+
logger.trace("Context not defined, returning.");
3736
return;
3837
}
39-
for (Hostnames.HostnameEntry hostnameEntry : ThreadCache.get().getHostnames().asArray()) {
38+
for (Hostnames.HostnameEntry hostnameEntry : Context.get().getHostnames().asArray()) {
4039
if (!hostnameEntry.getHostname().equals(hostname)) {
4140
continue;
4241
}

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

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package dev.aikido.agent_api.collectors;
22

3+
import dev.aikido.agent_api.context.Context;
4+
import dev.aikido.agent_api.context.ContextObject;
35
import dev.aikido.agent_api.thread_cache.ThreadCache;
46
import dev.aikido.agent_api.thread_cache.ThreadCacheObject;
57
import dev.aikido.agent_api.helpers.logging.LogManager;
@@ -14,15 +16,29 @@ public final class URLCollector {
1416

1517
private URLCollector() {}
1618
public static void report(URL url) {
17-
ThreadCacheObject threadCache = ThreadCache.get();
18-
if(threadCache != null && url != null) {
19+
if(url != null) {
1920
if (!url.getProtocol().startsWith("http")) {
2021
return; // Non-HTTP(S) URL
2122
}
2223
logger.trace("Adding a new URL to the cache: %s", url);
2324
int port = getPortFromURL(url);
24-
threadCache.getHostnames().add(url.getHost(), port);
25-
ThreadCache.set(threadCache);
25+
26+
// We store hostname and port in two places, Thread Cache and Context. Thread Cache is for reporting
27+
// outbound domains. Context is to have a map of hostnames with used port numbers to detect SSRF attacks.
28+
29+
// Add to thread cache :
30+
ThreadCacheObject threadCache = ThreadCache.get();
31+
if (threadCache != null) {
32+
threadCache.getHostnames().add(url.getHost(), port);
33+
ThreadCache.set(threadCache);
34+
}
35+
36+
// Add to context :
37+
ContextObject context = Context.get();
38+
if (context != null) {
39+
context.getHostnames().add(url.getHost(), port);
40+
Context.set(context);
41+
}
2642
}
2743
}
2844
}

agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.google.gson.Gson;
44
import com.google.gson.GsonBuilder;
5+
import dev.aikido.agent_api.storage.Hostnames;
56
import dev.aikido.agent_api.storage.RedirectNode;
67

78
import java.util.ArrayList;
@@ -26,6 +27,10 @@ public class ContextObject {
2627
protected transient ArrayList<RedirectNode> redirectStartNodes;
2728
protected transient Map<String, Map<String, String>> cache = new HashMap<>();
2829

30+
// We store hostnames in the context object so we can match a given hostname (by DNS request)
31+
// with its port number (which we know by instrumenting the URLs that get requested).
32+
protected transient Hostnames hostnames = new Hostnames(1000); // max 1000 entries
33+
2934
public boolean middlewareExecuted() {return executedMiddleware; }
3035
public void setExecutedMiddleware(boolean value) { executedMiddleware = value; }
3136

@@ -69,6 +74,7 @@ public HashMap<String, List<String>> getCookies() {
6974
return cookies;
7075
}
7176
public Map<String, Map<String, String>> getCache() { return cache; }
77+
public Hostnames getHostnames() { return hostnames; }
7278

7379
public String toJson() {
7480
Gson gson = new GsonBuilder().setPrettyPrinting().create();

agent_api/src/test/java/collectors/HostnameCollectorTest.java renamed to agent_api/src/test/java/collectors/DNSRecordCollectorTest.java

Lines changed: 25 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
package collectors;
22

3-
import dev.aikido.agent_api.collectors.HostnameCollector;
3+
import dev.aikido.agent_api.collectors.DNSRecordCollector;
44
import dev.aikido.agent_api.context.Context;
5+
import dev.aikido.agent_api.context.ContextObject;
56
import dev.aikido.agent_api.storage.Hostnames;
67
import dev.aikido.agent_api.thread_cache.ThreadCache;
7-
import dev.aikido.agent_api.thread_cache.ThreadCacheObject;
88
import dev.aikido.agent_api.vulnerabilities.ssrf.SSRFException;
99
import org.junit.jupiter.api.*;
1010
import org.junitpioneer.jupiter.SetEnvironmentVariable;
@@ -14,12 +14,11 @@
1414
import java.net.UnknownHostException;
1515
import java.util.List;
1616

17-
import static dev.aikido.agent_api.helpers.UnixTimeMS.getUnixTimeMS;
1817
import static org.junit.jupiter.api.Assertions.assertEquals;
1918
import static org.junit.jupiter.api.Assertions.assertThrows;
2019
import static org.mockito.Mockito.*;
2120

22-
public class HostnameCollectorTest {
21+
public class DNSRecordCollectorTest {
2322
InetAddress inetAddress1;
2423
InetAddress inetAddress2;
2524
@BeforeEach
@@ -32,52 +31,48 @@ void setup() throws UnknownHostException {
3231
@SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "token")
3332
@Test
3433
public void testThreadCacheNull() {
35-
// Early return because of Thread Cache being null :
36-
HostnameCollector.report("dev.aikido", new InetAddress[]{
34+
// Early return because of Context being null :
35+
DNSRecordCollector.report("dev.aikido", new InetAddress[]{
3736
inetAddress1, inetAddress2
3837
});
3938
}
4039

4140
@SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "token")
4241
@Test
4342
public void testThreadCacheHostnames() {
44-
ThreadCacheObject myThreadCache = mock(ThreadCacheObject.class);
45-
when(myThreadCache.getLastRenewedAtMS()).thenReturn(getUnixTimeMS());
46-
ThreadCache.set(myThreadCache);
47-
HostnameCollector.report("dev.aikido", new InetAddress[]{
43+
ContextObject myContextObject = mock(ContextObject.class);
44+
Context.set(myContextObject);
45+
DNSRecordCollector.report("dev.aikido", new InetAddress[]{
4846
inetAddress1, inetAddress2
4947
});
50-
verify(myThreadCache).getHostnames();
51-
52-
myThreadCache = mock(ThreadCacheObject.class);
53-
when(myThreadCache.getLastRenewedAtMS()).thenReturn(getUnixTimeMS());
48+
verify(myContextObject).getHostnames();
5449

50+
myContextObject = mock(ContextObject.class);
5551
Hostnames hostnames = new Hostnames(20);
56-
when(myThreadCache.getHostnames()).thenReturn(hostnames);
52+
when(myContextObject.getHostnames()).thenReturn(hostnames);
53+
54+
Context.set(myContextObject);
5755

58-
ThreadCache.set(myThreadCache);
59-
HostnameCollector.report("dev.aikido", new InetAddress[]{
56+
DNSRecordCollector.report("dev.aikido", new InetAddress[]{
6057
inetAddress1, inetAddress2
6158
});
62-
verify(myThreadCache, times(2)).getHostnames();
59+
verify(myContextObject, times(2)).getHostnames();
6360
}
6461

6562
@SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "token")
6663
@Test
6764
public void testHostnameSame() {
68-
ThreadCacheObject myThreadCache = mock(ThreadCacheObject.class);
69-
when(myThreadCache.getLastRenewedAtMS()).thenReturn(getUnixTimeMS());
70-
65+
ContextObject myContextObject = mock(ContextObject.class);
7166
Hostnames hostnames = new Hostnames(20);
7267
hostnames.add("dev.aikido.not", 80);
7368
hostnames.add("dev.aikido", 80);
74-
when(myThreadCache.getHostnames()).thenReturn(hostnames);
69+
when(myContextObject.getHostnames()).thenReturn(hostnames);
7570

76-
ThreadCache.set(myThreadCache);
77-
HostnameCollector.report("dev.aikido", new InetAddress[]{
71+
Context.set(myContextObject);
72+
DNSRecordCollector.report("dev.aikido", new InetAddress[]{
7873
inetAddress1, inetAddress2
7974
});
80-
verify(myThreadCache, times(2)).getHostnames();
75+
verify(myContextObject, times(2)).getHostnames();
8176
}
8277

8378
public static class SampleContextObject extends EmptySampleContextObject {
@@ -92,22 +87,16 @@ public SampleContextObject() {
9287
@SetEnvironmentVariable(key = "AIKIDO_BLOCK", value = "1")
9388
@Test
9489
public void testHostnameSameWithContextAsAttack() {
95-
ThreadCacheObject myThreadCache = mock(ThreadCacheObject.class);
96-
when(myThreadCache.getLastRenewedAtMS()).thenReturn(getUnixTimeMS());
97-
98-
Hostnames hostnames = new Hostnames(20);
99-
hostnames.add("dev.aikido.not", 80);
100-
hostnames.add("dev.aikido", 80);
101-
when(myThreadCache.getHostnames()).thenReturn(hostnames);
90+
ContextObject myContextObject = new SampleContextObject();
91+
myContextObject.getHostnames().add("dev.aikido.not", 80);
92+
myContextObject.getHostnames().add("dev.aikido", 80);
93+
Context.set(myContextObject);
10294

103-
ThreadCache.set(myThreadCache);
104-
Context.set(new SampleContextObject());
10595
Exception exception = assertThrows(SSRFException.class, () -> {
106-
HostnameCollector.report("dev.aikido", new InetAddress[]{
96+
DNSRecordCollector.report("dev.aikido", new InetAddress[]{
10797
inetAddress1, inetAddress2
10898
});
10999
});
110-
verify(myThreadCache, times(2)).getHostnames();
111100
assertEquals("Aikido Zen has blocked a server-side request forgery", exception.getMessage());
112101
}
113102

agent_api/src/test/java/collectors/URLCollectorTest.java

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import java.net.URL;
1515

1616
import static org.junit.jupiter.api.Assertions.assertEquals;
17+
import static org.junit.jupiter.api.Assertions.assertNull;
1718
import static utils.EmtpyThreadCacheObject.getEmptyThreadCacheObject;
1819

1920
public class URLCollectorTest {
@@ -55,6 +56,11 @@ public void testNewUrlConnectionWithHttp() throws IOException {
5556
assertEquals(1, hostnameArray.length);
5657
assertEquals(80, hostnameArray[0].getPort());
5758
assertEquals("app.local.aikido.io", hostnameArray[0].getHostname());
59+
60+
Hostnames.HostnameEntry[] hostnameArray2 = Context.get().getHostnames().asArray();
61+
assertEquals(1, hostnameArray2.length);
62+
assertEquals(80, hostnameArray2[0].getPort());
63+
assertEquals("app.local.aikido.io", hostnameArray2[0].getHostname());
5864
}
5965

6066
@SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "invalid-token")
@@ -67,6 +73,11 @@ public void testNewUrlConnectionHttps() throws IOException {
6773
assertEquals(1, hostnameArray.length);
6874
assertEquals(443, hostnameArray[0].getPort());
6975
assertEquals("aikido.dev", hostnameArray[0].getHostname());
76+
77+
Hostnames.HostnameEntry[] hostnameArray2 = Context.get().getHostnames().asArray();
78+
assertEquals(1, hostnameArray2.length);
79+
assertEquals(443, hostnameArray2[0].getPort());
80+
assertEquals("aikido.dev", hostnameArray2[0].getHostname());
7081
}
7182

7283
@SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "invalid-token")
@@ -77,5 +88,47 @@ public void testNewUrlConnectionFaultyProtocol() throws IOException {
7788
URLCollector.report(new URL("ftp://localhost:8080"));
7889
Hostnames.HostnameEntry[] hostnameArray = ThreadCache.get().getHostnames().asArray();
7990
assertEquals(0, hostnameArray.length);
91+
Hostnames.HostnameEntry[] hostnameArray2 = Context.get().getHostnames().asArray();
92+
assertEquals(0, hostnameArray2.length);
93+
}
94+
95+
@SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "invalid-token")
96+
@SetEnvironmentVariable(key = "AIKIDO_BLOCK", value = "true")
97+
@Test
98+
public void testWithNullURL() throws IOException {
99+
setContextAndLifecycle("");
100+
URLCollector.report(null);
101+
Hostnames.HostnameEntry[] hostnameArray = ThreadCache.get().getHostnames().asArray();
102+
assertEquals(0, hostnameArray.length);
103+
Hostnames.HostnameEntry[] hostnameArray2 = Context.get().getHostnames().asArray();
104+
assertEquals(0, hostnameArray2.length);
105+
}
106+
107+
@SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "invalid-token")
108+
@SetEnvironmentVariable(key = "AIKIDO_BLOCK", value = "true")
109+
@Test
110+
public void testWithNullContext() throws IOException {
111+
setContextAndLifecycle("");
112+
Context.reset();
113+
URLCollector.report(new URL("https://aikido.dev"));
114+
Hostnames.HostnameEntry[] hostnameArray = ThreadCache.get().getHostnames().asArray();
115+
assertEquals(1, hostnameArray.length);
116+
assertEquals(443, hostnameArray[0].getPort());
117+
assertEquals("aikido.dev", hostnameArray[0].getHostname());
118+
assertNull(Context.get());
119+
}
120+
121+
@SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "invalid-token")
122+
@SetEnvironmentVariable(key = "AIKIDO_BLOCK", value = "true")
123+
@Test
124+
public void testWithNullThreadCache() throws IOException {
125+
setContextAndLifecycle("");
126+
ThreadCache.reset();
127+
URLCollector.report(new URL("https://aikido.dev"));
128+
Hostnames.HostnameEntry[] hostnameArray = Context.get().getHostnames().asArray();
129+
assertEquals(1, hostnameArray.length);
130+
assertEquals(443, hostnameArray[0].getPort());
131+
assertEquals("aikido.dev", hostnameArray[0].getHostname());
132+
assertNull(ThreadCache.get());
80133
}
81134
}

agent_api/src/test/java/wrappers/ApacheHttpClientTest.java

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -88,17 +88,6 @@ public void testSSRFWithoutPortAndWithoutContext() throws Exception {
8888
});
8989
}
9090

91-
@SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "invalid-token-2")
92-
@SetEnvironmentVariable(key = "AIKIDO_BLOCK", value = "true")
93-
@Test
94-
public void testSSRFWithoutPortAndWithoutThreadCache() throws Exception {
95-
setContextAndLifecycle("http://localhost:80");
96-
ThreadCache.set(null);
97-
assertThrows(ConnectException.class, () -> {
98-
fetchResponse("http://localhost/weirdroute");
99-
});
100-
}
101-
10291
private void fetchResponse(String urlString) throws IOException {
10392
HttpGet request = new HttpGet(urlString);
10493
request.addHeader("Authorization", "Bearer invalid-token-2");

agent_api/src/test/java/wrappers/InetAddressTest.java

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -88,17 +88,6 @@ public void testSSRFWithoutPortAndWithoutContext() throws Exception {
8888
});
8989
}
9090

91-
@SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "invalid-token-2")
92-
@SetEnvironmentVariable(key = "AIKIDO_BLOCK", value = "true")
93-
@Test
94-
public void testSSRFWithoutPortAndWithoutThreadCache() throws Exception {
95-
setContextAndLifecycle("http://localhost:80");
96-
ThreadCache.set(null);
97-
assertThrows(ConnectException.class, () -> {
98-
fetchResponse("http://localhost/weirdroute");
99-
});
100-
}
101-
10291
@SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "invalid-token-2")
10392
@SetEnvironmentVariable(key = "AIKIDO_BLOCK", value = "true")
10493
@Test

0 commit comments

Comments
 (0)