Skip to content

Commit 72e7da0

Browse files
Merge pull request #247 from AikidoSec/dont-create-new-instances-for-vulnerabilities
re-use vulnerability type/detector instances to avoid unnecessary mem usage
2 parents 22b84d8 + 8a70bb6 commit 72e7da0

14 files changed

Lines changed: 99 additions & 94 deletions

File tree

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,7 @@ public static void report(Object command) {
2222
StatisticsStore.registerCall("runtime.Exec", OperationKind.EXEC_OP);
2323

2424
// scan
25-
Vulnerabilities.Vulnerability vulnerability = new Vulnerabilities.ShellInjectionVulnerability();
26-
Scanner.scanForGivenVulnerability(vulnerability, "runtime.Exec", new String[]{commandStr});
25+
Scanner.scanForGivenVulnerability(Vulnerabilities.SHELL_INJECTION, "runtime.Exec", new String[]{commandStr});
2726
}
2827
}
2928
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ public static void report(String hostname, InetAddress[] inetAddresses) {
4444
}
4545
logger.debug("Hostname: %s, Port: %s, IPs: %s", hostnameEntry.getHostname(), hostnameEntry.getPort(), ipAddresses);
4646

47-
Attack attack = new SSRFDetector().run(
47+
Attack attack = SSRFDetector.run(
4848
hostname, hostnameEntry.getPort(), ipAddresses, OPERATION_NAME
4949
);
5050
if (attack == null) {

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,18 +31,17 @@ public static void report(Object filePath, String operation, int depth) {
3131
logger.trace("Scan on %s for file: %s", operation, filePath);
3232
StatisticsStore.registerCall(operation, OperationKind.FS_OP);
3333

34-
Vulnerabilities.Vulnerability vulnerability = new Vulnerabilities.PathTraversalVulnerability();
3534
if (filePath instanceof String filePathString) {
36-
Scanner.scanForGivenVulnerability(vulnerability, operation, new String[]{filePathString});
35+
Scanner.scanForGivenVulnerability(Vulnerabilities.PATH_TRAVERSAL, operation, new String[]{filePathString});
3736
} else if (filePath instanceof URI filePathURI) {
3837
// File(...) Constructor also accepts URI objects, but path remains the same
3938
// So we just extract the path here : (i.e. new File("file:///../test.txt") --> "/../test.txt")
4039
String filePathString = filePathURI.getPath();
41-
Scanner.scanForGivenVulnerability(vulnerability, operation, new String[]{filePathString});
40+
Scanner.scanForGivenVulnerability(Vulnerabilities.PATH_TRAVERSAL, operation, new String[]{filePathString});
4241
} else if (filePath instanceof Path filePathAsPath) {
4342
// Some functions on Path also accept other paths
4443
String filePathString = filePathAsPath.toString();
45-
Scanner.scanForGivenVulnerability(vulnerability, operation, new String[]{filePathString});
44+
Scanner.scanForGivenVulnerability(Vulnerabilities.PATH_TRAVERSAL, operation, new String[]{filePathString});
4645
} else if (filePath instanceof Object[] filePaths) {
4746
// In Paths.get() sometimes you can have multiple paths provided, scan them individually :
4847
if (depth >= MAX_RECURSION_DEPTH) {

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ public static void report(String sql, String dialect, String operation) {
1515
StatisticsStore.registerCall(operation, OperationKind.SQL_OP);
1616

1717
// scan
18-
Vulnerabilities.Vulnerability vulnerability = new Vulnerabilities.SQLInjectionVulnerability();
19-
Scanner.scanForGivenVulnerability(vulnerability, operation, new String[]{sql, dialect});
18+
Scanner.scanForGivenVulnerability(Vulnerabilities.SQL_INJECTION, operation, new String[]{sql, dialect});
2019
}
2120
}

agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/Vulnerabilities.java

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,56 +3,65 @@
33
import dev.aikido.agent_api.vulnerabilities.path_traversal.PathTraversalDetector;
44
import dev.aikido.agent_api.vulnerabilities.shell_injection.ShellInjectionDetector;
55
import dev.aikido.agent_api.vulnerabilities.sql_injection.SqlDetector;
6-
import dev.aikido.agent_api.vulnerabilities.ssrf.SSRFDetector;
76

87
public final class Vulnerabilities {
98
private Vulnerabilities() {}
109
public interface Vulnerability {
1110
Detector getDetector();
1211
String getKind();
1312
}
14-
public static final class SQLInjectionVulnerability implements Vulnerability {
13+
14+
private static final class SQLInjectionVulnerability implements Vulnerability {
1515
@Override
1616
public Detector getDetector() {
17-
return new SqlDetector();
17+
return SqlDetector.INSTANCE;
1818
}
1919
@Override
2020
public String getKind() {
2121
return "sql_injection";
2222
}
2323
}
24-
public static final class PathTraversalVulnerability implements Vulnerability {
24+
public static final Vulnerability SQL_INJECTION = new SQLInjectionVulnerability();
25+
26+
private static final class PathTraversalVulnerability implements Vulnerability {
2527
@Override
2628
public Detector getDetector() {
27-
return new PathTraversalDetector();
29+
return PathTraversalDetector.INSTANCE;
2830
}
2931
@Override
3032
public String getKind() {
3133
return "path_traversal";
3234
}
3335
}
34-
public static final class SSRFVulnerability implements Vulnerability {
36+
public static final Vulnerability PATH_TRAVERSAL = new PathTraversalVulnerability();
37+
38+
private static final class SSRFVulnerability implements Vulnerability {
3539
@Override
3640
public Detector getDetector() { return null; }
3741
@Override
3842
public String getKind() {
3943
return "ssrf";
4044
}
4145
}
42-
public static final class StoredSSRFVulnerability implements Vulnerability {
46+
public static final Vulnerability SSRF = new SSRFVulnerability();
47+
48+
private static final class StoredSSRFVulnerability implements Vulnerability {
4349
@Override
4450
public Detector getDetector() { return null; }
4551
@Override
4652
public String getKind() {
4753
return "stored_ssrf";
4854
}
4955
}
50-
public static final class ShellInjectionVulnerability implements Vulnerability {
56+
public static final Vulnerability STORED_SSRF = new StoredSSRFVulnerability();
57+
58+
private static final class ShellInjectionVulnerability implements Vulnerability {
5159
@Override
52-
public Detector getDetector() { return new ShellInjectionDetector(); }
60+
public Detector getDetector() { return ShellInjectionDetector.INSTANCE; }
5361
@Override
5462
public String getKind() {
5563
return "shell_injection";
5664
}
5765
}
66+
public static final Vulnerability SHELL_INJECTION = new ShellInjectionVulnerability();
5867
}

agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/path_traversal/PathTraversalDetector.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
import static dev.aikido.agent_api.vulnerabilities.path_traversal.UnsafePathPartsChecker.containsUnsafePathParts;
99

1010
public class PathTraversalDetector implements Detector {
11+
public static final PathTraversalDetector INSTANCE = new PathTraversalDetector();
12+
1113
/**
1214
* @param userInput, this is the user input which we will evaluate
1315
* @param arguments, Includes one element : filename

agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/shell_injection/ShellInjectionDetector.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99
import static dev.aikido.agent_api.vulnerabilities.shell_injection.ShellSyntaxChecker.containsShellSyntax;
1010

1111
public class ShellInjectionDetector implements Detector {
12+
public static final ShellInjectionDetector INSTANCE = new ShellInjectionDetector();
13+
1214
@Override
1315
public DetectorResult run(String userInput, String[] arguments) {
1416
if (userInput.isEmpty() || arguments == null || arguments.length == 0 || arguments[0] == null) {

agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/SqlDetector.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,9 @@
88
import java.util.regex.Pattern;
99

1010
public class SqlDetector implements Detector {
11+
public static final SqlDetector INSTANCE = new SqlDetector();
1112
private static final Logger logger = LogManager.getLogger(SqlDetector.class);
13+
1214
/**
1315
* @param userInput contains the user input which we want to scan
1416
* @param arguments contains: [query, dialect]

agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/ssrf/SSRFDetector.java

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,22 +3,18 @@
33
import dev.aikido.agent_api.context.Context;
44
import dev.aikido.agent_api.context.ContextObject;
55
import dev.aikido.agent_api.vulnerabilities.Attack;
6-
import dev.aikido.agent_api.vulnerabilities.Detector;
76
import dev.aikido.agent_api.vulnerabilities.Vulnerabilities;
87

9-
import java.util.HashSet;
108
import java.util.List;
119
import java.util.Map;
1210

13-
import static dev.aikido.agent_api.helpers.ShouldBlockHelper.shouldBlock;
1411
import static dev.aikido.agent_api.helpers.StackTrace.getCurrentStackTrace;
1512
import static dev.aikido.agent_api.vulnerabilities.ssrf.FindHostnameInContext.findHostnameInContext;
1613
import static dev.aikido.agent_api.vulnerabilities.ssrf.IsPrivateIP.containsPrivateIP;
1714
import static dev.aikido.agent_api.vulnerabilities.ssrf.PrivateIPRedirectFinder.isRedirectToPrivateIP;
18-
import static dev.aikido.agent_api.vulnerabilities.ssrf.imds.Resolver.resolvesToImdsIp;
1915

20-
public class SSRFDetector {
21-
public Attack run(String hostname, int port, List<String> ipAddresses, String operation) {
16+
public final class SSRFDetector {
17+
public static Attack run(String hostname, int port, List<String> ipAddresses, String operation) {
2218
if(hostname == null || hostname.isEmpty()) {
2319
return null;
2420
}
@@ -39,7 +35,7 @@ public Attack run(String hostname, int port, List<String> ipAddresses, String op
3935
if(attackFindings != null) {
4036
return new Attack(
4137
operation,
42-
new Vulnerabilities.SSRFVulnerability(),
38+
Vulnerabilities.SSRF,
4339
attackFindings.source(),
4440
attackFindings.pathToPayload(),
4541
/*metadata*/ Map.of(

agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/ssrf/StoredSSRFDetector.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ public Attack run(String hostname, List<String> ipAddresses, String operation) {
2323

2424
return new Attack(
2525
operation,
26-
new Vulnerabilities.StoredSSRFVulnerability(),
26+
Vulnerabilities.STORED_SSRF,
2727
null, // source is null for stored attacks
2828
"", // path is empty
2929
Map.of(

0 commit comments

Comments
 (0)